Support join with dynamic fields

Signed-off-by: Tomoyuki Morita <[email protected]>

Author: ykmr1224

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

@@ -1223,6 +1223,8 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
               context.relBuilder.peek(),
               context.relBuilder.literal(context.sysLimit.joinSubsearchLimit())));
     }
+    List<String> leftAllFields = context.fieldBuilder.getAllFieldNames(1);
+    List<String> rightAllFields = context.fieldBuilder.getAllFieldNames(0);
     if (node.getJoinCondition().isEmpty()) {
       // join-with-field-list grammar
       List<String> leftColumns = context.fieldBuilder.getStaticFieldNames(1);
@@ -1254,12 +1256,12 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
           || (node.getArgumentMap().get("overwrite").equals(Literal.TRUE))) {
         toBeRemovedFields =
             duplicatedFieldNames.stream()
-                .map(field -> JoinAndLookupUtils.analyzeFieldsForLookUp(field, true, context))
+                .map(field -> JoinAndLookupUtils.analyzeFieldsInRight(field, context))
                 .toList();
       } else {
         toBeRemovedFields =
             duplicatedFieldNames.stream()
-                .map(field -> JoinAndLookupUtils.analyzeFieldsForLookUp(field, false, context))
+                .map(field -> JoinAndLookupUtils.analyzeFieldsInLeft(field, context))
                 .toList();
       }
       Literal max = node.getArgumentMap().get("max");
@@ -1284,13 +1286,16 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
       if (!toBeRemovedFields.isEmpty()) {
         context.relBuilder.projectExcept(toBeRemovedFields);
       }
+      context.fieldBuilder.reorganizeDynamicFields(leftAllFields, rightAllFields);
+
       return context.relBuilder.peek();
     }
     // The join-with-criteria grammar doesn't allow empty join condition
     RexNode joinCondition =
         node.getJoinCondition()
             .map(c -> rexVisitor.analyzeJoinCondition(c, context))
             .orElse(context.relBuilder.literal(true));
