@@ -61,21 +61,19 @@ use serde_json::{Map, Value};
6161///
6262/// * `schema` - The schema of the record batch.
6363/// * `batch` - The record batch to modify.
64- /// * `indexes` - The indexes of the columns to replace.
65- /// * `arrays` - The new arrays to replace the columns with.
64+ /// * `indexed_arrays` - A list of indexes and arrays to replace the columns indexed with.
6665///
6766/// # Returns
6867///
6968/// The modified record batch with the columns replaced.
7069pub fn replace_columns (
7170 schema : Arc < Schema > ,
7271 batch : & RecordBatch ,
73- indexes : & [ usize ] ,
74- arrays : & [ Arc < dyn Array + ' static > ] ,
72+ indexed_arrays : & [ ( usize , Arc < dyn Array + ' static > ) ] ,
7573) -> RecordBatch {
7674 let mut batch_arrays = batch. columns ( ) . iter ( ) . map ( Arc :: clone) . collect_vec ( ) ;
77- for ( & index, arr) in indexes . iter ( ) . zip ( arrays . iter ( ) ) {
78- batch_arrays[ index] = Arc :: clone ( arr) ;
75+ for ( index, arr) in indexed_arrays {
76+ batch_arrays[ * index] = Arc :: clone ( arr) ;
7977 }
8078 RecordBatch :: try_new ( schema, batch_arrays) . unwrap ( )
8179}
@@ -180,7 +178,7 @@ mod tests {
180178
181179 let arr: Arc < dyn Array + ' static > = Arc :: new ( Int32Array :: from_value ( 0 , 3 ) ) ;
182180
183- let new_rb = replace_columns ( schema_ref. clone ( ) , & rb, & [ 2 ] , & [ arr] ) ;
181+ let new_rb = replace_columns ( schema_ref. clone ( ) , & rb, & [ ( 2 , arr) ] ) ;
184182
185183 assert_eq ! ( new_rb. schema( ) , schema_ref) ;
186184 assert_eq ! ( new_rb. num_columns( ) , 3 ) ;
0 commit comments