[MINOR]: Use take_arrays in repartition (#12657)

* use take_arrays in repartition * format
apache · Sep 28, 2024 · 792f07d · 792f07d
1 parent 322d835
commit 792f07d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
@@ -1001,7 +1001,7 @@ mod tests {
     }
 
     #[test]
-    fn test_get_arrayref_at_indices() -> Result<()> {
+    fn test_take_arrays() -> Result<()> {
         let arrays: Vec<ArrayRef> = vec![
             Arc::new(Float64Array::from(vec![5.0, 7.0, 8.0, 9., 10.])),
             Arc::new(Float64Array::from(vec![2.0, 3.0, 3.0, 4.0, 5.0])),

diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs
@@ -41,7 +41,7 @@ use crate::{DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Stat
 use arrow::datatypes::{SchemaRef, UInt32Type};
 use arrow::record_batch::RecordBatch;
 use arrow_array::{PrimitiveArray, RecordBatchOptions};
-use datafusion_common::utils::{get_arrayref_at_indices, transpose};
+use datafusion_common::utils::{take_arrays, transpose};
 use datafusion_common::{not_impl_err, DataFusionError, Result};
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_execution::memory_pool::MemoryConsumer;
@@ -299,8 +299,7 @@ impl BatchPartitioner {
                             let _timer = partitioner_timer.timer();
 
                             // Produce batches based on indices
-                            let columns =
-                                get_arrayref_at_indices(batch.columns(), &indices)?;
+                            let columns = take_arrays(batch.columns(), &indices)?;
 
                             let mut options = RecordBatchOptions::new();
                             options = options.with_row_count(Some(indices.len()));