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

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

.github/workflows/stalled.yml

@@ -21,8 +21,9 @@ jobs:
         with:
           repo-token: ${{ steps.github_app_token.outputs.token }}
           stale-pr-label: 'stalled'
-          stale-pr-message: 'This PR is stalled because it has been open for 30 days with no activity.'
-          days-before-pr-stale: 30
+          stale-pr-message: 'This PR is stalled because it has been open for 2 weeks with no activity.'
+          days-before-pr-stale: 14
           days-before-issue-stale: -1
           days-before-pr-close: -1
           days-before-issue-close: -1
+          exempt-draft-pr: true

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

@@ -33,6 +33,8 @@ public enum Key {
     PPL_REX_MAX_MATCH_LIMIT("plugins.ppl.rex.max_match.limit"),
     PPL_VALUES_MAX_LIMIT("plugins.ppl.values.max.limit"),
     PPL_SYNTAX_LEGACY_PREFERRED("plugins.ppl.syntax.legacy.preferred"),
+    PPL_SUBSEARCH_MAXOUT("plugins.ppl.subsearch.maxout"),
+    PPL_JOIN_SUBSEARCH_MAXOUT("plugins.ppl.join.subsearch_maxout"),
 
     /** Enable Calcite as execution engine */
     CALCITE_ENGINE_ENABLED("plugins.calcite.enabled"),

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

@@ -5,14 +5,43 @@
 
 package org.opensearch.sql.ast.tree;
 
+import static org.opensearch.sql.ast.dsl.AstDSL.aggregate;
+import static org.opensearch.sql.ast.dsl.AstDSL.doubleLiteral;
+import static org.opensearch.sql.ast.dsl.AstDSL.eval;
+import static org.opensearch.sql.ast.dsl.AstDSL.function;
+import static org.opensearch.sql.ast.dsl.AstDSL.stringLiteral;
+import static org.opensearch.sql.ast.expression.IntervalUnit.SECOND;
+import static org.opensearch.sql.ast.tree.Timechart.PerFunctionRateExprBuilder.sum;
+import static org.opensearch.sql.ast.tree.Timechart.PerFunctionRateExprBuilder.timestampadd;
+import static org.opensearch.sql.ast.tree.Timechart.PerFunctionRateExprBuilder.timestampdiff;
+import static org.opensearch.sql.calcite.plan.OpenSearchConstants.IMPLICIT_FIELD_TIMESTAMP;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DIVIDE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MULTIPLY;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.SUM;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.TIMESTAMPADD;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.TIMESTAMPDIFF;
+
 import com.google.common.collect.ImmutableList;
 import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Optional;
 import lombok.AllArgsConstructor;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
+import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.dsl.AstDSL;
+import org.opensearch.sql.ast.expression.AggregateFunction;
+import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.IntervalUnit;
+import org.opensearch.sql.ast.expression.Let;
+import org.opensearch.sql.ast.expression.Span;
+import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.calcite.utils.PlanUtils;
 
 /** AST node represent Timechart operation. */
 @Getter
@@ -49,8 +78,9 @@ public Timechart useOther(Boolean useOther) {
   }
 
   @Override
-  public Timechart attach(UnresolvedPlan child) {
-    return toBuilder().child(child).build();
+  public UnresolvedPlan attach(UnresolvedPlan child) {
+    // Transform after child attached to avoid unintentionally overriding it
+    return toBuilder().child(child).build().transformPerFunction();
   }
 
   @Override
@@ -62,4 +92,112 @@ public List<UnresolvedPlan> getChild() {
   public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
     return nodeVisitor.visitTimechart(this, context);
   }
+
+  /**
+   * Transform per function to eval-based post-processing on sum result by timechart. Specifically,
+   * calculate how many seconds are in the time bucket based on the span option dynamically, then
+   * divide the aggregated sum value by the number of seconds to get the per-second rate.
+   *
+   * <p>For example, with span=5m per_second(field): per second rate = sum(field) / 300 seconds
+   *
+   * @return eval+timechart if per function present, or the original timechart otherwise.
+   */
+  private UnresolvedPlan transformPerFunction() {
+    Optional<PerFunction> perFuncOpt = PerFunction.from(aggregateFunction);
+    if (perFuncOpt.isEmpty()) {
+      return this;
+    }
+
+    PerFunction perFunc = perFuncOpt.get();
+    Span span = (Span) this.binExpression;
+    Field spanStartTime = AstDSL.field(IMPLICIT_FIELD_TIMESTAMP);
+    Function spanEndTime = timestampadd(span.getUnit(), span.getValue(), spanStartTime);
+    Function spanSeconds = timestampdiff(SECOND, spanStartTime, spanEndTime);
+
+    return eval(
+        timechart(AstDSL.alias(perFunc.aggName, sum(perFunc.aggArg))),
+        let(perFunc.aggName).multiply(perFunc.seconds).dividedBy(spanSeconds));
+  }
+
+  private Timechart timechart(UnresolvedExpression newAggregateFunction) {
+    return this.toBuilder().aggregateFunction(newAggregateFunction).build();
+  }
+
+  /** TODO: extend to support additional per_* functions */
+  @RequiredArgsConstructor
+  static class PerFunction {
+    private static final Map<String, Integer> UNIT_SECONDS = Map.of("per_second", 1);
+    private final String aggName;
+    private final UnresolvedExpression aggArg;
+    private final int seconds;
+
+    static Optional<PerFunction> from(UnresolvedExpression aggExpr) {
+      if (!(aggExpr instanceof AggregateFunction)) {
+        return Optional.empty();
+      }
+
+      AggregateFunction aggFunc = (AggregateFunction) aggExpr;
+      String aggFuncName = aggFunc.getFuncName().toLowerCase(Locale.ROOT);
+      if (!UNIT_SECONDS.containsKey(aggFuncName)) {
+        return Optional.empty();
+      }
+
+      String aggName = toAggName(aggFunc);
+      return Optional.of(
+          new PerFunction(aggName, aggFunc.getField(), UNIT_SECONDS.get(aggFuncName)));
+    }
+
+    private static String toAggName(AggregateFunction aggFunc) {
+      String fieldName =
+          (aggFunc.getField() instanceof Field)
+              ? ((Field) aggFunc.getField()).getField().toString()
+              : aggFunc.getField().toString();
+      return String.format(Locale.ROOT, "%s(%s)", aggFunc.getFuncName(), fieldName);
+    }
+  }
+
+  private PerFunctionRateExprBuilder let(String fieldName) {
+    return new PerFunctionRateExprBuilder(AstDSL.field(fieldName));
+  }
+
+  /** Fluent builder for creating Let expressions with mathematical operations. */
+  static class PerFunctionRateExprBuilder {
+    private final Field field;
+    private UnresolvedExpression expr;
+
+    PerFunctionRateExprBuilder(Field field) {
+      this.field = field;
+      this.expr = field;
+    }
+
+    PerFunctionRateExprBuilder multiply(Integer multiplier) {
+      // Promote to double literal to avoid integer division in downstream
+      this.expr =
+          function(
+              MULTIPLY.getName().getFunctionName(), expr, doubleLiteral(multiplier.doubleValue()));
+      return this;
+    }
+
+    Let dividedBy(UnresolvedExpression divisor) {
+      return AstDSL.let(field, function(DIVIDE.getName().getFunctionName(), expr, divisor));
+    }
+
+    static UnresolvedExpression sum(UnresolvedExpression field) {
+      return aggregate(SUM.getName().getFunctionName(), field);
+    }
+
+    static Function timestampadd(
+        SpanUnit unit, UnresolvedExpression value, UnresolvedExpression timestampField) {
+      UnresolvedExpression intervalUnit =
+          stringLiteral(PlanUtils.spanUnitToIntervalUnit(unit).toString());
+      return function(
+          TIMESTAMPADD.getName().getFunctionName(), intervalUnit, value, timestampField);
+    }
+
+    static Function timestampdiff(
+        IntervalUnit unit, UnresolvedExpression start, UnresolvedExpression end) {
+      return function(
+          TIMESTAMPDIFF.getName().getFunctionName(), stringLiteral(unit.toString()), start, end);
+    }
+  }
 }

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

@@ -35,7 +35,7 @@ public class CalcitePlanContext {
   public final ExtendedRexBuilder rexBuilder;
   public final FunctionProperties functionProperties;
   public final QueryType queryType;
-  public final Integer querySizeLimit;
+  public final SysLimit sysLimit;
 
   /** This thread local variable is only used to skip script encoding in script pushdown. */
   public static final ThreadLocal<Boolean> skipEncoding = ThreadLocal.withInitial(() -> false);
@@ -61,9 +61,9 @@ public class CalcitePlanContext {
 
   @Getter public Map<String, RexLambdaRef> rexLambdaRefMap;
 
-  private CalcitePlanContext(FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
+  private CalcitePlanContext(FrameworkConfig config, SysLimit sysLimit, QueryType queryType) {
     this.config = config;
-    this.querySizeLimit = querySizeLimit;
+    this.sysLimit = sysLimit;
     this.queryType = queryType;
     this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
     this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
@@ -102,12 +102,12 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
   }
 
   public CalcitePlanContext clone() {
-    return new CalcitePlanContext(config, querySizeLimit, queryType);
+    return new CalcitePlanContext(config, sysLimit, queryType);
   }
 
   public static CalcitePlanContext create(
-      FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
-    return new CalcitePlanContext(config, querySizeLimit, queryType);
+      FrameworkConfig config, SysLimit sysLimit, QueryType queryType) {
+    return new CalcitePlanContext(config, sysLimit, queryType);
   }
 
   /**

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

@@ -133,6 +133,8 @@
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
 import org.opensearch.sql.ast.tree.Window;
+import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
+import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
 import org.opensearch.sql.calcite.plan.OpenSearchConstants;
 import org.opensearch.sql.calcite.utils.BinUtils;
 import org.opensearch.sql.calcite.utils.JoinAndLookupUtils;
@@ -1136,6 +1138,15 @@ private Optional<RexLiteral> extractAliasLiteral(RexNode node) {
   public RelNode visitJoin(Join node, CalcitePlanContext context) {
     List<UnresolvedPlan> children = node.getChildren();
     children.forEach(c -> analyze(c, context));
+    // add join.subsearch_maxout limit to subsearch side, 0 and negative means unlimited.
+    if (context.sysLimit.joinSubsearchLimit() > 0) {
+      PlanUtils.replaceTop(
+          context.relBuilder,
+          LogicalSystemLimit.create(
+              SystemLimitType.JOIN_SUBSEARCH_MAXOUT,
+              context.relBuilder.peek(),
+              context.relBuilder.literal(context.sysLimit.joinSubsearchLimit())));
+    }
     if (node.getJoinCondition().isEmpty()) {
       // join-with-field-list grammar
       List<String> leftColumns = context.relBuilder.peek(1).getRowType().getFieldNames();
@@ -1914,6 +1925,9 @@ public RelNode visitFlatten(Flatten node, CalcitePlanContext context) {
 
   /** Helper method to get the function name for proper column naming */
   private String getValueFunctionName(UnresolvedExpression aggregateFunction) {
+    if (aggregateFunction instanceof Alias) {
+      return ((Alias) aggregateFunction).getName();
+    }
     if (!(aggregateFunction instanceof AggregateFunction)) {
       return "value";
     }

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

@@ -68,9 +68,13 @@
 import org.opensearch.sql.ast.expression.subquery.ExistsSubquery;
 import org.opensearch.sql.ast.expression.subquery.InSubquery;
 import org.opensearch.sql.ast.expression.subquery.ScalarSubquery;
+import org.opensearch.sql.ast.expression.subquery.SubqueryExpression;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
+import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
 import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
 import org.opensearch.sql.calcite.utils.PlanUtils;
+import org.opensearch.sql.calcite.utils.SubsearchUtils;
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
@@ -465,7 +469,7 @@ private RexNode extractRexNodeFromAlias(RexNode node) {
   public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) {
     List<RexNode> nodes = node.getChild().stream().map(child -> analyze(child, context)).toList();
     UnresolvedPlan subquery = node.getQuery();
-    RelNode subqueryRel = resolveSubqueryPlan(subquery, context);
+    RelNode subqueryRel = resolveSubqueryPlan(subquery, node, context);
     if (subqueryRel.getRowType().getFieldCount() != nodes.size()) {
       throw new SemanticCheckException(
           "The number of columns in the left hand side of an IN subquery does not match the number"
@@ -489,7 +493,7 @@ public RexNode visitScalarSubquery(ScalarSubquery node, CalcitePlanContext conte
     return context.relBuilder.scalarQuery(
         b -> {
           UnresolvedPlan subquery = node.getQuery();
-          return resolveSubqueryPlan(subquery, context);
+          return resolveSubqueryPlan(subquery, node, context);
         });
   }
 
@@ -498,21 +502,39 @@ public RexNode visitExistsSubquery(ExistsSubquery node, CalcitePlanContext conte
     return context.relBuilder.exists(
         b -> {
           UnresolvedPlan subquery = node.getQuery();
-          return resolveSubqueryPlan(subquery, context);
+          return resolveSubqueryPlan(subquery, node, context);
         });
   }
 
-  private RelNode resolveSubqueryPlan(UnresolvedPlan subquery, CalcitePlanContext context) {
+  private RelNode resolveSubqueryPlan(
+      UnresolvedPlan subquery, SubqueryExpression subqueryExpression, CalcitePlanContext context) {
     boolean isNestedSubquery = context.isResolvingSubquery();
     context.setResolvingSubquery(true);
     // clear and store the outer state
     boolean isResolvingJoinConditionOuter = context.isResolvingJoinCondition();
     if (isResolvingJoinConditionOuter) {
       context.setResolvingJoinCondition(false);
     }
-    RelNode subqueryRel = subquery.accept(planVisitor, context);
+    subquery.accept(planVisitor, context);
+    // add subsearch.maxout limit to exists-in subsearch, 0 and negative means unlimited
+    if (context.sysLimit.subsearchLimit() > 0 && !(subqueryExpression instanceof ScalarSubquery)) {
+      // Cannot add system limit to the top of subquery simply.
+      // Instead, add system limit under the correlated conditions.
+      SubsearchUtils.SystemLimitInsertionShuttle shuttle =
+          new SubsearchUtils.SystemLimitInsertionShuttle(context);
+      RelNode replacement = context.relBuilder.peek().accept(shuttle);
+      if (!shuttle.isCorrelatedConditionFound()) {
+        // If no correlated condition found, add system limit to the top of subquery.
+        replacement =
+            LogicalSystemLimit.create(
+                SystemLimitType.SUBSEARCH_MAXOUT,
+                replacement,
+                context.relBuilder.literal(context.sysLimit.subsearchLimit()));
+      }
+      PlanUtils.replaceTop(context.relBuilder, replacement);
+    }
     // pop the inner plan
-    context.relBuilder.build();
+    RelNode subqueryRel = context.relBuilder.build();
     // clear the exists subquery resolving state
     // restore to the previous state
     if (isResolvingJoinConditionOuter) {

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite;
+
+import org.opensearch.sql.common.setting.Settings;
+
+public record SysLimit(Integer querySizeLimit, Integer subsearchLimit, Integer joinSubsearchLimit) {
+  /** Create SysLimit from Settings. */
+  public static SysLimit fromSettings(Settings settings) {
+    return settings == null
+        ? UNLIMITED_SUBSEARCH
+        : new SysLimit(
+            settings.getSettingValue(Settings.Key.QUERY_SIZE_LIMIT),
+            settings.getSettingValue(Settings.Key.PPL_SUBSEARCH_MAXOUT),
+            settings.getSettingValue(Settings.Key.PPL_JOIN_SUBSEARCH_MAXOUT));
+  }
+
+  /** No limitation on subsearch */
+  public static SysLimit UNLIMITED_SUBSEARCH = new SysLimit(10000, 0, 0);
+
+  /** For testing only */
+  public static SysLimit DEFAULT = new SysLimit(10000, 10000, 50000);
+}

core/src/main/java/org/opensearch/sql/calcite/plan/LogicalSystemLimit.java

@@ -30,7 +30,11 @@ public enum SystemLimitType {
      *
      * <p>This type is used to indicate that the limit is applied to the system level.
      */
-    QUERY_SIZE_LIMIT
+    QUERY_SIZE_LIMIT,
+    /** The max output from subsearch to join against. */
+    JOIN_SUBSEARCH_MAXOUT,
+    /** Max output to return from a subsearch. */
+    SUBSEARCH_MAXOUT,
   }
 
   @Getter private final SystemLimitType type;

core/src/main/java/org/opensearch/sql/calcite/udf/udaf/PercentileApproxFunction.java

@@ -62,7 +62,7 @@ public Object result(PencentileApproAccumulator acc) {
         float floatRet = (float) retValue;
         return floatRet;
       default:
-        return acc.value();
+        return retValue;
     }
   }