Adding KNN codec that is based on Lucene92 codec (#444) (#446)

* Adding KNN codec that is based on Lucene92 codec Signed-off-by: Martin Gaievski <gaievski@amazon.com> (cherry picked from commit 7499375)
opensearch-project · Jul 14, 2022 · bb043e6 · bb043e6
1 parent 32395a3
commit bb043e6
Show file tree

Hide file tree

Showing 9 changed files with 126 additions and 94 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/codec/KNN920Codec/KNN920Codec.java b/src/main/java/org/opensearch/knn/index/codec/KNN920Codec/KNN920Codec.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.knn.index.codec.KNN920Codec;
+
+import lombok.Builder;
+import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.CompoundFormat;
+import org.apache.lucene.codecs.DocValuesFormat;
+import org.apache.lucene.codecs.FilterCodec;
+import org.opensearch.knn.index.codec.KNNFormatFacade;
+import org.opensearch.knn.index.codec.KNNFormatFactory;
+
+import static org.opensearch.knn.index.codec.KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate;
+
+/**
+ * KNN codec that is based on Lucene92 codec
+ */
+public final class KNN920Codec extends FilterCodec {
+
+    private static final String KNN920 = "KNN920Codec";
+    private final KNNFormatFacade knnFormatFacade;
+
+    /**
+     * No arg constructor that uses Lucene91 as the delegate
+     */
+    public KNN920Codec() {
+        this(createKNN92DefaultDelegate());
+    }
+
+    /**
+     * Constructor that takes a Codec delegate to delegate all methods this code does not implement to.
+     *
+     * @param delegate codec that will perform all operations this codec does not override
+     */
+    @Builder
+    public KNN920Codec(Codec delegate) {
+        super(KNN920, delegate);
+        knnFormatFacade = KNNFormatFactory.createKNN920Format(delegate);
+    }
+
+    @Override
+    public DocValuesFormat docValuesFormat() {
+        return knnFormatFacade.docValuesFormat();
+    }
+
+    @Override
+    public CompoundFormat compoundFormat() {
+        return knnFormatFacade.compoundFormat();
+    }
+}
diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNCodecFactory.java
@@ -4,42 +4,23 @@
  */
 package org.opensearch.knn.index.codec;
 
-import com.google.common.collect.ImmutableMap;
+import lombok.AllArgsConstructor;
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.backward_codecs.lucene91.Lucene91Codec;
+import org.apache.lucene.codecs.lucene92.Lucene92Codec;
 import org.opensearch.index.mapper.MapperService;
-import org.opensearch.knn.index.codec.util.CodecBuilder;
-
-import java.util.Map;
+import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;
 
 /**
  * Factory abstraction for KNN codec
  */
+@AllArgsConstructor
 public class KNNCodecFactory {
 
-    private final Map<KNNCodecVersion, CodecBuilder> codecByVersion;
-
-    private static final KNNCodecVersion LATEST_KNN_CODEC_VERSION = KNNCodecVersion.KNN910;
-
-    public KNNCodecFactory(MapperService mapperService) {
-        codecByVersion = ImmutableMap.of(KNNCodecVersion.KNN910, new CodecBuilder.KNN91CodecBuilder(mapperService));
-    }
+    private final MapperService mapperService;
 
     public Codec createKNNCodec(final Codec userCodec) {
-        return getCodec(LATEST_KNN_CODEC_VERSION, userCodec);
-    }
-
-    public Codec createKNNCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) {
-        return getCodec(knnCodecVersion, userCodec);
-    }
-
-    private Codec getCodec(final KNNCodecVersion knnCodecVersion, final Codec userCodec) {
-        try {
-            final CodecBuilder codecBuilder = codecByVersion.getOrDefault(knnCodecVersion, codecByVersion.get(LATEST_KNN_CODEC_VERSION));
-            return codecBuilder.userCodec(userCodec).build();
-        } catch (Exception ex) {
-            throw new RuntimeException("Cannot create instance of KNN codec", ex);
-        }
+        return KNN920Codec.builder().delegate(userCodec).build();
     }
 
     /**
@@ -50,12 +31,9 @@ public static class CodecDelegateFactory {
         public static Codec createKNN91DefaultDelegate() {
             return new Lucene91Codec();
         }
-    }
 
-    /**
-     * Collection of supported coded versions
-     */
-    enum KNNCodecVersion {
-        KNN910
+        public static Codec createKNN92DefaultDelegate() {
+            return new Lucene92Codec();
+        }
     }
 }
diff --git a/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java b/src/main/java/org/opensearch/knn/index/codec/KNNFormatFactory.java
@@ -13,11 +13,29 @@
  */
 public class KNNFormatFactory {
 
+    /**
+     * Return facade class that abstracts format specific to KNN910 codec
+     * @param delegate delegate codec that is wrapped by KNN codec
+     * @return
+     */
     public static KNNFormatFacade createKNN910Format(final Codec delegate) {
         final KNNFormatFacade knnFormatFacade = new KNNFormatFacade(
             new KNN80DocValuesFormat(delegate.docValuesFormat()),
             new KNN80CompoundFormat(delegate.compoundFormat())
         );
         return knnFormatFacade;
     }
+
+    /**
+     * Return facade class that abstracts format specific to KNN920 codec
+     * @param delegate delegate codec that is wrapped by KNN codec
+     * @return
+     */
+    public static KNNFormatFacade createKNN920Format(final Codec delegate) {
+        final KNNFormatFacade knnFormatFacade = new KNNFormatFacade(
+            new KNN80DocValuesFormat(delegate.docValuesFormat()),
+            new KNN80CompoundFormat(delegate.compoundFormat())
+        );
+        return knnFormatFacade;
+    }
 }
