Enabled Inner Product Space Type support for Lucene Engine

Signed-off-by: Navneet Verma <[email protected]>
navneet1v · Mar 17, 2024 · f225b43 · f225b43
1 parent 4734d88
commit f225b43
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 31 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Detect AVX2 Dynamically on the System [#1502](https://github.com/opensearch-project/k-NN/pull/1502)
 * Validate zero vector when using cosine metric [#1501](https://github.com/opensearch-project/k-NN/pull/1501)
 * Persist model definition in model metadata [#1527] (https://github.com/opensearch-project/k-NN/pull/1527)
+* Added Inner Product Space type support for Lucene Engine [#1551](https://github.com/opensearch-project/k-NN/pull/1551)
 ### Bug Fixes
 * Disable sdc table for HNSWPQ read-only indices [#1518](https://github.com/opensearch-project/k-NN/pull/1518)
 * Switch SpaceType.INNERPRODUCT's vector similarity function to MAXIMUM_INNER_PRODUCT [#1532](https://github.com/opensearch-project/k-NN/pull/1532)

diff --git a/src/main/java/org/opensearch/knn/index/util/Lucene.java b/src/main/java/org/opensearch/knn/index/util/Lucene.java
@@ -42,7 +42,7 @@ public class Lucene extends JVMLibrary {
                     )
                 )
                 .build()
-        ).addSpaces(SpaceType.L2, SpaceType.COSINESIMIL).build()
+        ).addSpaces(SpaceType.L2, SpaceType.COSINESIMIL, SpaceType.INNER_PRODUCT).build()
     );
 
     final static Lucene INSTANCE = new Lucene(METHODS, Version.LATEST.toString());

diff --git a/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java b/src/test/java/org/opensearch/knn/index/LuceneEngineIT.java
@@ -12,6 +12,7 @@
 import org.apache.commons.lang.math.RandomUtils;
 import org.apache.hc.core5.http.io.entity.EntityUtils;
 import org.apache.lucene.index.VectorSimilarityFunction;
+import org.apache.lucene.util.VectorUtil;
 import org.junit.After;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
@@ -80,36 +81,6 @@ public void testQuery_cosine() throws Exception {
         baseQueryTest(SpaceType.COSINESIMIL);
     }
 
-    public void testQuery_innerProduct_notSupported() throws Exception {
-        XContentBuilder builder = XContentFactory.jsonBuilder()
-            .startObject()
-            .startObject(PROPERTIES_FIELD_NAME)
-            .startObject(FIELD_NAME)
-            .field(TYPE_FIELD_NAME, KNN_VECTOR_TYPE)
-            .field(DIMENSION_FIELD_NAME, DIMENSION)
-            .startObject(KNNConstants.KNN_METHOD)
-            .field(KNNConstants.NAME, KNNEngine.LUCENE.getMethod(METHOD_HNSW).getMethodComponent().getName())
-            .field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, SpaceType.INNER_PRODUCT.getValue())
-            .field(KNNConstants.KNN_ENGINE, KNNEngine.LUCENE.getName())
-            .startObject(KNNConstants.PARAMETERS)
-            .field(KNNConstants.METHOD_PARAMETER_M, M)
-            .field(KNNConstants.METHOD_PARAMETER_EF_CONSTRUCTION, EF_CONSTRUCTION)
-            .endObject()
-            .endObject()
-            .endObject()
-            .endObject()
-            .endObject();
-
-        String mapping = builder.toString();
-
-        createIndex(INDEX_NAME, getKNNDefaultIndexSettings());
-
-        Request request = new Request("PUT", "/" + INDEX_NAME + "/_mapping");
-        request.setJsonEntity(mapping);
-
-        expectThrows(ResponseException.class, () -> client().performRequest(request));
-    }
-
     public void testQuery_invalidVectorDimensionInQuery() throws Exception {
 
         createKnnIndexMappingWithLuceneEngine(DIMENSION, SpaceType.L2, VectorDataType.FLOAT);
@@ -449,4 +420,53 @@ private void validateQueryResultsWithFilters(
                 .containsAll(expectedDocIdsKLimitsFilterResult)
         );
     }
