[arrow-rs] branch master updated: Support `concat_batches` for 0 columns (#4662)

[jakevin]

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

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


The following commit(s) were added to refs/heads/master by this push:
     new 696cbdbb72 Support `concat_batches` for 0 columns (#4662)
696cbdbb72 is described below

commit 696cbdbb72e821613f44a09e2b547d2f85c06089
Author: Daniël Heres <danielhe...@gmail.com>
AuthorDate: Tue Aug 8 08:01:52 2023 +0200

    Support `concat_batches` for 0 columns (#4662)
    
    * Support `concat_batches` for 0 columns
    
    * Add doc comment
---
 arrow-select/src/concat.rs | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs
index 0bf4c97ff8..31846ee1fd 100644
--- a/arrow-select/src/concat.rs
+++ b/arrow-select/src/concat.rs
@@ -97,6 +97,14 @@ pub fn concat_batches<'a>(
     schema: &SchemaRef,
     input_batches: impl IntoIterator<Item = &'a RecordBatch>,
 ) -> Result<RecordBatch, ArrowError> {
+    // When schema is empty, sum the number of the rows of all batches
+    if schema.fields().is_empty() {
+        let num_rows: usize = 
input_batches.into_iter().map(RecordBatch::num_rows).sum();
+        let mut options = RecordBatchOptions::default();
+        options.row_count = Some(num_rows);
+        return RecordBatch::try_new_with_options(schema.clone(), vec![], 
&options);
+    }
+
     let batches: Vec<&RecordBatch> = input_batches.into_iter().collect();
     if batches.is_empty() {
         return Ok(RecordBatch::new_empty(schema.clone()));
@@ -142,6 +150,21 @@ mod tests {
         assert!(re.is_err());
     }
 
+    #[test]
+    fn test_concat_batches_no_columns() {
+        // Test concat using empty schema / batches without columns
+        let schema = Arc::new(Schema::empty());
+
+        let mut options = RecordBatchOptions::default();
+        options.row_count = Some(100);
+        let batch =
+            RecordBatch::try_new_with_options(schema.clone(), vec![], 
&options).unwrap();
+        // put in 2 batches of 100 rows each
+        let re = concat_batches(&schema, &[batch.clone(), batch]).unwrap();
+
+        assert_eq!(re.num_rows(), 200);
+    }
+
     #[test]
     fn test_concat_one_element_vec() {
         let arr = Arc::new(PrimitiveArray::<Int64Type>::from(vec![