Vector formats in segments supported for all codecs

tteofili · Apr 24, 2024 · 5b73540 · 5b73540
1 parent f65fd7d
commit 5b73540
Showing 1 changed file with 16 additions and 3 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/engine/Engine.java b/server/src/main/java/org/elasticsearch/index/engine/Engine.java
@@ -58,7 +58,8 @@
 import org.elasticsearch.core.UpdateForV9;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.VersionType;
-import org.elasticsearch.index.codec.PerFieldMapperCodec;
+import org.elasticsearch.index.codec.Elasticsearch814Codec;
+import org.elasticsearch.index.codec.LegacyPerFieldMapperCodec;
 import org.elasticsearch.index.mapper.DocumentParser;
 import org.elasticsearch.index.mapper.IdFieldMapper;
 import org.elasticsearch.index.mapper.LuceneDocument;
@@ -1096,14 +1097,14 @@ private void fillSegmentInfo(
         segment.attributes.putAll(info.info.getAttributes());
         Codec codec = info.info.getCodec();
         Map<String, List<String>> knnFormats = null;
-        if (includeVectorFormatsInfo && codec instanceof PerFieldMapperCodec) {
+        if (includeVectorFormatsInfo) {
             try {
                 FieldInfos fieldInfos = segmentReader.getFieldInfos();
                 if (fieldInfos.hasVectorValues()) {
                     for (FieldInfo fieldInfo : fieldInfos) {
                         String name = fieldInfo.getName();
                         if (fieldInfo.hasVectorValues()) {
-                            KnnVectorsFormat knnVectorsFormatForField = ((PerFieldMapperCodec) codec).getKnnVectorsFormatForField(name);
+                            KnnVectorsFormat knnVectorsFormatForField = getKnnVectorsFormatForField(codec, name);
                             if (knnFormats == null) {
                                 knnFormats = new HashMap<>();
                             }
@@ -1130,6 +1131,18 @@ private void fillSegmentInfo(
         segments.put(info.info.name, segment);
     }
 
+    private static KnnVectorsFormat getKnnVectorsFormatForField(Codec codec, String name) {
+        KnnVectorsFormat format;
+        if (codec instanceof Elasticsearch814Codec) {
+            format = ((Elasticsearch814Codec) codec).getKnnVectorsFormatForField(name);
+        } else if (codec instanceof LegacyPerFieldMapperCodec) {
+            format = ((LegacyPerFieldMapperCodec) codec).getKnnVectorsFormatForField(name);
+        } else {
+            format = codec.knnVectorsFormat();
+        }
+        return format;
+    }
+
     /**
      * The list of segments in the engine.
      */