Merge remote-tracking branch 'upstream/main' into fix/background-io-thread-context

Author: Swiddis

common/build.gradle

@@ -45,7 +45,7 @@ dependencies {
     implementation "com.github.seancfoley:ipaddress:5.4.2"
 
     testImplementation group: 'junit', name: 'junit', version: '4.13.2'
-    testImplementation group: 'org.assertj', name: 'assertj-core', version: '3.9.1'
+    testImplementation group: 'org.assertj', name: 'assertj-core', version: '3.27.7'
     testImplementation group: 'com.google.guava', name: 'guava', version: "${guava_version}"
     testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: "${hamcrest_version}"
     testImplementation('org.junit.jupiter:junit-jupiter:5.9.3')

core/src/main/java/org/opensearch/sql/analysis/Analyzer.java

@@ -80,6 +80,7 @@
 import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.ML;
 import org.opensearch.sql.ast.tree.Multisearch;
+import org.opensearch.sql.ast.tree.MvCombine;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Patterns;
@@ -535,6 +536,11 @@ public LogicalPlan visitAddColTotals(AddColTotals node, AnalysisContext context)
     throw getOnlyForCalciteException("addcoltotals");
   }
 
+  @Override
+  public LogicalPlan visitMvCombine(MvCombine node, AnalysisContext context) {
+    throw getOnlyForCalciteException("mvcombine");
+  }
+
   /** Build {@link ParseExpression} to context and skip to child nodes. */
   @Override
   public LogicalPlan visitParse(Parse node, AnalysisContext context) {

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -68,6 +68,7 @@
 import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.ML;
 import org.opensearch.sql.ast.tree.Multisearch;
+import org.opensearch.sql.ast.tree.MvCombine;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Patterns;
@@ -466,4 +467,8 @@ public T visitAddTotals(AddTotals node, C context) {
   public T visitAddColTotals(AddColTotals node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitMvCombine(MvCombine node, C context) {
+    return visitChildren(node, context);
+  }
 }

core/src/main/java/org/opensearch/sql/ast/analysis/FieldResolutionVisitor.java

@@ -44,6 +44,7 @@
 import org.opensearch.sql.ast.tree.Join;
 import org.opensearch.sql.ast.tree.Lookup;
 import org.opensearch.sql.ast.tree.Multisearch;
+import org.opensearch.sql.ast.tree.MvCombine;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Patterns;
 import org.opensearch.sql.ast.tree.Project;
@@ -634,6 +635,22 @@ public Node visitExpand(Expand node, FieldResolutionContext context) {
     return node;
   }
 
+  @Override
+  public Node visitMvCombine(MvCombine node, FieldResolutionContext context) {
+    Set<String> mvCombineFields = extractFieldsFromExpression(node.getField());
+
+    FieldResolutionResult current = context.getCurrentRequirements();
+
+    Set<String> regularFields = new HashSet<>(current.getRegularFields());
+    regularFields.addAll(mvCombineFields);
+
+    context.pushRequirements(new FieldResolutionResult(regularFields, Set.of(ALL_FIELDS)));
+
+    visitChildren(node, context);
+    context.popRequirements();
+    return node;
+  }
+
   private Set<String> extractFieldsFromAggregation(UnresolvedExpression expr) {
     Set<String> fields = new HashSet<>();
     if (expr instanceof Alias alias) {

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -62,6 +62,7 @@
 import org.opensearch.sql.ast.tree.Head;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.MinSpanBin;
+import org.opensearch.sql.ast.tree.MvCombine;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Patterns;
 import org.opensearch.sql.ast.tree.Project;
@@ -468,6 +469,14 @@ public static List<Argument> defaultDedupArgs() {
         argument("consecutive", booleanLiteral(false)));
   }
 
+  public static MvCombine mvcombine(Field field) {
+    return new MvCombine(field, null);
+  }
+
+  public static MvCombine mvcombine(Field field, String delim) {
+    return new MvCombine(field, delim);
+  }
+
   public static List<Argument> sortOptions() {
     return exprList(argument("desc", booleanLiteral(false)));
   }

core/src/main/java/org/opensearch/sql/ast/tree/MvCombine.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import javax.annotation.Nullable;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Field;
+
+@Getter
+@ToString(callSuper = true)
+@EqualsAndHashCode(callSuper = false)
+public class MvCombine extends UnresolvedPlan {
+
+  private final Field field;
+  private final String delim;
+  @Nullable private UnresolvedPlan child;
+
+  public MvCombine(Field field, @Nullable String delim) {
+    this.field = field;
+    this.delim = (delim == null) ? " " : delim;
+  }
+
+  public MvCombine attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return child == null ? ImmutableList.of() : ImmutableList.of(child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitMvCombine(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -63,6 +63,7 @@
 import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.rex.RexWindowBounds;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.fun.SqlLibraryOperators;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.fun.SqlTrimFunction;
 import org.apache.calcite.sql.type.ArraySqlType;
@@ -124,6 +125,7 @@
 import org.opensearch.sql.ast.tree.Lookup.OutputStrategy;
 import org.opensearch.sql.ast.tree.ML;
 import org.opensearch.sql.ast.tree.Multisearch;
+import org.opensearch.sql.ast.tree.MvCombine;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.Parse;
 import org.opensearch.sql.ast.tree.Patterns;
@@ -3169,6 +3171,174 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
+  /**
+   * mvcombine command visitor to collapse rows that are identical on all non-target fields.
+   *
+   * <p>Grouping semantics:
+   *
+   * <ul>
+   *   <li>The target field is always excluded from the GROUP BY keys.
+   *   <li>Metadata fields (for example {@code _id}, {@code _index}, {@code _score}) are excluded
+   *       from GROUP BY keys <strong>unless</strong> they were explicitly projected earlier (for
+   *       example, via {@code fields}).
+   * </ul>
+   *
+   * <p>The target field values are aggregated using {@code ARRAY_AGG}, with {@code NULL} values
+   * filtered out. The aggregation result replaces the original target column and produces an {@code
+   * ARRAY<T>} output.
+   *
+   * <p>The original output column order is preserved. Metadata fields are projected as typed {@code
+   * NULL} literals after aggregation only when they are not part of grouping (since they were
+   * skipped).
+   *
+   * @param node mvcombine command to be visited
+   * @param context CalcitePlanContext containing the RelBuilder and planning context
+   * @return RelNode representing collapsed records with the target combined into a multivalue array
+   * @throws SemanticCheckException if the mvcombine target is not a direct field reference
+   */
+  @Override
+  public RelNode visitMvCombine(MvCombine node, CalcitePlanContext context) {
+    // 1) Lower the child plan first so the RelBuilder has the input schema on the stack.
+    visitChildren(node, context);
+
+    final RelBuilder relBuilder = context.relBuilder;
+
+    final RelNode input = relBuilder.peek();
+    final List<String> inputFieldNames = input.getRowType().getFieldNames();
+    final List<RelDataType> inputFieldTypes =
+        input.getRowType().getFieldList().stream().map(RelDataTypeField::getType).toList();
+
+    // If true, we should NOT auto-skip meta fields (because user explicitly projected them)
+    final boolean includeMetaFields = context.isProjectVisited();
+
+    // 2) Resolve the mvcombine target to an input column index (must be a direct field reference).
+    final Field targetField = node.getField();
+    final int targetIndex = resolveTargetIndex(targetField, context);
+    final String targetName = inputFieldNames.get(targetIndex);
+
+    // 3) Group by all non-target fields, skipping meta fields unless explicitly projected.
+    final List<RexNode> groupExprs =
+        buildGroupExpressionsExcludingTarget(
+            targetIndex, inputFieldNames, relBuilder, includeMetaFields);
+
+    // 4) Aggregate target values using ARRAY_AGG, filtering out NULLs.
+    performArrayAggAggregation(relBuilder, targetIndex, targetName, groupExprs);
+
+    // 5) Restore original output column order (ARRAY_AGG already returns ARRAY<T>).
+    restoreColumnOrderAfterArrayAgg(
+        relBuilder, inputFieldNames, inputFieldTypes, targetIndex, groupExprs, includeMetaFields);
+
+    return relBuilder.peek();
+  }
+
+  /** Resolves the mvcombine target expression to an input field index. */
+  private int resolveTargetIndex(Field targetField, CalcitePlanContext context) {
+    final RexNode targetRex;
+    try {
+      targetRex = rexVisitor.analyze(targetField, context);
+    } catch (IllegalArgumentException e) {
+      // Make missing-field behavior deterministic (and consistently mapped to 4xx)
+      // instead of leaking RelBuilder/rexVisitor exception wording.
+      throw new SemanticCheckException(
+          "mvcombine target field not found: " + targetField.getField().toString(), e);
+    }
+
+    if (!isInputRef(targetRex)) {
+      throw new SemanticCheckException(
+          "mvcombine target must be a direct field reference, but got: " + targetField);
+    }
+
+    final int index = ((RexInputRef) targetRex).getIndex();
+
+    final RelDataType fieldType =
+        context.relBuilder.peek().getRowType().getFieldList().get(index).getType();
+
+    if (SqlTypeUtil.isArray(fieldType) || SqlTypeUtil.isMultiset(fieldType)) {
+      throw new SemanticCheckException(
+          "mvcombine target cannot be an array/multivalue type, but got: " + fieldType);
+    }
+
+    return index;
+  }
+
+  /**
+   * Builds group-by expressions for mvcombine: all non-target input fields; meta fields are skipped
+   * unless includeMetaFields is true.
+   */
+  private List<RexNode> buildGroupExpressionsExcludingTarget(
+      int targetIndex,
+      List<String> inputFieldNames,
+      RelBuilder relBuilder,
+      boolean includeMetaFields) {
+
+    final List<RexNode> groupExprs = new ArrayList<>(Math.max(0, inputFieldNames.size() - 1));
+    for (int i = 0; i < inputFieldNames.size(); i++) {
+      if (i == targetIndex) {
+        continue;
+      }
+      if (isMetadataField(inputFieldNames.get(i)) && !includeMetaFields) {
+        continue;
+      }
+      groupExprs.add(relBuilder.field(i));
+    }
+    return groupExprs;
+  }
+
+  /** Applies mvcombine aggregation. */
+  private void performArrayAggAggregation(
+      RelBuilder relBuilder, int targetIndex, String targetName, List<RexNode> groupExprs) {
+
+    final RexNode targetRef = relBuilder.field(targetIndex);
+    final RexNode notNullTarget = relBuilder.isNotNull(targetRef);
+
+    final RelBuilder.AggCall aggCall =
+        relBuilder
+            .aggregateCall(SqlLibraryOperators.ARRAY_AGG, targetRef)
+            .filter(notNullTarget)
+            .as(targetName);
+
+    relBuilder.aggregate(relBuilder.groupKey(groupExprs), aggCall);
+  }
+
+  /**
+   * Restores the original output column order after the aggregate step. Meta fields are set to
+   * typed NULL only when they were skipped from grouping (includeMetaFields=false).
+   */
+  private void restoreColumnOrderAfterArrayAgg(
+      RelBuilder relBuilder,
+      List<String> inputFieldNames,
+      List<RelDataType> inputFieldTypes,
+      int targetIndex,
+      List<RexNode> groupExprs,
+      boolean includeMetaFields) {
+
+    final int aggregatedTargetPos = groupExprs.size();
+
+    final List<RexNode> projections = new ArrayList<>(inputFieldNames.size());
+    final List<String> projectionNames = new ArrayList<>(inputFieldNames.size());
+
+    int groupPos = 0;
+    for (int i = 0; i < inputFieldNames.size(); i++) {
+      final String fieldName = inputFieldNames.get(i);
+      projectionNames.add(fieldName);
+
+      if (i == targetIndex) {
+        // aggregated target is always the last field in the aggregate output
+        projections.add(relBuilder.field(aggregatedTargetPos));
+      } else if (isMetadataField(fieldName) && !includeMetaFields) {
+        // meta fields were skipped from grouping => not present in aggregate output => keep schema
+        // stable
+        projections.add(relBuilder.getRexBuilder().makeNullLiteral(inputFieldTypes.get(i)));
+      } else {
+        // grouped field (including meta fields when includeMetaFields=true)
+        projections.add(relBuilder.field(groupPos));
+        groupPos++;
+      }
+    }
+
+    relBuilder.project(projections, projectionNames, /* force= */ true);
+  }
+
   @Override
   public RelNode visitValues(Values values, CalcitePlanContext context) {
     if (values.getValues() == null || values.getValues().isEmpty()) {