From 676550899e019c377f467c6ca4c158e98f7aac4c Mon Sep 17 00:00:00 2001 From: DouPache Date: Sat, 28 Sep 2024 12:58:17 +0800 Subject: [PATCH 1/2] use take_arrays in repartition --- datafusion/common/src/utils/mod.rs | 2 +- datafusion/physical-plan/src/repartition/mod.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 355d699721a3a..116dab316bf59 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 = 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 10f898b26aa0b..1a4bb797c5f13 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; @@ -300,7 +300,7 @@ impl BatchPartitioner { // Produce batches based on indices let columns = - get_arrayref_at_indices(batch.columns(), &indices)?; + take_arrays(batch.columns(), &indices)?; let mut options = RecordBatchOptions::new(); options = options.with_row_count(Some(indices.len())); From 3299f16e89f9c579099e558c328d94b4486ea3e6 Mon Sep 17 00:00:00 2001 From: DouPache Date: Sat, 28 Sep 2024 13:09:50 +0800 Subject: [PATCH 2/2] format --- datafusion/physical-plan/src/repartition/mod.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 1a4bb797c5f13..5b25d582d20c3 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -299,8 +299,7 @@ impl BatchPartitioner { let _timer = partitioner_timer.timer(); // Produce batches based on indices - let columns = - take_arrays(batch.columns(), &indices)?; + let columns = take_arrays(batch.columns(), &indices)?; let mut options = RecordBatchOptions::new(); options = options.with_row_count(Some(indices.len()));