Enhance semantic field to allow to enable/disable chunking. (opensearch-project#1337)

* Implement the query logic for the semantic field.

Signed-off-by: Bo Zhang <[email protected]>

* Enhance semantic field to allow to enable/disable chunking.

Signed-off-by: Bo Zhang <[email protected]>

---------

Signed-off-by: Bo Zhang <[email protected]>

Author: bzhangam

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Semantic Field] Add semantic mapping transformer. ([#1276](https://github.com/opensearch-project/neural-search/pull/1276))
 - [Semantic Field] Add semantic ingest processor. ([#1309](https://github.com/opensearch-project/neural-search/pull/1309))
 - [Semantic Field] Implement the query logic for the semantic field. ([#1315](https://github.com/opensearch-project/neural-search/pull/1315))
+- [Semantic Field] Enhance semantic field to allow to enable/disable chunking. ([#1337](https://github.com/opensearch-project/neural-search/pull/1337))
 
 ### Enhancements
 - [Performance Improvement] Add custom bulk scorer for hybrid query (2-3x faster) ([#1289](https://github.com/opensearch-project/neural-search/pull/1289))

src/main/java/org/opensearch/neuralsearch/constants/SemanticFieldConstants.java

@@ -35,4 +35,10 @@ public class SemanticFieldConstants {
      * Default suffix for semantic info field name. It will be used to construct the field name of the semantic info.
      */
     public static final String DEFAULT_SEMANTIC_INFO_FIELD_NAME_SUFFIX = "_semantic_info";
+
+    /**
+     * Name of the field to control if we should do chunking for the semantic field. By default, the chunking is
+     * disabled to not downgrade the search performance.
+     */
+    public static final String CHUNKING = "chunking";
 }

src/main/java/org/opensearch/neuralsearch/constants/SemanticInfoFieldConstants.java

@@ -23,7 +23,7 @@ public class SemanticInfoFieldConstants {
 
     public static final String CHUNKS_FIELD_NAME = "chunks";
     public static final String CHUNKS_TEXT_FIELD_NAME = "text";
-    public static final String CHUNKS_EMBEDDING_FIELD_NAME = "embedding";
+    public static final String EMBEDDING_FIELD_NAME = "embedding";
 
     public static final String MODEL_FIELD_NAME = "model";
     public static final String MODEL_ID_FIELD_NAME = "id";

src/main/java/org/opensearch/neuralsearch/mapper/SemanticFieldMapper.java

@@ -30,6 +30,9 @@
 import java.util.Map;
 import java.util.Set;
 
+import static org.opensearch.neuralsearch.constants.MappingConstants.PATH_SEPARATOR;
+import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.CHUNKING;
+import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.DEFAULT_SEMANTIC_INFO_FIELD_NAME_SUFFIX;
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.MODEL_ID;
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.RAW_FIELD_TYPE;
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.SEARCH_MODEL_ID;
@@ -127,6 +130,14 @@ public static class Builder extends ParametrizedFieldMapper.Builder {
             null
         );
 
+        @Getter
+        protected final Parameter<Boolean> chunkingEnabled = Parameter.boolParam(
+            CHUNKING,
+            false,
+            m -> ((SemanticFieldMapper) m).semanticParameters.getChunkingEnabled(),
+            false
+        );
+
         @Setter
         protected ParametrizedFieldMapper.Builder delegateBuilder;
 
@@ -136,7 +147,7 @@ protected Builder(String name) {
 
         @Override
         protected List<Parameter<?>> getParameters() {
-            return List.of(modelId, searchModelId, rawFieldType, semanticInfoFieldName);
+            return List.of(modelId, searchModelId, rawFieldType, semanticInfoFieldName, chunkingEnabled);
         }
 
         @Override
@@ -157,12 +168,13 @@ public SemanticFieldMapper build(BuilderContext context) {
         }
 
         public SemanticParameters getSemanticParameters() {
-            return new SemanticParameters(
-                modelId.getValue(),
-                searchModelId.getValue(),
-                rawFieldType.getValue(),
-                semanticInfoFieldName.getValue()
-            );
+            return SemanticParameters.builder()
+                .modelId(modelId.getValue())
+                .searchModelId(searchModelId.getValue())
+                .rawFieldType(rawFieldType.getValue())
+                .semanticInfoFieldName(semanticInfoFieldName.getValue())
+                .chunkingEnabled(chunkingEnabled.getValue())
+                .build();
         }
     }
 
@@ -249,6 +261,15 @@ public SemanticFieldType(@NonNull final MappedFieldType delegate, @NonNull final
         public String typeName() {
             return SemanticFieldMapper.CONTENT_TYPE;
         }
+
+        public String getSemanticInfoFieldPath() {
+            final String[] paths = name().split("\\.");
+            final String semanticInfoFieldName = semanticParameters.getSemanticInfoFieldName();
+            paths[paths.length - 1] = semanticInfoFieldName == null
+                ? paths[paths.length - 1] + DEFAULT_SEMANTIC_INFO_FIELD_NAME_SUFFIX
+                : semanticInfoFieldName;
+            return String.join(PATH_SEPARATOR, paths);
+        }
     }
 
     @Override

src/main/java/org/opensearch/neuralsearch/mapper/dto/SemanticParameters.java

@@ -4,22 +4,18 @@
  */
 package org.opensearch.neuralsearch.mapper.dto;
 
+import lombok.Builder;
 import lombok.Getter;
 
 /**
  * A DTO to hold all the semantic parameters.
  */
 @Getter
+@Builder
 public class SemanticParameters {
     private final String modelId;
     private final String searchModelId;
     private final String rawFieldType;
     private final String semanticInfoFieldName;
-
-    public SemanticParameters(String modelId, String searchModelId, String rawFieldType, String semanticInfoFieldName) {
-        this.modelId = modelId;
-        this.searchModelId = searchModelId;
-        this.rawFieldType = rawFieldType;
-        this.semanticInfoFieldName = semanticInfoFieldName;
-    }
+    private final Boolean chunkingEnabled;
 }

src/main/java/org/opensearch/neuralsearch/mappingtransformer/SemanticInfoConfigBuilder.java

@@ -5,7 +5,6 @@
 package org.opensearch.neuralsearch.mappingtransformer;
 
 import org.opensearch.core.xcontent.NamedXContentRegistry;
-import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.ObjectMapper;
 import org.opensearch.index.mapper.RankFeaturesFieldMapper;
 import org.opensearch.index.mapper.TextFieldMapper;
@@ -22,7 +21,7 @@
 
 import static org.opensearch.neuralsearch.constants.MappingConstants.PROPERTIES;
 import static org.opensearch.neuralsearch.constants.MappingConstants.TYPE;
-import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.CHUNKS_EMBEDDING_FIELD_NAME;
+import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.EMBEDDING_FIELD_NAME;
 import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.CHUNKS_FIELD_NAME;
 import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.CHUNKS_TEXT_FIELD_NAME;
 import static org.opensearch.neuralsearch.constants.SemanticInfoFieldConstants.DEFAULT_MODEL_CONFIG;
@@ -45,6 +44,7 @@ public class SemanticInfoConfigBuilder {
     private String spaceType;
     private String knnMethodName = KNN_VECTOR_METHOD_DEFAULT_NAME;
     private Integer embeddingDimension;
+    private Boolean chunkingEnabled;
 
     public SemanticInfoConfigBuilder(@NonNull final NamedXContentRegistry xContentRegistry) {
         this.xContentRegistry = xContentRegistry;
@@ -72,7 +72,7 @@ public SemanticInfoConfigBuilder(@NonNull final NamedXContentRegistry xContentRe
      * @return Config of the semantic info fields.
      */
     public Map<String, Object> build() {
-        Map<String, Object> embeddingFieldConfig = switch (embeddingFieldType) {
+        final Map<String, Object> embeddingFieldConfig = switch (embeddingFieldType) {
             case KNNVectorFieldMapper.CONTENT_TYPE -> buildKnnFieldConfig();
             case RankFeaturesFieldMapper.CONTENT_TYPE -> buildRankFeaturesFieldConfig();
             default -> throw new IllegalArgumentException(
@@ -84,14 +84,17 @@ public Map<String, Object> build() {
             );
         };
 
-        final Map<String, Object> chunksConfig = Map.of(
-            TYPE,
-            ObjectMapper.NESTED_CONTENT_TYPE,
-            PROPERTIES,
-            Map.of(CHUNKS_TEXT_FIELD_NAME, Map.of(TYPE, TextFieldMapper.CONTENT_TYPE), CHUNKS_EMBEDDING_FIELD_NAME, embeddingFieldConfig)
-        );
-
-        return Map.of(PROPERTIES, Map.of(CHUNKS_FIELD_NAME, chunksConfig, MODEL_FIELD_NAME, DEFAULT_MODEL_CONFIG));
+        if (chunkingEnabled) {
+            final Map<String, Object> chunksConfig = Map.of(
+                TYPE,
+                ObjectMapper.NESTED_CONTENT_TYPE,
+                PROPERTIES,
+                Map.of(CHUNKS_TEXT_FIELD_NAME, Map.of(TYPE, TextFieldMapper.CONTENT_TYPE), EMBEDDING_FIELD_NAME, embeddingFieldConfig)
+            );
+            return Map.of(PROPERTIES, Map.of(CHUNKS_FIELD_NAME, chunksConfig, MODEL_FIELD_NAME, DEFAULT_MODEL_CONFIG));
+        } else {
+            return Map.of(PROPERTIES, Map.of(EMBEDDING_FIELD_NAME, embeddingFieldConfig, MODEL_FIELD_NAME, DEFAULT_MODEL_CONFIG));
+        }
     }
 
     private Map<String, Object> buildKnnFieldConfig() {
@@ -158,26 +161,22 @@ private void extractInfoForTextEmbeddingModel(@NonNull final MLModel mlModel, @N
 
         this.embeddingDimension = textEmbeddingModelConfig.getEmbeddingDimension();
 
-        final Map<String, Object> allConfigMap;
-        try {
-            allConfigMap = MapperService.parseMapping(xContentRegistry, textEmbeddingModelConfig.getAllConfig());
-        } catch (Exception e) {
-            throw new IllegalArgumentException(
-                String.format(
-                    Locale.ROOT,
-                    "Failed to parse the all_config of the model %s. Invalid all_config: %s",
-                    modelId,
-                    textEmbeddingModelConfig.getAllConfig()
-                )
-            );
+        final Map<String, Object> additionalConfig = textEmbeddingModelConfig.getAdditionalConfig();
+        if (additionalConfig != null) {
+            final Object spaceTypeObject = additionalConfig.get(KNN_VECTOR_METHOD_SPACE_TYPE_FIELD_NAME);
+            if (spaceTypeObject instanceof String == false) {
+                throw createInvalidSpaceTypeException(modelId);
+            }
+            this.spaceType = (String) spaceTypeObject;
+        } else {
+            throw createInvalidSpaceTypeException(modelId);
         }
-        final Object spaceTypeObject = allConfigMap.get(KNN_VECTOR_METHOD_SPACE_TYPE_FIELD_NAME);
-        if (spaceTypeObject instanceof String == false) {
-            throw new IllegalArgumentException(
-                String.format(Locale.ROOT, "space_type is not defined or not a string in the all_config of the model %s.", modelId)
-            );
-        }
-        this.spaceType = (String) spaceTypeObject;
+    }
+
+    private IllegalArgumentException createInvalidSpaceTypeException(@NonNull final String modelId) {
+        return new IllegalArgumentException(
+            String.format(Locale.ROOT, "space_type is not defined or not a string in the additional_config of the model %s.", modelId)
+        );
     }
 
     private void extractInfoForSparseModel() {
@@ -213,4 +212,9 @@ private String getUnsupportedRemoteModelError(final String modelType, @lombok.No
             String.join(",", SUPPORTED_REMOTE_MODEL_TYPES)
         );
     }
+
+    public SemanticInfoConfigBuilder chunkingEnabled(final Boolean chunkingEnabled) {
+        this.chunkingEnabled = chunkingEnabled;
+        return this;
+    }
 }

src/main/java/org/opensearch/neuralsearch/mappingtransformer/SemanticMappingTransformer.java

@@ -27,6 +27,7 @@
 import static org.opensearch.neuralsearch.constants.SemanticFieldConstants.SEMANTIC_INFO_FIELD_NAME;
 import static org.opensearch.neuralsearch.util.SemanticMappingUtils.collectSemanticField;
 import static org.opensearch.neuralsearch.util.SemanticMappingUtils.extractModelIdToFieldPathMap;
+import static org.opensearch.neuralsearch.util.SemanticMappingUtils.isChunkingEnabled;
 import static org.opensearch.neuralsearch.util.SemanticMappingUtils.getProperties;
 import static org.opensearch.neuralsearch.util.SemanticMappingUtils.validateModelId;
 import static org.opensearch.neuralsearch.util.SemanticMappingUtils.validateSemanticInfoFieldName;
@@ -190,8 +191,8 @@ private void modifyMappings(
             final List<String> fieldPathList = modelIdToFieldPathMap.get(modelId);
             for (String fieldPath : fieldPathList) {
                 try {
-                    final Map<String, Object> semanticInfoConfig = createSemanticInfoField(mlModel, modelId);
                     final Map<String, Object> fieldConfig = semanticFieldPathToConfigMap.get(fieldPath);
+                    final Map<String, Object> semanticInfoConfig = createSemanticInfoField(mlModel, modelId, fieldConfig, fieldPath);
                     setSemanticInfoField(mappings, fieldPath, fieldConfig.get(SEMANTIC_INFO_FIELD_NAME), semanticInfoConfig);
                 } catch (IllegalArgumentException e) {
                     throw new IllegalArgumentException(getModifyMappingErrorMessage(fieldPath, e.getMessage()), e);
@@ -207,9 +208,16 @@ private String getModifyMappingErrorMessage(@NonNull final String fieldPath, fin
     }
 
     @VisibleForTesting
-    private Map<String, Object> createSemanticInfoField(final @NonNull MLModel modelConfig, String modelId) {
-        SemanticInfoConfigBuilder builder = new SemanticInfoConfigBuilder(xContentRegistry);
-        return builder.mlModel(modelConfig, modelId).build();
+    private Map<String, Object> createSemanticInfoField(
+        final @NonNull MLModel modelConfig,
+        final String modelId,
+        @NonNull final Map<String, Object> fieldConfig,
+        String fieldPath
+    ) {
+        final SemanticInfoConfigBuilder builder = new SemanticInfoConfigBuilder(xContentRegistry);
+        builder.mlModel(modelConfig, modelId);
+        builder.chunkingEnabled(isChunkingEnabled(fieldConfig, fieldPath));
+        return builder.build();
     }
 
     @SuppressWarnings("unchecked")

src/main/java/org/opensearch/neuralsearch/processor/TextChunkingProcessor.java

@@ -33,7 +33,6 @@
 import static org.opensearch.neuralsearch.processor.chunker.Chunker.DEFAULT_MAX_CHUNK_LIMIT;
 import static org.opensearch.neuralsearch.processor.chunker.Chunker.DISABLED_MAX_CHUNK_LIMIT;
 import static org.opensearch.neuralsearch.processor.chunker.Chunker.CHUNK_STRING_COUNT_FIELD;
-import static org.opensearch.neuralsearch.processor.chunker.ChunkerParameterParser.parseInteger;
 import static org.opensearch.neuralsearch.processor.chunker.ChunkerParameterParser.parseIntegerWithDefault;
 
 /**
@@ -273,31 +272,11 @@ private void chunkMapType(
         }
     }
 
-    /**
-     * Chunk the content, update the runtime max_chunk_limit and return the result
-     */
-    private List<String> chunkString(final String content, final Map<String, Object> runTimeParameters) {
-        // return an empty list for empty string
-        if (StringUtils.isEmpty(content)) {
-            return List.of();
-        }
-        List<String> contentResult = chunker.chunk(content, runTimeParameters);
-        // update chunk_string_count for each string
-        int chunkStringCount = parseInteger(runTimeParameters, CHUNK_STRING_COUNT_FIELD);
-        runTimeParameters.put(CHUNK_STRING_COUNT_FIELD, chunkStringCount - 1);
-        // update runtime max_chunk_limit if not disabled
-        int runtimeMaxChunkLimit = parseInteger(runTimeParameters, MAX_CHUNK_LIMIT_FIELD);
-        if (runtimeMaxChunkLimit != DISABLED_MAX_CHUNK_LIMIT) {
-            runTimeParameters.put(MAX_CHUNK_LIMIT_FIELD, runtimeMaxChunkLimit - contentResult.size());
-        }
-        return contentResult;
-    }
-
     private List<String> chunkList(final List<String> contentList, final Map<String, Object> runTimeParameters) {
         // flatten original output format from List<List<String>> to List<String>
         List<String> result = new ArrayList<>();
         for (String content : contentList) {
-            result.addAll(chunkString(content, runTimeParameters));
+            result.addAll(chunker.chunkString(content, runTimeParameters));
         }
         return result;
     }
@@ -314,7 +293,7 @@ private List<String> chunkLeafType(final Object value, final Map<String, Object>
             if (StringUtils.isBlank(String.valueOf(value))) {
                 return result;
             }
-            result = chunkString(value.toString(), runTimeParameters);
+            result = chunker.chunkString(value.toString(), runTimeParameters);
         } else if (isListOfString(value)) {
             result = chunkList((List<String>) value, runTimeParameters);
         }