Merge branch 'main' into ad-operator

Author: songkant-aws

async-query/src/main/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunner.java

@@ -12,7 +12,6 @@
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
 import org.opensearch.plugins.Plugin;
-import org.opensearch.sql.legacy.executor.AsyncRestExecutor;
 import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService;
 import org.opensearch.sql.spark.asyncquery.model.NullAsyncQueryRequestContext;
 import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
@@ -21,6 +20,8 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.Client;
 
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
+
 /**
  * The job runner class for scheduling async query.
  *
@@ -37,7 +38,7 @@
 public class ScheduledAsyncQueryJobRunner implements ScheduledJobRunner {
   // Share SQL plugin thread pool
   private static final String ASYNC_QUERY_THREAD_POOL_NAME =
-      AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME;
+      SQL_WORKER_THREAD_POOL_NAME;
   private static final Logger LOGGER = LogManager.getLogger(ScheduledAsyncQueryJobRunner.class);
 
   private static final ScheduledAsyncQueryJobRunner INSTANCE = new ScheduledAsyncQueryJobRunner();

async-query/src/test/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunnerTest.java

@@ -15,6 +15,7 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
 
 import java.time.Instant;
 import org.apache.logging.log4j.LogManager;
@@ -87,7 +88,7 @@ public void testRunJobWithCorrectParameter() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+    verify(threadPool.executor(SQL_WORKER_THREAD_POOL_NAME))
         .submit(captor.capture());
 
     Runnable runnable = captor.getValue();
@@ -145,7 +146,7 @@ public void testDoRefreshThrowsException() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+    verify(threadPool.executor(SQL_WORKER_THREAD_POOL_NAME))
         .submit(captor.capture());
 
     Runnable runnable = captor.getValue();

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

@@ -97,16 +97,14 @@ private UnresolvedPlan transformPerFunction() {
 
     PerFunction perFunc = perFuncOpt.get();
     // For chart, the rowSplit should contain the span information
-    UnresolvedExpression spanExpr = rowSplit;
-    if (rowSplit instanceof Alias) {
-      spanExpr = ((Alias) rowSplit).getDelegated();
-    }
+    UnresolvedExpression spanExpr =
+        rowSplit instanceof Alias ? ((Alias) rowSplit).getDelegated() : rowSplit;
     if (!(spanExpr instanceof Span)) {
       return this; // Cannot transform without span information
     }
 
     Span span = (Span) spanExpr;
-    Field spanStartTime = AstDSL.implicitTimestampField();
+    Field spanStartTime = (Field) span.getField();
     Function spanEndTime = timestampadd(span.getUnit(), span.getValue(), spanStartTime);
     Function spanMillis = timestampdiff(MILLISECOND, spanStartTime, spanEndTime);
     final int SECOND_IN_MILLISECOND = 1000;

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

@@ -9,50 +9,27 @@
 import java.util.List;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
+import lombok.RequiredArgsConstructor;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 
 @Getter
 @ToString
 @EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
 public class StreamWindow extends UnresolvedPlan {
 
   private final List<UnresolvedExpression> windowFunctionList;
   private final List<UnresolvedExpression> groupList;
   private final boolean current;
   private final int window;
   private final boolean global;
+  private final boolean bucketNullable;
   private final UnresolvedExpression resetBefore;
   private final UnresolvedExpression resetAfter;
   @ToString.Exclude private UnresolvedPlan child;
 
-  /** StreamWindow Constructor. */
-  public StreamWindow(
-      List<UnresolvedExpression> windowFunctionList,
-      List<UnresolvedExpression> groupList,
-      boolean current,
-      int window,
-      boolean global,
-      UnresolvedExpression resetBefore,
-      UnresolvedExpression resetAfter) {
-    this.windowFunctionList = windowFunctionList;
-    this.groupList = groupList;
-    this.current = current;
-    this.window = window;
-    this.global = global;
-    this.resetBefore = resetBefore;
-    this.resetAfter = resetAfter;
-  }
-
-  public boolean isCurrent() {
-    return current;
-  }
-
-  public boolean isGlobal() {
-    return global;
-  }
-
   @Override
   public StreamWindow attach(UnresolvedPlan child) {
     this.child = child;

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

@@ -15,6 +15,7 @@
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_DEDUP;
+import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_MAIN;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_RARE_TOP;
 import static org.opensearch.sql.calcite.utils.PlanUtils.ROW_NUMBER_COLUMN_FOR_STREAMSTATS;
@@ -50,9 +51,6 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.Aggregate;
 import org.apache.calcite.rel.core.JoinRelType;
-import org.apache.calcite.rel.hint.HintStrategyTable;
-import org.apache.calcite.rel.hint.RelHint;
-import org.apache.calcite.rel.logical.LogicalAggregate;
 import org.apache.calcite.rel.logical.LogicalValues;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFamily;
@@ -147,6 +145,7 @@
 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.AliasFieldsWrappable;
 import org.opensearch.sql.calcite.plan.LogicalAD;
 import org.opensearch.sql.calcite.plan.LogicalSystemLimit;
 import org.opensearch.sql.calcite.plan.LogicalSystemLimit.SystemLimitType;
@@ -201,7 +200,11 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) {
       throw new CalciteUnsupportedException("information_schema is unsupported in Calcite");
     }
     context.relBuilder.scan(node.getTableQualifiedName().getParts());
