Fix PostingsFormat in KNN Codec (#236) (#248)

* postings list fix
opendistro-for-elasticsearch · Oct 14, 2020 · a496778 · a496778
1 parent ed9c7dc
commit a496778
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/...ain/java/com/amazon/opendistroforelasticsearch/knn/index/codec/KNN86Codec/KNN86Codec.java b/...ain/java/com/amazon/opendistroforelasticsearch/knn/index/codec/KNN86Codec/KNN86Codec.java
@@ -44,6 +44,7 @@ public final class KNN86Codec extends Codec {
     private final DocValuesFormat perFieldDocValuesFormat;
     private final CompoundFormat compoundFormat;
     private Codec lucene86Codec;
+    private PostingsFormat postingsFormat = null;
 
     public static final String KNN_86 = "KNN86Codec";
     public static final String LUCENE_86 = "Lucene86"; // Lucene Codec to be used
@@ -82,9 +83,17 @@ public DocValuesFormat docValuesFormat() {
      * approach of manually overriding.
      */
 
+
+    public void setPostingsFormat(PostingsFormat postingsFormat) {
+        this.postingsFormat = postingsFormat;
+    }
+
     @Override
     public PostingsFormat postingsFormat() {
-        return getDelegatee().postingsFormat();
+        if (this.postingsFormat == null) {
+            return getDelegatee().postingsFormat();
+        }
+        return this.postingsFormat;
     }
 
     @Override

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/KNNCodecService.java b/src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/KNNCodecService.java
@@ -16,7 +16,9 @@
 
 package com.amazon.opendistroforelasticsearch.knn.plugin;
 
+import com.amazon.opendistroforelasticsearch.knn.index.codec.KNN86Codec.KNN86Codec;
 import org.apache.lucene.codecs.Codec;
+import org.apache.lucene.codecs.PostingsFormat;
 import org.elasticsearch.index.codec.CodecService;
 
 import static com.amazon.opendistroforelasticsearch.knn.index.codec.KNN86Codec.KNN86Codec.KNN_86;
@@ -35,7 +37,7 @@ class KNNCodecService extends CodecService {
      * return the KNN Codec
      *
      * @param name dummy name
-     * @return KNN80Codec
+     * @return KNN86Codec
      */
     @Override
     public Codec codec(String name) {
@@ -45,4 +47,8 @@ public Codec codec(String name) {
         }
         return codec;
     }
+
+    public void setPostingsFormat(PostingsFormat postingsFormat) {
+        ((KNN86Codec)codec("")).setPostingsFormat(postingsFormat);
+    }
 }
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/KNNEngineFactory.java b/src/main/java/com/amazon/opendistroforelasticsearch/knn/plugin/KNNEngineFactory.java
@@ -16,7 +16,6 @@
 
 package com.amazon.opendistroforelasticsearch.knn.plugin;
 
-import org.elasticsearch.index.codec.CodecService;
 import org.elasticsearch.index.engine.Engine;
 import org.elasticsearch.index.engine.EngineConfig;
 import org.elasticsearch.index.engine.EngineFactory;
@@ -27,10 +26,11 @@
  */
 class KNNEngineFactory implements EngineFactory {
 
-    private static CodecService codecService = new KNNCodecService();
+    private static KNNCodecService codecService = new KNNCodecService();
 
     @Override
     public Engine newReadWriteEngine(EngineConfig config) {
+        codecService.setPostingsFormat(config.getCodec().postingsFormat());
         EngineConfig engineConfig = new EngineConfig(config.getShardId(), config.getAllocationId(),
                 config.getThreadPool(), config.getIndexSettings(), config.getWarmer(), config.getStore(),
                 config.getMergePolicy(), config.getAnalyzer(), config.getSimilarity(), codecService,