From cb8eaf2fc6d2f82b0f190f6378f7c112018f9bfd Mon Sep 17 00:00:00 2001
From: Navneet Verma <navneev@amazon.com>
Date: Sun, 17 Mar 2024 01:08:59 -0700
Subject: [PATCH] Enabled Inner Product Space Type support for Lucene Engine

Signed-off-by: Navneet Verma <navneev@amazon.com>
---
 .../org/opensearch/knn/index/util/Lucene.java |  2 +-
 .../opensearch/knn/index/LuceneEngineIT.java  | 50 +++++++++++++++++++
 .../knn/index/util/LuceneTests.java           | 14 ++++++
 3 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/opensearch/knn/index/util/Lucene.java b/src/main/java/org/opensearch/knn/index/util/Lucene.java
index bfa6cb040b..63642ae2c0 100644
--- 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
index 562765e0c8..8b57e3af61 100644
--- 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;
@@ -449,4 +450,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
index 38cacffa4d..6de46b52d9 100644
--- 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() {