IGNITE-27433 SQL Calcite: Improve unique index rows count estimation

Author: alex-plekhanov

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rel/AbstractIndexScan.java

@@ -24,6 +24,8 @@
 import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelTraitSet;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.RelInput;
 import org.apache.calcite.rel.RelWriter;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
@@ -34,6 +36,7 @@
 import org.apache.calcite.util.ImmutableBitSet;
 import org.apache.ignite.internal.processors.query.calcite.externalize.RelInputEx;
 import org.apache.ignite.internal.processors.query.calcite.metadata.cost.IgniteCost;
+import org.apache.ignite.internal.processors.query.calcite.prepare.bounds.MultiBounds;
 import org.apache.ignite.internal.processors.query.calcite.prepare.bounds.SearchBounds;
 import org.apache.ignite.internal.processors.query.calcite.schema.IgniteIndex;
 import org.apache.ignite.internal.processors.query.calcite.schema.IgniteTable;
@@ -108,6 +111,14 @@ public boolean isInlineScan() {
         return false;
     }
 
+    /** {@inheritDoc} */
+    @Override public double estimateRowCount(RelMetadataQuery mq) {
+        IgniteTable tbl = table.unwrap(IgniteTable.class);
+        IgniteIndex idx = tbl.getIndex(idxName);
+
+        return adjustRowCountByUniqueBoundsScan(super.estimateRowCount(mq), idx.collation(), searchBounds);
+    }
+
     /** {@inheritDoc} */
     @Override public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) {
         double rows = table.getRowCount();
@@ -131,12 +142,14 @@ public boolean isInlineScan() {
 
             if (searchBounds != null) {
                 selectivity = mq.getSelectivity(this, RexUtil.composeConjunction(builder,
-                        Commons.transform(searchBounds, b -> b == null ? null : b.condition())));
+                    Commons.transform(searchBounds, b -> b == null ? null : b.condition())));
 
                 cost = Math.log(rows) * IgniteCost.ROW_COMPARISON_COST;
-            }
 
-            rows *= selectivity;
+                rows *= selectivity;
+
+                rows = adjustRowCountByUniqueBoundsScan(rows, idx.collation(), searchBounds);
+            }
 
             if (rows <= 0)
                 rows = 1;
@@ -148,6 +161,30 @@ public boolean isInlineScan() {
         return planner.getCostFactory().makeCost(rows, cost, 0).plus(planner.getCostFactory().makeTinyCost());
     }
 
+    /** Rows count can't exceed count of exact bounds on unique index. */
+    private static double adjustRowCountByUniqueBoundsScan(double rows, RelCollation collation, List<SearchBounds> bounds) {
+        if (bounds == null)
+            return rows;
+
+        long exactBounds = 1L;
+
+        for (RelFieldCollation fldCol : collation.getFieldCollations()) {
+            SearchBounds bound = bounds.get(fldCol.getFieldIndex());
+
+            if (bound == null || bound.type() == SearchBounds.Type.RANGE)
+                return rows;
+
+            if (bound.type() == SearchBounds.Type.MULTI) {
+                if (((MultiBounds)bound).bounds().stream().anyMatch(b -> b.type() != SearchBounds.Type.EXACT))
+                    return rows;
+                else
+                    exactBounds *= ((MultiBounds)bound).bounds().size();
+            }
+        }
+
+        return Math.min(rows, exactBounds);
+    }
+
     /** */
     public List<SearchBounds> searchBounds() {
         return searchBounds;

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/CorrelatedNestedLoopJoinPlannerTest.java

@@ -19,7 +19,6 @@
 
 import java.util.List;
 import java.util.function.Predicate;
-
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.JoinRelType;
@@ -155,6 +154,31 @@ public void testJoinPushExpressionRule() throws Exception {
         );
     }
 
+    /** Check that CNLJ is used for unique scans (when left hand return no more than 1 row). */
+    @Test
+    public void testJoinForUniqueScans() throws Exception {
+        IgniteSchema publicSchema = createSchema(
+            createTable("EMP", 2 * DEFAULT_TBL_SIZE,
+                IgniteDistributions.affinity(1, "default", "hash"),
+                "EMPNO", INTEGER, "DEPTNO", INTEGER, "NAME", VARCHAR)
+                .addIndex("PK", 0)
+                .addIndex("AFF", 1),
+            createTable("DEPT", DEFAULT_TBL_SIZE,
+                IgniteDistributions.affinity(0, "default", "hash"),
+                "DEPTNO", INTEGER, "NAME", VARCHAR)
+                .addIndex("PK", 0)
+        );
+
+        // TODO https://issues.apache.org/jira/browse/IGNITE-16334
+        // For query SELECT * FROM ... plan still not optimal and contains other join type instead of CNLJ.
+        String sql = "SELECT count(*) FROM emp e JOIN dept d ON e.deptno = d.deptno WHERE e.empno = ?";
+
+        assertPlan(sql, publicSchema, hasChildThat(isInstanceOf(IgniteCorrelatedNestedLoopJoin.class)
+            .and(input(0, isIndexScan("EMP", "PK")))
+            .and(input(1, isIndexScan("DEPT", "PK")))
+        ));
+    }
+
     /** */
     private TestTable testTable(String name) {
         return createTable(name, IgniteDistributions.broadcast(), "ID", Integer.class, "JID", Integer.class, "VAL", String.class);

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/planner/IndexSearchBoundsPlannerTest.java

@@ -22,6 +22,8 @@
 import java.util.function.Predicate;
 import java.util.stream.IntStream;
 import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexLiteral;
@@ -461,6 +463,53 @@ public void testBoundsComplex() throws Exception {
         );
     }
 
