accessors for custom compression modes

Signed-off-by: Sarthak Aggarwal <sarthagg@amazon.com>
opensearch-project · Dec 13, 2023 · d77955a · d77955a
1 parent 82a4fea
commit d77955a
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 11 deletions.
diff --git a/src/main/java/org/opensearch/index/codec/customcodecs/Lucene95CustomStoredFieldsFormat.java b/src/main/java/org/opensearch/index/codec/customcodecs/Lucene95CustomStoredFieldsFormat.java
@@ -19,6 +19,8 @@
 import org.apache.lucene.store.IOContext;
 
 import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
 /** Stored field format used by pluggable codec */
@@ -27,15 +29,13 @@ public class Lucene95CustomStoredFieldsFormat extends StoredFieldsFormat {
     /** A key that we use to map to a mode */
     public static final String MODE_KEY = Lucene95CustomStoredFieldsFormat.class.getSimpleName() + ".mode";
 
-    private static final int ZSTD_BLOCK_LENGTH = 10 * 48 * 1024;
-    private static final int ZSTD_MAX_DOCS_PER_BLOCK = 4096;
-    private static final int ZSTD_BLOCK_SHIFT = 10;
-
-    private final CompressionMode zstdCompressionMode;
-    private final CompressionMode zstdNoDictCompressionMode;
+    public static final int ZSTD_BLOCK_LENGTH = 10 * 48 * 1024;
+    public static final int ZSTD_MAX_DOCS_PER_BLOCK = 4096;
+    public static final int ZSTD_BLOCK_SHIFT = 10;
 
     private final Lucene95CustomCodec.Mode mode;
     private final int compressionLevel;
+    Map<Lucene95CustomCodec.Mode, CompressionMode> compressionModeMap;
 
     /** default constructor */
     public Lucene95CustomStoredFieldsFormat() {
@@ -60,8 +60,8 @@ public Lucene95CustomStoredFieldsFormat(Lucene95CustomCodec.Mode mode) {
     public Lucene95CustomStoredFieldsFormat(Lucene95CustomCodec.Mode mode, int compressionLevel) {
         this.mode = Objects.requireNonNull(mode);
         this.compressionLevel = compressionLevel;
-        zstdCompressionMode = new ZstdCompressionMode(compressionLevel);
-        zstdNoDictCompressionMode = new ZstdNoDictCompressionMode(compressionLevel);
+        this.compressionModeMap = new HashMap<>();
+        populateCompressionModeMap(compressionLevel);
     }
 
     /**
@@ -98,21 +98,21 @@ public StoredFieldsWriter fieldsWriter(Directory directory, SegmentInfo si, IOCo
         return impl(mode).fieldsWriter(directory, si, context);
     }
 
-    StoredFieldsFormat impl(Lucene95CustomCodec.Mode mode) {
+    public StoredFieldsFormat impl(Lucene95CustomCodec.Mode mode) {
         switch (mode) {
             case ZSTD:
             case ZSTD_DEPRECATED:
                 return new Lucene90CompressingStoredFieldsFormat(
                     "CustomStoredFieldsZstd",
-                    zstdCompressionMode,
+                    compressionModeMap.get(Lucene95CustomCodec.Mode.ZSTD),
                     ZSTD_BLOCK_LENGTH,
                     ZSTD_MAX_DOCS_PER_BLOCK,
                     ZSTD_BLOCK_SHIFT
                 );
             case ZSTD_NO_DICT:
                 return new Lucene90CompressingStoredFieldsFormat(
                     "CustomStoredFieldsZstdNoDict",
-                    zstdNoDictCompressionMode,
+                    compressionModeMap.get(Lucene95CustomCodec.Mode.ZSTD_NO_DICT),
                     ZSTD_BLOCK_LENGTH,
                     ZSTD_MAX_DOCS_PER_BLOCK,
                     ZSTD_BLOCK_SHIFT
@@ -132,4 +132,14 @@ Lucene95CustomCodec.Mode getMode() {
     public int getCompressionLevel() {
         return compressionLevel;
     }
+
+    public Map<Lucene95CustomCodec.Mode, CompressionMode> getCompressionModes(){
+        return this.compressionModeMap;
+    }
+
+    private void populateCompressionModeMap(int compressionLevel) {
+        this.compressionModeMap.put(Lucene95CustomCodec.Mode.ZSTD, new ZstdCompressionMode(compressionLevel));
+        this.compressionModeMap.put(Lucene95CustomCodec.Mode.ZSTD_NO_DICT, new ZstdNoDictCompressionMode(compressionLevel));
+    }
+
 }
diff --git a/...t/java/org/opensearch/index/codec/customcodecs/Lucene95CustomStoredFieldsFormatTests.java b/...t/java/org/opensearch/index/codec/customcodecs/Lucene95CustomStoredFieldsFormatTests.java
@@ -44,4 +44,9 @@ public void testZstdNoDictLucene95CustomCodecModeWithCompressionLevel() {
         assertEquals(randomCompressionLevel, lucene95CustomStoredFieldsFormat.getCompressionLevel());
     }
 
+    public void testCompressionModes(){
+        Lucene95CustomStoredFieldsFormat lucene95CustomStoredFieldsFormat = new Lucene95CustomStoredFieldsFormat();
+        assertTrue(lucene95CustomStoredFieldsFormat.getCompressionModes().get(Lucene95CustomCodec.Mode.ZSTD) instanceof ZstdCompressionMode);
+        assertTrue(lucene95CustomStoredFieldsFormat.getCompressionModes().get(Lucene95CustomCodec.Mode.ZSTD_NO_DICT) instanceof ZstdNoDictCompressionMode);
+    }
 }