Merge origin/main into pushdown-limit

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

Author: yuancu

.github/workflows/maven-publish.yml

@@ -29,6 +29,30 @@ jobs:
         with:
           role-to-assume: ${{ secrets.PUBLISH_SNAPSHOTS_ROLE }}
           aws-region: us-east-1
+
+      # Create the initial direct-query directory structure
+      - name: Create direct-query directory structure in repository
+        run: |
+          # Get credentials for publishing
+          export SONATYPE_USERNAME=$(aws secretsmanager get-secret-value --secret-id maven-snapshots-username --query SecretString --output text)
+          export SONATYPE_PASSWORD=$(aws secretsmanager get-secret-value --secret-id maven-snapshots-password --query SecretString --output text)
+          echo "::add-mask::$SONATYPE_USERNAME"
+          echo "::add-mask::$SONATYPE_PASSWORD"
+          
+          # Create a placeholder file
+          TEMP_DIR=$(mktemp -d)
+          echo "Directory placeholder - $(date)" > "${TEMP_DIR}/.placeholder"
+          
+          # Upload the placeholder file to create the directory structure
+          echo "Creating initial directory structure..."
+          curl -X PUT -u "${SONATYPE_USERNAME}:${SONATYPE_PASSWORD}" \
+            --upload-file "${TEMP_DIR}/.placeholder" \
+            "https://aws.oss.sonatype.org/content/repositories/snapshots/org/opensearch/direct-query/.placeholder"
+          
+          # Clean up
+          rm -rf "${TEMP_DIR}"
+          echo "Directory structure created"
+
       - name: publish snapshots to maven
         run: |
           export SONATYPE_USERNAME=$(aws secretsmanager get-secret-value --secret-id maven-snapshots-username --query SecretString --output text)

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

