Fix UnindexedField infra to handle lookup&range index requirements

agourlay · agourlay · commit af248be7ab3c · 2025-05-16T08:28:42.000+02:00
diff --git a/lib/collection/src/problems/unindexed_field.rs b/lib/collection/src/problems/unindexed_field.rs
@@ -309,6 +309,8 @@ impl<'a> IssueExtractor<'a> {
 pub struct Extractor<'a> {
     payload_schema: &'a HashMap<PayloadKeyType, PayloadFieldSchema>,
     unindexed_schema: HashMap<PayloadKeyType, Vec<PayloadFieldSchema>>,
+    needs_match_index: bool,
+    needs_range_index: bool,
 }
 
 impl<'a> Extractor<'a> {
@@ -320,6 +322,8 @@ impl<'a> Extractor<'a> {
         let mut extractor = Self {
             payload_schema,
             unindexed_schema: HashMap::new(),
+            needs_match_index: false,
+            needs_range_index: false,
         };
 
         extractor.update_from_filter(None, filter);
@@ -332,6 +336,8 @@ impl<'a> Extractor<'a> {
         Self {
             payload_schema,
             unindexed_schema: HashMap::new(),
+            needs_match_index: false,
+            needs_range_index: false,
         }
     }
 
@@ -364,6 +370,8 @@ impl<'a> Extractor<'a> {
         match condition {
             Condition::Field(field_condition) => {
                 key = &field_condition.key;
+                self.update_needs_match(field_condition.r#match.is_some());
+                self.update_needs_range(field_condition.range.is_some());
                 inferred = infer_schema_from_field_condition(field_condition);
             }
             Condition::Filter(filter) => {
@@ -402,25 +410,36 @@ impl<'a> Extractor<'a> {
         }
     }
 
+    /// Capture that an index with Match capability is required.
+    /// Does not toggle off if has been set once already.
+    fn update_needs_match(&mut self, is_match: bool) {
+        if !self.needs_match_index {
+            self.needs_match_index = is_match
+        }
+    }
+
+    /// Capture that an index with Match capability is required.
+    /// Does not toggle off if has been set once already.
+    fn update_needs_range(&mut self, is_range: bool) {
+        if !self.needs_range_index {
+            self.needs_range_index = is_range
+        }
+    }
+
     fn needs_index(&self, key: &JsonPath, inferred: &[PayloadFieldSchema]) -> bool {
         match self.payload_schema.get(key) {
             Some(index_info) => {
-                let index_info_kind = index_info.kind();
+                if self.needs_match_index && !index_info.supports_match() {
+                    return true;
+                }
 
+                if self.needs_range_index && !index_info.supports_range() {
+                    return true;
+                }
+
+                let index_info_kind = index_info.kind();
                 let already_indexed = inferred
                     .iter()
-                    // TODO(strict-mode):
-                    // Use better comparisons for parametrized indexes. An idea is to make the inferring step
-                    // also output valid parametrized indexes and compare those instead of just the kind (index type)
-                    //
-                    // The only reason why it would be needed is because integer index can be parametrized
-                    // with just lookup or just range, so it is possible to make a false negative here. E.g.
-                    //
-                    // condition: MatchValue
-                    // inferred: FieldType(Integer)
-                    // index_info: FieldParams(IntegerIndex(range))
-                    //
-                    // In this case, we would assume that the field is indexed correctly when it is not
                     .map(PayloadFieldSchema::kind)
                     .any(|inferred| inferred == index_info_kind);
 
diff --git a/lib/segment/src/types.rs b/lib/segment/src/types.rs
@@ -1898,6 +1898,34 @@ impl PayloadFieldSchema {
             },
         }
     }
+
+    /// Check if this type supports a `range` condition
+    pub fn supports_range(&self) -> bool {
+        match self {
+            PayloadFieldSchema::FieldType(payload_schema_type) => match payload_schema_type {
+                PayloadSchemaType::Float => true,
+                PayloadSchemaType::Integer => true,
+                PayloadSchemaType::Keyword => false,
+                PayloadSchemaType::Uuid => false,
+                PayloadSchemaType::Bool => false,
+                PayloadSchemaType::Geo => false,
+                PayloadSchemaType::Text => false,
+                PayloadSchemaType::Datetime => false,
+            },
+            PayloadFieldSchema::FieldParams(payload_schema_params) => match payload_schema_params {
+                PayloadSchemaParams::Float(_) => true,
+                PayloadSchemaParams::Integer(integer_index_params) => {
+                    integer_index_params.range == Some(true)
+                }
+                PayloadSchemaParams::Keyword(_) => false,
+                PayloadSchemaParams::Uuid(_) => false,
+                PayloadSchemaParams::Bool(_) => false,
+                PayloadSchemaParams::Geo(_) => false,
+                PayloadSchemaParams::Text(_) => false,
+                PayloadSchemaParams::Datetime(_) => false,
+            },
+        }
+    }
 }
 
 impl From<PayloadSchemaType> for PayloadFieldSchema {
diff --git a/tests/openapi/test_strictmode.py b/tests/openapi/test_strictmode.py
@@ -162,7 +162,7 @@ def search_request_with_timeout(timeout):
     assert not search_fail.ok
 
 
-def test_strict_mode_unindexed_filter_read_validation(collection_name):
+def test_strict_mode_unindexed_filter_keyword_read_validation(collection_name):
     def search_request_with_filter():
         return request_with_validation(
             api='/collections/{collection_name}/points/search',
@@ -217,6 +217,84 @@ def search_request_with_filter():
     search_request_with_filter().raise_for_status()
 
 
+def test_strict_mode_unindexed_filter_integer_read_validation(collection_name):
+    def search_request_with_filter():
+        return 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": "count",
+                            "match": {
+                                "value": 1
+                            }
+                        }
+                    ]
+                },
+            }
+        )
+
+    search_request_with_filter().raise_for_status()
+
+    set_strict_mode(collection_name, {
+        "enabled": True,
+        "unindexed_filtering_retrieve": True,
+    })
+
+    search_request_with_filter().raise_for_status()
+
+    set_strict_mode(collection_name, {
+        "unindexed_filtering_retrieve": False,
+    })
+
+    search_fail = search_request_with_filter()
+
+    assert "count" in search_fail.json()['status']['error']
+    assert not search_fail.ok
+
+    request_with_validation(
+        api='/collections/{collection_name}/index',
+        method="PUT",
+        path_params={'collection_name': collection_name},
+        query_params={'wait': 'true'},
+        body={
+            "field_name": "count",
+            "field_schema": {
+                "type": "integer",
+                "lookup": False,
+                "range": True,
+            }
+        }
+    ).raise_for_status()
+
+    # the integer index does not support `lookup` on our match condition
+    assert "count" in search_fail.json()['status']['error']
+    assert not search_fail.ok
+
+    request_with_validation(
+        api='/collections/{collection_name}/index',
+        method="PUT",
+        path_params={'collection_name': collection_name},
+        query_params={'wait': 'true'},
+        body={
+            "field_name": "count",
+            "field_schema": {
+                "type": "integer",
+                "lookup": True,
+                "range": True,
+            }
+        }
+    ).raise_for_status()
+    
+    # We created an index with `lookup` on this field so it should work now
+    search_request_with_filter().raise_for_status()
+ 
+
 def test_strict_mode_unindexed_filter_write_validation(collection_name):
     def update_request_with_filter():
         return request_with_validation(
@@ -1013,7 +1091,7 @@ def upsert_points(ids: list[int]):
     })
 
     for i in range(32):
-        print(upsert_points([3, 4]).json())
+        upsert_points([3, 4]).json()
 
     # Max limit has been reached and one of the next requests must fail. Due to cache it might not be the first call!
     for i in range(32):
@@ -1368,6 +1446,26 @@ def query_request():
     assert "discount_price" in query_fail.json()['status']['error']
     assert "formula expression" in query_fail.json()['status']['error']
 
+    # Create index on `discount_price`
+    request_with_validation(
+        api='/collections/{collection_name}/index',
+        method="PUT",
+        path_params={'collection_name': collection_name},
+        query_params={'wait': 'true'},
+        body={
+            "field_name": "discount_price",
+            "field_schema": {
+                "type": "integer",
+                "lookup": True,
+                "range": False,
+            }
+        }
+    ).raise_for_status()
+
+    # Query succeeds with the index
+    query_ok = query_request()
+    assert query_ok.ok
+
 
 def test_strict_mode_read_rate_limiting_small_replenish(collection_name):
     """