diff --git a/src/main/java/org/opensearch/knn/index/codec/util/CodecBuilder.java b/src/main/java/org/opensearch/knn/index/codec/util/CodecBuilder.java
diff --git a/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec b/src/main/resources/META-INF/services/org.apache.lucene.codecs.Codec
@@ -2,4 +2,5 @@ org.opensearch.knn.index.codec.KNN80Codec.KNN80Codec
 org.opensearch.knn.index.codec.KNN84Codec.KNN84Codec
 org.opensearch.knn.index.codec.KNN86Codec.KNN86Codec
 org.opensearch.knn.index.codec.KNN87Codec.KNN87Codec
-org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec
+org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec
+org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec
diff --git a/src/test/java/org/opensearch/knn/index/codec/KNN920Codec/KNN920CodecTests.java b/src/test/java/org/opensearch/knn/index/codec/KNN920Codec/KNN920CodecTests.java
@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.index.codec.KNN920Codec;
+
+import org.opensearch.knn.index.codec.KNNCodecTestCase;
+
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
+
+public class KNN920CodecTests extends KNNCodecTestCase {
+
+    public void testMultiFieldsKnnIndex() throws Exception {
+        testMultiFieldsKnnIndex(new KNN920Codec());
+    }
+
+    public void testBuildFromModelTemplate() throws InterruptedException, ExecutionException, IOException {
+        testBuildFromModelTemplate(new KNN920Codec());
+    }
+}
diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecFactoryTests.java
@@ -7,9 +7,10 @@
 
 import org.apache.lucene.codecs.Codec;
 import org.apache.lucene.backward_codecs.lucene91.Lucene91Codec;
+import org.apache.lucene.codecs.lucene92.Lucene92Codec;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.knn.KNNTestCase;
-import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec;
+import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;
 
 import static org.mockito.Mockito.mock;
 
@@ -21,21 +22,18 @@ public void testKNN91DefaultDelegate() {
         assertTrue(knn91DefaultDelegate instanceof Lucene91Codec);
     }
 
-    public void testKNN91DefaultCodec() {
-        Lucene91Codec lucene91CodecDelegate = new Lucene91Codec();
-        MapperService mapperService = mock(MapperService.class);
-        KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
-        Codec knnCodec = knnCodecFactory.createKNNCodec(lucene91CodecDelegate);
-        assertNotNull(knnCodec);
-        assertTrue(knnCodec instanceof KNN910Codec);
+    public void testKNN92DefaultDelegate() {
+        Codec knn92DefaultDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate();
+        assertNotNull(knn92DefaultDelegate);
+        assertTrue(knn92DefaultDelegate instanceof Lucene92Codec);
     }
 
-    public void testKNN91CodecByVersion() {
+    public void testKNNDefaultCodec() {
         Lucene91Codec lucene91CodecDelegate = new Lucene91Codec();
         MapperService mapperService = mock(MapperService.class);
         KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
-        Codec knnCodec = knnCodecFactory.createKNNCodec(KNNCodecFactory.KNNCodecVersion.KNN910, lucene91CodecDelegate);
+        Codec knnCodec = knnCodecFactory.createKNNCodec(lucene91CodecDelegate);
         assertNotNull(knnCodec);
-        assertTrue(knnCodec instanceof KNN910Codec);
+        assertTrue(knnCodec instanceof KNN920Codec);
     }
 }
diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java b/src/test/java/org/opensearch/knn/index/codec/KNNCodecTestCase.java
@@ -11,7 +11,7 @@
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.knn.KNNTestCase;
 import org.opensearch.knn.common.KNNConstants;
-import org.opensearch.knn.index.codec.KNN910Codec.KNN910Codec;
+import org.opensearch.knn.index.codec.KNN920Codec.KNN920Codec;
 import org.opensearch.knn.jni.JNIService;
 import org.opensearch.knn.index.KNNQuery;
 import org.opensearch.knn.index.KNNSettings;
@@ -61,7 +61,7 @@
  */
 public class KNNCodecTestCase extends KNNTestCase {
 
-    private static final KNN910Codec ACTUAL_CODEC = new KNN910Codec();
+    private static final KNN920Codec ACTUAL_CODEC = new KNN920Codec();
     private static FieldType sampleFieldType;
     static {
         sampleFieldType = new FieldType(KNNVectorFieldMapper.Defaults.FIELD_TYPE);

diff --git a/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java b/src/test/java/org/opensearch/knn/index/codec/KNNFormatFactoryTests.java
@@ -24,4 +24,16 @@ public void testKNN91Format() {
         assertNotNull(knnFormatFacade.compoundFormat());
         assertNotNull(knnFormatFacade.docValuesFormat());
     }
+
+    public void testKNN92Format() {
+        final Codec lucene92CodecDelegate = KNNCodecFactory.CodecDelegateFactory.createKNN92DefaultDelegate();
+        MapperService mapperService = mock(MapperService.class);
+        KNNCodecFactory knnCodecFactory = new KNNCodecFactory(mapperService);
+        final Codec knnCodec = knnCodecFactory.createKNNCodec(lucene92CodecDelegate);
+        KNNFormatFacade knnFormatFacade = KNNFormatFactory.createKNN920Format(knnCodec);
+
+        assertNotNull(knnFormatFacade);
+        assertNotNull(knnFormatFacade.compoundFormat());
+        assertNotNull(knnFormatFacade.docValuesFormat());
+    }
 }