@@ -382,7 +382,7 @@ public LogicalPlan visitRareTopN(RareTopN node, AnalysisContext context) {
     fields.forEach(
         field -> newEnv.define(new Symbol(Namespace.FIELD_NAME, field.toString()), field.type()));
 
-    List<Argument> options = node.getNoOfResults();
+    List<Argument> options = node.getArguments();
     Integer noOfResults = (Integer) options.get(0).getValue().getValue();
 
     return new LogicalRareTopN(child, node.getCommandType(), noOfResults, fields, groupBys);
@@ -593,20 +593,22 @@ public LogicalPlan visitFillNull(final FillNull node, final AnalysisContext cont
 
     ImmutableList.Builder<Pair<ReferenceExpression, Expression>> expressionsBuilder =
         new Builder<>();
-    for (FillNull.NullableFieldFill fieldFill : node.getNullableFieldFills()) {
-      Expression fieldExpr =
-          expressionAnalyzer.analyze(fieldFill.getNullableFieldReference(), context);
+    for (Pair<Field, UnresolvedExpression> fieldFill : node.getReplacementPairs()) {
+      Expression fieldExpr = expressionAnalyzer.analyze(fieldFill.getLeft(), context);
       ReferenceExpression ref =
-          DSL.ref(fieldFill.getNullableFieldReference().getField().toString(), fieldExpr.type());
+          DSL.ref(fieldFill.getLeft().getField().toString(), fieldExpr.type());
       FunctionExpression ifNullFunction =
-          DSL.ifnull(ref, expressionAnalyzer.analyze(fieldFill.getReplaceNullWithMe(), context));
+          DSL.ifnull(ref, expressionAnalyzer.analyze(fieldFill.getRight(), context));
       expressionsBuilder.add(new ImmutablePair<>(ref, ifNullFunction));
       TypeEnvironment typeEnvironment = context.peek();
       // define the new reference in type env.
       typeEnvironment.define(ref);
     }
-
-    return new LogicalEval(child, expressionsBuilder.build());
+    List<Pair<ReferenceExpression, Expression>> expressions = expressionsBuilder.build();
+    if (expressions.isEmpty()) {
+      throw new SemanticCheckException("At least one field is required for fillnull in V2.");
+    }
+    return new LogicalEval(child, expressions);
   }
 
   /** Build {@link LogicalML} for ml command. */

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

@@ -13,7 +13,6 @@
 import java.util.Optional;
 import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
-import org.apache.commons.lang3.tuple.ImmutablePair;
 import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.ast.expression.AggregateFunction;
 import org.opensearch.sql.ast.expression.Alias;
@@ -524,21 +523,23 @@ public static Patterns patterns(
         input);
   }
 
-  public static FillNull fillNull(UnresolvedExpression replaceNullWithMe, Field... fields) {
-    return new FillNull(
-        FillNull.ContainNullableFieldFill.ofSameValue(
-            replaceNullWithMe, ImmutableList.copyOf(fields)));
+  public static FillNull fillNull(UnresolvedPlan input, UnresolvedExpression replacement) {
+    return FillNull.ofSameValue(replacement, ImmutableList.of()).attach(input);
   }
 
   public static FillNull fillNull(
-      List<ImmutablePair<Field, UnresolvedExpression>> fieldAndReplacements) {
-    ImmutableList.Builder<FillNull.NullableFieldFill> replacementsBuilder = ImmutableList.builder();
-    for (ImmutablePair<Field, UnresolvedExpression> fieldAndReplacement : fieldAndReplacements) {
+      UnresolvedPlan input, UnresolvedExpression replacement, Field... fields) {
+    return FillNull.ofSameValue(replacement, ImmutableList.copyOf(fields)).attach(input);
+  }
+
+  public static FillNull fillNull(
+      UnresolvedPlan input, List<Pair<Field, UnresolvedExpression>> fieldAndReplacements) {
+    ImmutableList.Builder<Pair<Field, UnresolvedExpression>> replacementsBuilder =
+        ImmutableList.builder();
+    for (Pair<Field, UnresolvedExpression> fieldAndReplacement : fieldAndReplacements) {
       replacementsBuilder.add(
-          new FillNull.NullableFieldFill(
-              fieldAndReplacement.getLeft(), fieldAndReplacement.getRight()));
+          Pair.of(fieldAndReplacement.getLeft(), fieldAndReplacement.getRight()));
     }
-    return new FillNull(
-        FillNull.ContainNullableFieldFill.ofVariousValue(replacementsBuilder.build()));
+    return FillNull.ofVariousValue(replacementsBuilder.build()).attach(input);
   }
 }

core/src/main/java/org/opensearch/sql/ast/expression/Argument.java

@@ -7,6 +7,7 @@
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Map;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -32,4 +33,29 @@ public List<UnresolvedExpression> getChild() {
   public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
     return nodeVisitor.visitArgument(this, context);
   }
+
+  /** ArgumentMap is a helper class to get argument value by name. */
+  public static class ArgumentMap {
+    private final Map<String, Literal> map;
+
+    public ArgumentMap(List<Argument> arguments) {
+      this.map =
+          arguments.stream()
+              .collect(java.util.stream.Collectors.toMap(Argument::getArgName, Argument::getValue));
+    }
+
+    public static ArgumentMap of(List<Argument> arguments) {
+      return new ArgumentMap(arguments);
+    }
+
+    /**
+     * Get argument value by name.
+     *
+     * @param name argument name
+     * @return argument value
+     */
+    public Literal get(String name) {
+      return map.get(name);
+    }
+  }
 }

core/src/main/java/org/opensearch/sql/ast/expression/WindowFrame.java

@@ -10,6 +10,7 @@
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import lombok.ToString;
+import org.opensearch.sql.ast.dsl.AstDSL;
 
 @EqualsAndHashCode(callSuper = false)
 @Getter
@@ -25,12 +26,25 @@ public enum FrameType {
     ROWS
   }
 
