Adds a mapper for context aware segments grouping criteria

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 3.x]
 ### Added
 - Use Lucene `pack` method for `half_float` and `usigned_long` when using `ApproximatePointRangeQuery`.
+- Add a mapper for context aware segments grouping criteria ([#19233](https://github.com/opensearch-project/OpenSearch/pull/19233))
 
 ### Changed
 - Refactor to move prepareIndex and prepareDelete methods to Engine class ([#19551](https://github.com/opensearch-project/OpenSearch/pull/19551))

modules/lang-painless/src/main/java/org/opensearch/painless/PainlessModulePlugin.java

@@ -59,6 +59,7 @@
 import org.opensearch.repositories.RepositoriesService;
 import org.opensearch.rest.RestController;
 import org.opensearch.rest.RestHandler;
+import org.opensearch.script.ContextAwareGroupingScript;
 import org.opensearch.script.DerivedFieldScript;
 import org.opensearch.script.IngestScript;
 import org.opensearch.script.ScoreScript;
@@ -120,6 +121,9 @@ public final class PainlessModulePlugin extends Plugin implements ScriptPlugin,
         derived.add(AllowlistLoader.loadFromResourceFiles(Allowlist.class, "org.opensearch.derived.txt"));
         map.put(DerivedFieldScript.CONTEXT, derived);
 
+        // Only basic painless support for ContextAwareGrouping script
+        map.put(ContextAwareGroupingScript.CONTEXT, new ArrayList<>(Allowlist.BASE_ALLOWLISTS));
+
         allowlists = map;
     }
 

modules/lang-painless/src/test/java/org/opensearch/painless/ContextAwareGroupingScriptTests.java

@@ -0,0 +1,67 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.painless;
+
+import org.opensearch.common.settings.Settings;
+import org.opensearch.painless.spi.Allowlist;
+import org.opensearch.script.ContextAwareGroupingScript;
+import org.opensearch.script.ScriptContext;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class ContextAwareGroupingScriptTests extends ScriptTestCase {
+
+    private static PainlessScriptEngine SCRIPT_ENGINE;
+
+    @Override
+    public void setUp() throws Exception {
+        super.setUp();
+
+        Map<ScriptContext<?>, List<Allowlist>> contexts = newDefaultContexts();
+        List<Allowlist> allowlists = new ArrayList<>(Allowlist.BASE_ALLOWLISTS);
+        contexts.put(ContextAwareGroupingScript.CONTEXT, allowlists);
+
+        SCRIPT_ENGINE = new PainlessScriptEngine(Settings.EMPTY, contexts);
+    }
+
+    @Override
+    public void tearDown() throws Exception {
+        super.tearDown();
+        SCRIPT_ENGINE = null;
+    }
+
+    @Override
+    protected PainlessScriptEngine getEngine() {
+        return SCRIPT_ENGINE;
+    }
+
+    public void testContextAwareGroupingScript() {
+        String stringConcat = "ctx.value + \"-context-aware\"";
+        ContextAwareGroupingScript script = compile(stringConcat);
+
+        assertEquals("value-context-aware", script.execute(Map.of("value", "value")));
+
+        String integerAddition = "String.valueOf(ctx.value / 100)";
+        ContextAwareGroupingScript integerAdditionScript = compile(integerAddition);
+        assertEquals("2", integerAdditionScript.execute(Map.of("value", 200)));
+    }
+
+    private ContextAwareGroupingScript compile(String expression) {
+        ContextAwareGroupingScript.Factory factory = getEngine().compile(
+            expression,
+            expression,
+            ContextAwareGroupingScript.CONTEXT,
+            Collections.emptyMap()
+        );
+        return factory.newInstance();
+    }
+}

modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/71_context_api.yml

@@ -2,7 +2,7 @@
     - do:
         scripts_painless_context: {}
     - match: { contexts.0: aggregation_selector}
-    - match: { contexts.24: update}
+    - match: { contexts.25: update}
 ---
 
 "Action to get all API values for score context":

server/src/main/java/org/opensearch/index/mapper/CompositeMappedFieldType.java

@@ -50,7 +50,8 @@ public CompositeMappedFieldType(String name, List<String> fields, CompositeField
      */
     @ExperimentalApi
     public enum CompositeFieldType {
-        STAR_TREE("star_tree");
+        STAR_TREE("star_tree"),
+        CONTEXT_AWARE_GROUPING("context_aware_grouping");
 
         private final String name;
 

server/src/main/java/org/opensearch/index/mapper/ContextAwareGroupingFieldMapper.java

@@ -0,0 +1,184 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.mapper;
+
+import org.opensearch.script.ContextAwareGroupingScript;
+import org.opensearch.script.Script;
+
+import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.opensearch.script.Script.DEFAULT_SCRIPT_LANG;
+
+/**
+ * A field mapper to specify context aware grouping mapper creation
+ *
+ * @opensearch.internal
+ */
+public class ContextAwareGroupingFieldMapper extends ParametrizedFieldMapper {
+
+    public static final String CONTENT_TYPE = "context_aware_grouping";
+
+    public static final Mapper.TypeParser PARSER = new TypeParser();
+
+    private static class TypeParser implements Mapper.TypeParser {
+
+        @Override
+        public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserContext context) throws MapperParsingException {
+            throw new IllegalStateException("ContextAwareGroupingFieldMapper needs objbuilder to validate node");
+        }
+
+        @Override
+        public Mapper.Builder<?> parse(String name, Map<String, Object> node, ParserContext context, ObjectMapper.Builder objBuilder)
+            throws MapperParsingException {
+            Builder builder = new Builder(name);
+            builder.parse(name, context, node);
+
+            if (builder.fields.isConfigured() == false) {
+                throw new MapperParsingException("[fields] in context_aware_grouping is required");
+            }
+
+            Set<String> propertyFieldNames = (Set<String>) objBuilder.mappersBuilders.stream()
+                .map(b -> ((Mapper.Builder) b).name())
+                .collect(Collectors.toSet());
+
+            if (propertyFieldNames.containsAll(builder.fields.getValue()) == false) {
+                throw new MapperParsingException(
+                    "[fields] should be from properties: [" + propertyFieldNames + "] but found [" + builder.fields.getValue() + "]"
+                );
+            }
+
+            final Script s = builder.script.getValue();
+            if (s != null) {
+                ContextAwareGroupingScript.Factory factory = context.scriptService()
+                    .compile(builder.script.get(), ContextAwareGroupingScript.CONTEXT);
+                builder.compiledScript = factory.newInstance();
+            }
+            return builder;
+        }
+    }
+
+    /**
+     * Builder for this field mapper
+     *
+     * @opensearch.internal
+     */
+    public static class Builder extends ParametrizedFieldMapper.Builder {
+
+        private static ContextAwareGroupingFieldMapper toType(FieldMapper in) {
+            return (ContextAwareGroupingFieldMapper) in;
+        }
+
+        private final Parameter<List<String>> fields = new Parameter<>("fields", true, Collections::emptyList, (n, c, o) -> {
+            if (!(o instanceof List)) {
+                throw new MapperParsingException("Expected [fields] to be a list of strings but got [" + o + "]");
+            }
+
+            List<String> fields = (List<String>) o;
+            if (fields.isEmpty()) {
+                throw new MapperParsingException("Expected [fields] in context_aware_grouping to have one value");
+            }
+
+            if (fields.size() > 1) {
+                throw new MapperParsingException("Currently [fields] in context_aware_grouping does not support multiple values");
+            }
+
+            return fields;
+        }, m -> toType(m).fields);
+
+        private final Parameter<Script> script = new Parameter<>("script", true, () -> null, (n, c, o) -> {
+            if (o == null) {
+                return null;
+            }
+
+            Script s = Script.parse(o);
+            if (!s.getLang().equals(DEFAULT_SCRIPT_LANG)) {
+                throw new MapperParsingException("context_aware_grouping only supports painless script");
+            }
+            return s;
+        }, m -> toType(m).script).acceptsNull();
+
+        private ContextAwareGroupingScript compiledScript;
+
+        /**
+         * Creates a new Builder with a field name
+         *
+         * @param name
+         */
+        protected Builder(String name) {
+            super(name);
+        }
+
+        protected Builder(String name, List<String> fields, Script script, ContextAwareGroupingScript contextAwareGroupingScript) {
+            super(name);
+            this.fields.setValue(fields);
+            this.script.setValue(script);
+            this.compiledScript = contextAwareGroupingScript;
+        }
+
+        @Override
+        protected List<Parameter<?>> getParameters() {
+            return List.of(fields, script);
+        }
+
+        @Override
+        public ParametrizedFieldMapper build(BuilderContext context) {
+            final ContextAwareGroupingFieldType contextAwareGroupingFieldType = new ContextAwareGroupingFieldType(
+                this.fields.getValue(),
+                this.compiledScript
+            );
+            return new ContextAwareGroupingFieldMapper(name, contextAwareGroupingFieldType, this);
+        }
+    }
+
+    private final List<String> fields;
+    private final Script script;
+    private final ContextAwareGroupingScript compiledScript;
+
+    /**
+     * Creates a new ParametrizedFieldMapper
+     *
+     * @param simpleName
+     * @param mappedFieldType
+     * @param builder
+     */
+    protected ContextAwareGroupingFieldMapper(
+        String simpleName,
+        ContextAwareGroupingFieldType mappedFieldType,
+        ContextAwareGroupingFieldMapper.Builder builder
+    ) {
+        super(simpleName, mappedFieldType, MultiFields.empty(), CopyTo.empty());
+        this.fields = builder.fields.getValue();
+        this.script = builder.script.getValue();
+        this.compiledScript = builder.compiledScript;
+    }
+
+    @Override
+    public Builder getMergeBuilder() {
+        return new Builder(CONTENT_TYPE, this.fields, this.script, this.compiledScript);
+    }
+
+    @Override
+    protected void parseCreateField(ParseContext context) throws IOException {
+        throw new MapperParsingException("context_aware_grouping cannot be ingested in the document");
+    }
+
+    public ContextAwareGroupingFieldType fieldType() {
+        return (ContextAwareGroupingFieldType) mappedFieldType;
+    }
+
+    @Override
+    protected String contentType() {
+        return CONTENT_TYPE;
+    }
+}

server/src/main/java/org/opensearch/index/mapper/ContextAwareGroupingFieldType.java

@@ -0,0 +1,60 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.index.mapper;
+
+import org.apache.lucene.search.Query;
+import org.opensearch.index.query.QueryShardContext;
+import org.opensearch.script.ContextAwareGroupingScript;
+import org.opensearch.search.lookup.SearchLookup;
+
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Field type for context_aware_grouping field mapper
+ *
+ * @opensearch.internal
+ */
+public class ContextAwareGroupingFieldType extends CompositeMappedFieldType {
+
+    private ContextAwareGroupingScript compiledScript;
+
+    public ContextAwareGroupingFieldType(final List<String> fields, final ContextAwareGroupingScript compiledScript) {
+        super(
+            ContextAwareGroupingFieldMapper.CONTENT_TYPE,
+            false,
+            false,
+            false,
+            TextSearchInfo.NONE,
+            Collections.emptyMap(),
+            fields,
+            CompositeFieldType.CONTEXT_AWARE_GROUPING
+        );
+        this.compiledScript = compiledScript;
+    }
+
+    public ContextAwareGroupingScript compiledScript() {
+        return compiledScript;
+    }
+
+    @Override
+    public ValueFetcher valueFetcher(QueryShardContext context, SearchLookup searchLookup, String format) {
+        throw new UnsupportedOperationException("valueFetcher is not supported for context_aware_grouping field");
+    }
+
+    @Override
+    public String typeName() {
+        return ContextAwareGroupingFieldMapper.CONTENT_TYPE;
+    }
+
+    @Override
+    public Query termQuery(Object value, QueryShardContext context) {
+        throw new UnsupportedOperationException("Term query is not supported for context_aware_grouping field");
+    }
+}

server/src/main/java/org/opensearch/index/mapper/DateFieldMapper.java

@@ -814,18 +814,7 @@ protected DateFieldMapper clone() {
 
     @Override
     protected void parseCreateField(ParseContext context) throws IOException {
-        String dateAsString;
-        if (context.externalValueSet()) {
-            Object dateAsObject = context.externalValue();
-            if (dateAsObject == null) {
-                dateAsString = null;
-            } else {
-                dateAsString = dateAsObject.toString();
-            }
-        } else {
-            dateAsString = context.parser().textOrNull();
-        }
-
+        String dateAsString = getFieldValue(context);
         long timestamp;
         if (dateAsString == null) {
             if (nullValue == null) {
@@ -875,6 +864,20 @@ boolean isSkiplistDefaultEnabled(IndexSortConfig indexSortConfig, String fieldNa
         return false;
     }
 
+    @Override
+    protected String getFieldValue(ParseContext context) throws IOException {
+        if (context.externalValueSet()) {
+            Object dateAsObject = context.externalValue();
+            if (dateAsObject == null) {
+                return null;
+            } else {
+                return dateAsObject.toString();
+            }
+        } else {
+            return context.parser().textOrNull();
+        }
+    }
+
     public Long getNullValue() {
         return nullValue;
     }