Merge pull request #32487 from vespa-engine/toregge/allow-more-than-o…

…ne-mapped-dimension-for-tensor-attribute-with-hnsw-index Allow more than one mapped dimension for tensor attribute with HNSW i…
vespa-engine · Sep 30, 2024 · 1d739a2 · 1d739a2
2 parents fa393bb + 8fc1031
commit 1d739a2
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 24 deletions.
diff --git a/config-model/src/main/java/com/yahoo/schema/processing/TensorFieldProcessor.java b/config-model/src/main/java/com/yahoo/schema/processing/TensorFieldProcessor.java
@@ -60,20 +60,8 @@ private boolean isTensorTypeThatSupportsHnswIndex(ImmutableSDField field) {
      */
     public static boolean isTensorTypeThatSupportsHnswIndex(TensorType type) {
         // Tensors with 1 indexed dimension support hnsw index (used for approximate nearest neighbor search).
-        if ((type.dimensions().size() == 1) &&
-                type.dimensions().get(0).isIndexed()) {
-            return true;
-        }
-        // Tensors with 1 mapped + 1 indexed dimension support hnsw index (aka multiple vectors per document).
-        if (type.dimensions().size() == 2) {
-            var a = type.dimensions().get(0);
-            var b = type.dimensions().get(1);
-            if ((a.isMapped() && b.isIndexed()) ||
-                    (a.isIndexed() && b.isMapped())) {
-                return true;
-            }
-        }
-        return false;
+        var indexedSubtype = type.indexedSubtype();
+        return (indexedSubtype.rank() == 1 && indexedSubtype.hasOnlyIndexedBoundDimensions());
     }
 
     private boolean isTensorTypeThatSupportsDirectStore(ImmutableSDField field) {

diff --git a/config-model/src/test/java/com/yahoo/schema/processing/TensorFieldTestCase.java b/config-model/src/test/java/com/yahoo/schema/processing/TensorFieldTestCase.java
@@ -91,6 +91,13 @@ void tensor_with_one_mapped_and_one_indexed_dimension_can_have_hnsw_index() thro
         assertHnswIndexParams("tensor(x[64],y{})", "", 16, 200);
     }
 
+    @Test
+    void tensor_with_two_mapped_and_one_indexed_dimension_can_have_hnsw_index() throws ParseException {
+        assertHnswIndexParams("tensor(x{},y{},z[64])", "", 16, 200);
+        assertHnswIndexParams("tensor(x{},y[64],z{})", "", 16, 200);
+        assertHnswIndexParams("tensor(x[64],y{},z{})", "", 16, 200);
+    }
+
     @Test
     void hnsw_index_parameters_can_be_specified() throws ParseException {
         assertHnswIndexParams("index { hnsw { max-links-per-node: 32 } }", 32, 200);

diff --git a/...iner-search/src/main/java/com/yahoo/search/searchers/ValidateNearestNeighborSearcher.java b/...iner-search/src/main/java/com/yahoo/search/searchers/ValidateNearestNeighborSearcher.java
@@ -104,16 +104,8 @@ private static boolean badQueryTensorType(TensorType queryTensorType) {
         }
 
         private static boolean isTensorTypeThatSupportsHnswIndex(TensorType tt) {
-            List<TensorType.Dimension> dims = tt.dimensions();
-            if (dims.size() == 1) {
-                return dims.get(0).isIndexed();
-            }
-            if (dims.size() == 2) {
-                var dims0 = dims.get(0);
-                var dims1 = dims.get(1);
-                return ((dims0.isMapped() && dims1.isIndexed()) || (dims0.isIndexed() && dims1.isMapped()));
-            }
-            return false;
+            var indexedSubtype = tt.indexedSubtype();
+            return (indexedSubtype.rank() == 1 && indexedSubtype.hasOnlyIndexedBoundDimensions());
         }
 
         /** Returns an error message if this is invalid, or null if it is valid */