-  public static WindowFrame defaultFrame() {
-    return new WindowFrame(
-        FrameType.ROWS, createBound("UNBOUNDED PRECEDING"), createBound("UNBOUNDED FOLLOWING"));
+  public static WindowFrame rowsUnbounded() {
+    return WindowFrame.of(
+        FrameType.ROWS,
+        AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
+        AstDSL.stringLiteral("UNBOUNDED FOLLOWING"));
   }
 
-  public static WindowFrame create(FrameType type, Literal lower, Literal upper) {
+  public static WindowFrame toCurrentRow() {
+    return WindowFrame.of(
+        FrameType.ROWS,
+        AstDSL.stringLiteral("UNBOUNDED PRECEDING"),
+        AstDSL.stringLiteral("CURRENT ROW"));
+  }
+
+  public static WindowFrame of(FrameType type, String lower, String upper) {
+    return WindowFrame.of(type, AstDSL.stringLiteral(lower), AstDSL.stringLiteral(upper));
+  }
+
+  public static WindowFrame of(FrameType type, Literal lower, Literal upper) {
     WindowBound lowerBound = null;
     WindowBound upperBound = null;
     if (lower != null) {

core/src/main/java/org/opensearch/sql/ast/expression/WindowFunction.java

@@ -26,7 +26,7 @@ public class WindowFunction extends UnresolvedExpression {
   private final UnresolvedExpression function;
   @Setter private List<UnresolvedExpression> partitionByList = new ArrayList<>();
   @Setter private List<Pair<SortOption, UnresolvedExpression>> sortList = new ArrayList<>();
-  @Setter private WindowFrame windowFrame = WindowFrame.defaultFrame();
+  @Setter private WindowFrame windowFrame = WindowFrame.rowsUnbounded();
 
   public WindowFunction(
       UnresolvedExpression function,

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

@@ -6,70 +6,53 @@
 package org.opensearch.sql.ast.tree;
 
 import java.util.List;
-import java.util.Objects;
-import lombok.AllArgsConstructor;
+import java.util.Optional;
+import lombok.EqualsAndHashCode;
 import lombok.Getter;
-import lombok.NonNull;
-import lombok.RequiredArgsConstructor;
+import lombok.ToString;
+import org.apache.commons.lang3.tuple.Pair;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.Node;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 
 /** AST node represent FillNull operation. */
-@RequiredArgsConstructor
-@AllArgsConstructor
+@Getter
+@EqualsAndHashCode(callSuper = false)
+@ToString
 public class FillNull extends UnresolvedPlan {
 
-  @Getter
-  @RequiredArgsConstructor
-  public static class NullableFieldFill {
-    @NonNull private final Field nullableFieldReference;
-    @NonNull private final UnresolvedExpression replaceNullWithMe;
+  public static FillNull ofVariousValue(List<Pair<Field, UnresolvedExpression>> replacements) {
+    return new FillNull(replacements);
   }
 
-  public interface ContainNullableFieldFill {
-    List<NullableFieldFill> getNullFieldFill();
-
-    static ContainNullableFieldFill ofVariousValue(List<NullableFieldFill> replacements) {
-      return new VariousValueNullFill(replacements);
-    }
-
-    static ContainNullableFieldFill ofSameValue(
-        UnresolvedExpression replaceNullWithMe, List<Field> nullableFieldReferences) {
-      return new SameValueNullFill(replaceNullWithMe, nullableFieldReferences);
+  public static FillNull ofSameValue(UnresolvedExpression replacement, List<Field> fieldList) {
+    List<Pair<Field, UnresolvedExpression>> replacementPairs =
+        fieldList.stream().map(f -> Pair.of(f, replacement)).toList();
+    FillNull instance = new FillNull(replacementPairs);
+    if (replacementPairs.isEmpty()) {
+      // no field specified, the replacement value will be applied to all fields.
+      instance.replacementForAll = Optional.of(replacement);
     }
+    return instance;
   }
 
-  private static class SameValueNullFill implements ContainNullableFieldFill {
-    @Getter private final List<NullableFieldFill> nullFieldFill;
+  private Optional<UnresolvedExpression> replacementForAll = Optional.empty();
 
-    public SameValueNullFill(
-        UnresolvedExpression replaceNullWithMe, List<Field> nullableFieldReferences) {
-      Objects.requireNonNull(replaceNullWithMe, "Null replacement is required");
-      this.nullFieldFill =
-          Objects.requireNonNull(nullableFieldReferences, "Nullable field reference is required")
-              .stream()
-              .map(nullableReference -> new NullableFieldFill(nullableReference, replaceNullWithMe))
-              .toList();
-    }
-  }
+  private final List<Pair<Field, UnresolvedExpression>> replacementPairs;
 
-  @RequiredArgsConstructor
-  private static class VariousValueNullFill implements ContainNullableFieldFill {
-    @NonNull @Getter private final List<NullableFieldFill> nullFieldFill;
+  FillNull(List<Pair<Field, UnresolvedExpression>> replacementPairs) {
+    this.replacementPairs = replacementPairs;
   }
 
   private UnresolvedPlan child;
 
-  @NonNull private final ContainNullableFieldFill containNullableFieldFill;
-
-  public List<NullableFieldFill> getNullableFieldFills() {
-    return containNullableFieldFill.getNullFieldFill();
+  public List<Field> getFields() {
+    return getReplacementPairs().stream().map(Pair::getLeft).toList();
   }
 
   @Override
-  public UnresolvedPlan attach(UnresolvedPlan child) {
+  public FillNull attach(UnresolvedPlan child) {
     this.child = child;
     return this;
   }

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

@@ -29,7 +29,8 @@ public class RareTopN extends UnresolvedPlan {
 
   private UnresolvedPlan child;
   private final CommandType commandType;
-  private final List<Argument> noOfResults;
+  // arguments: noOfResults: Integer, countField: String, showCount: Boolean
+  private final List<Argument> arguments;
   private final List<Field> fields;
   private final List<UnresolvedExpression> groupExprList;
 

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

@@ -37,7 +37,8 @@ public AggCall visitAlias(Alias node, CalcitePlanContext context) {
 
   @Override
   public AggCall visitAggregateFunction(AggregateFunction node, CalcitePlanContext context) {
-    RexNode field = rexNodeVisitor.analyze(node.getField(), context);
+    RexNode field =
+        node.getField() == null ? null : rexNodeVisitor.analyze(node.getField(), context);
     List<RexNode> argList = new ArrayList<>();
     for (UnresolvedExpression arg : node.getArgList()) {
       argList.add(rexNodeVisitor.analyze(arg, context));

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

@@ -31,6 +31,7 @@ public class CalcitePlanContext {
   public final ExtendedRexBuilder rexBuilder;
   public final FunctionProperties functionProperties;
   public final QueryType queryType;
+  public final Integer querySizeLimit;
 
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
@@ -46,8 +47,9 @@ public class CalcitePlanContext {
   private final Stack<RexCorrelVariable> correlVar = new Stack<>();
   private final Stack<List<RexNode>> windowPartitions = new Stack<>();
 
-  private CalcitePlanContext(FrameworkConfig config, QueryType queryType) {
+  private CalcitePlanContext(FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
     this.config = config;
+    this.querySizeLimit = querySizeLimit;
     this.queryType = queryType;
     this.connection = CalciteToolsHelper.connect(config, TYPE_FACTORY);
     this.relBuilder = CalciteToolsHelper.create(config, TYPE_FACTORY, connection);
@@ -84,7 +86,8 @@ public Optional<RexCorrelVariable> peekCorrelVar() {
     }
   }
 
-  public static CalcitePlanContext create(FrameworkConfig config, QueryType queryType) {
-    return new CalcitePlanContext(config, queryType);
+  public static CalcitePlanContext create(
+      FrameworkConfig config, Integer querySizeLimit, QueryType queryType) {
+    return new CalcitePlanContext(config, querySizeLimit, queryType);
   }
 }