+
+    @SneakyThrows
+    public void test_whenUsingIP_thenSuccess() {
+        XContentBuilder builder = XContentFactory.jsonBuilder()
+            .startObject()
+            .startObject("properties")
+            .startObject(FIELD_NAME)
+            .field("type", "knn_vector")
+            .field("dimension", 2)
+            .startObject(KNNConstants.KNN_METHOD)
+            .field(KNNConstants.NAME, KNNEngine.LUCENE.getMethod(KNNConstants.METHOD_HNSW).getMethodComponent().getName())
+            .field(KNNConstants.METHOD_PARAMETER_SPACE_TYPE, SpaceType.INNER_PRODUCT.getValue())
+            .field(KNNConstants.KNN_ENGINE, KNNEngine.LUCENE)
+            .endObject()
+            .endObject()
+            .endObject()
+            .endObject();
+        final String mapping = builder.toString();
+        createKnnIndex(INDEX_NAME, mapping);
+
+        final List<Float[]> dataVectors = Arrays.asList(new Float[] { -2.0f, 2.0f }, new Float[] { 2.0f, -2.0f });
+        final List<String> ids = Arrays.asList(DOC_ID, DOC_ID_2);
+
+        // Ingest all the documents
+        for (int i = 0; i < dataVectors.size(); i++) {
+            addKnnDoc(INDEX_NAME, ids.get(i), FIELD_NAME, dataVectors.get(i));
+        }
+        refreshIndex(INDEX_NAME);
+
+        float[] queryVector = new float[] { -2.0f, 2.0f };
+        int k = 2;
+        final Response response = searchKNNIndex(
+            INDEX_NAME,
+            new KNNQueryBuilder(FIELD_NAME, queryVector, k, QueryBuilders.matchAllQuery()),
+            k
+        );
+        final String responseBody = EntityUtils.toString(response.getEntity());
+        final List<Float> knnResults = parseSearchResponseScore(responseBody, FIELD_NAME);
+
+        // Check that the expected scores are returned
+        final List<Float> expectedScores = Arrays.asList(
+            VectorUtil.scaleMaxInnerProductScore(8.0f),
+            VectorUtil.scaleMaxInnerProductScore(-8.0f)
+        );
+        assertEquals(expectedScores.size(), knnResults.size());
+        for (int i = 0; i < expectedScores.size(); i++) {
+            assertEquals(expectedScores.get(i), knnResults.get(i), 0.0000001);
+        }
+    }
 }
diff --git a/src/test/java/org/opensearch/knn/index/util/LuceneTests.java b/src/test/java/org/opensearch/knn/index/util/LuceneTests.java
@@ -82,6 +82,20 @@ public void testLucenHNSWMethod() throws IOException {
         in = xContentBuilderToMap(xContentBuilder);
         KNNMethodContext knnMethodContext4 = KNNMethodContext.parse(in);
         assertNotNull(luceneHNSW.validate(knnMethodContext4));
+
+        // Check INNER_PRODUCT is supported with Lucene Engine
+        xContentBuilder = XContentFactory.jsonBuilder()
+            .startObject()
+            .field(NAME, METHOD_HNSW)
+            .field(METHOD_PARAMETER_SPACE_TYPE, SpaceType.INNER_PRODUCT.getValue())
+            .startObject(PARAMETERS)
+            .field(METHOD_PARAMETER_EF_CONSTRUCTION, efConstruction)
+            .field(METHOD_PARAMETER_M, m)
+            .endObject()
+            .endObject();
+        in = xContentBuilderToMap(xContentBuilder);
+        KNNMethodContext knnMethodContext5 = KNNMethodContext.parse(in);
+        assertNull(luceneHNSW.validate(knnMethodContext5));
     }
 
     public void testGetExtension() {