This is an automated email from the ASF dual-hosted git repository. dheres pushed a commit to branch support_concat_batches_no_columns2 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit e29aa252bbf53f4032e1ec1c37c6d99bcf83ce1c Author: Daniƫl Heres <[email protected]> AuthorDate: Mon Aug 7 22:26:56 2023 +0200 Support `concat_batches` for 0 columns --- arrow-select/src/concat.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/arrow-select/src/concat.rs b/arrow-select/src/concat.rs index 0bf4c97ff8..0d36271a73 100644 --- a/arrow-select/src/concat.rs +++ b/arrow-select/src/concat.rs @@ -97,6 +97,13 @@ pub fn concat_batches<'a>( schema: &SchemaRef, input_batches: impl IntoIterator<Item = &'a RecordBatch>, ) -> Result<RecordBatch, ArrowError> { + 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 +149,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![ @@ -718,4 +740,4 @@ mod tests { assert_eq!(data.buffers()[1].len(), 200); assert_eq!(data.buffers()[1].capacity(), 256); // Nearest multiple of 64 } -} +} \ No newline at end of file
