diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 1adac74ecb3c5..0aa22653b44ee 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -161,3 +161,7 @@ name = "bounded_window" [[bench]] harness = false name = "window_filter" + +[[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..6a7c1fc7adcb2 --- /dev/null +++ b/datafusion/physical-plan/benches/range_repartition.rs @@ -0,0 +1,418 @@ +// 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(); +} + +fn bench_range_partitioner_construction(c: &mut Criterion) { + let mut group = c.benchmark_group("range_partitioner_construction"); + + 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; + + 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, _| { + b.iter(|| { + let partitioner = BatchPartitioner::try_new_range_partitioner( + black_box(&range_part), + Time::default(), + ) + .unwrap(); + black_box(partitioner); + }); + }, + ); + } + 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, + bench_range_partitioner_construction +); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 8f4b8558a592b..b1425dca4e33f 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,18 @@ use crate::{ PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, }; -use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; +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 +88,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, @@ -558,6 +562,24 @@ impl RepartitionExecState { ); } + let range_router = if let Partitioning::Range(range_partitioning) = &partitioning + { + let ordering = range_partitioning.ordering(); + let sort_options: Vec = + ordering.iter().map(|e| e.options).collect(); + let data_types = ordering + .iter() + .map(|e| e.expr.data_type(input.schema().as_ref())) + .collect::>>()?; + Some(Arc::new(RangeRouter::try_new_with_data_types( + &sort_options, + range_partitioning.split_points(), + &data_types, + )?)) + } else { + None + }; + // launch one async task per *input* partition let mut spawned_tasks = Vec::with_capacity(num_input_partitions); for (i, (stream, metrics)) in @@ -591,6 +613,7 @@ impl RepartitionExecState { stream, txs, partitioning.clone(), + range_router.clone(), metrics, // preserve_order depends on partition index to start from 0 if preserve_order { 0 } else { i }, @@ -634,14 +657,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: Arc, /// 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 +672,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 +703,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 +732,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 +764,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 +778,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 +805,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 +848,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 +867,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,26 +1004,62 @@ 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. + #[deprecated(since = "55.0.0", note = "Use try_new_range_partitioner instead")] + 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 = Arc::new(RangeRouter::try_new( + &sort_options, + range_partitioning.split_points(), + )?); + + Ok(Self::new_range_partitioner_with_router( + ordering, + router, + num_partitions, + timer, + )) + } + /// Create a new [`BatchPartitioner`] for range-based repartitioning using a pre-constructed [`RangeRouter`]. + /// + /// # Parameters + /// - `ordering`: Lexical ordering expressions for range partitioning. + /// - `router`: Shared [`RangeRouter`] reference. + /// - `num_partitions`: Total number of output partitions. + /// - `timer`: Metric used to record time spent during repartitioning. + pub(crate) fn new_range_partitioner_with_router( + ordering: LexOrdering, + router: Arc, + num_partitions: usize, + timer: metrics::Time, + ) -> Self { Self { state: BatchPartitionerState::Range { - partition_buffer: Vec::with_capacity(ordering.len()), ordering, - sort_options, - split_points, + router, indices: vec![vec![]; num_partitions], }, timer, @@ -1053,7 +1099,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 +1191,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 +1209,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 +1225,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 +1273,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(); @@ -2139,16 +2160,29 @@ impl RepartitionExec { mut stream: SendableRecordBatchStream, mut output_channels: HashMap, partitioning: Partitioning, + range_router: Option>, metrics: RepartitionMetrics, input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = BatchPartitioner::try_new( - partitioning, - metrics.repartition_time.clone(), - input_partition, - num_input_partitions, - )?; + let mut partitioner = match (partitioning, range_router) { + (Partitioning::Range(range_partitioning), Some(router)) => { + let ordering = range_partitioning.ordering().clone(); + let num_partitions = range_partitioning.partition_count(); + BatchPartitioner::new_range_partitioner_with_router( + ordering, + router, + num_partitions, + metrics.repartition_time.clone(), + ) + } + (partitioning, _) => BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?, + }; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -2504,8 +2538,11 @@ mod tests { {collect, expressions::col}, }; - use arrow::array::{ArrayRef, StringArray, UInt32Array}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ + Array, ArrayRef, Decimal128Array, Int64Array, StringArray, + TimestampNanosecondArray, UInt32Array, + }; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datafusion_common::ScalarValue; use datafusion_common::cast::{as_string_array, as_uint32_array}; use datafusion_common::exec_err; @@ -2513,7 +2550,9 @@ mod tests { use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; - use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; + use datafusion_physical_expr::{ + LexOrdering, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; use insta::assert_snapshot; #[derive(Debug)] @@ -3014,6 +3053,101 @@ mod tests { Ok(()) } + /// Split points are not required to carry the key's exact type: `compare_rows`, + /// the routing function the range router replaced, compares `Decimal128` on + /// scale alone and ignores precision. Routing must not depend on the split + /// point's precision matching the column's. + #[tokio::test] + async fn range_repartition_routes_decimal_with_wider_column_precision() -> Result<()> + { + let schema = Arc::new(Schema::new(vec![Field::new( + "k", + DataType::Decimal128(20, 2), + true, + )])); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("k", &schema)?, + SortOptions::default(), + )]) + .unwrap(); + // Split point is Decimal128(10, 2); the column is Decimal128(20, 2). + let partitioning = Partitioning::Range(RangePartitioning::try_new( + ordering, + vec![SplitPoint::new(vec![ScalarValue::Decimal128( + Some(1000), + 10, + 2, + )])], + )?); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new( + Decimal128Array::from(vec![Some(500i128), Some(2000i128)]) + .with_precision_and_scale(20, 2)?, + )], + )?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!(1, partition_row_count(&output_partitions[0])); + assert_eq!(1, partition_row_count(&output_partitions[1])); + + Ok(()) + } + + /// Same contract for timestamps: `compare_rows` ignores the timezone, since + /// the underlying values are UTC either way. A tz-less split point must + /// still route a `Timestamp(ns, "UTC")` column — including when the key is + /// compound, where the single-column fast path does not apply. + #[tokio::test] + async fn range_repartition_routes_compound_timestamp_key_ignoring_timezone() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "t", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + true, + ), + Field::new("i", DataType::Int64, true), + ])); + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new(col("t", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("i", &schema)?, SortOptions::default()), + ]) + .unwrap(); + // Split point timestamp carries no timezone; the column carries "UTC". + let partitioning = Partitioning::Range(RangePartitioning::try_new( + ordering, + vec![SplitPoint::new(vec![ + ScalarValue::TimestampNanosecond(Some(100), None), + ScalarValue::Int64(Some(0)), + ])], + )?); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new( + TimestampNanosecondArray::from(vec![Some(50i64), Some(200)]) + .with_timezone("UTC"), + ), + Arc::new(Int64Array::from(vec![Some(1i64), Some(2)])), + ], + )?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!(1, partition_row_count(&output_partitions[0])); + assert_eq!(1, partition_row_count(&output_partitions[1])); + + Ok(()) + } + #[test] fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { // Three columns so the projection both narrows the schema (required for diff --git a/datafusion/physical-plan/src/repartition/range.rs b/datafusion/physical-plan/src/repartition/range.rs new file mode 100644 index 0000000000000..81fb4c74f4210 --- /dev/null +++ b/datafusion/physical-plan/src/repartition/range.rs @@ -0,0 +1,907 @@ +// 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::{Row, RowConverter, Rows, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_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 { + data_types: Vec, + 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, + /// inferring key data types from the split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result { + let data_types: Option> = if !split_points.is_empty() { + Some( + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect(), + ) + } else { + None + }; + Self::try_new_with_optional_data_types( + sort_options, + split_points, + data_types.as_deref(), + ) + } + + /// Constructs the best router for the given sort options, split points, and target data types, + /// coercing split point values to the target data types. + pub(crate) fn try_new_with_data_types( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + data_types: &[DataType], + ) -> Result { + Self::try_new_with_optional_data_types( + sort_options, + split_points, + Some(data_types), + ) + } + + fn try_new_with_optional_data_types( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + data_types: Option<&[DataType]>, + ) -> Result { + validate_range_split_points(split_points, sort_options)?; + + let (data_types, split_points) = if let Some(target_dts) = data_types { + if target_dts.len() != sort_options.len() { + return plan_err!( + "Range partitioning expected {} data types for sort options, but got {}", + sort_options.len(), + target_dts.len() + ); + } + let coerced_split_points = split_points + .iter() + .map(|sp| { + let vals = sp + .values() + .iter() + .zip(target_dts) + .map(|(val, target_dt)| { + if val.data_type() == *target_dt { + Ok(val.clone()) + } else { + val.cast_to(target_dt) + } + }) + .collect::>>()?; + Ok(SplitPoint::new(vals)) + }) + .collect::>>()?; + (target_dts.to_vec(), coerced_split_points) + } else if !split_points.is_empty() { + let dts: Vec = (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect(); + (dts, split_points.to_vec()) + } else { + (vec![], 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 { + data_types, + split_points, + 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 { + data_types, + split_points, + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Data types configured in this router. + #[cfg(test)] + pub(crate) fn data_types(&self) -> &[DataType] { + &self.data_types + } + + /// 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(()); + } + + if arrays.len() != self.data_types.len() { + return exec_err!( + "Range partitioning expected {} columns, but got {}", + self.data_types.len(), + arrays.len() + ); + } + + for (i, (arr, expected_dt)) in arrays.iter().zip(&self.data_types).enumerate() { + if arr.data_type() != expected_dt { + return exec_err!( + "Range partitioning expected column {i} to be of type {expected_dt:?}, but got {:?}", + arr.data_type() + ); + } + } + + 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: Option, +} + +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() { + None + } 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::>>()?; + + Some(row_converter.convert_columns(&split_point_arrays)?) + }; + + 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)?; + if let Some(sp_rows) = &self.split_point_rows { + for (row_idx, row) in rows.iter().enumerate() { + let partition = partition_point_rows(sp_rows, &row); + emit(row_idx, partition); + } + } else { + for row_idx in 0..rows.num_rows() { + emit(row_idx, 0); + } + } + Ok(()) + } +} + +#[inline] +fn partition_point_rows(sp_rows: &Rows, row: &Row<'_>) -> usize { + let mut left = 0; + let mut right = sp_rows.num_rows(); + while left < right { + let mid = left + (right - left) / 2; + if sp_rows.row(mid) <= *row { + left = mid + 1; + } else { + right = mid; + } + } + left +} + +#[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]])) + } + + #[test] + fn test_router_decimal_precision_widening() -> Result<()> { + // Split point defined with precision 10, scale 2 + let split_points = vec![SplitPoint::new(vec![ScalarValue::Decimal128( + Some(1000), + 10, + 2, + )])]; + let sort_options = vec![SortOptions::default()]; + let target_data_types = vec![DataType::Decimal128(20, 2)]; + + let router = RangeRouter::try_new_with_data_types( + &sort_options, + &split_points, + &target_data_types, + )?; + assert_eq!(router.data_types(), &target_data_types); + + // Column array is Decimal128(20, 2) + let array = Arc::new( + Decimal128Array::from(vec![Some(500i128), Some(1000i128), Some(2000i128)]) + .with_precision_and_scale(20, 2)?, + ) as ArrayRef; + + assert_routing(&router, &[array], &[0, 1, 1], Some(&[vec![0], vec![1, 2]])) + } + + #[test] + fn test_router_timestamp_timezone_coercion() -> Result<()> { + // Split point defined without timezone + let split_points = vec![SplitPoint::new(vec![ + ScalarValue::TimestampNanosecond(Some(100), None), + ScalarValue::Int64(Some(0)), + ])]; + let sort_options = vec![SortOptions::default(), SortOptions::default()]; + let target_data_types = vec![ + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + DataType::Int64, + ]; + + let router = RangeRouter::try_new_with_data_types( + &sort_options, + &split_points, + &target_data_types, + )?; + assert_eq!(router.data_types(), &target_data_types); + + let col1 = Arc::new( + TimestampNanosecondArray::from(vec![Some(50i64), Some(200i64)]) + .with_timezone("UTC"), + ) as ArrayRef; + let col2 = Arc::new(Int64Array::from(vec![Some(1i64), Some(2i64)])) as ArrayRef; + + assert_routing(&router, &[col1, col2], &[0, 1], Some(&[vec![0], vec![1]])) + } + + #[test] + fn test_router_type_mismatch_error() -> Result<()> { + let split_points = make_split_points_1d(vec![ScalarValue::Int64(Some(10))]); + let sort_options = vec![SortOptions::default()]; + + let router = RangeRouter::try_new(&sort_options, &split_points)?; + + // Pass Float64Array instead of Int64Array + let invalid_input = Arc::new(Float64Array::from(vec![5.0, 15.0])) as ArrayRef; + let err = router.route_with(&[invalid_input], |_, _| {}).unwrap_err(); + assert!( + err.to_string() + .contains("Range partitioning expected column 0 to be of type Int64") + ); + + // Pass wrong column count + let err = router.route_with(&[], |_, _| {}).unwrap_err(); + assert!( + err.to_string() + .contains("Range partitioning expected 1 columns, but got 0") + ); + + Ok(()) + } +}