fix(providers): handle empty projection pushdown for count(*) (#97) (#102)

DataFusion's projection pushdown passes `Some([])` to `TableProvider::scan`
for queries like `SELECT count(*) FROM t` that only need the row count. The
three hand-rolled providers each mishandled this:

- SQLite's `build_sql` emitted `SELECT  FROM "t"` (empty column list), which
  SQLite rejected with a parse error. `execute` also lost the row count for
  zero-column batches.
- Mongo and Redis built zero-column RecordBatches via `RecordBatch::try_new`,
  which Arrow rejects with "must either specify a row count or at least one
  column".

SQLite now emits `SELECT 1 FROM t` for empty projections, and all three
providers pass the real row count through `RecordBatchOptions::with_row_count`
so aggregates see the correct input cardinality. Delegated providers
(Postgres / MySQL / SeekDB via `datafusion-table-providers`, plus Lance and
Iceberg) are unaffected — upstream already handles this.

Regression tests (integration, `#[ignore]`d) cover bare `count(*)`,
`count(*)` with a filter, and `count(*)` on an empty table.

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: BtXin

crates/skardi/src/sources/providers/mongo/mod.rs

@@ -3,7 +3,8 @@ pub mod fts_table_function;
 
 use anyhow::{Context, Result};
 use arrow::array::{
-    ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray,
+    ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, RecordBatchOptions,
+    StringArray,
 };
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use async_trait::async_trait;
@@ -423,7 +424,17 @@ impl TableProvider for MongoTableProvider {
         let batch = if let Some(proj) = projection {
             let projected_schema = Arc::new(self.schema.project(proj)?);
             let columns: Vec<ArrayRef> = proj.iter().map(|&i| batch.column(i).clone()).collect();
-            RecordBatch::try_new(projected_schema, columns)?
+            if proj.is_empty() {
+                // DataFusion pushes an empty projection for `count(*)`-style
+                // queries where only the row count matters. `RecordBatch::try_new`
+                // rejects a zero-column batch unless we supply the row count
+                // explicitly, so pass it through so aggregates see the real
+                // input cardinality.
+                let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+                RecordBatch::try_new_with_options(projected_schema, columns, &options)?
+            } else {
+                RecordBatch::try_new(projected_schema, columns)?
+            }
         } else {
             batch
         };
@@ -1989,6 +2000,49 @@ mod tests {
         assert!(total_rows(&batches) >= 2); // at least Electronics and Furniture
     }
 
+    /// Regression test for #97 (Mongo half): projection pushdown emits
+    /// `Some([])` for `count(*)` without a filter, and the Mongo provider
+    /// previously built a zero-column batch with `RecordBatch::try_new`, which
+    /// Arrow rejects with "must either specify a row count or at least one
+    /// column". Needs a dedicated collection so a bare `count(*)` has a known
+    /// value regardless of what other parallel tests do to `products`.
+    #[tokio::test]
+    #[ignore]
+    async fn test_count_star_pushdown() {
+        // Seed a scratch collection with a known row count, registering it
+        // *after* seeding so the provider's PRAGMA-equivalent picks up the rows.
+        let raw_uri = "mongodb://root:rootpass@127.0.0.1:27017";
+        let seed_client = Client::with_uri_str(raw_uri)
+            .await
+            .expect("connect to mongo");
+        let scratch = seed_client
+            .database("mydb")
+            .collection::<Document>("count_star_scratch");
+        scratch.drop().await.expect("drop scratch");
+        scratch
+            .insert_many(vec![
+                doc! { "_id": "A", "product_id": "A", "name": "a", "price": 1.0 },
+                doc! { "_id": "B", "product_id": "B", "name": "b", "price": 2.0 },
+                doc! { "_id": "C", "product_id": "C", "name": "c", "price": 3.0 },
+            ])
+            .await
+            .expect("seed scratch");
+
+        let mut ctx = SessionContext::new();
+        register_ci_collection(&mut ctx, "count_star_scratch", "product_id").await;
+
+        let batches = query_all(&ctx, "SELECT count(*) FROM count_star_scratch").await;
+        assert_eq!(total_rows(&batches), 1);
+        let counts = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(counts.value(0), 3);
+
+        scratch.drop().await.expect("drop scratch");
+    }
+
     // ─── Filter pushdown tests (integration) ────────────────────────────
 
     #[tokio::test]

crates/skardi/src/sources/providers/redis/datasource.rs

@@ -8,8 +8,8 @@ use std::{
 use anyhow::Result;
 use arrow::{
     array::{
-        ArrayRef, RecordBatch, StringBuilder, UInt64Array, as_boolean_array, as_largestring_array,
-        as_string_array,
+        ArrayRef, RecordBatch, RecordBatchOptions, StringBuilder, UInt64Array, as_boolean_array,
+        as_largestring_array, as_string_array,
     },
     datatypes::{DataType, Field, Schema, SchemaRef},
 };
@@ -822,8 +822,21 @@ where
             .map(|mut b| Arc::new(b.finish()) as ArrayRef)
             .collect();
 
-        RecordBatch::try_new(self.projected_schema.clone(), arrays)
-            .map_err(|e| DataFusionError::Execution(format!("Error building RecordBatch: {}", e)))
+        if arrays.is_empty() {
+            // DataFusion pushes an empty projection for `count(*)`-style queries
+            // where only the row count matters. `RecordBatch::try_new` rejects a
+            // zero-column batch unless we supply the row count explicitly, so
+            // pass `count` through so aggregates see the real input cardinality.
+            let options = RecordBatchOptions::new().with_row_count(Some(count));
+            RecordBatch::try_new_with_options(self.projected_schema.clone(), arrays, &options)
+                .map_err(|e| {
+                    DataFusionError::Execution(format!("Error building RecordBatch: {}", e))
+                })
+        } else {
+            RecordBatch::try_new(self.projected_schema.clone(), arrays).map_err(|e| {
+                DataFusionError::Execution(format!("Error building RecordBatch: {}", e))
+            })
+        }
     }
 }
 
@@ -2102,6 +2115,45 @@ mod tests {
         assert!(ci_total_rows(&batches) >= 2); // at least Electronics and Furniture
     }
 
+    /// Regression test for #97 (Redis half): projection pushdown emits
+    /// `Some([])` for `count(*)`, and `fetch_partition` previously built a
+    /// zero-column batch via `RecordBatch::try_new`, which Arrow rejects with
+    /// "must either specify a row count or at least one column". Uses a
+    /// dedicated table so the bare `count(*)` has a known value regardless of
+    /// what other parallel tests do to `products`.
+    #[tokio::test]
+    #[ignore]
+    async fn test_count_star_pushdown_live() {
+        let table_name = "count_star_scratch_live";
+        clear_ci_table(table_name);
+
+        let mut ctx = SessionContext::new();
+        let mut extra_options = HashMap::new();
+        extra_options.insert("columns".to_string(), "product_id,name,price".to_string());
+        register_ci_table_with_options(&mut ctx, table_name, Some(&extra_options));
+
+        ctx.sql(&format!(
+            "INSERT INTO {table_name} (product_id, name, price)
+             VALUES ('A', 'a', '1.0'), ('B', 'b', '2.0'), ('C', 'c', '3.0')"
+        ))
+        .await
+        .expect("parse insert")
+        .collect()
+        .await
+        .expect("execute insert");
+
+        let batches = ci_query_all(&ctx, &format!("SELECT count(*) FROM {table_name}")).await;
+        assert_eq!(ci_total_rows(&batches), 1);
+        let counts = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<arrow::array::Int64Array>()
+            .unwrap();
+        assert_eq!(counts.value(0), 3);
+
+        clear_ci_table(table_name);
+    }
+
     #[tokio::test]
     #[ignore]
     async fn test_empty_table_declared_schema_live() {

crates/skardi/src/sources/providers/sqlite/mod.rs

@@ -11,7 +11,7 @@ pub use vec_to_binary::register_vec_to_binary_udf;
 use anyhow::{Context, Result};
 use arrow::array::{
     ArrayRef, BinaryArray, BooleanArray, FixedSizeListArray, Float32Array, Float64Array,
-    Int64Array, ListArray, RecordBatch, StringArray, UInt64Array,
+    Int64Array, ListArray, RecordBatch, RecordBatchOptions, StringArray, UInt64Array,
 };
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use async_trait::async_trait;
@@ -877,9 +877,18 @@ impl SqliteScanExec {
             .map(|n| format!(" LIMIT {n}"))
             .unwrap_or_default();
 
-        Ok(format!(
-            "SELECT {} FROM {table}{where_clause}{limit_clause}",
+        // DataFusion pushes down an empty projection for queries like `count(*)`
+        // where only the row count is needed. SQLite rejects `SELECT  FROM t`, so
+        // emit a constant projection that preserves row count (`SELECT 1 FROM t`)
+        // and let `execute` return a zero-column batch with the correct row count.
+        let projection_clause = if columns.is_empty() {
+            "1".to_string()
+        } else {
             columns.join(", ")
+        };
+
+        Ok(format!(
+            "SELECT {projection_clause} FROM {table}{where_clause}{limit_clause}"
         ))
     }
 }
@@ -940,12 +949,13 @@ impl ExecutionPlan for SqliteScanExec {
             .collect();
 
         let future = async move {
-            let batch: Vec<Vec<tokio_rusqlite::rusqlite::types::Value>> = conn
-                .call(
+            let (batch, row_count): (Vec<Vec<tokio_rusqlite::rusqlite::types::Value>>, usize) =
+                conn.call(
                     move |conn| -> std::result::Result<_, tokio_rusqlite::rusqlite::Error> {
                         let mut stmt = conn.prepare(&sql)?;
                         let mut col_values: Vec<Vec<tokio_rusqlite::rusqlite::types::Value>> =
                             (0..num_cols).map(|_| Vec::new()).collect();
+                        let mut row_count: usize = 0;
 
                         let mut rows = stmt.query([])?;
                         while let Some(row) = rows.next()? {
@@ -954,9 +964,10 @@ impl ExecutionPlan for SqliteScanExec {
                                     row.get(col_idx)?;
                                 col_values[col_idx].push(val);
                             }
+                            row_count += 1;
                         }
 
-                        Ok(col_values)
+                        Ok((col_values, row_count))
                     },
                 )
                 .await
@@ -969,7 +980,16 @@ impl ExecutionPlan for SqliteScanExec {
                 .map(|(values, data_type)| sqlite_values_to_arrow(&values, data_type))
                 .collect();
 
-            RecordBatch::try_new(output_schema, arrays).map_err(DataFusionError::from)
+            if num_cols == 0 {
+                // Zero-column batch (e.g. from `count(*)` projection pushdown).
+                // RecordBatch::try_new would return a 0-row batch; pass the row
+                // count explicitly so aggregates see the real input cardinality.
+                let options = RecordBatchOptions::new().with_row_count(Some(row_count));
+                RecordBatch::try_new_with_options(output_schema, arrays, &options)
+                    .map_err(DataFusionError::from)
+            } else {
+                RecordBatch::try_new(output_schema, arrays).map_err(DataFusionError::from)
+            }
         };
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
@@ -1732,6 +1752,75 @@ mod tests {
         assert_eq!(total_rows(&batches), 2);
     }
 
+    /// Regression test for #97: `count(*)` was rewritten to an empty projection
+    /// (`SELECT  FROM "t"`), which SQLite rejects with a syntax error.
+    #[tokio::test]
+    #[ignore]
+    async fn test_count_star_pushdown() {
+        let db_path = create_test_db().await;
+        let db = db_path.to_str().unwrap();
+        let mut ctx = SessionContext::new();
+        register_test_table(&mut ctx, db).await;
+
+        let batches = query_all(&ctx, "SELECT count(*) FROM test_items").await;
+        assert_eq!(total_rows(&batches), 1);
+
+        let counts = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(counts.value(0), 3);
+    }
+
+    /// `count(*)` combined with a WHERE clause still needs the projection pushdown
+    /// to produce the correct row count after filtering.
+    #[tokio::test]
+    #[ignore]
+    async fn test_count_star_with_filter() {
+        let db_path = create_test_db().await;
+        let db = db_path.to_str().unwrap();
+        let mut ctx = SessionContext::new();
+        register_test_table(&mut ctx, db).await;
+
+        let batches = query_all(&ctx, "SELECT count(*) FROM test_items WHERE id > 1").await;
+        assert_eq!(total_rows(&batches), 1);
+
+        let counts = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(counts.value(0), 2);
+    }
+
+    /// `count(*)` over an empty table must return 0, not a SQLite syntax error.
+    #[tokio::test]
+    #[ignore]
+    async fn test_count_star_empty_table() {
+        let db_path = create_test_db().await;
+        let db = db_path.to_str().unwrap();
+        let mut ctx = SessionContext::new();
+        register_test_table(&mut ctx, db).await;
+
+        ctx.sql("DELETE FROM test_items")
+            .await
+            .expect("parse delete")
+            .collect()
+            .await
+            .expect("execute delete");
+
+        let batches = query_all(&ctx, "SELECT count(*) FROM test_items").await;
+        assert_eq!(total_rows(&batches), 1);
+
+        let counts = batches[0]
+            .column(0)
+            .as_any()
+            .downcast_ref::<Int64Array>()
+            .unwrap();
+        assert_eq!(counts.value(0), 0);
+    }
+
     // ─── Insert test ────────────────────────────────────────────────────
 
     #[tokio::test]