+    /** Tests row count estimation by bounds scan. */
+    @Test
+    public void testEstimateRowCount() throws Exception {
+        int t1Size = 1_000;
+        int t2Size = 5;
+
+        IgniteSchema schema = createSchema(createTable("T1", t1Size, IgniteDistributions.broadcast(),
+            "ID1", SqlTypeName.INTEGER, "ID2", SqlTypeName.INTEGER, "ID3", SqlTypeName.INTEGER)
+            .addIndex("IDX1", 0)
+            .addIndex("IDX2", 1, 2),
+            createTable("T2", t2Size, IgniteDistributions.broadcast(),
+                "ID1", SqlTypeName.INTEGER, "ID2", SqlTypeName.INTEGER, "ID3", SqlTypeName.INTEGER)
+                .addIndex("IDX1", 0, 1)
+        );
+
+        assertPlan("SELECT * FROM t1 WHERE id1 = ?", schema, isIndexScan("T1", "IDX1")
+            .and(estimatedRowCountBetween(1, 1)));
+
+        assertPlan("SELECT * FROM t1 WHERE id1 in (?, ?, ?)", schema, isIndexScan("T1", "IDX1")
+            .and(estimatedRowCountBetween(3, 3)));
+
+        // Can't estimate row count by bounds containing range.
+        assertPlan("SELECT * FROM t1 WHERE id1 > ?", schema, isIndexScan("T1", "IDX1")
+            .and(estimatedRowCountBetween(2, t1Size)));
+
+        // Index scan on IDX2 for column ID2 is not unique, can't estimate row count by bounds.
+        assertPlan("SELECT * FROM t1 WHERE id2 = ?", schema, isIndexScan("T1", "IDX2")
+            .and(estimatedRowCountBetween(2, t1Size)));
+
+        assertPlan("SELECT * FROM t1 WHERE id2 = ? AND id3 = ?", schema, isIndexScan("T1", "IDX2")
+            .and(estimatedRowCountBetween(1, 1)));
+
+        assertPlan("SELECT * FROM t1 WHERE id2 = ? AND id3 in (?, ?, ?)", schema, isIndexScan("T1", "IDX2")
+            .and(estimatedRowCountBetween(3, 3)));
+
+        assertPlan("SELECT * FROM t1 WHERE id2 in (?, ?, ?) AND id3 in (?, ?, ?)", schema, isIndexScan("T1", "IDX2")
+            .and(estimatedRowCountBetween(9, 9)));
+
+        // Can't estimate row count by bounds containing range.
+        assertPlan("SELECT * FROM t1 WHERE id2 in (?, ?, ?) AND id3 between ? and ?", schema, isIndexScan("T1", "IDX2")
+            .and(estimatedRowCountBetween(10, t1Size)));
+
+        // Row count for small table is limited by estimation by condition.
+        assertPlan("SELECT * FROM t2 WHERE id1 in (?, ?, ?) and id2 in (?, ?, ?)", schema, isIndexScan("T2", "IDX1")
+            .and(estimatedRowCountBetween(0, t2Size)));
+    }
+
     /** */
     private void assertBounds(String sql, Predicate<SearchBounds>... predicates) throws Exception {
         assertPlan(sql, publicSchema, nodeOrAnyChild(isInstanceOf(IgniteIndexScan.class)
@@ -516,4 +565,21 @@ private static boolean matchValue(Object val, RexNode bound) {
         return Objects.toString(val).equals(Objects.toString(
             bound instanceof RexLiteral ? ((RexLiteral)bound).getValueAs(val.getClass()) : bound));
     }
+
+    /** */
+    protected <T extends RelNode> Predicate<T> estimatedRowCountBetween(double min, double max) {
+        return node -> {
+            RelMetadataQuery mq = node.getCluster().getMetadataQuery();
+
+            double rowCnt = node.estimateRowCount(mq);
+
+            if (rowCnt >= min && rowCnt <= max)
+                return true;
+
+            lastErrorMsg = "Unexpected estimated row count [node=" + node + ", rowCnt=" + rowCnt + ']';
+
+            return false;
+        };
+    }
+
 }