ARROW-8255: [Rust] [DataFusion] Bug fix for COUNT(*)

This fixes a bug where `SELECT COUNT(1) FROM table` or `SELECT COUNT(*) FROM table` would fail because the projection push down rule would generate a plan where no columns would be read from the table.

Closes #6755 from andygrove/ARROW-8255

Authored-by: Andy Grove <andygrove73@gmail.com>
Signed-off-by: Andy Grove <andygrove73@gmail.com>

rust/datafusion/src/optimizer/projection_push_down.rs

@@ -114,6 +114,12 @@ impl ProjectionPushDown {
                 let mut projection: Vec<usize> = Vec::with_capacity(accum.len());
                 accum.iter().for_each(|i| projection.push(*i));
 
+                // Ensure that we are reading at least one column from the table in case the query
+                // does not reference any columns directly such as "SELECT COUNT(1) FROM table"
+                if projection.is_empty() {
+                    projection.push(0);
+                }
+
                 // sort the projection otherwise we get non-deterministic behavior
                 projection.sort();
 
@@ -123,6 +129,7 @@ impl ProjectionPushDown {
                 for i in &projection {
                     projected_fields.push(table_schema.fields()[*i].clone());
                 }
+
                 let projected_schema = Schema::new(projected_fields);
 
                 // now that the table scan is returning a different schema we need to

rust/datafusion/tests/sql.rs

@@ -309,6 +309,26 @@ fn csv_query_external_table_count() {
     assert_eq!(expected, actual);
 }
 
+#[test]
+fn csv_query_count_star() {
+    let mut ctx = ExecutionContext::new();
+    register_aggregate_csv_by_sql(&mut ctx);
+    let sql = "SELECT COUNT(*) FROM aggregate_test_100";
+    let actual = execute(&mut ctx, sql).join("\n");
+    let expected = "100".to_string();
+    assert_eq!(expected, actual);
+}
+
+#[test]
+fn csv_query_count_one() {
+    let mut ctx = ExecutionContext::new();
+    register_aggregate_csv_by_sql(&mut ctx);
+    let sql = "SELECT COUNT(1) FROM aggregate_test_100";
+    let actual = execute(&mut ctx, sql).join("\n");
+    let expected = "100".to_string();
+    assert_eq!(expected, actual);
+}
+
 fn aggr_test_schema() -> Arc<Schema> {
     Arc::new(Schema::new(vec![
         Field::new("c1", DataType::Utf8, false),