Merge remote-tracking branch 'origin/main' into type-checker

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

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

@@ -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/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/calcite/CalciteRelNodeVisitor.java

@@ -26,6 +26,7 @@
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.ViewExpanders;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexCorrelVariable;
 import org.apache.calcite.rex.RexInputRef;
@@ -37,7 +38,7 @@
 import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.tools.RelBuilder.AggCall;
 import org.apache.calcite.util.Holder;
-import org.apache.calcite.util.Pair;
+import org.apache.commons.lang3.tuple.Pair;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
 import org.opensearch.sql.ast.Node;
@@ -411,16 +412,16 @@ public RelNode visitAggregation(Aggregation node, CalcitePlanContext context) {
     //                   \- Scan t
     Pair<List<AggCall>, List<RexNode>> resolved = resolveAggCallAndGroupBy(node, context);
     List<RexInputRef> trimmedRefs = new ArrayList<>();
-    trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.right)); // group-by keys first
-    trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.left));
+    trimmedRefs.addAll(PlanUtils.getInputRefs(resolved.getRight())); // group-by keys first
+    trimmedRefs.addAll(PlanUtils.getInputRefsFromAggCall(resolved.getLeft()));
     context.relBuilder.project(trimmedRefs);
 
     // Re-resolve aggCalls and group-by list based on adding trimmed Project.
     // Using re-resolving rather than Calcite Mapping (ref Calcite ProjectTableScanRule)
     // because that Mapping only works for RexNode, but we need both AggCall and RexNode list.
     Pair<List<AggCall>, List<RexNode>> reResolved = resolveAggCallAndGroupBy(node, context);
-    List<AggCall> aggList = reResolved.left;
-    List<RexNode> groupByList = reResolved.right;
+    List<AggCall> aggList = reResolved.getLeft();
+    List<RexNode> groupByList = reResolved.getRight();
     context.relBuilder.aggregate(context.relBuilder.groupKey(groupByList), aggList);
 
     // schema reordering
@@ -670,6 +671,42 @@ public RelNode visitWindow(Window node, CalcitePlanContext context) {
     return context.relBuilder.peek();
   }
 
+  @Override
+  public RelNode visitFillNull(FillNull node, CalcitePlanContext context) {
+    visitChildren(node, context);
+    if (node.getFields().size()
+        != new HashSet<>(node.getFields().stream().map(f -> f.getField().toString()).toList())
+            .size()) {
+      throw new IllegalArgumentException("The field list cannot be duplicated in fillnull");
+    }
+    List<RexNode> projects = new ArrayList<>();
+    List<RelDataTypeField> fieldsList = context.relBuilder.peek().getRowType().getFieldList();
+    for (RelDataTypeField field : fieldsList) {
+      RexNode fieldRef = context.rexBuilder.makeInputRef(field.getType(), field.getIndex());
+      boolean toReplace = false;
+      for (Pair<Field, UnresolvedExpression> pair : node.getReplacementPairs()) {
+        if (field.getName().equalsIgnoreCase(pair.getLeft().getField().toString())) {
+          RexNode replacement = rexVisitor.analyze(pair.getRight(), context);
+          RexNode coalesce = context.rexBuilder.coalesce(fieldRef, replacement);
+          RexNode coalesceWithAlias = context.relBuilder.alias(coalesce, field.getName());
+          projects.add(coalesceWithAlias);
+          toReplace = true;
+          break;
+        }
+      }
+      if (!toReplace && node.getReplacementForAll().isEmpty()) {
+        projects.add(fieldRef);
+      } else if (node.getReplacementForAll().isPresent()) {
+        RexNode replacement = rexVisitor.analyze(node.getReplacementForAll().get(), context);
+        RexNode coalesce = context.rexBuilder.coalesce(fieldRef, replacement);
+        RexNode coalesceWithAlias = context.relBuilder.alias(coalesce, field.getName());
+        projects.add(coalesceWithAlias);
+      }
+    }
+    context.relBuilder.project(projects);
+    return context.relBuilder.peek();
+  }
+
   /*
    * Unsupported Commands of PPL with Calcite for OpenSearch 3.0.0-beta
    */
@@ -703,11 +740,6 @@ public RelNode visitKmeans(Kmeans node, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("Kmeans command is unsupported in Calcite");
   }
 
-  @Override
-  public RelNode visitFillNull(FillNull fillNull, CalcitePlanContext context) {
-    throw new CalciteUnsupportedException("FillNull command is unsupported in Calcite");
-  }
-
   @Override
   public RelNode visitRareTopN(RareTopN node, CalcitePlanContext context) {
     throw new CalciteUnsupportedException("Rare and Top commands are unsupported in Calcite");

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

@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.plan;
+
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.schema.TranslatableTable;
+import org.apache.calcite.schema.impl.AbstractTable;
+import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory;
+
+/**
+ * Abstract class to map the {@link org.opensearch.sql.storage.Table} and {@link
+ * org.apache.calcite.schema.Table}
+ */
+public abstract class AbstractOpenSearchTable extends AbstractTable
+    implements TranslatableTable, org.opensearch.sql.storage.Table {
+
+  @Override
+  public RelDataType getRowType(RelDataTypeFactory relDataTypeFactory) {
+    return OpenSearchTypeFactory.convertSchema(this);
+  }
+}