[CALCITE-6236] Fix EnumerableBatchNestedLoopJoin::estimateRowCount

sap-contributions · Jan 31, 2024 · c0315af · c0315af
1 parent 75511b8
commit c0315af
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableBatchNestedLoopJoin.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableBatchNestedLoopJoin.java
@@ -31,8 +31,10 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelWriter;
 import org.apache.calcite.rel.core.CorrelationId;
+import org.apache.calcite.rel.core.Filter;
 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.rel.metadata.RelMdCollation;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rex.RexNode;
@@ -122,11 +124,12 @@ public static EnumerableBatchNestedLoopJoin create(
   @Override public @Nullable RelOptCost computeSelfCost(
       final RelOptPlanner planner,
       final RelMetadataQuery mq) {
-    double rowCount = mq.getRowCount(this);
+    double rowCount = estimateRowCount(mq);
 
     final double rightRowCount = right.estimateRowCount(mq);
     final double leftRowCount = left.estimateRowCount(mq);
-    if (Double.isInfinite(leftRowCount) || Double.isInfinite(rightRowCount)) {
+    if (Double.isInfinite(leftRowCount) || Double.isInfinite(rightRowCount)
+        || Double.isInfinite(rowCount)) {
       return planner.getCostFactory().makeInfiniteCost();
     }
 
@@ -144,6 +147,27 @@ public static EnumerableBatchNestedLoopJoin create(
         rowCount + leftRowCount, 0, 0).plus(rescanCost);
   }
 
+  @Override public  double estimateRowCount(final RelMetadataQuery mq) {
+    // See CALCITE-6236
+    // There is a filter inserted by EnumerableBatchNestedLoopJoinRule on the right side.
+    // This filter reduces the row count of the right side by it's selectivity.
+    // To get the correct row count, we need to build a join based on the original right side
+    RelNode filter = right.stripped();
+    while (!(filter instanceof Filter)) {
+      if (filter.getInputs().isEmpty()) {
+        return Double.POSITIVE_INFINITY;
+      }
+      filter = filter.getInput(0).stripped();
+    }
+    if (filter.getInputs().isEmpty()) {
+      return Double.POSITIVE_INFINITY;
+    }
+    LogicalJoin join =
+        new LogicalJoin(getCluster(), traitSet, hints, left, filter.getInput(0).stripped(),
+            condition, variablesSet, joinType, false, ImmutableList.of());
+    return mq.getRowCount(join);
+  }
+
   @Override public RelWriter explainTerms(RelWriter pw) {
     super.explainTerms(pw);
     return pw.item("batchSize", variablesSet.size());

diff --git a/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableBatchNestedLoopJoinTest.java b/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableBatchNestedLoopJoinTest.java
@@ -226,6 +226,7 @@ class EnumerableBatchNestedLoopJoinTest {
             + "join locations l on e.empid <> l.empid and d.deptno = l.empid")
         .withHook(Hook.PLANNER, (Consumer<RelOptPlanner>) planner -> {
           planner.removeRule(EnumerableRules.ENUMERABLE_CORRELATE_RULE);
+          planner.removeRule(EnumerableRules.ENUMERABLE_JOIN_RULE);
           // Use a small batch size, otherwise we will run into Janino's
           // "InternalCompilerException: Code of method grows beyond 64 KB".
           planner.addRule(