[CALCITE-6236] EnumerableBatchNestedLoopJoin::estimateRowCount returns wrong value

Author: kramerul

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

@@ -55,6 +55,7 @@
 public class EnumerableBatchNestedLoopJoin extends Join implements EnumerableRel {
 
   private final ImmutableBitSet requiredColumns;
+  private final @Nullable Double originRowCount;
   protected EnumerableBatchNestedLoopJoin(
       RelOptCluster cluster,
       RelTraitSet traits,
@@ -63,18 +64,32 @@ protected EnumerableBatchNestedLoopJoin(
       RexNode condition,
       Set<CorrelationId> variablesSet,
       ImmutableBitSet requiredColumns,
-      JoinRelType joinType) {
+      JoinRelType joinType,
+      @Nullable Double originRowCount) {
     super(cluster, traits, ImmutableList.of(), left, right, condition, variablesSet, joinType);
     this.requiredColumns = requiredColumns;
+    this.originRowCount = originRowCount;
   }
 
+  @Deprecated
   public static EnumerableBatchNestedLoopJoin create(
       RelNode left,
       RelNode right,
       RexNode condition,
       ImmutableBitSet requiredColumns,
       Set<CorrelationId> variablesSet,
       JoinRelType joinType) {
+    return create(left, right, condition, requiredColumns, variablesSet, joinType, null);
+  }
+
+  public static EnumerableBatchNestedLoopJoin create(
+      RelNode left,
+      RelNode right,
+      RexNode condition,
+      ImmutableBitSet requiredColumns,
+      Set<CorrelationId> variablesSet,
+      JoinRelType joinType,
+      @Nullable Double originRowCount) {
     final RelOptCluster cluster = left.getCluster();
     final RelMetadataQuery mq = cluster.getMetadataQuery();
     final RelTraitSet traitSet =
@@ -89,9 +104,11 @@ public static EnumerableBatchNestedLoopJoin create(
         condition,
         variablesSet,
         requiredColumns,
-        joinType);
+        joinType,
+        originRowCount);
   }
 
+
   @Override public @Nullable Pair<RelTraitSet, List<RelTraitSet>> passThroughTraits(
       final RelTraitSet required) {
     return EnumerableTraitsUtils.passThroughTraitsForJoin(
@@ -115,8 +132,8 @@ public static EnumerableBatchNestedLoopJoin create(
   @Override public EnumerableBatchNestedLoopJoin copy(RelTraitSet traitSet,
       RexNode condition, RelNode left, RelNode right, JoinRelType joinType,
       boolean semiJoinDone) {
-    return new EnumerableBatchNestedLoopJoin(getCluster(), traitSet,
-        left, right, condition, variablesSet, requiredColumns, joinType);
+    return new EnumerableBatchNestedLoopJoin(getCluster(), traitSet, left, right, condition,
+        variablesSet, requiredColumns, joinType, originRowCount);
   }
 
   @Override public @Nullable RelOptCost computeSelfCost(
@@ -149,6 +166,10 @@ public static EnumerableBatchNestedLoopJoin create(
     return pw.item("batchSize", variablesSet.size());
   }
 
+  public @Nullable Double getOriginRowCount() {
+    return originRowCount;
+  }
+
   @Override public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
     final BlockBuilder builder = new BlockBuilder();
     final Result leftResult =

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

@@ -32,6 +32,7 @@
 import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.tools.RelBuilderFactory;
 import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.Util;
 
 import org.immutables.value.Value;
 
@@ -137,6 +138,8 @@ public EnumerableBatchNestedLoopJoinRule(RelBuilderFactory relBuilderFactory,
     // Push a filter with batchSize disjunctions
     relBuilder.push(join.getRight()).filter(relBuilder.or(conditionList));
     final RelNode right = relBuilder.build();
+    final double originRowCount =
+        Util.first(call.getMetadataQuery().getRowCount(join), Double.MAX_VALUE);
 
     call.transformTo(
         EnumerableBatchNestedLoopJoin.create(
@@ -147,7 +150,8 @@ public EnumerableBatchNestedLoopJoinRule(RelBuilderFactory relBuilderFactory,
             join.getCondition(),
             requiredColumns.build(),
             correlationIds,
-            join.getJoinType()));
+            join.getJoinType(),
+            originRowCount));
   }
 
   /** Rule configuration. */

core/src/main/java/org/apache/calcite/rel/metadata/RelMdRowCount.java

@@ -16,6 +16,7 @@
  */
 package org.apache.calcite.rel.metadata;
 
+import org.apache.calcite.adapter.enumerable.EnumerableBatchNestedLoopJoin;
 import org.apache.calcite.adapter.enumerable.EnumerableLimit;
 import org.apache.calcite.plan.RelOptPredicateList;
 import org.apache.calcite.plan.volcano.RelSubset;
@@ -183,6 +184,12 @@ public Double getRowCount(Calc rel, RelMetadataQuery mq) {
   }
 
   public @Nullable Double getRowCount(Join rel, RelMetadataQuery mq) {
+    if (rel instanceof EnumerableBatchNestedLoopJoin) {
+      Double originRowCount = ((EnumerableBatchNestedLoopJoin) rel).getOriginRowCount();
+      if (originRowCount != null) {
+        return originRowCount;
+      }
+    }
     return RelMdUtil.getJoinRowCount(mq, rel, rel.getCondition());
   }
 

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(