Merge remote-tracking branch 'upstream/main' into issues/4789

Author: LantaoJin

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

@@ -57,6 +57,8 @@
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.expression.WindowFunction;
 import org.opensearch.sql.ast.tree.AD;
+import org.opensearch.sql.ast.tree.AddColTotals;
+import org.opensearch.sql.ast.tree.AddTotals;
 import org.opensearch.sql.ast.tree.Aggregation;
 import org.opensearch.sql.ast.tree.Append;
 import org.opensearch.sql.ast.tree.AppendCol;
@@ -521,6 +523,16 @@ public LogicalPlan visitEval(Eval node, AnalysisContext context) {
     return new LogicalEval(child, expressionsBuilder.build());
   }
 
+  @Override
+  public LogicalPlan visitAddTotals(AddTotals node, AnalysisContext context) {
+    throw getOnlyForCalciteException("addtotals");
+  }
+
+  @Override
+  public LogicalPlan visitAddColTotals(AddColTotals node, AnalysisContext context) {
+    throw getOnlyForCalciteException("addcoltotals");
+  }
+
   /** 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

@@ -45,6 +45,8 @@
 import org.opensearch.sql.ast.statement.Query;
 import org.opensearch.sql.ast.statement.Statement;
 import org.opensearch.sql.ast.tree.AD;
+import org.opensearch.sql.ast.tree.AddColTotals;
+import org.opensearch.sql.ast.tree.AddTotals;
 import org.opensearch.sql.ast.tree.Aggregation;
 import org.opensearch.sql.ast.tree.Append;
 import org.opensearch.sql.ast.tree.AppendCol;
@@ -451,4 +453,12 @@ public T visitAppend(Append node, C context) {
   public T visitMultisearch(Multisearch node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitAddTotals(AddTotals node, C context) {
+    return visitChildren(node, context);
+  }
+
+  public T visitAddColTotals(AddColTotals node, C context) {
+    return visitChildren(node, context);
+  }
 }

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

@@ -0,0 +1,47 @@
+/*
+ * 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.Map;
+import lombok.*;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.Literal;
+
+/**
+ * AST node representing the PPL addcoltotals command. Computes column-wise totals across events and
+ * optionally appends a summary event.
+ *
+ * @see AddTotals for row-wise totals
+ */
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class AddColTotals extends UnresolvedPlan {
+  private final List<Field> fieldList;
+  private final Map<String, Literal> options;
+  private UnresolvedPlan child;
+
+  @Override
+  public AddColTotals 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> visitor, C context) {
+    return visitor.visitAddColTotals(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/ast/tree/AddTotals.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 java.util.Map;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.Literal;
+
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class AddTotals extends UnresolvedPlan {
+  private final List<Field> fieldList;
+  private final Map<String, Literal> options;
+  private UnresolvedPlan child;
+
+  @Override
+  public AddTotals 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> visitor, C context) {
+    return visitor.visitAddTotals(this, context);
+  }
+}

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

@@ -8,6 +8,7 @@
 import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
 
 import java.sql.Connection;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -61,6 +62,16 @@ public class CalcitePlanContext {
 
   @Getter public Map<String, RexLambdaRef> rexLambdaRefMap;
 
+  /**
+   * List of captured variables from outer scope for lambda functions. When a lambda body references
+   * a field that is not a lambda parameter, it gets captured and stored here. The captured
+   * variables are passed as additional arguments to the transform function.
+   */
+  @Getter private List<RexNode> capturedVariables;
+
+  /** Whether we're currently inside a lambda context. */
+  @Getter @Setter private boolean inLambdaContext = false;
+
   private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType queryType) {
     this.config = config;
     this.sysLimit = sysLimit;
@@ -70,6 +81,24 @@ private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType
     this.rexBuilder = new ExtendedRexBuilder(relBuilder.getRexBuilder());
     this.functionProperties = new FunctionProperties(QueryType.PPL);
     this.rexLambdaRefMap = new HashMap<>();
+    this.capturedVariables = new ArrayList<>();
+  }
+
+  /**
+   * Private constructor for creating a context that shares relBuilder with parent. Used by clone()
+   * to create lambda contexts that can resolve fields from the parent context.
+   */
+  private CalcitePlanContext(CalcitePlanContext parent) {
+    this.config = parent.config;
+    this.sysLimit = parent.sysLimit;
+    this.queryType = parent.queryType;
+    this.connection = parent.connection;
+    this.relBuilder = parent.relBuilder; // Share the same relBuilder
+    this.rexBuilder = parent.rexBuilder; // Share the same rexBuilder
+    this.functionProperties = parent.functionProperties;
+    this.rexLambdaRefMap = new HashMap<>(); // New map for lambda variables
+    this.capturedVariables = new ArrayList<>(); // New list for captured variables
+    this.inLambdaContext = true; // Mark that we're inside a lambda
   }
 
   public RexNode resolveJoinCondition(
@@ -101,8 +130,13 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
     }
   }
 
+  /**
+   * Creates a clone of this context that shares the relBuilder with the parent. This allows lambda
+   * expressions to reference fields from the current row while having their own lambda variable
+   * mappings.
+   */
   public CalcitePlanContext clone() {
-    return new CalcitePlanContext(config, sysLimit, queryType);
+    return new CalcitePlanContext(this);
   }
 
   public static CalcitePlanContext create(
@@ -134,4 +168,42 @@ public static boolean isLegacyPreferred() {
   public void putRexLambdaRefMap(Map<String, RexLambdaRef> candidateMap) {
     this.rexLambdaRefMap.putAll(candidateMap);
   }
+
+  /**
+   * Captures an external variable for use inside a lambda. Returns a RexLambdaRef that references
+   * the captured variable by its index in the captured variables list. The actual RexNode value is
+   * stored in capturedVariables and will be passed as additional arguments to the transform
+   * function.
+   *
+   * @param fieldRef The RexInputRef representing the external field
+   * @param fieldName The name of the field being captured
+   * @return A RexLambdaRef that can be used inside the lambda to reference the captured value
+   */
+  public RexLambdaRef captureVariable(RexNode fieldRef, String fieldName) {
+    // Check if this variable is already captured
+    for (int i = 0; i < capturedVariables.size(); i++) {
+      if (capturedVariables.get(i).equals(fieldRef)) {
+        // Return existing reference - offset by number of lambda params (1 for array element)
+        return rexLambdaRefMap.get("__captured_" + i);
+      }
+    }
+
+    // Add to captured variables list
+    int captureIndex = capturedVariables.size();
+    capturedVariables.add(fieldRef);
+
+    // Create a lambda ref for this captured variable
+    // The index is offset by the number of lambda parameters (1 for single-param lambda)
+    // Count only actual lambda parameters, not captured variables
+    int lambdaParamCount =
+        (int)
+            rexLambdaRefMap.keySet().stream().filter(key -> !key.startsWith("__captured_")).count();
+    RexLambdaRef lambdaRef =
+        new RexLambdaRef(lambdaParamCount + captureIndex, fieldName, fieldRef.getType());
+
+    // Store it so we can find it again if the same field is referenced multiple times
+    rexLambdaRefMap.put("__captured_" + captureIndex, lambdaRef);
+
+    return lambdaRef;
+  }
 }