+    JoinAndLookupUtils.verifyJoinConditionNotUseAnyType(joinCondition, context);
     if (node.getJoinType() == SEMI || node.getJoinType() == ANTI) {
       // semi and anti join only return left table outputs
       context.relBuilder.join(
@@ -1301,7 +1306,7 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
       // when a new project add to stack. To avoid `id0`, we will rename the `id0` to `alias.id`
       // or `tableIdentifier.id`:
       List<String> leftColumns = context.fieldBuilder.getStaticFieldNames(1);
-      List<String> rightColumns = context.fieldBuilder.getStaticFieldNames();
+      List<String> rightColumns = context.fieldBuilder.getStaticFieldNames(0);
       List<String> rightTableName =
           PlanUtils.findTable(context.relBuilder.peek()).getQualifiedName();
       // Using `table.column` instead of `catalog.database.table.column` as column prefix because
@@ -1336,6 +1341,8 @@ public RelNode visitJoin(Join node, CalcitePlanContext context) {
       }
       context.relBuilder.join(
           JoinAndLookupUtils.translateJoinType(node.getJoinType()), joinCondition);
+
+      context.fieldBuilder.reorganizeDynamicFields(leftAllFields, rightAllFields);
       JoinAndLookupUtils.renameToExpectedFields(
           rightColumnsWithAliasIfConflict, leftColumns.size(), context);
     }
@@ -1368,8 +1375,13 @@ public Void visitInputRef(RexInputRef inputRef) {
 
   private static RexNode buildJoinConditionByFieldName(
       CalcitePlanContext context, String fieldName) {
-    RexNode lookupKey = JoinAndLookupUtils.analyzeFieldsForLookUp(fieldName, false, context);
-    RexNode sourceKey = JoinAndLookupUtils.analyzeFieldsForLookUp(fieldName, true, context);
+    RexNode lookupKey = JoinAndLookupUtils.analyzeFieldsInRight(fieldName, context);
+    RexNode sourceKey = JoinAndLookupUtils.analyzeFieldsInLeft(fieldName, context);
+    if (context.fieldBuilder.isAnyType(sourceKey)) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Source key `%s` needs to be specific type. Please cast explicitly.", fieldName));
+    }
     return context.rexBuilder.equals(sourceKey, lookupKey);
   }
 
@@ -1396,6 +1408,10 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
     // Get lookupColumns from top of stack (after above potential projection).
     List<String> lookupTableFieldNames = context.fieldBuilder.getStaticFieldNames();
 
+    // For merging with dynamic fields later
+    List<String> leftAllFields = context.fieldBuilder.getAllFieldNames(1);
+    List<String> rightAllFields = context.fieldBuilder.getAllFieldNames(0);
+
     // 3. Find fields which should be removed in lookup-table.
     // For lookup table, the mapping fields should be dropped after join
     // unless they are explicitly put in the output fields
@@ -1409,6 +1425,7 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
             .toList();
     List<RexNode> toBeRemovedLookupFields =
         toBeRemovedLookupFieldNames.stream()
+            .filter(d -> lookupTableFieldNames.contains(d))
             .map(d -> (RexNode) context.fieldBuilder.staticField(2, 1, d))
             .toList();
     List<RexNode> toBeRemovedFields = new ArrayList<>(toBeRemovedLookupFields);
@@ -1420,7 +1437,7 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
 
     List<RexNode> duplicatedSourceFields =
         duplicatedFieldNamesMap.keySet().stream()
-            .map(field -> JoinAndLookupUtils.analyzeFieldsForLookUp(field, true, context))
+            .map(field -> JoinAndLookupUtils.analyzeFieldsInLeft(field, context))
             .toList();
     // Duplicated fields in source-field should always be removed.
     toBeRemovedFields.addAll(duplicatedSourceFields);
@@ -1432,7 +1449,7 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
     if (!duplicatedFieldNamesMap.isEmpty() && node.getOutputStrategy() == OutputStrategy.APPEND) {
       List<RexNode> duplicatedProvidedFields =
           duplicatedFieldNamesMap.values().stream()
-              .map(field -> JoinAndLookupUtils.analyzeFieldsForLookUp(field, false, context))
+              .map(field -> JoinAndLookupUtils.analyzeFieldsInRight(field, context))
               .toList();
       for (int i = 0; i < duplicatedProvidedFields.size(); ++i) {
         newCoalesceList.add(
@@ -1469,7 +1486,7 @@ public RelNode visitLookup(Lookup node, CalcitePlanContext context) {
       context.relBuilder.projectExcept(toBeRemovedFields);
     }
 
-    // TODO: dedupe dynamic fields
+    context.fieldBuilder.reorganizeDynamicFields(leftAllFields, rightAllFields);
 
     // 7. Rename the fields to the expected names.
     JoinAndLookupUtils.renameToExpectedFields(

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

@@ -21,7 +21,6 @@
 import org.apache.calcite.rel.core.Project;
 import org.apache.calcite.rel.logical.LogicalAggregate;
 import org.apache.calcite.rel.logical.LogicalProject;
-import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
@@ -34,6 +33,7 @@
 import org.apache.calcite.util.mapping.Mappings;
 import org.apache.commons.lang3.tuple.Pair;
 import org.immutables.value.Value;
+import org.opensearch.sql.calcite.ExtendedRexBuilder;
 import org.opensearch.sql.calcite.rel.RelBuilderWrapper;
 import org.opensearch.sql.calcite.rel.RelFieldBuilder;
 
@@ -78,7 +78,8 @@ public void apply(RelOptRuleCall call, LogicalAggregate aggregate, LogicalProjec
 
     final RelBuilder rawRelBuilder = call.builder();
     final RelBuilderWrapper relBuilder = new RelBuilderWrapper(rawRelBuilder);
-    final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
+    final ExtendedRexBuilder rexBuilder =
+        (ExtendedRexBuilder) aggregate.getCluster().getRexBuilder();
     final RelFieldBuilder fieldBuilder = new RelFieldBuilder(rawRelBuilder, rexBuilder);
     relBuilder.push(project.getInput());
 

core/src/main/java/org/opensearch/sql/calcite/rel/QualifiedNameResolver.java

@@ -38,18 +38,29 @@ public static Optional<RexNode> resolveField(
     if (inputFieldNames.contains(fieldName)) {
       return Optional.of(context.fieldBuilder.staticField(inputCount, inputOrdinal, fieldName));
     } else if (context.fieldBuilder.isDynamicFieldsExist()) {
-      return Optional.of(context.fieldBuilder.dynamicField(fieldName));
+      return Optional.of(context.fieldBuilder.dynamicField(inputCount, inputOrdinal, fieldName));
     }
     return Optional.empty();
   }
 
+  /** Resolve field in the top of the stack */
+  public static Optional<RexNode> resolveField(String fieldName, CalcitePlanContext context) {
+    return resolveField(1, 0, fieldName, context);
+  }
+
+  /** Resolve field in the specified input. Throw exception if not found. */
   public static RexNode resolveFieldOrThrow(
       int inputCount, int inputOrdinal, String fieldName, CalcitePlanContext context) {
     return resolveField(inputCount, inputOrdinal, fieldName, context)
         .orElseThrow(
             () -> new IllegalArgumentException(String.format("Field [%s] not found.", fieldName)));
   }
 
+  /** Resolve field in the top of the stack. Throw exception if not found. */
+  public static RexNode resolveFieldOrThrow(String fieldName, CalcitePlanContext context) {
+    return resolveFieldOrThrow(1, 0, fieldName, context);
+  }
+
   /**
    * Resolves a qualified name to a RexNode based on the current context.
    *
@@ -78,6 +89,8 @@ private static RexNode resolveInJoinCondition(
 
     return resolveFieldWithAlias(nameNode, context, 2)
         .or(() -> resolveFieldWithoutAlias(nameNode, context, 2))
+        .or(() -> resolveDynamicFieldsWithAlias(nameNode, context, 2))
+        .or(() -> resolveDynamicFields(nameNode, context, 2))
         .orElseThrow(() -> getNotFoundException(nameNode));
   }
 
@@ -139,6 +152,28 @@ private static Optional<RexNode> resolveFieldWithAlias(
     return Optional.empty();
   }
 
+  private static Optional<RexNode> resolveDynamicFieldsWithAlias(
+      QualifiedName nameNode, CalcitePlanContext context, int inputCount) {
+    List<String> parts = nameNode.getParts();
+    log.debug(
+        "resolveDynamicFieldsWithAlias() called with nameNode={}, parts={}, inputCount={}",
+        nameNode,
+        parts,
+        inputCount);
+
+    if (parts.size() >= 2) {
+      // Consider first part as table alias
+      String alias = parts.get(0);
+
+      String fieldName = String.join(".", parts.subList(1, parts.size()));
+      Optional<RexNode> dynamicField =
+          tryToResolveField(alias, DYNAMIC_FIELDS_MAP, context, inputCount);
+      return dynamicField.map(field -> createItemAccess(field, fieldName, context));
+    }
+
+    return Optional.empty();
+  }
+
   private static Optional<RexNode> resolveDynamicFields(
       QualifiedName nameNode, CalcitePlanContext context, int inputCount) {
     List<String> parts = nameNode.getParts();