[arrow-rs] 01/01: Filter empty record batch:

[dh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dheres pushed a commit to branch filter_empty_record_batch
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git

commit e17eeb19c548ef1d8149ac37cabae51fbde31913
Author: Daniël Heres <da...@gmail.com>
AuthorDate: Fri Aug 4 21:45:33 2023 +0200

    Filter empty record batch:
---
 arrow-select/src/filter.rs | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs
index c89491944a..7675a54241 100644
--- a/arrow-select/src/filter.rs
+++ b/arrow-select/src/filter.rs
@@ -187,8 +187,8 @@ pub fn filter_record_batch(
         .iter()
         .map(|a| filter_array(a, &filter))
         .collect::<Result<Vec<_>, _>>()?;
-
-    RecordBatch::try_new(record_batch.schema(), filtered_arrays)
+    let options = RecordBatchOptions::default().with_row_count(Some(filter.count()));
+    RecordBatch::try_new_with_options(record_batch.schema(), filtered_arrays, &options)
 }
 
 /// A builder to construct [`FilterPredicate`]
@@ -301,6 +301,10 @@ impl FilterPredicate {
     pub fn filter(&self, values: &dyn Array) -> Result<ArrayRef, ArrowError> {
         filter_array(values, self)
     }
+
+    pub fn count(&self) -> usize {
+        self.count
+    }
 }
 
 fn filter_array(
@@ -987,6 +991,21 @@ mod tests {
         assert_eq!(out.as_ref(), &a.slice(0, 2));
     }
 
+    #[test]
+    fn test_empty_record_batch() {
+        let pred = BooleanArray::from(vec![Some(true), Some(true), None]);
+        let options = RecordBatchOptions::default().with_row_count(Some(100));
+        let record_batch = RecordBatch::try_new_with_options(
+            Arc::new(Schema::empty()),
+            vec![],
+            &options,
+        )
+        .unwrap();
+        let out = filter_record_batch(&record_batch, &pred).unwrap();
+
+        assert_eq!(out.num_rows(), 2);
+    }
+
     #[test]
     fn test_fast_path() {
         let a: PrimitiveArray<Int64Type> =