Merge remote-tracking branch 'origin/main' into issues/4201

Signed-off-by: Yuanchun Shen <[email protected]>

Author: yuancu

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -50,6 +50,8 @@ public enum Key {
     /** Common Settings for SQL and PPL. */
     QUERY_MEMORY_LIMIT("plugins.query.memory_limit"),
     QUERY_SIZE_LIMIT("plugins.query.size_limit"),
+    QUERY_BUCKET_SIZE("plugins.query.buckets"),
+    SEARCH_MAX_BUCKETS("search.max_buckets"),
     ENCYRPTION_MASTER_KEY("plugins.query.datasources.encryption.masterkey"),
     DATASOURCES_URI_HOSTS_DENY_LIST("plugins.query.datasources.uri.hosts.denylist"),
     DATASOURCES_LIMIT("plugins.query.datasources.limit"),

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

@@ -85,6 +85,7 @@
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.RelationSubquery;
 import org.opensearch.sql.ast.tree.Rename;
+import org.opensearch.sql.ast.tree.Replace;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.SPath;
@@ -801,6 +802,11 @@ public LogicalPlan visitCloseCursor(CloseCursor closeCursor, AnalysisContext con
     return new LogicalCloseCursor(closeCursor.getChild().get(0).accept(this, context));
   }
 
+  @Override
+  public LogicalPlan visitReplace(Replace node, AnalysisContext context) {
+    throw getOnlyForCalciteException("Replace");
+  }
+
   @Override
   public LogicalPlan visitJoin(Join node, AnalysisContext context) {
     throw getOnlyForCalciteException("Join");

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

@@ -73,6 +73,7 @@
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.RelationSubquery;
 import org.opensearch.sql.ast.tree.Rename;
+import org.opensearch.sql.ast.tree.Replace;
 import org.opensearch.sql.ast.tree.Reverse;
 import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.SPath;
@@ -245,6 +246,10 @@ public T visitRename(Rename node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitReplace(Replace node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitEval(Eval node, C context) {
     return visitChildren(node, context);
   }

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

@@ -0,0 +1,58 @@
+/*
+ * 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 java.util.Set;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.Setter;
+import lombok.ToString;
+import org.jetbrains.annotations.Nullable;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Field;
+
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+public class Replace extends UnresolvedPlan {
+  private final List<ReplacePair> replacePairs;
+  private final Set<Field> fieldList;
+  @Nullable private UnresolvedPlan child;
+
+  /**
+   * Constructor with multiple pattern/replacement pairs.
+   *
+   * @param replacePairs List of pattern/replacement pairs
+   * @param fieldList Set of fields to apply replacements to
+   */
+  public Replace(List<ReplacePair> replacePairs, Set<Field> fieldList) {
+    this.replacePairs = replacePairs;
+    this.fieldList = fieldList;
+  }
+
+  @Override
+  public Replace attach(UnresolvedPlan child) {
+    if (null == this.child) {
+      this.child = child;
+    } else {
+      this.child.attach(child);
+    }
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return this.child == null ? ImmutableList.of() : ImmutableList.of(this.child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitReplace(this, context);
+  }
+}

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

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.ast.expression.Literal;
+
+/** A pair of pattern and replacement literals for the Replace command. */
+@Getter
+@AllArgsConstructor
+@EqualsAndHashCode
+@ToString
+public class ReplacePair {
+  private final Literal pattern;
+  private final Literal replacement;
+}

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

@@ -121,6 +121,8 @@
 import org.opensearch.sql.ast.tree.Regex;
 import org.opensearch.sql.ast.tree.Relation;
 import org.opensearch.sql.ast.tree.Rename;
+import org.opensearch.sql.ast.tree.Replace;
+import org.opensearch.sql.ast.tree.ReplacePair;
 import org.opensearch.sql.ast.tree.Rex;
 import org.opensearch.sql.ast.tree.SPath;
 import org.opensearch.sql.ast.tree.Search;
@@ -2411,6 +2413,51 @@ public RelNode visitValues(Values values, CalcitePlanContext context) {
     }
   }
 
+  @Override
+  public RelNode visitReplace(Replace node, CalcitePlanContext context) {
+    visitChildren(node, context);
+
+    List<String> fieldNames = context.relBuilder.peek().getRowType().getFieldNames();
+
+    // Create a set of field names to replace for quick lookup
+    Set<String> fieldsToReplace =
+        node.getFieldList().stream().map(f -> f.getField().toString()).collect(Collectors.toSet());
+
+    // Validate that all fields to replace exist by calling field() on each
+    // This leverages relBuilder.field()'s built-in validation which throws
+    // IllegalArgumentException if any field doesn't exist
+    for (String fieldToReplace : fieldsToReplace) {
+      context.relBuilder.field(fieldToReplace);
+    }
+
+    List<RexNode> projectList = new ArrayList<>();
+
+    // Project all fields, replacing specified ones in-place
+    for (String fieldName : fieldNames) {
+      if (fieldsToReplace.contains(fieldName)) {
+        // Replace this field in-place with all pattern/replacement pairs applied sequentially
+        RexNode fieldRef = context.relBuilder.field(fieldName);
+
+        // Apply all replacement pairs sequentially (nested REPLACE calls)
+        for (ReplacePair pair : node.getReplacePairs()) {
+          RexNode patternNode = rexVisitor.analyze(pair.getPattern(), context);
+          RexNode replacementNode = rexVisitor.analyze(pair.getReplacement(), context);
+          fieldRef =
+              context.relBuilder.call(
+                  SqlStdOperatorTable.REPLACE, fieldRef, patternNode, replacementNode);
+        }
+
+        projectList.add(fieldRef);
+      } else {
+        // Keep original field unchanged
+        projectList.add(context.relBuilder.field(fieldName));
+      }
+    }
+
+    context.relBuilder.project(projectList, fieldNames);
+    return context.relBuilder.peek();
+  }
+
   private void buildParseRelNode(Parse node, CalcitePlanContext context) {
     RexNode sourceField = rexVisitor.analyze(node.getSourceField(), context);
     ParseMethod parseMethod = node.getParseMethod();

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -62,7 +62,9 @@ public enum BuiltinFunctionName {
   /** Collection functions */
   ARRAY(FunctionName.of("array")),
   ARRAY_LENGTH(FunctionName.of("array_length")),
+  MAP_APPEND(FunctionName.of("map_append"), true),
   MAP_CONCAT(FunctionName.of("map_concat"), true),
+  MAP_REMOVE(FunctionName.of("map_remove"), true),
   MVAPPEND(FunctionName.of("mvappend")),
   MVJOIN(FunctionName.of("mvjoin")),
   FORALL(FunctionName.of("forall")),

core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MVAppendCore.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.CollectionUDF;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/** Core logic for `mvappend` command to collect elements from list of args */
+public class MVAppendCore {
+
+  /**
+   * Collect non-null elements from `args`. If an item is a list, it will collect non-null elements
+   * of the list. See {@ref MVAppendFunctionImplTest} for detailed behavior.
+   */
+  public static List<Object> collectElements(Object... args) {
+    List<Object> elements = new ArrayList<>();
+
+    for (Object arg : args) {
+      if (arg == null) {
+        continue;
+      } else if (arg instanceof List) {
+        addListElements((List<?>) arg, elements);
+      } else {
+        elements.add(arg);
+      }
+    }
+
+    return elements.isEmpty() ? null : elements;
+  }
+
+  private static void addListElements(List<?> list, List<Object> elements) {
+    for (Object item : list) {
+      if (item != null) {
+        elements.add(item);
+      }
+    }
+  }
+}

core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MVAppendFunctionImpl.java

@@ -7,7 +7,6 @@
 
 import static org.apache.calcite.sql.type.SqlTypeUtil.createArrayType;
 
-import java.util.ArrayList;
 import java.util.List;
 import org.apache.calcite.adapter.enumerable.NotNullImplementor;
 import org.apache.calcite.adapter.enumerable.NullPolicy;
@@ -98,33 +97,6 @@ public Expression implement(
   }
 
   public static Object mvappend(Object... args) {
-    List<Object> elements = collectElements(args);
-    return elements.isEmpty() ? null : elements;
-  }
-
-  private static List<Object> collectElements(Object... args) {
-    List<Object> elements = new ArrayList<>();
-
-    for (Object arg : args) {
-      if (arg == null) {
-        continue;
-      }
-
-      if (arg instanceof List) {
-        addListElements((List<?>) arg, elements);
-      } else {
-        elements.add(arg);
-      }
-    }
-
-    return elements;
-  }
-
-  private static void addListElements(List<?> list, List<Object> elements) {
-    for (Object item : list) {
-      if (item != null) {
-        elements.add(item);
-      }
-    }
+    return MVAppendCore.collectElements(args);
   }
 }