From 36ed5633b60d88a651f64b7e6667694be87981bb Mon Sep 17 00:00:00 2001
From: Dooyong Kim <kdooyong@amazon.com>
Date: Mon, 2 Dec 2024 16:14:50 -0800
Subject: [PATCH] Fixed a bug that does not get a space type from index setting
 when it is empty for compatibility.

 Signed-off-by: Dooyong Kim <kdooyong@amazon.com>
---
 .../knn/index/engine/SpaceTypeResolver.java   | 20 +++++++++++--------
 .../index/mapper/KNNVectorFieldMapper.java    |  7 +++++--
 .../mapper/KNNVectorFieldMapperUtil.java      |  4 +---
 .../plugin/rest/RestTrainModelHandler.java    |  1 +
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/main/java/org/opensearch/knn/index/engine/SpaceTypeResolver.java b/src/main/java/org/opensearch/knn/index/engine/SpaceTypeResolver.java
index a12ffbc7b..dbfec33d0 100644
--- a/src/main/java/org/opensearch/knn/index/engine/SpaceTypeResolver.java
+++ b/src/main/java/org/opensearch/knn/index/engine/SpaceTypeResolver.java
@@ -42,7 +42,7 @@ public SpaceType resolveSpaceType(
         SpaceType topLevelSpaceType = getSpaceTypeFromString(topLevelSpaceTypeString);
 
         if (isSpaceTypeConfigured(methodSpaceType) == false && isSpaceTypeConfigured(topLevelSpaceType) == false) {
-            return getSpaceTypeFromVectorDataType(vectorDataType);
+            return SpaceType.UNDEFINED;
         }
 
         if (isSpaceTypeConfigured(methodSpaceType) == false) {
@@ -75,13 +75,6 @@ private SpaceType getSpaceTypeFromMethodContext(final KNNMethodContext knnMethod
         return knnMethodContext.getSpaceType();
     }
 
-    private SpaceType getSpaceTypeFromVectorDataType(final VectorDataType vectorDataType) {
-        if (vectorDataType == VectorDataType.BINARY) {
-            return SpaceType.DEFAULT_BINARY;
-        }
-        return SpaceType.DEFAULT;
-    }
-
     private SpaceType getSpaceTypeFromString(final String spaceType) {
         if (Strings.isEmpty(spaceType)) {
             return SpaceType.UNDEFINED;
@@ -93,4 +86,15 @@ private SpaceType getSpaceTypeFromString(final String spaceType) {
     private boolean isSpaceTypeConfigured(final SpaceType spaceType) {
         return spaceType != null && spaceType != SpaceType.UNDEFINED;
     }
+
+    public SpaceType pickDefaultSpaceTypeWhenEmpty(SpaceType resolvedSpaceType, VectorDataType vectorDataType) {
+        if (resolvedSpaceType != SpaceType.UNDEFINED) {
+            return resolvedSpaceType;
+        } else {
+            if (vectorDataType == VectorDataType.BINARY) {
+                return SpaceType.DEFAULT_BINARY;
+            }
+            return SpaceType.DEFAULT;
+        }
+    }
 }
diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java
index 6e5138a56..0f9fa0165 100644
--- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java
+++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapper.java
@@ -375,12 +375,15 @@ public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserCont
                 // If the original knnMethodContext is not null, resolve its space type and engine from the rest of the
                 // configuration. This is consistent with the existing behavior for space type in 2.16 where we modify the
                 // parsed value
-                SpaceType resolvedSpaceType = SpaceTypeResolver.INSTANCE.resolveSpaceType(
+                final SpaceType resolvedSpaceType = SpaceTypeResolver.INSTANCE.resolveSpaceType(
                     builder.originalParameters.getKnnMethodContext(),
                     builder.vectorDataType.get(),
                     builder.topLevelSpaceType.get()
                 );
-                setSpaceType(builder.originalParameters.getKnnMethodContext(), resolvedSpaceType);
+                setSpaceType(
+                    builder.originalParameters.getKnnMethodContext(),
+                    SpaceTypeResolver.INSTANCE.pickDefaultSpaceTypeWhenEmpty(resolvedSpaceType, builder.vectorDataType.get())
+                );
                 validateSpaceType(builder);
                 resolveKNNMethodComponents(builder, parserContext, resolvedSpaceType);
                 validateFromKNNMethod(builder);
diff --git a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java
index b3727f2ef..607bc476a 100644
--- a/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java
+++ b/src/main/java/org/opensearch/knn/index/mapper/KNNVectorFieldMapperUtil.java
@@ -199,9 +199,7 @@ static KNNMethodContext createKNNMethodContextFromLegacy(
         SpaceType topLevelSpaceType
     ) {
         // If top level spaceType is set then use that spaceType otherwise default to spaceType from index-settings
-        final SpaceType finalSpaceToSet = topLevelSpaceType != SpaceType.UNDEFINED
-            ? topLevelSpaceType
-            : KNNVectorFieldMapperUtil.getSpaceType(indexSettings);
+        final SpaceType finalSpaceToSet = topLevelSpaceType != SpaceType.UNDEFINED ? topLevelSpaceType : getSpaceType(indexSettings);
         return new KNNMethodContext(
             KNNEngine.NMSLIB,
             finalSpaceToSet,
diff --git a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java
index 71f7201de..b90adcf31 100644
--- a/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java
+++ b/src/main/java/org/opensearch/knn/plugin/rest/RestTrainModelHandler.java
@@ -165,6 +165,7 @@ private TrainingModelRequest createTransportRequest(RestRequest restRequest) thr
             vectorDataType,
             topLevelSpaceType.getValue()
         );
+        resolvedSpaceType = SpaceTypeResolver.INSTANCE.pickDefaultSpaceTypeWhenEmpty(resolvedSpaceType, vectorDataType);
         setSpaceType(knnMethodContext, resolvedSpaceType);
         TrainingModelRequest trainingModelRequest = new TrainingModelRequest(
             modelId,