From 4bd2a6dd7420d8e66d9cfd3ec397cc6c91ed6f1d Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 28 Sep 2024 13:50:22 +0800 Subject: [PATCH 1/2] fix Signed-off-by: jayzhan211 --- datafusion/physical-plan/src/repartition/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 d874367e2eb68958e92e52a72dceb3d57d82c57c Mon Sep 17 00:00:00 2001 From: jayzhan211 Date: Sat, 28 Sep 2024 13:53:13 +0800 Subject: [PATCH 2/2] fmt Signed-off-by: jayzhan211 --- 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()));