[CALCITE-7755] Support IEJoin for inequality joins

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

@@ -0,0 +1,287 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.adapter.enumerable;
+
+import org.apache.calcite.adapter.java.JavaTypeFactory;
+import org.apache.calcite.linq4j.function.Function1;
+import org.apache.calcite.linq4j.tree.BlockBuilder;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.linq4j.tree.ExpressionType;
+import org.apache.calcite.linq4j.tree.Expressions;
+import org.apache.calcite.linq4j.tree.ParameterExpression;
+import org.apache.calcite.plan.RelOptCluster;
+import org.apache.calcite.plan.RelOptCost;
+import org.apache.calcite.plan.RelOptPlanner;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.RelNodes;
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.metadata.RelMdUtil;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.sql.type.SqlTypeUtil;
+import org.apache.calcite.util.BuiltInMethod;
+import org.apache.calcite.util.Util;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.lang.reflect.Type;
+import java.util.ArrayList;
+import java.util.List;
+
+import static java.util.Objects.requireNonNull;
+
+/** Implementation of an inner IEJoin with two inequality predicates in
+ * {@link EnumerableConvention enumerable calling convention}. */
+public class EnumerableIEJoin extends Join implements EnumerableRel {
+  private final ImmutableList<Condition> conditions;
+
+  protected EnumerableIEJoin(RelOptCluster cluster, RelTraitSet traitSet,
+      RelNode left, RelNode right, RexNode condition) {
+    super(cluster, traitSet, ImmutableList.of(), left, right, condition,
+        ImmutableSet.of(), JoinRelType.INNER);
+    final List<RexNode> conjunctions = RelOptUtil.conjunctions(condition);
+    if (conjunctions.size() != 2) {
+      throw new IllegalArgumentException(
+          "condition must contain exactly two supported cross-input inequalities");
+    }
+    final int leftFieldCount = left.getRowType().getFieldCount();
+    final Condition first =
+        analyzeConjunction(conjunctions.get(0), leftFieldCount);
+    final Condition second =
+        analyzeConjunction(conjunctions.get(1), leftFieldCount);
+    if (first == null || second == null) {
+      throw new IllegalArgumentException(
+          "condition must contain supported cross-input inequalities");
+    }
+    final ImmutableList<Condition> inequalities = ImmutableList.of(first, second);
+    for (Condition inequality : inequalities) {
+      if (!supportsKeyTypes(left, right, inequality)) {
+        throw new IllegalArgumentException("unsupported IEJoin key types: left "
+            + left.getRowType().getFieldList().get(inequality.leftKey).getType()
+            + ", right "
+            + right.getRowType().getFieldList().get(inequality.rightKey).getType());
+      }
+    }
+    this.conditions = inequalities;
+  }
+
+  /** Creates an EnumerableIEJoin. */
+  public static EnumerableIEJoin create(RelNode left, RelNode right,
+      RexNode condition) {
+    return new EnumerableIEJoin(left.getCluster(),
+        left.getCluster().traitSetOf(EnumerableConvention.INSTANCE),
+        left, right, condition);
+  }
+
+  @Override public EnumerableIEJoin copy(RelTraitSet traitSet,
+      RexNode condition, RelNode left, RelNode right, JoinRelType joinType,
+      boolean semiJoinDone) {
+    if (joinType != JoinRelType.INNER) {
+      throw new IllegalArgumentException("EnumerableIEJoin only supports inner joins");
+    }
+    return new EnumerableIEJoin(getCluster(), traitSet, left, right,
+        condition);
+  }
+
+  @Override public @Nullable RelOptCost computeSelfCost(RelOptPlanner planner,
+      RelMetadataQuery mq) {
+    final Double leftRows = mq.getRowCount(left);
+    final Double rightRows = mq.getRowCount(right);
+    final Double joinRows = mq.getRowCount(this);
+    if (leftRows == null || rightRows == null || joinRows == null) {
+      return null;
+    }
+    double outputRows = joinRows;
+    if (RelNodes.COMPARATOR.compare(left, right) > 0) {
+      outputRows = RelMdUtil.addEpsilon(outputRows);
+    }
+    final double inputRows = leftRows + rightRows;
+    // Sort the combined inputs by each inequality key, then scan and emit pairs.
+    final double cost =
+        2D * Util.nLogN(inputRows) + inputRows + outputRows;
+    return planner.getCostFactory().makeCost(cost, 0, 0);
+  }
+
+  @Override public Result implement(EnumerableRelImplementor implementor,
+      Prefer pref) {
+    final BlockBuilder builder = new BlockBuilder();
+    final Result leftResult =
+        implementor.visitChild(this, 0, (EnumerableRel) left, pref);
+    final Expression leftExpression =
+        builder.append("left", leftResult.block);
+    final Result rightResult =
+        implementor.visitChild(this, 1, (EnumerableRel) right, pref);
+    final Expression rightExpression =
+        builder.append("right", rightResult.block);
+    final ParameterExpression leftParameter =
+        Expressions.parameter(leftResult.physType.getJavaRowType(), "leftRow");
+    final ParameterExpression rightParameter =
+        Expressions.parameter(rightResult.physType.getJavaRowType(), "rightRow");
+    final JavaTypeFactory typeFactory = implementor.getTypeFactory();
+    final List<Expression> keySelectors = new ArrayList<>();
+    final List<Expression> comparators = new ArrayList<>();
+
+    for (Condition condition : conditions) {
+      final RelDataType leftType =
+          left.getRowType().getFieldList().get(condition.leftKey).getType();
+      final RelDataType rightType =
+          right.getRowType().getFieldList().get(condition.rightKey).getType();
+      // Use SQL storage types so timestamp comparisons use millisecond precision.
+      final RelDataType keyType =
+          typeFactory.toSql(
+              requireNonNull(typeFactory.leastRestrictive(ImmutableList.of(leftType, rightType))));
+      final Type keyClass = typeFactory.getJavaClass(keyType);
+      // For nullable INTEGER keys in array rows:
+      // leftRow -> (Integer) leftRow[leftKey]
+      // rightRow -> (Integer) rightRow[rightKey]
+      keySelectors.add(
+          Expressions.lambda(
+          Function1.class,
+          EnumUtils.convert(
+              leftResult.physType.fieldReference(
+              leftParameter, condition.leftKey), keyClass), leftParameter));
+      keySelectors.add(
+          Expressions.lambda(
+          Function1.class,
+          EnumUtils.convert(
+              rightResult.physType.fieldReference(
+              rightParameter, condition.rightKey), keyClass), rightParameter));
+      // PhysType generates comparators for row fields, so wrap the key in a
+      // scalar row type.
+      // For nullable INTEGER keys: (a, b) -> Utilities.compareNullsLast(a, b)
+      final RelDataType keyRowType =
+          typeFactory.builder().add("key", keyType).build();
+      final PhysType keyPhysType =
+          PhysTypeImpl.of(typeFactory, keyRowType, JavaRowFormat.SCALAR);
+      comparators.add(
+          keyPhysType.generateComparator(
+              RelCollations.of(
+                  new RelFieldCollation(0,
+                      RelFieldCollation.Direction.ASCENDING,
+                      RelFieldCollation.NullDirection.LAST))));
+    }
+
+    final PhysType physType =
+        PhysTypeImpl.of(typeFactory, getRowType(), pref.preferArray());
+    final List<Expression> arguments = new ArrayList<>();
+    arguments.add(leftExpression);
+    arguments.add(rightExpression);
+    arguments.addAll(keySelectors);
+    arguments.addAll(comparators);
+    arguments.add(Expressions.constant(conditions.get(0).operator));
+    arguments.add(Expressions.constant(conditions.get(1).operator));
+    // For two-field array rows, the result selector is:
+    // (leftRow, rightRow) -> new Object[] {
+    //     leftRow[0], leftRow[1], rightRow[0], rightRow[1] }
+    arguments.add(
+        EnumUtils.joinSelector(joinType, physType,
+        ImmutableList.of(leftResult.physType, rightResult.physType)));
+
+    // return EnumerableDefaults.ieJoin(left, right,
+    //     leftKey1, rightKey1, leftKey2, rightKey2,
+    //     comparator1, comparator2, operator1, operator2, resultSelector);
+    return implementor.result(physType,
+        builder.append(
+            Expressions.call(BuiltInMethod.IE_JOIN.method,
+            arguments)).toBlock());
+  }
+
+  static @Nullable Condition analyzeConjunction(RexNode node,
+      int leftFieldCount) {
+    if (!(node instanceof RexCall) || ((RexCall) node).operands.size() != 2) {
+      return null;
+    }
+    final RexCall call = (RexCall) node;
+    if (!(call.operands.get(0) instanceof RexInputRef)
+        || !(call.operands.get(1) instanceof RexInputRef)) {
+      return null;
+    }
+    final int first = ((RexInputRef) call.operands.get(0)).getIndex();
+    final int second = ((RexInputRef) call.operands.get(1)).getIndex();
+    final boolean firstIsLeft = first < leftFieldCount;
+    final boolean secondIsLeft = second < leftFieldCount;
+    if (firstIsLeft == secondIsLeft) {
+      return null;
+    }
+    final ExpressionType operator;
+    switch (firstIsLeft ? call.getKind() : call.getKind().reverse()) {
+    case LESS_THAN:
+      operator = ExpressionType.LessThan;
+      break;
+    case LESS_THAN_OR_EQUAL:
+      operator = ExpressionType.LessThanOrEqual;
+      break;
+    case GREATER_THAN:
+      operator = ExpressionType.GreaterThan;
+      break;
+    case GREATER_THAN_OR_EQUAL:
+      operator = ExpressionType.GreaterThanOrEqual;
+      break;
+    default:
+      return null;
+    }
+    return firstIsLeft
+        ? new Condition(first, second - leftFieldCount, operator)
+        : new Condition(second, first - leftFieldCount, operator);
+  }
+
+  static boolean supportsKeyTypes(RelNode left, RelNode right,
+      Condition condition) {
+    final RelDataType leftType =
+        left.getRowType().getFieldList().get(condition.leftKey).getType();
+    final RelDataType rightType =
+        right.getRowType().getFieldList().get(condition.rightKey).getType();
+    final SqlTypeName typeName = leftType.getSqlTypeName();
+    // Floating-point sort order disagrees with <, <=, > and >= for NaN
+    // and signed zero.
+    return SqlTypeUtil.equalSansNullability(
+        left.getCluster().getTypeFactory(), leftType, rightType)
+        && (SqlTypeUtil.isBoolean(leftType)
+            || (SqlTypeUtil.isExactNumeric(leftType)
+                && !SqlTypeName.UNSIGNED_TYPES.contains(typeName))
+            || SqlTypeUtil.isCharacter(leftType)
+            || SqlTypeUtil.isBinary(leftType)
+            || SqlTypeUtil.isDatetime(leftType)
+            || SqlTypeUtil.isInterval(leftType));
+  }
+
+  /** A normalized IEJoin condition. */
+  static final class Condition {
+    final int leftKey;
+    final int rightKey;
+    final ExpressionType operator;
+
+    private Condition(int leftKey, int rightKey,
+        ExpressionType operator) {
+      this.leftKey = leftKey;
+      this.rightKey = rightKey;
+      this.operator = operator;
+    }
+  }
+}

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

