@@ -61,19 +61,21 @@ use serde_json::{Map, Value};
61
61
///
62
62
/// * `schema` - The schema of the record batch.
63
63
/// * `batch` - The record batch to modify.
64
- /// * `indexed_arrays` - A list of indexes and arrays to replace the columns indexed with.
64
+ /// * `indexes` - The indexes of the columns to replace.
65
+ /// * `arrays` - The new arrays to replace the columns with.
65
66
///
66
67
/// # Returns
67
68
///
68
69
/// The modified record batch with the columns replaced.
69
70
pub fn replace_columns (
70
71
schema : Arc < Schema > ,
71
72
batch : & RecordBatch ,
72
- indexed_arrays : & [ ( usize , Arc < dyn Array + ' static > ) ] ,
73
+ indexes : & [ usize ] ,
74
+ arrays : & [ Arc < dyn Array + ' static > ] ,
73
75
) -> RecordBatch {
74
76
let mut batch_arrays = batch. columns ( ) . iter ( ) . map ( Arc :: clone) . collect_vec ( ) ;
75
- for ( index, arr) in indexed_arrays {
76
- batch_arrays[ * index] = Arc :: clone ( arr) ;
77
+ for ( & index, arr) in indexes . iter ( ) . zip ( arrays . iter ( ) ) {
78
+ batch_arrays[ index] = Arc :: clone ( arr) ;
77
79
}
78
80
RecordBatch :: try_new ( schema, batch_arrays) . unwrap ( )
79
81
}
@@ -176,7 +178,7 @@ mod tests {
176
178
177
179
let arr: Arc < dyn Array + ' static > = Arc :: new ( Int32Array :: from_value ( 0 , 3 ) ) ;
178
180
179
- let new_rb = replace_columns ( schema_ref. clone ( ) , & rb, & [ ( 2 , arr) ] ) ;
181
+ let new_rb = replace_columns ( schema_ref. clone ( ) , & rb, & [ 2 ] , & [ arr] ) ;
180
182
181
183
assert_eq ! ( new_rb. schema( ) , schema_ref) ;
182
184
assert_eq ! ( new_rb. num_columns( ) , 3 ) ;
0 commit comments