[CALCITE-7327] Support IS NOT DISTINCT FROM as equi condition of hash join

Author: silundong

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableHashJoin.java

@@ -218,8 +218,10 @@ private Result implementHashSemiJoin(EnumerableRelImplementor implementor, Prefe
                 Expressions.list(
                     leftExpression,
                     rightExpression,
-                    leftResult.physType.generateAccessorWithoutNulls(joinInfo.leftKeys),
-                    rightResult.physType.generateAccessorWithoutNulls(joinInfo.rightKeys),
+                    leftResult.physType.generateNullAwareAccessor(
+                        joinInfo.leftKeys, joinInfo.nullExclusionFlags),
+                    rightResult.physType.generateNullAwareAccessor(
+                        joinInfo.rightKeys, joinInfo.nullExclusionFlags),
                     Util.first(keyPhysType.comparer(),
                         Expressions.constant(null)),
                     predicate)))
@@ -264,8 +266,10 @@ private Result implementHashJoin(EnumerableRelImplementor implementor, Prefer pr
                 BuiltInMethod.HASH_JOIN.method,
                 Expressions.list(
                     rightExpression,
-                    leftResult.physType.generateAccessorWithoutNulls(joinInfo.leftKeys),
-                    rightResult.physType.generateAccessorWithoutNulls(joinInfo.rightKeys),
+                    leftResult.physType.generateNullAwareAccessor(
+                        joinInfo.leftKeys, joinInfo.nullExclusionFlags),
+                    rightResult.physType.generateNullAwareAccessor(
+                        joinInfo.rightKeys, joinInfo.nullExclusionFlags),
                     EnumUtils.joinSelector(joinType,
                         physType,
                         ImmutableList.of(

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeJoin.java

@@ -34,6 +34,7 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.CorrelationId;
 import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinInfo;
 import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.metadata.RelMdCollation;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
@@ -71,6 +72,9 @@
  * {@link EnumerableConvention enumerable calling convention} using
  * a merge algorithm. */
 public class EnumerableMergeJoin extends Join implements EnumerableRel {
+  @SuppressWarnings("HidingField")
+  private final JoinInfo joinInfo;
+
   protected EnumerableMergeJoin(
       RelOptCluster cluster,
       RelTraitSet traits,
@@ -80,6 +84,12 @@ protected EnumerableMergeJoin(
       Set<CorrelationId> variablesSet,
       JoinRelType joinType) {
     super(cluster, traits, ImmutableList.of(), left, right, condition, variablesSet, joinType);
+    // TODO: support IS NOT DISTINCT FROM condition as join keys of MergeJoin
+    // EnumerableMergeJoin cannot use IS NOT DISTINCT FROM condition as join keys
+    // (In the algorithm of MergeJoin in Enumerable convention, it will stop
+    // when leftKey or rightKey is NULL), so we create a new JoinInfo that only
+    // considers EQUALS.
+    this.joinInfo = JoinInfo.createWithStrictEquality(left, right, condition);
     assert getConvention() instanceof EnumerableConvention;
     final List<RelCollation> leftCollations = getCollations(left.getTraitSet());
     final List<RelCollation> rightCollations = getCollations(right.getTraitSet());

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeJoinRule.java

@@ -60,7 +60,10 @@ protected EnumerableMergeJoinRule(Config config) {
 
   @Override public @Nullable RelNode convert(RelNode rel) {
     Join join = (Join) rel;
-    final JoinInfo info = join.analyzeCondition();
+    // EnumerableMergeJoin cannot use IS NOT DISTINCT FROM condition as join keys. More details
+    // in EnumerableMergeJoin.java.
+    final JoinInfo info =
+        JoinInfo.createWithStrictEquality(join.getLeft(), join.getRight(), join.getCondition());
     if (!EnumerableMergeJoin.isMergeJoinSupported(join.getJoinType())) {
       // EnumerableMergeJoin only supports certain join types.
       return null;

core/src/main/java/org/apache/calcite/adapter/enumerable/PhysType.java

@@ -134,6 +134,15 @@ Expression fieldReference(Expression expression, int field,
    */
   Expression generateAccessorWithoutNulls(List<Integer> fields);
 
+  /**
+   * Similar to {@link #generateAccessor(List)} and {@link #generateAccessorWithoutNulls(List)},
+   * but it's null-aware. It returns a Expression which evaluates to null (if one of
+   * field is null and it isn't null-safe) or a list of
+   * fields that may contain null (no field is null, or there are fields with null but they are
+   * null-safe) at runtime.
+   */
+  Expression generateNullAwareAccessor(List<Integer> fields, List<Boolean> nullExclusionFlags);
+
   /** Generates a selector for the given fields from an expression, with the
    * default row format. */
   Expression generateSelector(

core/src/main/java/org/apache/calcite/adapter/enumerable/PhysTypeImpl.java

@@ -642,6 +642,21 @@ private List<Expression> fieldReferences(
     }
   }
 
+  private static Expression getListExpressionAllowSingleElement(
+      Expressions.FluentList<Expression> list) {
+    assert list.size() > 0;
+
+    if (list.size() == 1) {
+      return Expressions.call(
+          List.class,
+          null,
+          BuiltInMethod.LIST1.method,
+          list);
+    } else {
+      return getListExpression(list);
+    }
+  }
+
   private static Expression getListExpression(Expressions.FluentList<Expression> list) {
     assert list.size() >= 2;
 
@@ -713,6 +728,41 @@ private static Expression getListExpression(Expressions.FluentList<Expression> l
     return Expressions.lambda(Function1.class, exp, v1);
   }
 
+  @Override public Expression generateNullAwareAccessor(
+      List<Integer> fields,
+      List<Boolean> nullExclusionFlags) {
+    assert fields.size() == nullExclusionFlags.size();
+    ParameterExpression v1 = Expressions.parameter(javaRowClass, "v1");
+    if (fields.isEmpty()) {
+      return Expressions.lambda(
+          Function1.class,
+          Expressions.field(
+              null,
+              BuiltInMethod.COMPARABLE_EMPTY_LIST.field),
+          v1);
+    }
+    Expressions.FluentList<Expression> list = Expressions.list();
+    for (int field : fields) {
+      list.add(fieldReference(v1, field));
+    }
+
+    // in the HashJoin key selector scenario, when there is exactly one join key and it is
+    // null-safe, a row whose join key is null must still be correctly recognized and extracted.
+    // Therefore, when list.size() == 1, this method returns a list containing a single
+    // element (which may be null) rather than returning the element directly.
+    Expression exp = getListExpressionAllowSingleElement(list);
+    for (int i = list.size() - 1; i >= 0; i--) {
+      if (nullExclusionFlags.get(i)) {
+        exp =
+            Expressions.condition(
+                Expressions.equal(list.get(i), Expressions.constant(null)),
+                Expressions.constant(null),
+                exp);
+      }
+    }
+    return Expressions.lambda(Function1.class, exp, v1);
+  }
+
   @Override public Expression fieldReference(
       Expression expression, int field) {
     return fieldReference(expression, field, null);

core/src/main/java/org/apache/calcite/plan/RelOptUtil.java

@@ -1470,11 +1470,24 @@ private static void splitJoinCondition(
     nonEquiList.add(condition);
   }
 
-  /** Builds an equi-join condition from a set of left and right keys. */
+  /** Builds an equi-join condition by conjoining EQUALS operator for each corresponding pair of
+   * leftKeys and rightKeys. */
   public static RexNode createEquiJoinCondition(
       final RelNode left, final List<Integer> leftKeys,
       final RelNode right, final List<Integer> rightKeys,
       final RexBuilder rexBuilder) {
+    List<Boolean> filterNulls = Collections.nCopies(leftKeys.size(), Boolean.TRUE);
+    return createHashJoinCondition(left, leftKeys, right, rightKeys,
+        filterNulls, rexBuilder);
+  }
+
+  /** Builds an equi-join condition by conjoining operators for each corresponding pair of
+   * leftKeys and rightKeys. The operator is EQUALS if filterNulls is true for that
+   * position, otherwise IS NOT DISTINCT FROM. */
+  public static RexNode createHashJoinCondition(
+      final RelNode left, final List<Integer> leftKeys,
+      final RelNode right, final List<Integer> rightKeys,
+      final List<Boolean> filterNulls, final RexBuilder rexBuilder) {
     final List<RelDataType> leftTypes =
         RelOptUtil.getFieldTypeList(left.getRowType());
     final List<RelDataType> rightTypes =
@@ -1484,7 +1497,11 @@ public static RexNode createEquiJoinCondition(
           @Override public RexNode get(int index) {
             final int leftKey = leftKeys.get(index);
             final int rightKey = rightKeys.get(index);
-            return rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
+            final SqlOperator operator =
+                filterNulls.get(index)
+                  ? SqlStdOperatorTable.EQUALS
+                  : SqlStdOperatorTable.IS_NOT_DISTINCT_FROM;
+            return rexBuilder.makeCall(operator,
                 rexBuilder.makeInputRef(leftTypes.get(leftKey), leftKey),
                 rexBuilder.makeInputRef(rightTypes.get(rightKey),
                     leftTypes.size() + rightKey));

core/src/main/java/org/apache/calcite/rel/core/Join.java

@@ -104,7 +104,7 @@ protected Join(
     this.condition = requireNonNull(condition, "condition");
     this.variablesSet = ImmutableSet.copyOf(variablesSet);
     this.joinType = requireNonNull(joinType, "joinType");
-    this.joinInfo = JoinInfo.createWithStrictEquality(left, right, condition);
+    this.joinInfo = JoinInfo.of(left, right, condition);
     this.hints = ImmutableList.copyOf(hints);
   }
 

core/src/main/java/org/apache/calcite/rel/core/JoinInfo.java

@@ -29,45 +29,55 @@
 import com.google.common.collect.ImmutableList;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 import static java.util.Objects.requireNonNull;
 
 /** An analyzed join condition.
  *
  * <p>It is useful for the many algorithms that care whether a join has an
- * equi-join condition.
+ * equi-join condition (contains EQUALS and IS NOT DISTINCT FROM).
  *
- * <p>You can create one using {@link #createWithStrictEquality}, or call
+ * <p>You can create one using {@link #of(RelNode, RelNode, RexNode)},
+ * {@link #createWithStrictEquality}, or call
  * {@link Join#analyzeCondition()}; many kinds of join cache their
  * join info, especially those that are equi-joins.
  *
  * @see Join#analyzeCondition() */
 public class JoinInfo {
   public final ImmutableIntList leftKeys;
   public final ImmutableIntList rightKeys;
+  // for each join key, whether it filters out nulls. If TRUE, the join key uses EQUALS semantics
+  // (not null-safe); if FALSE, it uses IS NOT DISTINCT FROM semantics (null-safe).
+  public final ImmutableList<Boolean> nullExclusionFlags;
+  // non-equi parts of join condition.
+  // after CALCITE-7327, IS NOT DISTINCT FROM can be treated as a hash join key and is no longer
+  // part of nonEquiConditions.
   public final ImmutableList<RexNode> nonEquiConditions;
 
   /** Creates a JoinInfo. */
   protected JoinInfo(ImmutableIntList leftKeys, ImmutableIntList rightKeys,
-      ImmutableList<RexNode> nonEquiConditions) {
+      ImmutableList<Boolean> nullExclusionFlags, ImmutableList<RexNode> nonEquiConditions) {
     this.leftKeys = requireNonNull(leftKeys, "leftKeys");
     this.rightKeys = requireNonNull(rightKeys, "rightKeys");
+    this.nullExclusionFlags = requireNonNull(nullExclusionFlags, "nullExclusionFlags");
     this.nonEquiConditions =
         requireNonNull(nonEquiConditions, "nonEquiConditions");
-    assert leftKeys.size() == rightKeys.size();
+    assert leftKeys.size() == rightKeys.size() && leftKeys.size() == nullExclusionFlags.size();
   }
 
   /** Creates a {@code JoinInfo} by analyzing a condition. */
   public static JoinInfo of(RelNode left, RelNode right, RexNode condition) {
     final List<Integer> leftKeys = new ArrayList<>();
     final List<Integer> rightKeys = new ArrayList<>();
-    final List<Boolean> filterNulls = new ArrayList<>();
+    final List<Boolean> nullExclusionFlags = new ArrayList<>();
     final List<RexNode> nonEquiList = new ArrayList<>();
     RelOptUtil.splitJoinCondition(left, right, condition, leftKeys, rightKeys,
-        filterNulls, nonEquiList);
+        nullExclusionFlags, nonEquiList);
     return new JoinInfo(ImmutableIntList.copyOf(leftKeys),
-        ImmutableIntList.copyOf(rightKeys), ImmutableList.copyOf(nonEquiList));
+        ImmutableIntList.copyOf(rightKeys), ImmutableList.copyOf(nullExclusionFlags),
+        ImmutableList.copyOf(nonEquiList));
   }
 
   /** Creates a {@code JoinInfo} by analyzing a condition.
@@ -82,14 +92,18 @@ public static JoinInfo createWithStrictEquality(RelNode left,
     final List<RexNode> nonEquiList = new ArrayList<>();
     RelOptUtil.splitJoinCondition(left, right, condition, leftKeys, rightKeys,
         null, nonEquiList);
+    List<Boolean> nullExclusionFlags = Collections.nCopies(leftKeys.size(), Boolean.TRUE);
     return new JoinInfo(ImmutableIntList.copyOf(leftKeys),
-        ImmutableIntList.copyOf(rightKeys), ImmutableList.copyOf(nonEquiList));
+        ImmutableIntList.copyOf(rightKeys), ImmutableList.copyOf(nullExclusionFlags),
+        ImmutableList.copyOf(nonEquiList));
   }
 
-  /** Creates an equi-join. */
+  /** Creates an equi-join (only considers EQUALS operations). */
   public static JoinInfo of(ImmutableIntList leftKeys,
       ImmutableIntList rightKeys) {
-    return new JoinInfo(leftKeys, rightKeys, ImmutableList.of());
+    List<Boolean> nullExclusionFlags = Collections.nCopies(leftKeys.size(), Boolean.TRUE);
+    return new JoinInfo(leftKeys, rightKeys,
+        ImmutableList.copyOf(nullExclusionFlags), ImmutableList.of());
   }
 
   /** Returns whether this is an equi-join. */
@@ -117,7 +131,7 @@ public RexNode getRemaining(RexBuilder rexBuilder) {
 
   public RexNode getEquiCondition(RelNode left, RelNode right,
       RexBuilder rexBuilder) {
-    return RelOptUtil.createEquiJoinCondition(left, leftKeys, right, rightKeys,
+    return RelOptUtil.createHashJoinCondition(left, leftKeys, right, rightKeys, nullExclusionFlags,
         rexBuilder);
   }
 

core/src/main/java/org/apache/calcite/rel/rules/LoptOptimizeJoinRule.java

@@ -2053,7 +2053,7 @@ public static boolean isRemovableSelfJoin(Join joinRel) {
    */
   private static boolean areSelfJoinKeysUnique(RelMetadataQuery mq,
       RelNode leftRel, RelNode rightRel, RexNode joinFilters) {
-    final JoinInfo joinInfo = JoinInfo.createWithStrictEquality(leftRel, rightRel, joinFilters);
+    final JoinInfo joinInfo = JoinInfo.of(leftRel, rightRel, joinFilters);
 
     // Make sure each key on the left maps to the same simple column as the
     // corresponding key on the right

core/src/main/java/org/apache/calcite/rel/rules/LoptSemiJoinOptimizer.java

@@ -264,7 +264,7 @@ private static int isSuitableFilter(
 
     RelNode factRel = multiJoin.getJoinFactor(factIdx);
     RelNode dimRel = multiJoin.getJoinFactor(dimIdx);
-    final JoinInfo joinInfo = JoinInfo.createWithStrictEquality(factRel, dimRel, semiJoinCondition);
+    final JoinInfo joinInfo = JoinInfo.of(factRel, dimRel, semiJoinCondition);
     assert !joinInfo.leftKeys.isEmpty();
 
     // mutable copies