Fix condition check for values_count (qdrant#1502)

jankieAI · Mar 15, 2023 · a93a55e · a93a55e
1 parent 1275264
commit a93a55e
Show file tree

Hide file tree

Showing 2 changed files with 135 additions and 0 deletions.
diff --git a/lib/segment/src/payload_storage/condition_checker.rs b/lib/segment/src/payload_storage/condition_checker.rs
@@ -95,6 +95,10 @@ impl ValueChecker for ValuesCount {
     fn check_match(&self, payload: &Value) -> bool {
         self.check_count(payload)
     }
+
+    fn check(&self, payload: &Value) -> bool {
+        self.check_count(payload)
+    }
 }
 
 #[cfg(test)]
@@ -135,4 +139,40 @@ mod tests {
         assert!(near_berlin_query.check(&berlin_and_moscow));
         assert!(!miss_geo_query.check(&berlin_and_moscow));
     }
+
+    #[test]
+    fn test_value_count() {
+        let countries = json!([
+            {
+                "country": "Germany",
+            },
+            {
+                "country": "France",
+            }
+        ]);
+
+        let gt_one_country_query = ValuesCount {
+            lt: None,
+            gt: Some(1),
+            gte: None,
+            lte: None,
+        };
+        assert!(gt_one_country_query.check(&countries));
+
+        let gt_two_countries_query = ValuesCount {
+            lt: None,
+            gt: Some(2),
+            gte: None,
+            lte: None,
+        };
+        assert!(!gt_two_countries_query.check(&countries));
+
+        let gte_two_countries_query = ValuesCount {
+            lt: None,
+            gt: None,
+            gte: Some(2),
+            lte: None,
+        };
+        assert!(gte_two_countries_query.check(&countries));
+    }
 }
diff --git a/openapi/tests/openapi_integration/test_filter_values_count.py b/openapi/tests/openapi_integration/test_filter_values_count.py
@@ -0,0 +1,95 @@
+import random
+
+import pytest
+
+from .helpers.helpers import request_with_validation
+from .helpers.collection_setup import basic_collection_setup, drop_collection
+
+collection_name = 'test_collection_filter_values_count'
+
+
+@pytest.fixture(autouse=True)
+def setup():
+    basic_collection_setup(collection_name=collection_name)
+    yield
+    drop_collection(collection_name=collection_name)
+
+
+def test_filter_values_count():
+    response = request_with_validation(
+        api='/collections/{collection_name}/points/search',
+        method="POST",
+        path_params={'collection_name': collection_name},
+        body={
+            "vector": [0.2, 0.1, 0.9, 0.7],
+            "limit": 3,
+            "filter": {
+                "must": [
+                    {
+                        "key": "city",
+                        "values_count": {
+                            "gt": 2
+                        }
+                    }
+                ]
+            }
+        }
+    )
+    assert response.ok
+
+    json = response.json()
+    assert len(json['result']) == 0
+
+    response = request_with_validation(
+        api='/collections/{collection_name}/points/search',
+        method="POST",
+        path_params={'collection_name': collection_name},
+        body={
+            "vector": [0.2, 0.1, 0.9, 0.7],
+            "limit": 3,
+            "filter": {
+                "must": [
+                    {
+                        "key": "city",
+                        "values_count": {
+                            "gte": 2
+                        }
+                    }
+                ]
+            }
+        }
+    )
+    assert response.ok
+
+    json = response.json()
+    assert len(json['result']) == 3
+
+    ids = [x['id'] for x in json['result']]
+    assert 2 in ids
+    assert 3 in ids
+    assert 4 in ids
+
+    response = request_with_validation(
+        api='/collections/{collection_name}/points/search',
+        method="POST",
+        path_params={'collection_name': collection_name},
+        body={
+            "vector": [0.2, 0.1, 0.9, 0.7],
+            "limit": 3,
+            "filter": {
+                "must": [
+                    {
+                        "key": "city",
+                        "values_count": {
+                            "lt": 2
+                        }
+                    }
+                ]
+            }
+        }
+    )
+    assert response.ok
+
+    json = response.json()
+    assert len(json['result']) == 1
+    assert json['result'][0]['id'] == 1