diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index cc30651b44951..061d14477d0a2 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -152,3 +152,7 @@ required-features = ["test_utils"] [[bench]] harness = false name = "bounded_window" + +[[bench]] +harness = false +name = "range_repartition" diff --git a/datafusion/physical-plan/benches/range_repartition.rs b/datafusion/physical-plan/benches/range_repartition.rs new file mode 100644 index 0000000000000..7b2b0d5cc5afd --- /dev/null +++ b/datafusion/physical-plan/benches/range_repartition.rs @@ -0,0 +1,373 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{ + LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning, SplitPoint, +}; +use datafusion_physical_plan::metrics::Time; +use datafusion_physical_plan::repartition::{BatchPartitioner, RangeExpr}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +const BATCH_SIZE: usize = 8192; +const PARTITION_COUNTS: [usize; 7] = [8, 16, 32, 64, 128, 256, 512]; +const SEED: u64 = 42; + +fn create_i64_uniform_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(SEED); + let key_values: Vec = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let payload_values: Vec = (0..BATCH_SIZE).map(|i| i as i64).collect(); + + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(key_values)) as ArrayRef, + Arc::new(Int64Array::from(payload_values)) as ArrayRef, + ], + ) + .unwrap() +} + +fn create_i64_sequential_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch { + let key_values: Vec = (0..BATCH_SIZE) + .map(|i| ((i as i64) * max_val) / (BATCH_SIZE as i64)) + .collect(); + let payload_values: Vec = (0..BATCH_SIZE).map(|i| i as i64).collect(); + + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(key_values)) as ArrayRef, + Arc::new(Int64Array::from(payload_values)) as ArrayRef, + ], + ) + .unwrap() +} + +fn create_utf8_uniform_batch(schema: &SchemaRef, max_val: usize) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(SEED); + let key_strings: Vec = (0..BATCH_SIZE) + .map(|_| format!("key_{:010}", rng.random_range(0..max_val))) + .collect(); + let payload_values: Vec = (0..BATCH_SIZE).map(|i| i as i64).collect(); + + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(StringArray::from_iter_values( + key_strings.iter().map(String::as_str), + )) as ArrayRef, + Arc::new(Int64Array::from(payload_values)) as ArrayRef, + ], + ) + .unwrap() +} + +fn create_composite_i64_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(SEED); + let key1_values: Vec = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let key2_values: Vec = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let payload_values: Vec = (0..BATCH_SIZE).map(|i| i as i64).collect(); + + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(key1_values)) as ArrayRef, + Arc::new(Int64Array::from(key2_values)) as ArrayRef, + Arc::new(Int64Array::from(payload_values)) as ArrayRef, + ], + ) + .unwrap() +} + +fn bench_range_repartition_i64_uniform(c: &mut Criterion) { + let mut group = c.benchmark_group("range_repartition_i64_uniform"); + group.throughput(Throughput::Elements(BATCH_SIZE as u64)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Int64, false), + ])); + + let max_val = 1_000_000i64; + let batch = create_i64_uniform_batch(&schema, max_val); + + for &num_partitions in &PARTITION_COUNTS { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("key", &schema).unwrap(), + SortOptions::default(), + )]) + .unwrap(); + + let split_points: Vec = (1..num_partitions) + .map(|i| { + let val = (i as i64 * max_val) / (num_partitions as i64); + SplitPoint::new(vec![ScalarValue::Int64(Some(val))]) + }) + .collect(); + + let range_part = RangePartitioning::try_new(ordering, split_points).unwrap(); + + group.bench_with_input( + BenchmarkId::new("partitions", num_partitions), + &num_partitions, + |b, _| { + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); + b.iter(|| { + partitioner + .partition(batch.clone(), |p, b| { + black_box((p, b)); + Ok(()) + }) + .unwrap(); + }); + }, + ); + } + group.finish(); +} + +fn bench_range_repartition_i64_sequential(c: &mut Criterion) { + let mut group = c.benchmark_group("range_repartition_i64_sequential"); + group.throughput(Throughput::Elements(BATCH_SIZE as u64)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Int64, false), + ])); + + let max_val = 1_000_000i64; + let batch = create_i64_sequential_batch(&schema, max_val); + + for &num_partitions in &PARTITION_COUNTS { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("key", &schema).unwrap(), + SortOptions::default(), + )]) + .unwrap(); + + let split_points: Vec = (1..num_partitions) + .map(|i| { + let val = (i as i64 * max_val) / (num_partitions as i64); + SplitPoint::new(vec![ScalarValue::Int64(Some(val))]) + }) + .collect(); + + let range_part = RangePartitioning::try_new(ordering, split_points).unwrap(); + + group.bench_with_input( + BenchmarkId::new("partitions", num_partitions), + &num_partitions, + |b, _| { + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); + b.iter(|| { + partitioner + .partition(batch.clone(), |p, b| { + black_box((p, b)); + Ok(()) + }) + .unwrap(); + }); + }, + ); + } + group.finish(); +} + +fn bench_range_repartition_utf8_uniform(c: &mut Criterion) { + let mut group = c.benchmark_group("range_repartition_utf8_uniform"); + group.throughput(Throughput::Elements(BATCH_SIZE as u64)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("payload", DataType::Int64, false), + ])); + + let max_val = 1_000_000usize; + let batch = create_utf8_uniform_batch(&schema, max_val); + + for &num_partitions in &PARTITION_COUNTS { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("key", &schema).unwrap(), + SortOptions::default(), + )]) + .unwrap(); + + let split_points: Vec = (1..num_partitions) + .map(|i| { + let val = (i * max_val) / num_partitions; + SplitPoint::new(vec![ScalarValue::Utf8(Some(format!("key_{val:010}")))]) + }) + .collect(); + + let range_part = RangePartitioning::try_new(ordering, split_points).unwrap(); + + group.bench_with_input( + BenchmarkId::new("partitions", num_partitions), + &num_partitions, + |b, _| { + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); + b.iter(|| { + partitioner + .partition(batch.clone(), |p, b| { + black_box((p, b)); + Ok(()) + }) + .unwrap(); + }); + }, + ); + } + group.finish(); +} + +fn bench_range_repartition_composite_i64(c: &mut Criterion) { + let mut group = c.benchmark_group("range_repartition_composite_i64"); + group.throughput(Throughput::Elements(BATCH_SIZE as u64)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("key1", DataType::Int64, false), + Field::new("key2", DataType::Int64, false), + Field::new("payload", DataType::Int64, false), + ])); + + let max_val = 1_000_000i64; + let batch = create_composite_i64_batch(&schema, max_val); + + for &num_partitions in &PARTITION_COUNTS { + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new(col("key1", &schema).unwrap(), SortOptions::default()), + PhysicalSortExpr::new(col("key2", &schema).unwrap(), SortOptions::default()), + ]) + .unwrap(); + + let split_points: Vec = (1..num_partitions) + .map(|i| { + let val1 = (i as i64 * max_val) / (num_partitions as i64); + let val2 = 0i64; + SplitPoint::new(vec![ + ScalarValue::Int64(Some(val1)), + ScalarValue::Int64(Some(val2)), + ]) + }) + .collect(); + + let range_part = RangePartitioning::try_new(ordering, split_points).unwrap(); + + group.bench_with_input( + BenchmarkId::new("partitions", num_partitions), + &num_partitions, + |b, _| { + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); + b.iter(|| { + partitioner + .partition(batch.clone(), |p, b| { + black_box((p, b)); + Ok(()) + }) + .unwrap(); + }); + }, + ); + } + group.finish(); +} + +fn bench_range_expr_routing_i64(c: &mut Criterion) { + let mut group = c.benchmark_group("range_expr_routing_i64"); + group.throughput(Throughput::Elements(BATCH_SIZE as u64)); + + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", DataType::Int64, false), + ])); + + let max_val = 1_000_000i64; + let batch = create_i64_uniform_batch(&schema, max_val); + + for &num_partitions in &PARTITION_COUNTS { + let col_expr = col("key", &schema).unwrap(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(&col_expr), + SortOptions::default(), + )]) + .unwrap(); + + let split_points: Vec = (1..num_partitions) + .map(|i| { + let val = (i as i64 * max_val) / (num_partitions as i64); + SplitPoint::new(vec![ScalarValue::Int64(Some(val))]) + }) + .collect(); + + let range_part = RangePartitioning::try_new(ordering, split_points).unwrap(); + let range_expr = RangeExpr::try_new(vec![col_expr], &range_part).unwrap(); + + group.bench_with_input( + BenchmarkId::new("partitions", num_partitions), + &num_partitions, + |b, _| { + b.iter(|| { + let res = range_expr.evaluate(&batch).unwrap(); + black_box(res); + }); + }, + ); + } + group.finish(); +} + +criterion_group!( + benches, + bench_range_repartition_i64_uniform, + bench_range_repartition_i64_sequential, + bench_range_repartition_utf8_uniform, + bench_range_repartition_composite_i64, + bench_range_expr_routing_i64 +); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 8f4b8558a592b..f97cc5ed98c50 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -19,7 +19,6 @@ //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. -use std::cmp::Ordering; use std::fmt::{Debug, Display, Formatter}; use std::pin::Pin; use std::sync::Arc; @@ -47,18 +46,20 @@ use crate::{ PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, }; -use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; +#[cfg(test)] +use arrow::array::Array; +use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; use arrow::compute::take_arrays; use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type}; use arrow_schema::SortOptions; +use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; +use datafusion_common::utils::transpose; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, - assert_or_internal_err, internal_datafusion_err, internal_err, - validate_range_split_points, + ColumnStatistics, DataFusionError, HashMap, SplitPoint, assert_or_internal_err, + internal_datafusion_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; @@ -89,6 +90,11 @@ use log::trace; use parking_lot::Mutex; mod distributor_channels; +mod range; + +use range::RangeRouter; +use std::hash::{Hash, Hasher}; + use crate::repartition::distributor_channels::SendError; use distributor_channels::{ DistributionReceiver, DistributionSender, channels, partition_aware_channels, @@ -634,14 +640,10 @@ enum BatchPartitionerState { Range { /// Ordered partitioning key. ordering: LexOrdering, - /// Sort options from the `LexOrdering` - sort_options: Vec, - /// Boundaries between adjacent partitions. - split_points: Vec, + /// Router for partition assignment. + router: RangeRouter, /// Row indices grouped by output partition indices: Vec>, - /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points - partition_buffer: Vec, }, } @@ -653,11 +655,28 @@ pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_ /// /// This uses the same routing function as [`BatchPartitioner`], so dynamic /// filtering and repartitioning agree for every [`ScalarValue`] comparison. -#[derive(Debug, Hash, PartialEq, Eq)] +#[derive(Debug, Clone)] pub struct RangeExpr { on_columns: Vec, - split_points: Vec, - sort_options: Vec, + router: RangeRouter, +} + +impl PartialEq for RangeExpr { + fn eq(&self, other: &Self) -> bool { + self.on_columns == other.on_columns + && self.router.split_points() == other.router.split_points() + && self.router.sort_options() == other.router.sort_options() + } +} + +impl Eq for RangeExpr {} + +impl Hash for RangeExpr { + fn hash(&self, state: &mut H) { + self.on_columns.hash(state); + self.router.split_points().hash(state); + self.router.sort_options().hash(state); + } } impl RangeExpr { @@ -667,34 +686,26 @@ impl RangeExpr { on_columns: Vec, range_partitioning: &RangePartitioning, ) -> Result { - let sort_options = range_partitioning + let sort_options: Vec = range_partitioning .ordering() .iter() .map(|expr| expr.options) .collect(); - Self::try_new_parts( - on_columns, - range_partitioning.split_points().to_vec(), - sort_options, - ) + Self::try_new_parts(on_columns, range_partitioning.split_points(), &sort_options) } fn try_new_parts( on_columns: Vec, - split_points: Vec, - sort_options: Vec, + split_points: &[SplitPoint], + sort_options: &[SortOptions], ) -> Result { assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key"); assert_or_internal_err!( on_columns.len() == sort_options.len(), "RangeExpr key count must match sort options" ); - validate_range_split_points(&split_points, &sort_options)?; - Ok(Self { - on_columns, - split_points, - sort_options, - }) + let router = RangeRouter::try_new(sort_options, split_points)?; + Ok(Self { on_columns, router }) } /// Get the columns used to compute Range partition IDs. @@ -704,12 +715,12 @@ impl RangeExpr { /// Returns the Range split points used for routing. pub fn split_points(&self) -> &[SplitPoint] { - &self.split_points + self.router.split_points() } /// Returns the per-key sort options used for routing. pub fn sort_options(&self) -> &[SortOptions] { - &self.sort_options + self.router.sort_options() } } @@ -736,8 +747,8 @@ impl PhysicalExpr for RangeExpr { ); Ok(Arc::new(Self::try_new_parts( children, - self.split_points.clone(), - self.sort_options.clone(), + self.router.split_points(), + self.router.sort_options(), )?)) } @@ -750,17 +761,14 @@ impl PhysicalExpr for RangeExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { + if self.router.split_points().is_empty() { + return Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(0)))); + } + let arrays = evaluate_expressions_to_arrays(self.on_columns.iter(), batch)?; - let mut row_key_buffer = Vec::with_capacity(arrays.len()); let mut partition_ids = Vec::with_capacity(batch.num_rows()); - for row_idx in 0..batch.num_rows() { - extract_row_at_idx_to_buf(&arrays, row_idx, &mut row_key_buffer)?; - partition_ids.push(range_partition_id( - &row_key_buffer, - &self.split_points, - &self.sort_options, - )? as u64); - } + self.router + .route_partition_ids(&arrays, &mut partition_ids)?; Ok(ColumnarValue::Array(Arc::new(UInt64Array::from( partition_ids, )))) @@ -780,12 +788,13 @@ impl PhysicalExpr for RangeExpr { let sort_exprs = self .on_columns .iter() - .zip(&self.sort_options) + .zip(self.router.sort_options()) .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options)) .collect::>(); let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?; let split_point = self - .split_points + .router + .split_points() .iter() .map(|split_point| { let value = split_point @@ -822,10 +831,11 @@ impl RangeExpr { return internal_err!("PhysicalExprNode is not a RangeExpr"); }; let sort_exprs = sort_exprs_try_from_proto(&range_expr.sort_expr, ctx)?; - let (on_columns, sort_options) = sort_exprs - .into_iter() - .map(|sort_expr| (sort_expr.expr, sort_expr.options)) - .unzip(); + let (on_columns, sort_options): (Vec, Vec) = + sort_exprs + .into_iter() + .map(|sort_expr| (sort_expr.expr, sort_expr.options)) + .unzip(); let split_points = range_expr .split_point .iter() @@ -840,29 +850,12 @@ impl RangeExpr { .collect::>>()?; Ok(Arc::new(Self::try_new_parts( on_columns, - split_points, - sort_options, + &split_points, + &sort_options, )?)) } } -fn range_partition_id( - row_key: &[ScalarValue], - split_points: &[SplitPoint], - sort_options: &[SortOptions], -) -> Result { - let mut low = 0; - let mut high = split_points.len(); - while low < high { - let mid = low + (high - low) / 2; - match compare_rows(row_key, split_points[mid].values(), sort_options)? { - Ordering::Less => high = mid, - Ordering::Equal | Ordering::Greater => low = mid + 1, - } - } - Ok(low) -} - /// Computes `value % divisor` without division in the hot loop when `divisor` /// is fixed for many values. /// @@ -994,30 +987,42 @@ impl BatchPartitioner { } } + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Panics + /// Panics if the range partitioning is invalid or cannot construct a range router. + /// Prefer [`Self::try_new_range_partitioner`] for fallible construction. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + Self::try_new_range_partitioner(range_partitioning, timer) + .expect("valid range partitioning") + } + /// Create a new [`BatchPartitioner`] for range-based repartitioning. /// /// # Parameters /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions /// - `timer`: Metric used to record time spent during repartitioning. - pub fn new_range_partitioner( + pub fn try_new_range_partitioner( range_partitioning: &RangePartitioning, timer: metrics::Time, - ) -> Self { + ) -> Result { let ordering = range_partitioning.ordering().clone(); - let split_points = range_partitioning.split_points().to_vec(); let num_partitions = range_partitioning.partition_count(); let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); + let router = + RangeRouter::try_new(&sort_options, range_partitioning.split_points())?; - Self { + Ok(Self { state: BatchPartitionerState::Range { - partition_buffer: Vec::with_capacity(ordering.len()), ordering, - sort_options, - split_points, + router, indices: vec![vec![]; num_partitions], }, timer, - } + }) } /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. @@ -1053,7 +1058,7 @@ impl BatchPartitioner { )) } Partitioning::Range(range_repartitioning) => { - Ok(Self::new_range_partitioner(&range_repartitioning, timer)) + Self::try_new_range_partitioner(&range_repartitioning, timer) } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") @@ -1145,14 +1150,12 @@ impl BatchPartitioner { } BatchPartitionerState::Range { ordering, - sort_options, - split_points, + router, indices, - partition_buffer, } => { // Tracking time required for distributing indexes across output partitions let timer = self.timer.timer(); - if split_points.is_empty() { + if router.num_split_points() == 0 { timer.done(); Box::new(std::iter::once(Ok((0, batch)))) } else { @@ -1165,13 +1168,7 @@ impl BatchPartitioner { v.clear(); } - Self::partition_range_indices( - &arrays, - split_points, - sort_options, - partition_buffer, - indices, - )?; + router.route_indices(&arrays, indices)?; // Finished building index-arrays for output partitions timer.done(); @@ -1187,28 +1184,6 @@ impl BatchPartitioner { Ok(it) } - /// Groups input row indices by range partition. This populates `indices[p]` with the - /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. - fn partition_range_indices( - arrays: &[Arc], - split_points: &[SplitPoint], - sort_options: &[SortOptions], - row_key_buffer: &mut Vec, - indices: &mut [Vec], - ) -> Result<()> { - let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); - for row_idx in 0..num_rows { - // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row - extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; - - let partition = - range_partition_id(row_key_buffer, split_points, sort_options)?; - indices[partition].push(row_idx as u32) - } - - Ok(()) - } - // return the number of output partitions fn num_partitions(&self) -> usize { match &self.state { @@ -1257,6 +1232,11 @@ impl BatchPartitioner { return Ok(vec![]); } + if partition_ranges.len() == 1 && partition_ranges[0].2 == batch.num_rows() { + let (partition, _, _) = partition_ranges[0]; + return Ok(vec![Ok((partition, batch.clone()))]); + } + let batches = { let _timer = timer.timer(); let indices_array: PrimitiveArray = reordered_indices.into(); diff --git a/datafusion/physical-plan/src/repartition/range.rs b/datafusion/physical-plan/src/repartition/range.rs new file mode 100644 index 0000000000000..0bfa23c59a729 --- /dev/null +++ b/datafusion/physical-plan/src/repartition/range.rs @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec, + sort_options: Vec, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec = if !split_points.is_empty() { + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + + // Try single-column primitive fast path + if data_types.len() == 1 + && !sort_options.is_empty() + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(&data_types, sort_options, split_points)?; + Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Split points configured in this router. + pub(crate) fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Sort options configured in this router. + pub(crate) fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + self.split_points.len() + } + + /// Generic routing entry point that calls `emit(row_idx, partition)` for every row. + pub(crate) fn route_with(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + if self.split_points.is_empty() { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + emit(row_idx, 0); + } + return Ok(()); + } + + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_with(first_col.as_ref(), emit) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_with(arrays, emit), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec], + ) -> Result<()> { + self.route_with(arrays, |row_idx, partition| { + indices[partition].push(row_idx as u32); + }) + } + + /// Appends output partition IDs to `partition_ids`. + pub(crate) fn route_partition_ids( + &self, + arrays: &[ArrayRef], + partition_ids: &mut Vec, + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + partition_ids + .try_reserve(num_rows) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + self.route_with(arrays, |_row_idx, partition| { + partition_ids.push(partition as u64); + }) + } +} + +macro_rules! define_primitive_router { + ($( ($variant:ident, $type:ty, $arrow_type:ty, $array:ident) ),* $(,)?) => { + /// Specialized router for primitive scalar types. + #[derive(Debug, Clone)] + enum PrimitiveRangeRouter { + $( $variant(PrimitiveValuesRouter<$type>), )* + Float32(FloatValuesRouter), + Float64(FloatValuesRouter), + } + + impl PrimitiveRangeRouter { + fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option { + if split_points.is_empty() { + return None; + } + + let scalars = split_points.iter().map(|sp| sp.values()[0].clone()); + let split_array = ScalarValue::iter_to_array(scalars).ok()?; + if split_array.null_count() > 0 { + return None; + } + + macro_rules! make_primitive { + ($target_arrow_type:ty, $target_variant:ident) => {{ + let arr = split_array + .as_any() + .downcast_ref::>()?; + let vals = arr.values().to_vec(); + Some(Self::$target_variant(PrimitiveValuesRouter::new(vals, sort_options))) + }}; + } + + match split_array.data_type() { + DataType::Int8 => make_primitive!(Int8Type, Int8), + DataType::Int16 => make_primitive!(Int16Type, Int16), + DataType::Int32 => make_primitive!(Int32Type, Int32), + DataType::Int64 => make_primitive!(Int64Type, Int64), + DataType::UInt8 => make_primitive!(UInt8Type, UInt8), + DataType::UInt16 => make_primitive!(UInt16Type, UInt16), + DataType::UInt32 => make_primitive!(UInt32Type, UInt32), + DataType::UInt64 => make_primitive!(UInt64Type, UInt64), + DataType::Date32 => make_primitive!(Date32Type, Date32), + DataType::Date64 => make_primitive!(Date64Type, Date64), + DataType::Time32(TimeUnit::Second) => make_primitive!(Time32SecondType, Time32Second), + DataType::Time32(TimeUnit::Millisecond) => make_primitive!(Time32MillisecondType, Time32Millisecond), + DataType::Time64(TimeUnit::Microsecond) => make_primitive!(Time64MicrosecondType, Time64Microsecond), + DataType::Time64(TimeUnit::Nanosecond) => make_primitive!(Time64NanosecondType, Time64Nanosecond), + DataType::Timestamp(TimeUnit::Second, _) => make_primitive!(TimestampSecondType, TimestampSecond), + DataType::Timestamp(TimeUnit::Millisecond, _) => make_primitive!(TimestampMillisecondType, TimestampMillisecond), + DataType::Timestamp(TimeUnit::Microsecond, _) => make_primitive!(TimestampMicrosecondType, TimestampMicrosecond), + DataType::Timestamp(TimeUnit::Nanosecond, _) => make_primitive!(TimestampNanosecondType, TimestampNanosecond), + DataType::Float32 => { + let arr = split_array.as_any().downcast_ref::()?; + let vals = arr.values().to_vec(); + Some(Self::Float32(FloatValuesRouter::new(vals, sort_options))) + } + DataType::Float64 => { + let arr = split_array.as_any().downcast_ref::()?; + let vals = arr.values().to_vec(); + Some(Self::Float64(FloatValuesRouter::new(vals, sort_options))) + } + _ => None, + } + } + + fn route_with(&self, array: &dyn Array, emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + match self { + $( + Self::$variant(r) => { + let arr = array.as_any().downcast_ref::<$array>().ok_or_else(|| { + DataFusionError::Internal(format!("Expected {}", stringify!($array))) + })?; + r.route_with(arr, emit); + Ok(()) + } + )* + Self::Float32(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float32Array".to_string()) + })?; + r.route_with(arr, emit); + Ok(()) + } + Self::Float64(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float64Array".to_string()) + })?; + r.route_with(arr, emit); + Ok(()) + } + } + } + } + }; +} + +define_primitive_router!( + (Int8, i8, Int8Type, Int8Array), + (Int16, i16, Int16Type, Int16Array), + (Int32, i32, Int32Type, Int32Array), + (Int64, i64, Int64Type, Int64Array), + (UInt8, u8, UInt8Type, UInt8Array), + (UInt16, u16, UInt16Type, UInt16Array), + (UInt32, u32, UInt32Type, UInt32Array), + (UInt64, u64, UInt64Type, UInt64Array), + (Date32, i32, Date32Type, Date32Array), + (Date64, i64, Date64Type, Date64Array), + (Time32Second, i32, Time32SecondType, Time32SecondArray), + ( + Time32Millisecond, + i32, + Time32MillisecondType, + Time32MillisecondArray + ), + ( + Time64Microsecond, + i64, + Time64MicrosecondType, + Time64MicrosecondArray + ), + ( + Time64Nanosecond, + i64, + Time64NanosecondType, + Time64NanosecondArray + ), + ( + TimestampSecond, + i64, + TimestampSecondType, + TimestampSecondArray + ), + ( + TimestampMillisecond, + i64, + TimestampMillisecondType, + TimestampMillisecondArray + ), + ( + TimestampMicrosecond, + i64, + TimestampMicrosecondType, + TimestampMicrosecondArray + ), + ( + TimestampNanosecond, + i64, + TimestampNanosecondType, + TimestampNanosecondArray + ), +); + +/// Generic router for primitive integer and temporal types. +#[derive(Debug, Clone)] +struct PrimitiveValuesRouter { + split_points: Vec, + sort_options: SortOptions, +} + +impl PrimitiveValuesRouter { + fn new(split_points: Vec, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } + + fn route_with, E: FnMut(usize, usize)>( + &self, + array: &PrimitiveArray, + mut emit: E, + ) { + let split_points = &self.split_points; + let descending = self.sort_options.descending; + let nulls_first = self.sort_options.nulls_first; + + if array.null_count() == 0 { + let values = array.values().as_ref(); + if !descending { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| sp <= val); + emit(idx, p); + } + } else { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| sp >= val); + emit(idx, p); + } + } + } else { + let null_partition = if nulls_first { 0 } else { split_points.len() }; + for idx in 0..array.len() { + if array.is_null(idx) { + emit(idx, null_partition); + } else { + let val = array.value(idx); + let p = if !descending { + split_points.partition_point(|&sp| sp <= val) + } else { + split_points.partition_point(|&sp| sp >= val) + }; + emit(idx, p); + } + } + } + } +} + +/// Generic router for floating point values using total ordering. +#[derive(Debug, Clone)] +struct FloatValuesRouter { + split_points: Vec, + sort_options: SortOptions, +} + +impl FloatValuesRouter { + fn new(split_points: Vec, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } +} + +macro_rules! impl_float_values_router { + ($t:ty, $arr:ty) => { + impl FloatValuesRouter<$t> { + fn route_with(&self, array: &$arr, mut emit: E) { + let split_points = &self.split_points; + let descending = self.sort_options.descending; + let nulls_first = self.sort_options.nulls_first; + + if array.null_count() == 0 { + let values = array.values().as_ref(); + if !descending { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Greater + }); + emit(idx, p); + } + } else { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Less + }); + emit(idx, p); + } + } + } else { + let null_partition = if nulls_first { 0 } else { split_points.len() }; + for idx in 0..array.len() { + if array.is_null(idx) { + emit(idx, null_partition); + } else { + let val = array.value(idx); + let p = if !descending { + split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Greater + }) + } else { + split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Less + }) + }; + emit(idx, p); + } + } + } + } + } + }; +} + +impl_float_values_router!(f32, Float32Array); +impl_float_values_router!(f64, Float64Array); + +/// Router backed by Arrow's RowConverter. +#[derive(Debug, Clone)] +struct RowConverterRangeRouter { + converter: Arc, + split_point_rows: Vec, +} + +impl RowConverterRangeRouter { + fn try_new( + data_types: &[DataType], + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result { + let sort_fields = data_types + .iter() + .zip(sort_options) + .map(|(dt, opt)| SortField::new_with_options(dt.clone(), *opt)) + .collect::>(); + + if !RowConverter::supports_fields(&sort_fields) { + return not_impl_err!( + "Range partitioning is not supported for data types: {:?}", + data_types + ); + } + + let row_converter = RowConverter::new(sort_fields)?; + let num_cols = data_types.len(); + + let split_point_rows = if split_points.is_empty() { + vec![] + } else { + let split_point_arrays = (0..num_cols) + .map(|col_idx| { + let col_scalars = + split_points.iter().map(|sp| sp.values()[col_idx].clone()); + ScalarValue::iter_to_array(col_scalars) + }) + .collect::>>()?; + + row_converter + .convert_columns(&split_point_arrays)? + .iter() + .map(|r| r.owned()) + .collect() + }; + + Ok(Self { + converter: Arc::new(row_converter), + split_point_rows, + }) + } + + fn route_with( + &self, + arrays: &[ArrayRef], + mut emit: E, + ) -> Result<()> { + let rows = self.converter.convert_columns(arrays)?; + let sp_rows = &self.split_point_rows; + + for (row_idx, row) in rows.iter().enumerate() { + let partition = sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()); + emit(row_idx, partition); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + fn make_split_points_1d(scalars: Vec) -> Vec { + scalars + .into_iter() + .map(|s| SplitPoint::new(vec![s])) + .collect() + } + + fn assert_routing( + router: &RangeRouter, + arrays: &[ArrayRef], + expected_partition_ids: &[u64], + expected_indices: Option<&[Vec]>, + ) -> Result<()> { + let mut partition_ids = Vec::new(); + router.route_partition_ids(arrays, &mut partition_ids)?; + assert_eq!(partition_ids, expected_partition_ids); + + if let Some(expected) = expected_indices { + let mut indices = vec![vec![]; expected.len()]; + router.route_indices(arrays, &mut indices)?; + assert_eq!(indices, expected); + } + + Ok(()) + } + + #[test] + fn test_primitive_router_i64_asc() -> Result<()> { + let split_points = make_split_points_1d(vec![ + ScalarValue::Int64(Some(10)), + ScalarValue::Int64(Some(20)), + ScalarValue::Int64(Some(30)), + ]); + let sort_options = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); + + let input = Arc::new(Int64Array::from(vec![ + Some(5), + Some(10), + Some(15), + Some(20), + Some(25), + Some(30), + Some(35), + None, + ])) as ArrayRef; + + assert_routing( + &router, + &[input], + &[0, 1, 1, 2, 2, 3, 3, 0], + Some(&[vec![0, 7], vec![1, 2], vec![3, 4], vec![5, 6]]), + ) + } + + #[test] + fn test_primitive_router_i64_desc() -> Result<()> { + let split_points = make_split_points_1d(vec![ + ScalarValue::Int64(Some(30)), + ScalarValue::Int64(Some(20)), + ScalarValue::Int64(Some(10)), + ]); + let sort_options = vec![SortOptions { + descending: true, + nulls_first: false, + }]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); + + let input = Arc::new(Int64Array::from(vec![ + Some(35), + Some(30), + Some(25), + Some(20), + Some(15), + Some(10), + Some(5), + None, + ])) as ArrayRef; + + assert_routing( + &router, + &[input], + &[0, 1, 1, 2, 2, 3, 3, 3], + Some(&[vec![0], vec![1, 2], vec![3, 4], vec![5, 6, 7]]), + ) + } + + #[test] + fn test_float_router() -> Result<()> { + let split_points = make_split_points_1d(vec![ + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(100.0)), + ]); + let sort_options = vec![SortOptions { + descending: false, + nulls_first: false, + }]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); + + let input = Arc::new(Float64Array::from(vec![ + Some(-10.0), + Some(0.0), + Some(50.0), + Some(100.0), + Some(200.0), + None, + ])) as ArrayRef; + + assert_routing( + &router, + &[input], + &[0, 1, 1, 2, 2, 2], + Some(&[vec![0], vec![1, 2], vec![3, 4, 5]]), + ) + } + + #[test] + fn test_row_converter_strings() -> Result<()> { + let split_points = make_split_points_1d(vec![ + ScalarValue::Utf8(Some("d".to_string())), + ScalarValue::Utf8(Some("m".to_string())), + ScalarValue::Utf8(Some("s".to_string())), + ]); + let sort_options = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert!(matches!(router.inner, RangeRouterInner::Row(_))); + + let input = Arc::new(StringArray::from(vec![ + Some("apple"), + Some("d"), + Some("frog"), + Some("m"), + Some("orange"), + Some("s"), + Some("zebra"), + None, + ])) as ArrayRef; + + assert_routing( + &router, + &[input], + &[0, 1, 1, 2, 2, 3, 3, 0], + Some(&[vec![0, 7], vec![1, 2], vec![3, 4], vec![5, 6]]), + ) + } + + #[test] + fn test_row_converter_composite_keys() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(Some("b".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(Some("d".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(2)), + ScalarValue::Utf8(Some("a".to_string())), + ]), + ]; + let sort_options = vec![ + SortOptions { + descending: false, + nulls_first: false, + }, + SortOptions { + descending: false, + nulls_first: false, + }, + ]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert!(matches!(router.inner, RangeRouterInner::Row(_))); + + let col1 = Arc::new(Int64Array::from(vec![1, 1, 1, 1, 2, 2, 3])) as ArrayRef; + let col2 = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "a", "z", "a"])) + as ArrayRef; + + assert_routing( + &router, + &[col1, col2], + &[0, 1, 1, 2, 3, 3, 3], + Some(&[vec![0], vec![1, 2], vec![3], vec![4, 5, 6]]), + ) + } + + #[test] + fn test_router_empty_split_points() -> Result<()> { + let split_points = vec![]; + let sort_options = vec![SortOptions::default()]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + assert_eq!(router.num_split_points(), 0); + + let input = Arc::new(Int64Array::from(vec![10, 20, 30])) as ArrayRef; + assert_routing(&router, &[input], &[0, 0, 0], Some(&[vec![0, 1, 2]])) + } +}