@@ -104,7 +104,10 @@ pub fn concat(arrays: &[&dyn Array]) -> Result<ArrayRef> {
104
104
}
105
105
106
106
/// Concatenates `batches` together into a single record batch.
107
- pub fn concat_batches ( schema : & SchemaRef , batches : & [ RecordBatch ] ) -> Result < RecordBatch > {
107
+ pub fn concat_batches (
108
+ schema : & SchemaRef ,
109
+ batches : & [ RecordBatch ] ,
110
+ ) -> Result < RecordBatch > {
108
111
if batches. is_empty ( ) {
109
112
return Ok ( RecordBatch :: new_empty ( schema. clone ( ) ) ) ;
110
113
}
@@ -613,16 +616,16 @@ mod tests {
613
616
Arc :: new( StringArray :: from( vec![ "a" , "b" ] ) ) ,
614
617
] ,
615
618
)
616
- . unwrap ( ) ;
619
+ . unwrap ( ) ;
617
620
let batch2 = RecordBatch :: try_new (
618
621
schema. clone ( ) ,
619
622
vec ! [
620
623
Arc :: new( Int32Array :: from( vec![ 3 , 4 ] ) ) ,
621
624
Arc :: new( StringArray :: from( vec![ "c" , "d" ] ) ) ,
622
625
] ,
623
626
)
624
- . unwrap ( ) ;
625
- let new_batch = RecordBatch :: concat ( & schema, & [ batch1, batch2] ) . unwrap ( ) ;
627
+ . unwrap ( ) ;
628
+ let new_batch = concat_batches ( & schema, & [ batch1, batch2] ) . unwrap ( ) ;
626
629
assert_eq ! ( new_batch. schema( ) . as_ref( ) , schema. as_ref( ) ) ;
627
630
assert_eq ! ( 2 , new_batch. num_columns( ) ) ;
628
631
assert_eq ! ( 4 , new_batch. num_rows( ) ) ;
@@ -634,7 +637,7 @@ mod tests {
634
637
Field :: new( "a" , DataType :: Int32 , false ) ,
635
638
Field :: new( "b" , DataType :: Utf8 , false ) ,
636
639
] ) ) ;
637
- let batch = RecordBatch :: concat ( & schema, & [ ] ) . unwrap ( ) ;
640
+ let batch = concat_batches ( & schema, & [ ] ) . unwrap ( ) ;
638
641
assert_eq ! ( batch. schema( ) . as_ref( ) , schema. as_ref( ) ) ;
639
642
assert_eq ! ( 0 , batch. num_rows( ) ) ;
640
643
}
@@ -656,16 +659,16 @@ mod tests {
656
659
Arc :: new( StringArray :: from( vec![ "a" , "b" ] ) ) ,
657
660
] ,
658
661
)
659
- . unwrap ( ) ;
662
+ . unwrap ( ) ;
660
663
let batch2 = RecordBatch :: try_new (
661
664
schema2,
662
665
vec ! [
663
666
Arc :: new( Int32Array :: from( vec![ 3 , 4 ] ) ) ,
664
667
Arc :: new( StringArray :: from( vec![ "c" , "d" ] ) ) ,
665
668
] ,
666
669
)
667
- . unwrap ( ) ;
668
- let error = RecordBatch :: concat ( & schema1, & [ batch1, batch2] ) . unwrap_err ( ) ;
670
+ . unwrap ( ) ;
671
+ let error = concat_batches ( & schema1, & [ batch1, batch2] ) . unwrap_err ( ) ;
669
672
assert_eq ! (
670
673
error. to_string( ) ,
671
674
"Invalid argument error: batches[1] schema is different with argument schema." ,
0 commit comments