accessors for custom compression modes

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

Hide file tree

Showing 2 changed files with 19 additions and 4 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
@@ -27,9 +27,9 @@ 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;
+    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 CompressionMode zstdCompressionMode;
     private final CompressionMode zstdNoDictCompressionMode;
@@ -98,7 +98,7 @@ 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:
@@ -132,4 +132,9 @@ Lucene95CustomCodec.Mode getMode() {
     public int getCompressionLevel() {
         return compressionLevel;
     }
+
+    public CompressionMode getCompressionMode() {
+        return mode == Lucene95CustomCodec.Mode.ZSTD_NO_DICT ? zstdNoDictCompressionMode : zstdCompressionMode;
+    }
+
 }
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,14 @@ public void testZstdNoDictLucene95CustomCodecModeWithCompressionLevel() {
         assertEquals(randomCompressionLevel, lucene95CustomStoredFieldsFormat.getCompressionLevel());
     }
 
+    public void testCompressionModes(){
+        Lucene95CustomStoredFieldsFormat lucene95CustomStoredFieldsFormat = new Lucene95CustomStoredFieldsFormat();
+        assertTrue(lucene95CustomStoredFieldsFormat.getCompressionMode() instanceof ZstdCompressionMode);
+    }
+
+    public void testZstdNoDictCompressionModes(){
+        Lucene95CustomStoredFieldsFormat lucene95CustomStoredFieldsFormat = new Lucene95CustomStoredFieldsFormat(Lucene95CustomCodec.Mode.ZSTD_NO_DICT);
+        assertTrue(lucene95CustomStoredFieldsFormat.getCompressionMode() instanceof ZstdNoDictCompressionMode);
+    }
+
 }