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

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/integration/IndexRebuildIntegrationTest.java

@@ -87,23 +87,24 @@ public class IndexRebuildIntegrationTest extends AbstractBasicIntegrationTest {
 
         client.cluster().state(ClusterState.ACTIVE);
 
-        executeSql("CREATE TABLE tbl (id INT PRIMARY KEY, val VARCHAR, val2 VARCHAR) WITH CACHE_NAME=\"test\"");
+        executeSql("CREATE TABLE tbl (id INT, id2 INT, val VARCHAR, val2 VARCHAR, PRIMARY KEY(id, id2)) " +
+            "WITH CACHE_NAME=\"test\",affinity_key=id");
         executeSql("CREATE INDEX idx_id_val ON tbl (id DESC, val)");
         executeSql("CREATE INDEX idx_id_val2 ON tbl (id, val2 DESC)");
 
         for (int i = 0; i < 100; i++)
-            executeSql("INSERT INTO tbl VALUES (?, ?, ?)", i, "val" + i, "val" + i);
+            executeSql("INSERT INTO tbl VALUES (?, ?, ?, ?)", i, i, "val" + i, "val" + i);
 
-        executeSql("CREATE TABLE tbl2 (id INT PRIMARY KEY, val VARCHAR)");
+        executeSql("CREATE TABLE tbl2 (id INT, id2 INT, val VARCHAR, PRIMARY KEY(id, id2)) WITH affinity_key=id");
 
         for (int i = 0; i < 100; i++)
-            executeSql("INSERT INTO tbl2 VALUES (?, ?)", i, "val" + i);
+            executeSql("INSERT INTO tbl2 VALUES (?, ?, ?)", i, i, "val" + i);
     }
 
     /** */
     @Test
     public void testRebuildOnInitiatorNode() throws Exception {
-        String sql = "SELECT * FROM tbl WHERE id = 0 AND val='val0'";
+        String sql = "SELECT id, val, val2 FROM tbl WHERE id = 0 AND val='val0'";
 
         QueryChecker validChecker = assertQuery(grid(0), sql)
             .matches(QueryChecker.containsIndexScan("PUBLIC", "TBL", "IDX_ID_VAL"))
@@ -123,7 +124,7 @@ public void testRebuildOnRemoteNodeUncorrelated() throws Exception {
 
         // Uncorrelated, filter by indexed field, without projection (idenitity projection).
         for (int i = 0; i < 10; i++) {
-            QueryChecker checker = assertQuery(initNode, "SELECT * FROM tbl WHERE id = ? AND val = ?")
+            QueryChecker checker = assertQuery(initNode, "SELECT id, val, val2 FROM tbl WHERE id = ? AND val = ?")
                 .withParams(i, "val" + i)
                 .matches(CoreMatchers.not(QueryChecker.containsSubPlan("IgniteSort")))
                 .matches(QueryChecker.containsIndexScan("PUBLIC", "TBL", "IDX_ID_VAL"))
@@ -150,7 +151,7 @@ public void testRebuildOnRemoteNodeSorted() throws Exception {
         IgniteEx initNode = grid(0);
 
         // Order by part of index collation, without projection.
-        String sql = "SELECT * FROM tbl WHERE id >= 10 and id <= 15 ORDER BY id DESC";
+        String sql = "SELECT id, val, val2 FROM tbl WHERE id >= 10 and id <= 15 ORDER BY id DESC";
 
         QueryChecker checker = assertQuery(initNode, sql)
             .matches(CoreMatchers.not(QueryChecker.containsSubPlan("IndexSpool")))
@@ -192,7 +193,7 @@ public void testRebuildOnRemoteNodeSorted() throws Exception {
         executeSql("CREATE INDEX idx_val ON tbl (val DESC)");
 
         try {
-            sql = "SELECT * FROM tbl WHERE val BETWEEN 'val10' AND 'val15' ORDER BY val";
+            sql = "SELECT id, val, val2 FROM tbl WHERE val BETWEEN 'val10' AND 'val15' ORDER BY val";
 
             checker = assertQuery(initNode, sql)
                 .matches(QueryChecker.containsSubPlan("IgniteSort"))

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,20 @@ 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;
+        };
+    }
 }