-    return context.relBuilder.peek();
+    RelNode scan = context.relBuilder.peek();
+    if (scan instanceof AliasFieldsWrappable) {
+      return ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
+    }
+    return scan;
   }
 
   // This is a tool method to add an existed RelOptTable to builder stack, not used for now
@@ -1057,7 +1060,7 @@ private Pair<List<RexNode>, List<AggCall>> aggregateWithTrimming(
     List<String> intendedGroupKeyAliases = getGroupKeyNamesAfterAggregation(reResolved.getLeft());
     context.relBuilder.aggregate(
         context.relBuilder.groupKey(reResolved.getLeft()), reResolved.getRight());
-    if (hintBucketNonNull) addIgnoreNullBucketHintToAggregate(context);
+    if (hintBucketNonNull) PlanUtils.addIgnoreNullBucketHintToAggregate(context.relBuilder);
     // During aggregation, Calcite projects both input dependencies and output group-by fields.
     // When names conflict, Calcite adds numeric suffixes (e.g., "value0").
     // Apply explicit renaming to restore the intended aliases.
@@ -1128,8 +1131,7 @@ private Pair<List<RexNode>, List<AggCall>> resolveAttributesForAggregation(
   @Override
   public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     Argument.ArgumentMap statsArgs = Argument.ArgumentMap.of(node.getArgExprList());
-    Boolean bucketNullable =
-        (Boolean) statsArgs.getOrDefault(Argument.BUCKET_NULLABLE, Literal.TRUE).getValue();
+    Boolean bucketNullable = (Boolean) statsArgs.get(Argument.BUCKET_NULLABLE).getValue();
     int nGroup = node.getGroupExprList().size() + (Objects.nonNull(node.getSpan()) ? 1 : 0);
     BitSet nonNullGroupMask = new BitSet(nGroup);
     if (!bucketNullable) {
@@ -1320,7 +1322,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
                 : duplicatedFieldNames.stream()
                     .map(a -> (RexNode) context.relBuilder.field(a))
                     .toList();
-        buildDedupNotNull(context, dedupeFields, allowedDuplication);
+        buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
       }
       context.relBuilder.join(
           JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1376,7 +1378,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
         List<RexNode> dedupeFields =
             getRightColumnsInJoinCriteria(context.relBuilder, joinCondition);
 
-        buildDedupNotNull(context, dedupeFields, allowedDuplication);
+        buildDedupNotNull(context, dedupeFields, allowedDuplication, true);
       }
       context.relBuilder.join(
           JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
@@ -1541,24 +1543,20 @@ public RelNode visitDedupe(Dedupe node, CalcitePlanContext context) {
     if (keepEmpty) {
       buildDedupOrNull(context, dedupeFields, allowedDuplication);
     } else {
-      buildDedupNotNull(context, dedupeFields, allowedDuplication);
+      buildDedupNotNull(context, dedupeFields, allowedDuplication, false);
     }
     return context.relBuilder.peek();
   }
 
   private static void buildDedupOrNull(
       CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
     /*
-     * | dedup 2 a, b keepempty=false
-     * DropColumns('_row_number_dedup_)
-     * +- Filter ('_row_number_dedup_ <= n OR isnull('a) OR isnull('b))
-     *    +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
+     * | dedup 2 a, b keepempty=true
+     * LogicalProject(...)
+     * +- LogicalFilter(condition=[OR(IS NULL(a), IS NULL(b), <=(_row_number_dedup_, 1))])
+     *    +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
      *        +- ...
      */
-    // Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
-    // specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a
-    // ASC
-    // NULLS FIRST, 'b ASC NULLS FIRST]
     RexNode rowNumber =
         context
             .relBuilder
@@ -1581,16 +1579,21 @@ private static void buildDedupOrNull(
   }
 
   private static void buildDedupNotNull(
-      CalcitePlanContext context, List<RexNode> dedupeFields, Integer allowedDuplication) {
+      CalcitePlanContext context,
+      List<RexNode> dedupeFields,
+      Integer allowedDuplication,
+      boolean fromJoinMaxOption) {
     /*
      * | dedup 2 a, b keepempty=false
-     * DropColumns('_row_number_dedup_)
-     * +- Filter ('_row_number_dedup_ <= n)
-     *    +- Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST, specifiedwindowoundedpreceding$(), currentrow$())) AS _row_number_dedup_], ['a, 'b], ['a ASC NULLS FIRST, 'b ASC NULLS FIRST]
-     *       +- Filter (isnotnull('a) AND isnotnull('b))
-     *          +- ...
+     * LogicalProject(...)
+     * +- LogicalFilter(condition=[<=(_row_number_dedup_, n)]))
+     *    +- LogicalProject(..., _row_number_dedup_=[ROW_NUMBER() OVER (PARTITION BY a, b ORDER BY a, b)])
+     *        +- LogicalFilter(condition=[AND(IS NOT NULL(a), IS NOT NULL(b))])
+     *           +- ...
      */
     // Filter (isnotnull('a) AND isnotnull('b))
+    String rowNumberAlias =
+        fromJoinMaxOption ? ROW_NUMBER_COLUMN_FOR_JOIN_MAX_DEDUP : ROW_NUMBER_COLUMN_FOR_DEDUP;
     context.relBuilder.filter(
         context.relBuilder.and(dedupeFields.stream().map(context.relBuilder::isNotNull).toList()));
     // Window [row_number() windowspecdefinition('a, 'b, 'a ASC NULLS FIRST, 'b ASC NULLS FIRST,
@@ -1604,15 +1607,15 @@ private static void buildDedupNotNull(
             .partitionBy(dedupeFields)
             .orderBy(dedupeFields)
             .rowsTo(RexWindowBounds.CURRENT_ROW)
-            .as(ROW_NUMBER_COLUMN_FOR_DEDUP);
+            .as(rowNumberAlias);
     context.relBuilder.projectPlus(rowNumber);
-    RexNode _row_number_dedup_ = context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_DEDUP);
+    RexNode rowNumberField = context.relBuilder.field(rowNumberAlias);
     // Filter ('_row_number_dedup_ <= n)
     context.relBuilder.filter(
         context.relBuilder.lessThanOrEqual(
-            _row_number_dedup_, context.relBuilder.literal(allowedDuplication)));
+            rowNumberField, context.relBuilder.literal(allowedDuplication)));
     // DropColumns('_row_number_dedup_)
-    context.relBuilder.projectExcept(_row_number_dedup_);
+    context.relBuilder.projectExcept(rowNumberField);
   }
 
   @Override
@@ -1745,20 +1748,25 @@ public RelNode visitStreamWindow(StreamWindow node, CalcitePlanContext context)
               .as(ROW_NUMBER_COLUMN_FOR_STREAMSTATS);
       context.relBuilder.projectPlus(streamSeq);
 
-      // construct groupNotNull predicate
-      List<RexNode> groupByList =
-          groupList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
-      List<RexNode> notNullList =
-          PlanUtils.getSelectColumns(groupByList).stream()
-              .map(context.relBuilder::field)
-              .map(context.relBuilder::isNotNull)
-              .toList();
-      RexNode groupNotNull = context.relBuilder.and(notNullList);
+      if (!node.isBucketNullable()) {
+        // construct groupNotNull predicate
+        List<RexNode> groupByList =
+            groupList.stream().map(expr -> rexVisitor.analyze(expr, context)).toList();
+        List<RexNode> notNullList =
+            PlanUtils.getSelectColumns(groupByList).stream()
+                .map(context.relBuilder::field)
+                .map(context.relBuilder::isNotNull)
+                .toList();
+        RexNode groupNotNull = context.relBuilder.and(notNullList);
+
+        // wrap each expr: CASE WHEN groupNotNull THEN rawExpr ELSE CAST(NULL AS rawType) END
+        List<RexNode> wrappedOverExprs =
+            wrapWindowFunctionsWithGroupNotNull(overExpressions, groupNotNull, context);
+        context.relBuilder.projectPlus(wrappedOverExprs);
+      } else {
+        context.relBuilder.projectPlus(overExpressions);
+      }
 
-      // wrap each expr: CASE WHEN groupNotNull THEN rawExpr ELSE CAST(NULL AS rawType) END
-      List<RexNode> wrappedOverExprs =
-          wrapWindowFunctionsWithGroupNotNull(overExpressions, groupNotNull, context);
-      context.relBuilder.projectPlus(wrappedOverExprs);
       // resort when there is by condition
       context.relBuilder.sort(context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_STREAMSTATS));
       context.relBuilder.projectExcept(context.relBuilder.field(ROW_NUMBER_COLUMN_FOR_STREAMSTATS));
@@ -1814,11 +1822,11 @@ private RelNode buildStreamWindowJoinPlan(
       RexNode segRight = context.relBuilder.field(segmentCol);
       RexNode segOuter = context.relBuilder.field(v.get(), segmentCol);
       RexNode frame = buildResetFrameFilter(context, node, outerSeq, rightSeq, segOuter, segRight);
-      RexNode group = buildGroupFilter(context, groupList, v.get());
+      RexNode group = buildGroupFilter(context, node, groupList, v.get());
       filter = (group == null) ? frame : context.relBuilder.and(frame, group);
     } else { // global + window + by condition
       RexNode frame = buildFrameFilter(context, node, outerSeq, rightSeq);
-      RexNode group = buildGroupFilter(context, groupList, v.get());
+      RexNode group = buildGroupFilter(context, node, groupList, v.get());
       filter = context.relBuilder.and(frame, group);
     }
     context.relBuilder.filter(filter);
@@ -1968,7 +1976,10 @@ private RexNode buildResetFrameFilter(
   }
 
   private RexNode buildGroupFilter(
-      CalcitePlanContext context, List<UnresolvedExpression> groupList, RexCorrelVariable correl) {
+      CalcitePlanContext context,
+      StreamWindow node,
+      List<UnresolvedExpression> groupList,
+      RexCorrelVariable correl) {
     // build conjunctive equality filters: right.g_i = outer.g_i
     if (groupList.isEmpty()) {
       return null;
@@ -1980,7 +1991,17 @@ private RexNode buildGroupFilter(
                   String groupName = extractGroupFieldName(expr);
                   RexNode rightGroup = context.relBuilder.field(groupName);
                   RexNode outerGroup = context.relBuilder.field(correl, groupName);
-                  return context.relBuilder.equals(rightGroup, outerGroup);
+                  RexNode equalCondition = context.relBuilder.equals(rightGroup, outerGroup);
+                  // handle bucket_nullable case
+                  if (!node.isBucketNullable()) {
+                    return equalCondition;
+                  } else {
+                    RexNode bothNull =
+                        context.relBuilder.and(
+                            context.relBuilder.isNull(rightGroup),
+                            context.relBuilder.isNull(outerGroup));
+                    return context.relBuilder.or(equalCondition, bothNull);
+                  }
                 })
             .toList();
     return context.relBuilder.and(equalsList);
@@ -2395,25 +2416,6 @@ public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
-  private static void addIgnoreNullBucketHintToAggregate(CalcitePlanContext context) {
-    final RelHint statHits =
-        RelHint.builder("stats_args").hintOption(Argument.BUCKET_NULLABLE, "false").build();
-    assert context.relBuilder.peek() instanceof LogicalAggregate
-        : "Stats hits should be added to LogicalAggregate";
-    context.relBuilder.hints(statHits);
-    context
-        .relBuilder
-        .getCluster()
-        .setHintStrategies(
-            HintStrategyTable.builder()
-                .hintStrategy(
-                    "stats_args",
-                    (hint, rel) -> {
-                      return rel instanceof LogicalAggregate;
-                    })
-                .build());
-  }
-
   @Override
   public RelNode visitTableFunction(TableFunction node, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("Table function is unsupported in Calcite");