From 07fa010f1e456bc392f5be0a23f3ab1aee19b73e Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Sun, 23 Aug 2026 13:59:55 -0700 Subject: [PATCH 1/3] Range partition on primitives. --- datafusion/physical-plan/Cargo.toml | 4 + .../benches/range_repartition.rs | 361 +++++++ .../physical-plan/src/repartition/mod.rs | 138 ++- .../physical-plan/src/repartition/range.rs | 903 ++++++++++++++++++ 4 files changed, 1330 insertions(+), 76 deletions(-) create mode 100644 datafusion/physical-plan/benches/range_repartition.rs create mode 100644 datafusion/physical-plan/src/repartition/range.rs 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..4a55aec1daa00 --- /dev/null +++ b/datafusion/physical-plan/benches/range_repartition.rs @@ -0,0 +1,361 @@ +// 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::new_range_partitioner(&range_part, Time::default()); + 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::new_range_partitioner(&range_part, Time::default()); + 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::new_range_partitioner(&range_part, Time::default()); + 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::new_range_partitioner(&range_part, Time::default()); + 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..7ee7202825ece 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,21 @@ 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; +#[cfg(any(test, feature = "proto"))] +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, validate_range_split_points, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; @@ -89,6 +91,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 +641,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 +656,30 @@ 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.split_points == other.split_points + && self.sort_options == other.sort_options + } +} + +impl Eq for RangeExpr {} + +impl Hash for RangeExpr { + fn hash(&self, state: &mut H) { + self.on_columns.hash(state); + self.split_points.hash(state); + self.sort_options.hash(state); + } } impl RangeExpr { @@ -690,10 +712,19 @@ impl RangeExpr { "RangeExpr key count must match sort options" ); validate_range_split_points(&split_points, &sort_options)?; + let data_types: Vec = if !split_points.is_empty() { + (0..on_columns.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; Ok(Self { on_columns, split_points, sort_options, + router, }) } @@ -751,16 +782,9 @@ impl PhysicalExpr for RangeExpr { fn evaluate(&self, batch: &RecordBatch) -> Result { 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, )))) @@ -846,23 +870,6 @@ impl RangeExpr { } } -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. /// @@ -1004,16 +1011,25 @@ impl BatchPartitioner { timer: metrics::Time, ) -> Self { let ordering = range_partitioning.ordering().clone(); - let split_points = range_partitioning.split_points().to_vec(); + let split_points = range_partitioning.split_points(); let num_partitions = range_partitioning.partition_count(); let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); + let data_types: Vec = if !split_points.is_empty() { + (0..ordering.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + let router = RangeRouter::try_new(&data_types, &sort_options, split_points) + .unwrap_or_else(|_| { + RangeRouter::new_fallback(split_points.to_vec(), sort_options) + }); Self { state: BatchPartitionerState::Range { - partition_buffer: Vec::with_capacity(ordering.len()), ordering, - sort_options, - split_points, + router, indices: vec![vec![]; num_partitions], }, timer, @@ -1145,14 +1161,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 +1179,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 +1195,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 { diff --git a/datafusion/physical-plan/src/repartition/range.rs b/datafusion/physical-plan/src/repartition/range.rs new file mode 100644 index 0000000000000..202c816865d9b --- /dev/null +++ b/datafusion/physical-plan/src/repartition/range.rs @@ -0,0 +1,903 @@ +// 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 arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf}; +use datafusion_common::{DataFusionError, Result, ScalarValue}; +use datafusion_physical_expr::SplitPoint; + +/// An router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + 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), + /// Fallback for rare types not supported by RowConverter. + Fallback(FallbackRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given key types, split points, and sort options. + pub(crate) fn try_new( + data_types: &[DataType], + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result { + if split_points.is_empty() { + return Ok(Self::new_fallback( + split_points.to_vec(), + sort_options.to_vec(), + )); + } + + // Try single-column primitive fast path + if data_types.len() == 1 + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter fast path + if let Some(row_router) = + RowConverterRangeRouter::try_new(data_types, sort_options, split_points)? + { + return Ok(Self { + inner: RangeRouterInner::Row(row_router), + }); + } + + // Fallback + Ok(Self::new_fallback( + split_points.to_vec(), + sort_options.to_vec(), + )) + } + + /// Constructs a fallback router using dynamic row-by-row comparisons. + pub(crate) fn new_fallback( + split_points: Vec, + sort_options: Vec, + ) -> Self { + Self { + inner: RangeRouterInner::Fallback(FallbackRangeRouter { + split_points, + sort_options, + }), + } + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + match &self.inner { + RangeRouterInner::Primitive(r) => r.num_split_points(), + RangeRouterInner::Row(r) => r.num_split_points(), + RangeRouterInner::Fallback(r) => r.num_split_points(), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec], + ) -> Result<()> { + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_indices(first_col.as_ref(), indices) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_indices(arrays, indices), + RangeRouterInner::Fallback(r) => r.route_indices(arrays, indices), + } + } + + /// Appends output partition IDs to `partition_ids`. + pub(crate) fn route_partition_ids( + &self, + arrays: &[ArrayRef], + partition_ids: &mut Vec, + ) -> Result<()> { + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_partition_ids(first_col.as_ref(), partition_ids) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_partition_ids(arrays, partition_ids), + RangeRouterInner::Fallback(r) => r.route_partition_ids(arrays, partition_ids), + } + } +} + +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 num_split_points(&self) -> usize { + match self { + $( Self::$variant(r) => r.num_split_points(), )* + Self::Float32(r) => r.num_split_points(), + Self::Float64(r) => r.num_split_points(), + } + } + + fn route_indices(&self, array: &dyn Array, indices: &mut [Vec]) -> Result<()> { + match self { + $( + Self::$variant(r) => { + let arr = array.as_any().downcast_ref::<$array>().ok_or_else(|| { + DataFusionError::Internal(format!("Expected {}", stringify!($array))) + })?; + r.route_indices(arr, indices); + Ok(()) + } + )* + Self::Float32(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float32Array".to_string()) + })?; + r.route_indices(arr, indices); + Ok(()) + } + Self::Float64(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float64Array".to_string()) + })?; + r.route_indices(arr, indices); + Ok(()) + } + } + } + + fn route_partition_ids(&self, array: &dyn Array, partition_ids: &mut Vec) -> Result<()> { + match self { + $( + Self::$variant(r) => { + let arr = array.as_any().downcast_ref::<$array>().ok_or_else(|| { + DataFusionError::Internal(format!("Expected {}", stringify!($array))) + })?; + r.route_partition_ids(arr, partition_ids); + Ok(()) + } + )* + Self::Float32(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float32Array".to_string()) + })?; + r.route_partition_ids(arr, partition_ids); + Ok(()) + } + Self::Float64(r) => { + let arr = array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Internal("Expected Float64Array".to_string()) + })?; + r.route_partition_ids(arr, partition_ids); + 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 num_split_points(&self) -> usize { + self.split_points.len() + } + + fn route_indices>( + &self, + array: &PrimitiveArray, + indices: &mut [Vec], + ) { + 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); + indices[p].push(idx as u32); + } + } else { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| sp >= val); + indices[p].push(idx as u32); + } + } + } else { + let null_partition = if nulls_first { 0 } else { split_points.len() }; + for idx in 0..array.len() { + if array.is_null(idx) { + indices[null_partition].push(idx as u32); + } 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) + }; + indices[p].push(idx as u32); + } + } + } + } + + fn route_partition_ids>( + &self, + array: &PrimitiveArray, + partition_ids: &mut Vec, + ) { + 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 &val in values { + let p = split_points.partition_point(|&sp| sp <= val); + partition_ids.push(p as u64); + } + } else { + for &val in values { + let p = split_points.partition_point(|&sp| sp >= val); + partition_ids.push(p as u64); + } + } + } else { + let null_partition = + (if nulls_first { 0 } else { split_points.len() }) as u64; + for idx in 0..array.len() { + if array.is_null(idx) { + partition_ids.push(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) + }; + partition_ids.push(p as u64); + } + } + } + } +} + +/// 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, + } + } + + fn num_split_points(&self) -> usize { + self.split_points.len() + } +} + +macro_rules! impl_float_values_router { + ($t:ty, $arr:ty) => { + impl FloatValuesRouter<$t> { + fn route_indices(&self, array: &$arr, indices: &mut [Vec]) { + 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 + }); + indices[p].push(idx as u32); + } + } else { + for (idx, &val) in values.iter().enumerate() { + let p = split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Less + }); + indices[p].push(idx as u32); + } + } + } else { + let null_partition = if nulls_first { 0 } else { split_points.len() }; + for idx in 0..array.len() { + if array.is_null(idx) { + indices[null_partition].push(idx as u32); + } 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 + }) + }; + indices[p].push(idx as u32); + } + } + } + } + + fn route_partition_ids(&self, array: &$arr, partition_ids: &mut Vec) { + 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 &val in values { + let p = split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Greater + }); + partition_ids.push(p as u64); + } + } else { + for &val in values { + let p = split_points.partition_point(|&sp| { + sp.total_cmp(&val) != Ordering::Less + }); + partition_ids.push(p as u64); + } + } + } else { + let null_partition = + (if nulls_first { 0 } else { split_points.len() }) as u64; + for idx in 0..array.len() { + if array.is_null(idx) { + partition_ids.push(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 + }) + }; + partition_ids.push(p as u64); + } + } + } + } + } + }; +} + +impl_float_values_router!(f32, Float32Array); +impl_float_values_router!(f64, Float64Array); + +/// Router backed by Arrow's RowConverter. +#[derive(Debug, Clone)] +struct RowConverterRangeRouter { + sort_fields: Vec, + 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 Ok(None); + } + + let row_converter = RowConverter::new(sort_fields.clone())?; + let num_cols = data_types.len(); + + let mut split_point_arrays = Vec::with_capacity(num_cols); + for col_idx in 0..num_cols { + let col_scalars = split_points.iter().map(|sp| sp.values()[col_idx].clone()); + let col_array = ScalarValue::iter_to_array(col_scalars)?; + split_point_arrays.push(col_array); + } + + let rows = row_converter.convert_columns(&split_point_arrays)?; + let split_point_rows = (0..rows.num_rows()) + .map(|i| rows.row(i).owned()) + .collect::>(); + + Ok(Some(Self { + sort_fields, + split_point_rows, + })) + } + + fn num_split_points(&self) -> usize { + self.split_point_rows.len() + } + + fn route_indices(&self, arrays: &[ArrayRef], indices: &mut [Vec]) -> Result<()> { + let row_converter = RowConverter::new(self.sort_fields.clone())?; + let rows = row_converter.convert_columns(arrays)?; + let num_rows = rows.num_rows(); + let sp_rows = &self.split_point_rows; + + for row_idx in 0..num_rows { + let row = rows.row(row_idx); + let partition = sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()); + indices[partition].push(row_idx as u32); + } + Ok(()) + } + + fn route_partition_ids( + &self, + arrays: &[ArrayRef], + partition_ids: &mut Vec, + ) -> Result<()> { + let row_converter = RowConverter::new(self.sort_fields.clone())?; + let rows = row_converter.convert_columns(arrays)?; + let num_rows = rows.num_rows(); + let sp_rows = &self.split_point_rows; + + for row_idx in 0..num_rows { + let row = rows.row(row_idx); + let partition = sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()); + partition_ids.push(partition as u64); + } + Ok(()) + } +} + +/// Fallback router using dynamic row-by-row comparisons. +#[derive(Debug, Clone)] +struct FallbackRangeRouter { + split_points: Vec, + sort_options: Vec, +} + +impl FallbackRangeRouter { + fn num_split_points(&self) -> usize { + self.split_points.len() + } + + fn route_indices(&self, arrays: &[ArrayRef], indices: &mut [Vec]) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + let mut row_key_buffer = Vec::with_capacity(arrays.len()); + for row_idx in 0..num_rows { + extract_row_at_idx_to_buf(arrays, row_idx, &mut row_key_buffer)?; + let partition = range_partition_id_fallback( + &row_key_buffer, + &self.split_points, + &self.sort_options, + )?; + indices[partition].push(row_idx as u32); + } + Ok(()) + } + + fn route_partition_ids( + &self, + arrays: &[ArrayRef], + partition_ids: &mut Vec, + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + let mut row_key_buffer = Vec::with_capacity(arrays.len()); + for row_idx in 0..num_rows { + extract_row_at_idx_to_buf(arrays, row_idx, &mut row_key_buffer)?; + let partition = range_partition_id_fallback( + &row_key_buffer, + &self.split_points, + &self.sort_options, + )?; + partition_ids.push(partition as u64); + } + Ok(()) + } +} + +fn range_partition_id_fallback( + 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) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + #[test] + fn test_primitive_router_i64_asc() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(30))]), + ]; + let sort_options = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + let data_types = vec![DataType::Int64]; + + let router = RangeRouter::try_new(&data_types, &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; + + let mut partition_ids = Vec::new(); + router.route_partition_ids(&[Arc::clone(&input)], &mut partition_ids)?; + // Split points: 10, 20, 30. Partitions: 0 (<10), 1 (10..20), 2 (20..30), 3 (>=30). + // For 5: <10 -> 0 + // For 10: <=10 -> 1 (partition_point returns index where sp <= val is false, so sp=10 <= 10 is true -> idx 1) + // For 15: <=10 true, <=20 false -> 1 + // For 20: <=20 true, <=30 false -> 2 + // For 25: -> 2 + // For 30: -> 3 + // For 35: -> 3 + // For None (nulls_first = true): -> 0 + assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 0]); + + let mut indices = vec![vec![]; 4]; + router.route_indices(&[input], &mut indices)?; + assert_eq!(indices[0], vec![0, 7]); + assert_eq!(indices[1], vec![1, 2]); + assert_eq!(indices[2], vec![3, 4]); + assert_eq!(indices[3], vec![5, 6]); + + Ok(()) + } + + #[test] + fn test_primitive_router_i64_desc() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int64(Some(30))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + ]; + let sort_options = vec![SortOptions { + descending: true, + nulls_first: false, + }]; + let data_types = vec![DataType::Int64]; + + let router = RangeRouter::try_new(&data_types, &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; + + let mut partition_ids = Vec::new(); + router.route_partition_ids(&[Arc::clone(&input)], &mut partition_ids)?; + // DESC split points: 30, 20, 10. + // For 35: sp >= 35 is false for all -> 0 + // For 30: sp >= 30 is true for 30 (idx 0), false for rest -> 1 + // For 25: sp >= 25 is true for 30 -> 1 + // For 20: sp >= 20 is true for 30, 20 -> 2 + // For 15: -> 2 + // For 10: -> 3 + // For 5: -> 3 + // For None (nulls_first = false): -> 3 + assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 3]); + + Ok(()) + } + + #[test] + fn test_float_router() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Float64(Some(0.0))]), + SplitPoint::new(vec![ScalarValue::Float64(Some(100.0))]), + ]; + let sort_options = vec![SortOptions { + descending: false, + nulls_first: false, + }]; + let data_types = vec![DataType::Float64]; + + let router = RangeRouter::try_new(&data_types, &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; + + let mut partition_ids = Vec::new(); + router.route_partition_ids(&[input], &mut partition_ids)?; + assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 2]); + + Ok(()) + } + + #[test] + fn test_row_converter_strings() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Utf8(Some("d".to_string()))]), + SplitPoint::new(vec![ScalarValue::Utf8(Some("m".to_string()))]), + SplitPoint::new(vec![ScalarValue::Utf8(Some("s".to_string()))]), + ]; + let sort_options = vec![SortOptions { + descending: false, + nulls_first: true, + }]; + let data_types = vec![DataType::Utf8]; + + let router = RangeRouter::try_new(&data_types, &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; + + let mut partition_ids = Vec::new(); + router.route_partition_ids(&[input], &mut partition_ids)?; + assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 0]); + + Ok(()) + } + + #[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 data_types = vec![DataType::Int64, DataType::Utf8]; + + let router = RangeRouter::try_new(&data_types, &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; + + let mut partition_ids = Vec::new(); + router.route_partition_ids(&[col1, col2], &mut partition_ids)?; + assert_eq!(partition_ids, vec![0, 1, 1, 2, 3, 3, 3]); + + Ok(()) + } +} From ccaab1314ed029e3345171bde888e891d4957de9 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Mon, 24 Aug 2026 12:53:29 -0700 Subject: [PATCH 2/3] Review feedback. --- .../benches/range_repartition.rs | 28 +- .../physical-plan/src/repartition/mod.rs | 25 +- .../physical-plan/src/repartition/range.rs | 381 +++++++----------- 3 files changed, 189 insertions(+), 245 deletions(-) diff --git a/datafusion/physical-plan/benches/range_repartition.rs b/datafusion/physical-plan/benches/range_repartition.rs index 4a55aec1daa00..7b2b0d5cc5afd 100644 --- a/datafusion/physical-plan/benches/range_repartition.rs +++ b/datafusion/physical-plan/benches/range_repartition.rs @@ -141,8 +141,11 @@ fn bench_range_repartition_i64_uniform(c: &mut Criterion) { BenchmarkId::new("partitions", num_partitions), &num_partitions, |b, _| { - let mut partitioner = - BatchPartitioner::new_range_partitioner(&range_part, Time::default()); + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); b.iter(|| { partitioner .partition(batch.clone(), |p, b| { @@ -189,8 +192,11 @@ fn bench_range_repartition_i64_sequential(c: &mut Criterion) { BenchmarkId::new("partitions", num_partitions), &num_partitions, |b, _| { - let mut partitioner = - BatchPartitioner::new_range_partitioner(&range_part, Time::default()); + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); b.iter(|| { partitioner .partition(batch.clone(), |p, b| { @@ -237,8 +243,11 @@ fn bench_range_repartition_utf8_uniform(c: &mut Criterion) { BenchmarkId::new("partitions", num_partitions), &num_partitions, |b, _| { - let mut partitioner = - BatchPartitioner::new_range_partitioner(&range_part, Time::default()); + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); b.iter(|| { partitioner .partition(batch.clone(), |p, b| { @@ -290,8 +299,11 @@ fn bench_range_repartition_composite_i64(c: &mut Criterion) { BenchmarkId::new("partitions", num_partitions), &num_partitions, |b, _| { - let mut partitioner = - BatchPartitioner::new_range_partitioner(&range_part, Time::default()); + let mut partitioner = BatchPartitioner::try_new_range_partitioner( + &range_part, + Time::default(), + ) + .unwrap(); b.iter(|| { partitioner .partition(batch.clone(), |p, b| { diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 7ee7202825ece..ab799fda48527 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -52,7 +52,6 @@ use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array} use arrow::compute::take_arrays; use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type}; use arrow_schema::SortOptions; -#[cfg(any(test, feature = "proto"))] use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; @@ -781,6 +780,10 @@ impl PhysicalExpr for RangeExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { + if self.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 partition_ids = Vec::with_capacity(batch.num_rows()); self.router @@ -1006,10 +1009,10 @@ impl BatchPartitioner { /// # 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(); let num_partitions = range_partitioning.partition_count(); @@ -1021,19 +1024,16 @@ impl BatchPartitioner { } else { vec![] }; - let router = RangeRouter::try_new(&data_types, &sort_options, split_points) - .unwrap_or_else(|_| { - RangeRouter::new_fallback(split_points.to_vec(), sort_options) - }); + let router = RangeRouter::try_new(&data_types, &sort_options, split_points)?; - Self { + Ok(Self { state: BatchPartitionerState::Range { ordering, router, indices: vec![vec![]; num_partitions], }, timer, - } + }) } /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. @@ -1069,7 +1069,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:?}") @@ -1243,6 +1243,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 index 202c816865d9b..734a8392ae52a 100644 --- a/datafusion/physical-plan/src/repartition/range.rs +++ b/datafusion/physical-plan/src/repartition/range.rs @@ -18,13 +18,13 @@ //! 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::utils::{compare_rows, extract_row_at_idx_to_buf}; -use datafusion_common::{DataFusionError, Result, ScalarValue}; +use datafusion_common::{DataFusionError, Result, ScalarValue, not_impl_err}; use datafusion_physical_expr::SplitPoint; /// An router for assigning rows to range partitions. @@ -39,8 +39,6 @@ enum RangeRouterInner { Primitive(PrimitiveRangeRouter), /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. Row(RowConverterRangeRouter), - /// Fallback for rare types not supported by RowConverter. - Fallback(FallbackRangeRouter), } impl RangeRouter { @@ -50,15 +48,9 @@ impl RangeRouter { sort_options: &[SortOptions], split_points: &[SplitPoint], ) -> Result { - if split_points.is_empty() { - return Ok(Self::new_fallback( - split_points.to_vec(), - sort_options.to_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]) { @@ -67,33 +59,12 @@ impl RangeRouter { }); } - // Try RowConverter fast path - if let Some(row_router) = - RowConverterRangeRouter::try_new(data_types, sort_options, split_points)? - { - return Ok(Self { - inner: RangeRouterInner::Row(row_router), - }); - } - - // Fallback - Ok(Self::new_fallback( - split_points.to_vec(), - sort_options.to_vec(), - )) - } - - /// Constructs a fallback router using dynamic row-by-row comparisons. - pub(crate) fn new_fallback( - split_points: Vec, - sort_options: Vec, - ) -> Self { - Self { - inner: RangeRouterInner::Fallback(FallbackRangeRouter { - split_points, - sort_options, - }), - } + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(data_types, sort_options, split_points)?; + Ok(Self { + inner: RangeRouterInner::Row(row_router), + }) } /// Number of split points configured in this router. @@ -101,7 +72,6 @@ impl RangeRouter { match &self.inner { RangeRouterInner::Primitive(r) => r.num_split_points(), RangeRouterInner::Row(r) => r.num_split_points(), - RangeRouterInner::Fallback(r) => r.num_split_points(), } } @@ -120,7 +90,6 @@ impl RangeRouter { } } RangeRouterInner::Row(r) => r.route_indices(arrays, indices), - RangeRouterInner::Fallback(r) => r.route_indices(arrays, indices), } } @@ -139,7 +108,6 @@ impl RangeRouter { } } RangeRouterInner::Row(r) => r.route_partition_ids(arrays, partition_ids), - RangeRouterInner::Fallback(r) => r.route_partition_ids(arrays, partition_ids), } } } @@ -403,22 +371,24 @@ impl PrimitiveValuesR if array.null_count() == 0 { let values = array.values().as_ref(); if !descending { - for &val in values { - let p = split_points.partition_point(|&sp| sp <= val); - partition_ids.push(p as u64); - } + partition_ids.extend( + values + .iter() + .map(|&val| split_points.partition_point(|&sp| sp <= val) as u64), + ); } else { - for &val in values { - let p = split_points.partition_point(|&sp| sp >= val); - partition_ids.push(p as u64); - } + partition_ids.extend( + values + .iter() + .map(|&val| split_points.partition_point(|&sp| sp >= val) as u64), + ); } } else { let null_partition = (if nulls_first { 0 } else { split_points.len() }) as u64; - for idx in 0..array.len() { + partition_ids.extend((0..array.len()).map(|idx| { if array.is_null(idx) { - partition_ids.push(null_partition); + null_partition } else { let val = array.value(idx); let p = if !descending { @@ -426,9 +396,9 @@ impl PrimitiveValuesR } else { split_points.partition_point(|&sp| sp >= val) }; - partition_ids.push(p as u64); + p as u64 } - } + })); } } } @@ -508,26 +478,24 @@ macro_rules! impl_float_values_router { if array.null_count() == 0 { let values = array.values().as_ref(); if !descending { - for &val in values { - let p = split_points.partition_point(|&sp| { + partition_ids.extend(values.iter().map(|&val| { + split_points.partition_point(|&sp| { sp.total_cmp(&val) != Ordering::Greater - }); - partition_ids.push(p as u64); - } + }) as u64 + })); } else { - for &val in values { - let p = split_points.partition_point(|&sp| { + partition_ids.extend(values.iter().map(|&val| { + split_points.partition_point(|&sp| { sp.total_cmp(&val) != Ordering::Less - }); - partition_ids.push(p as u64); - } + }) as u64 + })); } } else { let null_partition = (if nulls_first { 0 } else { split_points.len() }) as u64; - for idx in 0..array.len() { + partition_ids.extend((0..array.len()).map(|idx| { if array.is_null(idx) { - partition_ids.push(null_partition); + null_partition } else { let val = array.value(idx); let p = if !descending { @@ -539,9 +507,9 @@ macro_rules! impl_float_values_router { sp.total_cmp(&val) != Ordering::Less }) }; - partition_ids.push(p as u64); + p as u64 } - } + })); } } } @@ -554,7 +522,7 @@ impl_float_values_router!(f64, Float64Array); /// Router backed by Arrow's RowConverter. #[derive(Debug, Clone)] struct RowConverterRangeRouter { - sort_fields: Vec, + converter: Arc, split_point_rows: Vec, } @@ -563,7 +531,7 @@ impl RowConverterRangeRouter { data_types: &[DataType], sort_options: &[SortOptions], split_points: &[SplitPoint], - ) -> Result> { + ) -> Result { let sort_fields = data_types .iter() .zip(sort_options) @@ -571,28 +539,37 @@ impl RowConverterRangeRouter { .collect::>(); if !RowConverter::supports_fields(&sort_fields) { - return Ok(None); + return not_impl_err!( + "Range partitioning is not supported for data types: {:?}", + data_types + ); } - let row_converter = RowConverter::new(sort_fields.clone())?; + let row_converter = RowConverter::new(sort_fields)?; let num_cols = data_types.len(); - let mut split_point_arrays = Vec::with_capacity(num_cols); - for col_idx in 0..num_cols { - let col_scalars = split_points.iter().map(|sp| sp.values()[col_idx].clone()); - let col_array = ScalarValue::iter_to_array(col_scalars)?; - split_point_arrays.push(col_array); - } - - let rows = row_converter.convert_columns(&split_point_arrays)?; - let split_point_rows = (0..rows.num_rows()) - .map(|i| rows.row(i).owned()) - .collect::>(); - - Ok(Some(Self { - sort_fields, + 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 num_split_points(&self) -> usize { @@ -600,13 +577,10 @@ impl RowConverterRangeRouter { } fn route_indices(&self, arrays: &[ArrayRef], indices: &mut [Vec]) -> Result<()> { - let row_converter = RowConverter::new(self.sort_fields.clone())?; - let rows = row_converter.convert_columns(arrays)?; - let num_rows = rows.num_rows(); + let rows = self.converter.convert_columns(arrays)?; let sp_rows = &self.split_point_rows; - for row_idx in 0..num_rows { - let row = rows.row(row_idx); + for (row_idx, row) in rows.iter().enumerate() { let partition = sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()); indices[partition].push(row_idx as u32); } @@ -618,96 +592,56 @@ impl RowConverterRangeRouter { arrays: &[ArrayRef], partition_ids: &mut Vec, ) -> Result<()> { - let row_converter = RowConverter::new(self.sort_fields.clone())?; - let rows = row_converter.convert_columns(arrays)?; - let num_rows = rows.num_rows(); + let rows = self.converter.convert_columns(arrays)?; let sp_rows = &self.split_point_rows; - for row_idx in 0..num_rows { - let row = rows.row(row_idx); - let partition = sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()); - partition_ids.push(partition as u64); - } + partition_ids.extend( + rows.iter().map(|row| { + sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()) as u64 + }), + ); Ok(()) } } -/// Fallback router using dynamic row-by-row comparisons. -#[derive(Debug, Clone)] -struct FallbackRangeRouter { - split_points: Vec, - sort_options: Vec, -} +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; -impl FallbackRangeRouter { - fn num_split_points(&self) -> usize { - self.split_points.len() + fn make_split_points_1d(scalars: Vec) -> Vec { + scalars + .into_iter() + .map(|s| SplitPoint::new(vec![s])) + .collect() } - fn route_indices(&self, arrays: &[ArrayRef], indices: &mut [Vec]) -> Result<()> { - let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); - let mut row_key_buffer = Vec::with_capacity(arrays.len()); - for row_idx in 0..num_rows { - extract_row_at_idx_to_buf(arrays, row_idx, &mut row_key_buffer)?; - let partition = range_partition_id_fallback( - &row_key_buffer, - &self.split_points, - &self.sort_options, - )?; - indices[partition].push(row_idx as u32); - } - Ok(()) - } - - fn route_partition_ids( - &self, + fn assert_routing( + router: &RangeRouter, arrays: &[ArrayRef], - partition_ids: &mut Vec, + expected_partition_ids: &[u64], + expected_indices: Option<&[Vec]>, ) -> Result<()> { - let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); - let mut row_key_buffer = Vec::with_capacity(arrays.len()); - for row_idx in 0..num_rows { - extract_row_at_idx_to_buf(arrays, row_idx, &mut row_key_buffer)?; - let partition = range_partition_id_fallback( - &row_key_buffer, - &self.split_points, - &self.sort_options, - )?; - partition_ids.push(partition as u64); - } - Ok(()) - } -} + let mut partition_ids = Vec::new(); + router.route_partition_ids(arrays, &mut partition_ids)?; + assert_eq!(partition_ids, expected_partition_ids); -fn range_partition_id_fallback( - 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, + if let Some(expected) = expected_indices { + let mut indices = vec![vec![]; expected.len()]; + router.route_indices(arrays, &mut indices)?; + assert_eq!(indices, expected); } - } - Ok(low) -} -#[cfg(test)] -mod tests { - use super::*; - use std::sync::Arc; + Ok(()) + } #[test] fn test_primitive_router_i64_asc() -> Result<()> { - let split_points = vec![ - SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), - SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), - SplitPoint::new(vec![ScalarValue::Int64(Some(30))]), - ]; + 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, @@ -728,36 +662,21 @@ mod tests { None, ])) as ArrayRef; - let mut partition_ids = Vec::new(); - router.route_partition_ids(&[Arc::clone(&input)], &mut partition_ids)?; - // Split points: 10, 20, 30. Partitions: 0 (<10), 1 (10..20), 2 (20..30), 3 (>=30). - // For 5: <10 -> 0 - // For 10: <=10 -> 1 (partition_point returns index where sp <= val is false, so sp=10 <= 10 is true -> idx 1) - // For 15: <=10 true, <=20 false -> 1 - // For 20: <=20 true, <=30 false -> 2 - // For 25: -> 2 - // For 30: -> 3 - // For 35: -> 3 - // For None (nulls_first = true): -> 0 - assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 0]); - - let mut indices = vec![vec![]; 4]; - router.route_indices(&[input], &mut indices)?; - assert_eq!(indices[0], vec![0, 7]); - assert_eq!(indices[1], vec![1, 2]); - assert_eq!(indices[2], vec![3, 4]); - assert_eq!(indices[3], vec![5, 6]); - - Ok(()) + 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 = vec![ - SplitPoint::new(vec![ScalarValue::Int64(Some(30))]), - SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), - SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), - ]; + 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, @@ -778,28 +697,20 @@ mod tests { None, ])) as ArrayRef; - let mut partition_ids = Vec::new(); - router.route_partition_ids(&[Arc::clone(&input)], &mut partition_ids)?; - // DESC split points: 30, 20, 10. - // For 35: sp >= 35 is false for all -> 0 - // For 30: sp >= 30 is true for 30 (idx 0), false for rest -> 1 - // For 25: sp >= 25 is true for 30 -> 1 - // For 20: sp >= 20 is true for 30, 20 -> 2 - // For 15: -> 2 - // For 10: -> 3 - // For 5: -> 3 - // For None (nulls_first = false): -> 3 - assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 3]); - - Ok(()) + 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 = vec![ - SplitPoint::new(vec![ScalarValue::Float64(Some(0.0))]), - SplitPoint::new(vec![ScalarValue::Float64(Some(100.0))]), - ]; + 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, @@ -818,20 +729,21 @@ mod tests { None, ])) as ArrayRef; - let mut partition_ids = Vec::new(); - router.route_partition_ids(&[input], &mut partition_ids)?; - assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 2]); - - Ok(()) + 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 = vec![ - SplitPoint::new(vec![ScalarValue::Utf8(Some("d".to_string()))]), - SplitPoint::new(vec![ScalarValue::Utf8(Some("m".to_string()))]), - SplitPoint::new(vec![ScalarValue::Utf8(Some("s".to_string()))]), - ]; + 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, @@ -852,11 +764,12 @@ mod tests { None, ])) as ArrayRef; - let mut partition_ids = Vec::new(); - router.route_partition_ids(&[input], &mut partition_ids)?; - assert_eq!(partition_ids, vec![0, 1, 1, 2, 2, 3, 3, 0]); - - Ok(()) + 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] @@ -894,10 +807,24 @@ mod tests { let col2 = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "a", "z", "a"])) as ArrayRef; - let mut partition_ids = Vec::new(); - router.route_partition_ids(&[col1, col2], &mut partition_ids)?; - assert_eq!(partition_ids, vec![0, 1, 1, 2, 3, 3, 3]); + assert_routing( + &router, + &[col1, col2], + &[0, 1, 1, 2, 3, 3, 3], + Some(&[vec![0], vec![1, 2], vec![3], vec![4, 5, 6]]), + ) + } - Ok(()) + #[test] + fn test_router_empty_split_points() -> Result<()> { + let split_points = vec![]; + let sort_options = vec![SortOptions::default()]; + let data_types = vec![DataType::Int64]; + + let router = RangeRouter::try_new(&data_types, &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]])) } } From 06849b5e05555845c13bbdd7a34ac5519b291d44 Mon Sep 17 00:00:00 2001 From: Stu Hood Date: Wed, 26 Aug 2026 12:17:23 -0700 Subject: [PATCH 3/3] Review feedback. --- .../physical-plan/src/repartition/mod.rs | 93 +++--- .../physical-plan/src/repartition/range.rs | 302 ++++++------------ 2 files changed, 137 insertions(+), 258 deletions(-) diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index ab799fda48527..f97cc5ed98c50 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -59,7 +59,7 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::transpose; use datafusion_common::{ ColumnStatistics, DataFusionError, HashMap, SplitPoint, assert_or_internal_err, - internal_datafusion_err, internal_err, validate_range_split_points, + internal_datafusion_err, internal_err, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; @@ -658,16 +658,14 @@ pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_ #[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.split_points == other.split_points - && self.sort_options == other.sort_options + && self.router.split_points() == other.router.split_points() + && self.router.sort_options() == other.router.sort_options() } } @@ -676,8 +674,8 @@ impl Eq for RangeExpr {} impl Hash for RangeExpr { fn hash(&self, state: &mut H) { self.on_columns.hash(state); - self.split_points.hash(state); - self.sort_options.hash(state); + self.router.split_points().hash(state); + self.router.sort_options().hash(state); } } @@ -688,43 +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)?; - let data_types: Vec = if !split_points.is_empty() { - (0..on_columns.len()) - .map(|col_idx| split_points[0].values()[col_idx].data_type()) - .collect() - } else { - vec![] - }; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; - Ok(Self { - on_columns, - split_points, - sort_options, - router, - }) + let router = RangeRouter::try_new(sort_options, split_points)?; + Ok(Self { on_columns, router }) } /// Get the columns used to compute Range partition IDs. @@ -734,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() } } @@ -766,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(), )?)) } @@ -780,7 +761,7 @@ impl PhysicalExpr for RangeExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { - if self.split_points.is_empty() { + if self.router.split_points().is_empty() { return Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(0)))); } @@ -807,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 @@ -849,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() @@ -867,8 +850,8 @@ impl RangeExpr { .collect::>>()?; Ok(Arc::new(Self::try_new_parts( on_columns, - split_points, - sort_options, + &split_points, + &sort_options, )?)) } } @@ -1004,6 +987,19 @@ 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 @@ -1014,17 +1010,10 @@ impl BatchPartitioner { timer: metrics::Time, ) -> Result { let ordering = range_partitioning.ordering().clone(); - let split_points = range_partitioning.split_points(); let num_partitions = range_partitioning.partition_count(); let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); - let data_types: Vec = if !split_points.is_empty() { - (0..ordering.len()) - .map(|col_idx| split_points[0].values()[col_idx].data_type()) - .collect() - } else { - vec![] - }; - let router = RangeRouter::try_new(&data_types, &sort_options, split_points)?; + let router = + RangeRouter::try_new(&sort_options, range_partitioning.split_points())?; Ok(Self { state: BatchPartitionerState::Range { diff --git a/datafusion/physical-plan/src/repartition/range.rs b/datafusion/physical-plan/src/repartition/range.rs index 734a8392ae52a..0bfa23c59a729 100644 --- a/datafusion/physical-plan/src/repartition/range.rs +++ b/datafusion/physical-plan/src/repartition/range.rs @@ -24,12 +24,16 @@ 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}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; use datafusion_physical_expr::SplitPoint; -/// An router for assigning rows to range partitions. +/// A router for assigning rows to range partitions. #[derive(Debug, Clone)] pub(crate) struct RangeRouter { + split_points: Vec, + sort_options: Vec, inner: RangeRouterInner, } @@ -42,12 +46,21 @@ enum RangeRouterInner { } impl RangeRouter { - /// Constructs the best router for the given key types, split points, and sort options. + /// Constructs the best router for the given sort options and split points. pub(crate) fn try_new( - data_types: &[DataType], 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() @@ -55,60 +68,86 @@ impl RangeRouter { 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)?; + 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 { - match &self.inner { - RangeRouterInner::Primitive(r) => r.num_split_points(), - RangeRouterInner::Row(r) => r.num_split_points(), - } + self.split_points.len() } - /// Groups row indices from `arrays` into partition index buckets. - pub(crate) fn route_indices( - &self, - arrays: &[ArrayRef], - indices: &mut [Vec], - ) -> Result<()> { + /// 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_indices(first_col.as_ref(), indices) + r.route_with(first_col.as_ref(), emit) } else { Ok(()) } } - RangeRouterInner::Row(r) => r.route_indices(arrays, indices), + 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<()> { - match &self.inner { - RangeRouterInner::Primitive(r) => { - if let Some(first_col) = arrays.first() { - r.route_partition_ids(first_col.as_ref(), partition_ids) - } else { - Ok(()) - } - } - RangeRouterInner::Row(r) => r.route_partition_ids(arrays, partition_ids), - } + 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); + }) } } @@ -177,22 +216,17 @@ macro_rules! define_primitive_router { } } - fn num_split_points(&self) -> usize { - match self { - $( Self::$variant(r) => r.num_split_points(), )* - Self::Float32(r) => r.num_split_points(), - Self::Float64(r) => r.num_split_points(), - } - } - - fn route_indices(&self, array: &dyn Array, indices: &mut [Vec]) -> Result<()> { + 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_indices(arr, indices); + r.route_with(arr, emit); Ok(()) } )* @@ -200,42 +234,14 @@ macro_rules! define_primitive_router { let arr = array.as_any().downcast_ref::().ok_or_else(|| { DataFusionError::Internal("Expected Float32Array".to_string()) })?; - r.route_indices(arr, indices); + 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_indices(arr, indices); - Ok(()) - } - } - } - - fn route_partition_ids(&self, array: &dyn Array, partition_ids: &mut Vec) -> Result<()> { - match self { - $( - Self::$variant(r) => { - let arr = array.as_any().downcast_ref::<$array>().ok_or_else(|| { - DataFusionError::Internal(format!("Expected {}", stringify!($array))) - })?; - r.route_partition_ids(arr, partition_ids); - Ok(()) - } - )* - Self::Float32(r) => { - let arr = array.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Internal("Expected Float32Array".to_string()) - })?; - r.route_partition_ids(arr, partition_ids); - Ok(()) - } - Self::Float64(r) => { - let arr = array.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Internal("Expected Float64Array".to_string()) - })?; - r.route_partition_ids(arr, partition_ids); + r.route_with(arr, emit); Ok(()) } } @@ -315,14 +321,10 @@ impl PrimitiveValuesR } } - fn num_split_points(&self) -> usize { - self.split_points.len() - } - - fn route_indices>( + fn route_with, E: FnMut(usize, usize)>( &self, array: &PrimitiveArray, - indices: &mut [Vec], + mut emit: E, ) { let split_points = &self.split_points; let descending = self.sort_options.descending; @@ -333,19 +335,19 @@ impl PrimitiveValuesR if !descending { for (idx, &val) in values.iter().enumerate() { let p = split_points.partition_point(|&sp| sp <= val); - indices[p].push(idx as u32); + emit(idx, p); } } else { for (idx, &val) in values.iter().enumerate() { let p = split_points.partition_point(|&sp| sp >= val); - indices[p].push(idx as u32); + 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) { - indices[null_partition].push(idx as u32); + emit(idx, null_partition); } else { let val = array.value(idx); let p = if !descending { @@ -353,54 +355,11 @@ impl PrimitiveValuesR } else { split_points.partition_point(|&sp| sp >= val) }; - indices[p].push(idx as u32); + emit(idx, p); } } } } - - fn route_partition_ids>( - &self, - array: &PrimitiveArray, - partition_ids: &mut Vec, - ) { - 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 { - partition_ids.extend( - values - .iter() - .map(|&val| split_points.partition_point(|&sp| sp <= val) as u64), - ); - } else { - partition_ids.extend( - values - .iter() - .map(|&val| split_points.partition_point(|&sp| sp >= val) as u64), - ); - } - } else { - let null_partition = - (if nulls_first { 0 } else { split_points.len() }) as u64; - partition_ids.extend((0..array.len()).map(|idx| { - if array.is_null(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) - }; - p as u64 - } - })); - } - } } /// Generic router for floating point values using total ordering. @@ -417,16 +376,12 @@ impl FloatValuesRouter { sort_options, } } - - fn num_split_points(&self) -> usize { - self.split_points.len() - } } macro_rules! impl_float_values_router { ($t:ty, $arr:ty) => { impl FloatValuesRouter<$t> { - fn route_indices(&self, array: &$arr, indices: &mut [Vec]) { + 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; @@ -438,21 +393,21 @@ macro_rules! impl_float_values_router { let p = split_points.partition_point(|&sp| { sp.total_cmp(&val) != Ordering::Greater }); - indices[p].push(idx as u32); + emit(idx, p); } } else { for (idx, &val) in values.iter().enumerate() { let p = split_points.partition_point(|&sp| { sp.total_cmp(&val) != Ordering::Less }); - indices[p].push(idx as u32); + 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) { - indices[null_partition].push(idx as u32); + emit(idx, null_partition); } else { let val = array.value(idx); let p = if !descending { @@ -464,54 +419,11 @@ macro_rules! impl_float_values_router { sp.total_cmp(&val) != Ordering::Less }) }; - indices[p].push(idx as u32); + emit(idx, p); } } } } - - fn route_partition_ids(&self, array: &$arr, partition_ids: &mut Vec) { - 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 { - partition_ids.extend(values.iter().map(|&val| { - split_points.partition_point(|&sp| { - sp.total_cmp(&val) != Ordering::Greater - }) as u64 - })); - } else { - partition_ids.extend(values.iter().map(|&val| { - split_points.partition_point(|&sp| { - sp.total_cmp(&val) != Ordering::Less - }) as u64 - })); - } - } else { - let null_partition = - (if nulls_first { 0 } else { split_points.len() }) as u64; - partition_ids.extend((0..array.len()).map(|idx| { - if array.is_null(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 - }) - }; - p as u64 - } - })); - } - } } }; } @@ -572,34 +484,18 @@ impl RowConverterRangeRouter { }) } - fn num_split_points(&self) -> usize { - self.split_point_rows.len() - } - - fn route_indices(&self, arrays: &[ArrayRef], indices: &mut [Vec]) -> 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()); - indices[partition].push(row_idx as u32); - } - Ok(()) - } - - fn route_partition_ids( + fn route_with( &self, arrays: &[ArrayRef], - partition_ids: &mut Vec, + mut emit: E, ) -> Result<()> { let rows = self.converter.convert_columns(arrays)?; let sp_rows = &self.split_point_rows; - partition_ids.extend( - rows.iter().map(|row| { - sp_rows.partition_point(|sp| sp.as_ref() <= row.as_ref()) as u64 - }), - ); + 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(()) } } @@ -646,9 +542,8 @@ mod tests { descending: false, nulls_first: true, }]; - let data_types = vec![DataType::Int64]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + let router = RangeRouter::try_new(&sort_options, &split_points)?; assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); let input = Arc::new(Int64Array::from(vec![ @@ -681,9 +576,8 @@ mod tests { descending: true, nulls_first: false, }]; - let data_types = vec![DataType::Int64]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + let router = RangeRouter::try_new(&sort_options, &split_points)?; assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); let input = Arc::new(Int64Array::from(vec![ @@ -715,9 +609,8 @@ mod tests { descending: false, nulls_first: false, }]; - let data_types = vec![DataType::Float64]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + let router = RangeRouter::try_new(&sort_options, &split_points)?; assert!(matches!(router.inner, RangeRouterInner::Primitive(_))); let input = Arc::new(Float64Array::from(vec![ @@ -748,9 +641,8 @@ mod tests { descending: false, nulls_first: true, }]; - let data_types = vec![DataType::Utf8]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + let router = RangeRouter::try_new(&sort_options, &split_points)?; assert!(matches!(router.inner, RangeRouterInner::Row(_))); let input = Arc::new(StringArray::from(vec![ @@ -798,9 +690,8 @@ mod tests { nulls_first: false, }, ]; - let data_types = vec![DataType::Int64, DataType::Utf8]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + 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; @@ -819,9 +710,8 @@ mod tests { fn test_router_empty_split_points() -> Result<()> { let split_points = vec![]; let sort_options = vec![SortOptions::default()]; - let data_types = vec![DataType::Int64]; - let router = RangeRouter::try_new(&data_types, &sort_options, &split_points)?; + 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;