@@ -0,0 +1,111 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.calcite.adapter.enumerable;
+
+import org.apache.calcite.plan.Convention;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.convert.ConverterRule;
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.logical.LogicalJoin;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgram;
+import org.apache.calcite.rex.RexUtil;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+import java.util.List;
+
+import static java.util.Objects.requireNonNull;
+
+/** Planner rule that converts an inner {@link LogicalJoin} whose condition
+ * consists of at least two cross-input field inequalities to an
+ * {@link EnumerableIEJoin}.
+ *
+ * <p>Based on Khayyat et al.,
+ * <a href="https://doi.org/10.14778/2831360.2831362">"Lightning Fast and Space
+ * Efficient Inequality Joins," PVLDB 8(13), 2015</a>. The first two
+ * inequalities drive the join and additional inequalities are evaluated by an
+ * {@link EnumerableCalc}.
+ *
+ * @see EnumerableRules#ENUMERABLE_IE_JOIN_RULE
+ */
+class EnumerableIEJoinRule extends ConverterRule {
+  /** Default configuration. */
+  static final Config DEFAULT_CONFIG = Config.INSTANCE
+      .withConversion(LogicalJoin.class, Convention.NONE,
+          EnumerableConvention.INSTANCE, "EnumerableIEJoinRule")
+      .withRuleFactory(EnumerableIEJoinRule::new);
+
+  /** Called from the Config. */
+  protected EnumerableIEJoinRule(Config config) {
+    super(config);
+  }
+
+  @Override public @Nullable RelNode convert(RelNode rel) {
+    final Join join = (Join) rel;
+    if (join.getJoinType() != JoinRelType.INNER
+        || !join.getVariablesSet().isEmpty()
+        || !join.getSystemFieldList().isEmpty()) {
+      return null;
+    }
+
+    final int leftFieldCount = join.getLeft().getRowType().getFieldCount();
+    final List<RexNode> conjunctions =
+        RelOptUtil.conjunctions(join.getCondition());
+    if (conjunctions.size() < 2) {
+      return null;
+    }
+    for (int i = 0; i < conjunctions.size(); i++) {
+      final EnumerableIEJoin.Condition condition =
+          EnumerableIEJoin.analyzeConjunction(conjunctions.get(i), leftFieldCount);
+      if (condition == null) {
+        return null;
+      }
+      if (i < 2
+          && !EnumerableIEJoin.supportsKeyTypes(
+              join.getLeft(), join.getRight(), condition)) {
+        return null;
+      }
+    }
+
+    final RelNode left = convert(join.getLeft(), join.getLeft().getTraitSet()
+        .replace(EnumerableConvention.INSTANCE));
+    final RelNode right = convert(join.getRight(), join.getRight().getTraitSet()
+        .replace(EnumerableConvention.INSTANCE));
+    final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
+    final RexNode ieCondition =
+        requireNonNull(RexUtil.composeConjunction(rexBuilder, conjunctions.subList(0, 2)));
+    final EnumerableIEJoin ieJoin =
+        EnumerableIEJoin.create(left, right, ieCondition);
+    if (conjunctions.size() == 2) {
+      return ieJoin;
+    }
+
+    final RexNode residual =
+        requireNonNull(
+            RexUtil.composeConjunction(rexBuilder,
+            conjunctions.subList(2, conjunctions.size())));
+    final RexProgram program =
+        RexProgram.create(ieJoin.getRowType(),
+            rexBuilder.identityProjects(ieJoin.getRowType()), residual,
+            ieJoin.getRowType(), rexBuilder);
+    return EnumerableCalc.create(ieJoin, program);
+  }
+}