IGNITE-23235 SQL Calcite: Fix not-null count calculation if index reb…

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/LogicalRelImplementor.java

@@ -447,8 +447,16 @@ private boolean hasExchange(RelNode rel) {
         else {
             CollectNode<Row> replacement = CollectNode.createCountCollector(ctx);
 
-            replacement.register(new ScanStorageNode<>(tbl.name(), ctx, rel.getTable().getRowType(), tbl.scan(ctx,
-                ctx.group(rel.sourceId()), ImmutableBitSet.of(0))));
+            replacement.register(
+                new ScanStorageNode<>(
+                    tbl.name(),
+                    ctx,
+                    rel.getTable().getRowType(),
+                    tbl.scan(ctx, ctx.group(rel.sourceId()), ImmutableBitSet.of(rel.fieldIndex())),
+                    rel.notNull() ? r -> ctx.rowHandler().get(0, r) != null : null,
+                    null
+                )
+            );
 
             return replacement;
         }

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

@@ -52,6 +52,9 @@ public class IgniteIndexCount extends AbstractRelNode implements SourceAwareIgni
     /** */
     private final boolean notNull;
 
+    /** */
+    private final int fieldIdx;
+
     /**
      * Constructor for deserialization.
      *
@@ -70,6 +73,7 @@ public IgniteIndexCount(RelInput input) {
             sourceId = -1L;
 
         notNull = input.getBoolean("notNull", false);
+        fieldIdx = ((Number)input.get("fieldIdx")).intValue();
     }
 
     /**
@@ -86,9 +90,10 @@ public IgniteIndexCount(
         RelTraitSet traits,
         RelOptTable tbl,
         String idxName,
-        boolean notNull
+        boolean notNull,
+        int fieldIdx
     ) {
-        this(-1, cluster, traits, tbl, idxName, notNull);
+        this(-1, cluster, traits, tbl, idxName, notNull, fieldIdx);
     }
 
     /**
@@ -107,14 +112,16 @@ private IgniteIndexCount(
         RelTraitSet traits,
         RelOptTable tbl,
         String idxName,
-        boolean notNull
+        boolean notNull,
+        int fieldIdx
     ) {
         super(cluster, traits);
 
         this.idxName = idxName;
         this.tbl = tbl;
         this.sourceId = sourceId;
         this.notNull = notNull;
+        this.fieldIdx = fieldIdx;
     }
 
     /** {@inheritDoc} */
@@ -156,7 +163,8 @@ public String indexName() {
             .item("index", idxName)
             .item("table", tbl.getQualifiedName())
             .itemIf("sourceId", sourceId, sourceId != -1L)
-            .item("notNull", notNull);
+            .item("notNull", notNull)
+            .item("fieldIdx", fieldIdx);
     }
 
     /** {@inheritDoc} */
@@ -166,16 +174,21 @@ public String indexName() {
 
     /** {@inheritDoc} */
     @Override public IgniteRel clone(RelOptCluster cluster, List<IgniteRel> inputs) {
-        return new IgniteIndexCount(sourceId, cluster, traitSet, tbl, idxName, notNull);
+        return new IgniteIndexCount(sourceId, cluster, traitSet, tbl, idxName, notNull, fieldIdx);
     }
 
     /** {@inheritDoc} */
     @Override public IgniteRel clone(long srcId) {
-        return new IgniteIndexCount(srcId, getCluster(), traitSet, tbl, idxName, notNull);
+        return new IgniteIndexCount(srcId, getCluster(), traitSet, tbl, idxName, notNull, fieldIdx);
     }
 
     /** */
     public boolean notNull() {
         return notNull;
     }
+
+    /** */
+    public int fieldIndex() {
+        return fieldIdx;
+    }
 }

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/rule/IndexCountRule.java

@@ -76,6 +76,7 @@ private IndexCountRule(IndexCountRule.Config cfg) {
 
         IgniteIndex idx = null;
         boolean notNull = false;
+        int fieldIdx = 0;
 
         if (argList.isEmpty())
             idx = table.getIndex(QueryUtils.PRIMARY_KEY_INDEX);
@@ -84,7 +85,7 @@ private IndexCountRule(IndexCountRule.Config cfg) {
                 return;
 
             notNull = true;
-            int fieldIdx = argList.get(0);
+            fieldIdx = argList.get(0);
 
             if (!scan.requiredColumns().isEmpty())
                 fieldIdx = scan.requiredColumns().nth(fieldIdx);
@@ -121,7 +122,8 @@ else if (table.distribution().getType() == RelDistribution.Type.HASH_DISTRIBUTED
             idxTraits,
             scan.getTable(),
             idx.name(),
-            notNull
+            notNull,
+            fieldIdx
         );
 
         RelBuilder b = call.builder();

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/exec/LogicalRelImplementorTest.java

@@ -200,9 +200,9 @@ private void checkIndexFirstOrLastRewriter(boolean first) {
     @Test
     public void testIndexCountRewriter() {
         IgniteIndexCount idxCnt = new IgniteIndexCount(cluster, cluster.traitSet(),
-            qctx.catalogReader().getTable(F.asList("PUBLIC", "TBL")), QueryUtils.PRIMARY_KEY_INDEX, false);
+            qctx.catalogReader().getTable(F.asList("PUBLIC", "TBL")), QueryUtils.PRIMARY_KEY_INDEX, false, 0);
 
-        checkCollectNode(relImplementor.visit(idxCnt));
+        checkCollectNode(relImplementor.visit(idxCnt), false);
 
         tbl.addIndex(QueryUtils.PRIMARY_KEY_INDEX, 2);
 
@@ -215,15 +215,22 @@ public void testIndexCountRewriter() {
 
         tbl.markIndexRebuildInProgress(true);
 
-        checkCollectNode(relImplementor.visit(idxCnt));
+        checkCollectNode(relImplementor.visit(idxCnt), false);
+
+        // Check not-null filter.
+        idxCnt = new IgniteIndexCount(cluster, cluster.traitSet(),
+            qctx.catalogReader().getTable(F.asList("PUBLIC", "TBL")), QueryUtils.PRIMARY_KEY_INDEX, true, 2);
+
+        checkCollectNode(relImplementor.visit(idxCnt), true);
     }
 
     /** */
-    private void checkCollectNode(Node<Object[]> node) {
+    private void checkCollectNode(Node<Object[]> node, boolean hasFilter) {
         assertTrue(node instanceof CollectNode);
         assertTrue(node.sources() != null && node.sources().size() == 1);
         assertTrue(node.sources().get(0) instanceof ScanNode);
         assertNull(node.sources().get(0).sources());
+        assertEquals(hasFilter, ((ScanNode<?>)node.sources().get(0)).filter() != null);
         assertEquals(tbl.getRowType(tf), node.sources().get(0).rowType());
     }
 

modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/IndexRebuildIntegrationTest.java

@@ -285,8 +285,10 @@ public void testIndexCountAtUnavailableIndex() throws IgniteCheckedException {
 
             executeSql(ddl);
 
+            executeSql("CREATE INDEX idx_val ON tbl3(val)");
+
             for (int i = 0; i < records; i++)
-                executeSql("INSERT INTO tbl3 VALUES (?, ?, ?)", i, "val" + i, "val" + i);
+                executeSql("INSERT INTO tbl3 VALUES (?, ?, ?)", i, i % 2 == 0 ? null : "val" + i, "val" + i);
 
             IgniteCacheTable tbl3 = (IgniteCacheTable)srvEngine.schemaHolder().schema("PUBLIC").getTable("TBL3");
 
@@ -300,8 +302,10 @@ public void testIndexCountAtUnavailableIndex() throws IgniteCheckedException {
             });
 
             try {
-                for (int i = 0; i < iterations; ++i)
+                for (int i = 0; i < iterations; ++i) {
                     assertQuery("select COUNT(*) from tbl3").returns((long)records).check();
+                    assertQuery("select COUNT(val) from tbl3").returns((long)records / 2L).check();
+                }
             }
             finally {
                 stop.set(true);