From 792f07d83f9ff379471cab8693a52b1f61d1cac7 Mon Sep 17 00:00:00 2001
From: doupache <143783826+doupache@users.noreply.github.com>
Date: Sat, 28 Sep 2024 15:24:03 +0800
Subject: [PATCH] [MINOR]:  Use take_arrays in repartition  (#12657)

* use take_arrays in repartition

* format
---
 datafusion/common/src/utils/mod.rs              | 2 +-
 datafusion/physical-plan/src/repartition/mod.rs | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs
index 355d699721a3..116dab316bf5 100644
--- 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
index 10f898b26aa0..5b25d582d20c 100644
--- 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()));