fix: execute "select count(*) from tbl" always getting zero (#114)

* fix: execute "select count(*) from tbl" always getting zero * docs: explain this special case in comment Signed-off-by: Ruihang Xia <[email protected]> --------- Signed-off-by: Ruihang Xia <[email protected]> Co-authored-by: Ruihang Xia <[email protected]>
datafusion-contrib · Aug 7, 2024 · 67db60d · 67db60d
1 parent 56c9497
commit 67db60d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/examples/datafusion_integration.rs b/examples/datafusion_integration.rs
@@ -20,6 +20,8 @@ async fn main() -> Result<()> {
         .show()
         .await?;
 
+    ctx.sql("select count(*) from table1").await?.show().await?;
+
     ctx.read_orc(
         "tests/basic/data/alltypes.snappy.orc",
         OrcReadOptions::default(),

diff --git a/src/array_decoder/mod.rs b/src/array_decoder/mod.rs
@@ -7,7 +7,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field};
 use arrow::datatypes::{
     Date32Type, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type, Int8Type, SchemaRef,
 };
-use arrow::record_batch::RecordBatch;
+use arrow::record_batch::{RecordBatch, RecordBatchOptions};
 use snafu::{ensure, ResultExt};
 
 use crate::column::{get_present_vec, Column};
@@ -569,7 +569,21 @@ impl NaiveStripeDecoder {
         let fields = self.inner_decode_next_batch(remaining)?;
 
         if fields.is_empty() {
-            Ok(None)
+            if remaining == 0 {
+                Ok(None)
+            } else {
+                // In case of empty projection, we need to create a RecordBatch with `row_count` only
+                // to reflect the row number
+                Ok(Some(
+                    RecordBatch::try_new_with_options(
+                        Arc::clone(&self.schema_ref),
+                        fields,
+                        &RecordBatchOptions::new()
+                            .with_row_count(Some(self.batch_size.min(remaining))),
+                    )
+                    .context(error::ConvertRecordBatchSnafu)?,
+                ))
+            }
         } else {
             //TODO(weny): any better way?
             let fields = self