Add 64x compression

Signed-off-by: John Mazanec <[email protected]>
opensearch-project · Oct 15, 2024 · 78915d1 · 78915d1
1 parent 8f76a48
commit 78915d1
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 15 deletions.
diff --git a/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java b/src/main/java/org/opensearch/knn/index/mapper/CompressionLevel.java
@@ -27,20 +27,10 @@ public enum CompressionLevel {
     x4(4, "4x", null, Collections.emptySet()),
     x8(8, "8x", new RescoreContext(2.0f, false), Set.of(Mode.ON_DISK)),
     x16(16, "16x", new RescoreContext(3.0f, false), Set.of(Mode.ON_DISK)),
-    x32(32, "32x", new RescoreContext(3.0f, false), Set.of(Mode.ON_DISK));
+    x32(32, "32x", new RescoreContext(3.0f, false), Set.of(Mode.ON_DISK)),
+    x64(64, "64x", new RescoreContext(5.0f, false), Set.of(Mode.ON_DISK));
 
-    public static final CompressionLevel MAX_COMPRESSION_LEVEL = CompressionLevel.x32;
-
-    // Internally, an empty string is easier to deal with them null. However, from the mapping,
-    // we do not want users to pass in the empty string and instead want null. So we make the conversion here
-    public static final String[] NAMES_ARRAY = new String[] {
-        NOT_CONFIGURED.getName(),
-        x1.getName(),
-        x2.getName(),
-        x4.getName(),
-        x8.getName(),
-        x16.getName(),
-        x32.getName() };
+    public static final CompressionLevel MAX_COMPRESSION_LEVEL = CompressionLevel.x64;
 
     /**
      * Default is set to 1x and is a noop

diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java
@@ -15,6 +15,7 @@
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
+import com.google.common.annotations.VisibleForTesting;
 import lombok.Getter;
 import lombok.Setter;
 import lombok.extern.log4j.Log4j2;
@@ -75,6 +76,17 @@ private static KNNVectorFieldMapper toType(FieldMapper in) {
         return (KNNVectorFieldMapper) in;
     }
 
+    // Supported compression levels for knn_vector field type
+    @VisibleForTesting
+    public static final String[] MAPPING_COMPRESSION_NAMES_ARRAY = new String[] {
+        CompressionLevel.NOT_CONFIGURED.getName(),
+        CompressionLevel.x1.getName(),
+        CompressionLevel.x2.getName(),
+        CompressionLevel.x4.getName(),
+        CompressionLevel.x8.getName(),
+        CompressionLevel.x16.getName(),
+        CompressionLevel.x32.getName() };
+
     /**
      * Builder for KNNVectorFieldMapper. This class defines the set of parameters that can be applied to the knn_vector
      * field type
@@ -161,7 +173,7 @@ public static class Builder extends ParametrizedFieldMapper.Builder {
             KNNConstants.COMPRESSION_LEVEL_PARAMETER,
             false,
             m -> toType(m).originalMappingParameters.getCompressionLevel(),
-            CompressionLevel.NAMES_ARRAY
+            MAPPING_COMPRESSION_NAMES_ARRAY
         ).acceptsNull();
 
         // A top level space Type field.

diff --git a/src/test/java/org/opensearch/knn/integ/ModeAndCompressionIT.java b/src/test/java/org/opensearch/knn/integ/ModeAndCompressionIT.java
@@ -39,6 +39,7 @@
 import static org.opensearch.knn.common.KNNConstants.TRAIN_FIELD_PARAMETER;
 import static org.opensearch.knn.common.KNNConstants.TRAIN_INDEX_PARAMETER;
 import static org.opensearch.knn.common.KNNConstants.VECTOR_DATA_TYPE_FIELD;
+import static org.opensearch.knn.index.mapper.KNNVectorFieldMapper.MAPPING_COMPRESSION_NAMES_ARRAY;
 
 public class ModeAndCompressionIT extends KNNRestTestCase {
 
@@ -253,7 +254,7 @@ public void testTraining_whenInvalid_thenFail() {
     public void testTraining_whenValid_thenSucceed() {
         setupTrainingIndex();
         XContentBuilder builder;
-        for (String compressionLevel : CompressionLevel.NAMES_ARRAY) {
+        for (String compressionLevel : MAPPING_COMPRESSION_NAMES_ARRAY) {
             if (compressionLevel.equals("4x")) {
                 continue;
             }