gene-bordegaray commented on code in PR #24598: URL: https://github.com/apache/datafusion/pull/24598#discussion_r3880807598
########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec<DataType> = if !split_points.is_empty() { Review Comment: here we infer the data type of the spilt pionts and this gets passed to try_new() ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; Review Comment: another point: that if there is only one split point, then we aren't really validating anything. So I think `RowConverter` should check the expected types and schemas and compare it with the inptu arraya before calling arrow ########## datafusion/physical-plan/src/repartition/mod.rs: ########## Review Comment: this is importatnt to note, maybe for me or a follow up. `pull_from_input` is called once for each partition and make a new `BatchPartitioner` which for range also makes a `RangeRouter`. With this if we have a lot of input / output partitions, for each of these ew are encoding a lot of split points repetitively. If this becomes visible we should have an `Arc<RangeRouter>` I think this is follow up but maybe document this comehwere? ########## datafusion/physical-plan/benches/range_repartition.rs: ########## @@ -0,0 +1,373 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{ + LexOrdering, PhysicalExpr, PhysicalSortExpr, RangePartitioning, SplitPoint, +}; +use datafusion_physical_plan::metrics::Time; +use datafusion_physical_plan::repartition::{BatchPartitioner, RangeExpr}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +const BATCH_SIZE: usize = 8192; +const PARTITION_COUNTS: [usize; 7] = [8, 16, 32, 64, 128, 256, 512]; +const SEED: u64 = 42; + +fn create_i64_uniform_batch(schema: &SchemaRef, max_val: i64) -> RecordBatch { + let mut rng = StdRng::seed_from_u64(SEED); + let key_values: Vec<i64> = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let payload_values: Vec<i64> = (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<i64> = (0..BATCH_SIZE) + .map(|i| ((i as i64) * max_val) / (BATCH_SIZE as i64)) + .collect(); + let payload_values: Vec<i64> = (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<String> = (0..BATCH_SIZE) + .map(|_| format!("key_{:010}", rng.random_range(0..max_val))) + .collect(); + let payload_values: Vec<i64> = (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<i64> = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let key2_values: Vec<i64> = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..max_val)) + .collect(); + let payload_values: Vec<i64> = (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<SplitPoint> = (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<SplitPoint> = (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( Review Comment: I think we woudl want to add the try_new_range_partitioner in the bench to ensure we are constructing efficiently with clones and what not 👍 ########## datafusion/physical-plan/src/repartition/mod.rs: ########## Review Comment: I think this might be worth to do in this PR ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -47,18 +46,20 @@ use crate::{ PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, }; -use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; +#[cfg(test)] Review Comment: probably shouldnt have ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec<DataType> = if !split_points.is_empty() { + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + + // Try single-column primitive fast path + if data_types.len() == 1 + && !sort_options.is_empty() + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(&data_types, sort_options, split_points)?; + Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Split points configured in this router. + pub(crate) fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Sort options configured in this router. + pub(crate) fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + self.split_points.len() + } + + /// Generic routing entry point that calls `emit(row_idx, partition)` for every row. + pub(crate) fn route_with<E>(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + if self.split_points.is_empty() { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + emit(row_idx, 0); + } + return Ok(()); + } + + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_with(first_col.as_ref(), emit) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_with(arrays, emit), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec<u32>], + ) -> 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<u64>, + ) -> 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<f32>), + Float64(FloatValuesRouter<f64>), + } + + impl PrimitiveRangeRouter { + fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option<Self> { + 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::<PrimitiveArray<$target_arrow_type>>()?; + 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::<Float32Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float32(FloatValuesRouter::new(vals, sort_options))) + } + DataType::Float64 => { + let arr = split_array.as_any().downcast_ref::<Float64Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float64(FloatValuesRouter::new(vals, sort_options))) + } + _ => None, + } + } + + fn route_with<E>(&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::<Float32Array>().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::<Float64Array>().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<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> PrimitiveValuesRouter<T> { + fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } + + fn route_with<A: ArrowPrimitiveType<Native = T>, E: FnMut(usize, usize)>( + &self, + array: &PrimitiveArray<A>, + 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<T: Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: Copy + Send + Sync + 'static> FloatValuesRouter<T> { + fn new(split_points: Vec<T>, 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<E: FnMut(usize, usize)>(&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<RowConverter>, + split_point_rows: Vec<OwnedRow>, +} + +impl RowConverterRangeRouter { + fn try_new( + data_types: &[DataType], + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + let sort_fields = data_types + .iter() + .zip(sort_options) + .map(|(dt, opt)| SortField::new_with_options(dt.clone(), *opt)) + .collect::<Vec<_>>(); + + if !RowConverter::supports_fields(&sort_fields) { + return not_impl_err!( + "Range partitioning is not supported for data types: {:?}", + data_types + ); + } + + let row_converter = RowConverter::new(sort_fields)?; + let num_cols = data_types.len(); + + let split_point_rows = if split_points.is_empty() { + vec![] + } else { + let split_point_arrays = (0..num_cols) + .map(|col_idx| { + let col_scalars = + split_points.iter().map(|sp| sp.values()[col_idx].clone()); + ScalarValue::iter_to_array(col_scalars) + }) + .collect::<Result<Vec<_>>>()?; + + row_converter + .convert_columns(&split_point_arrays)? + .iter() + .map(|r| r.owned()) + .collect() + }; + + Ok(Self { + converter: Arc::new(row_converter), + split_point_rows, + }) + } + + fn route_with<E: FnMut(usize, usize)>( + &self, + arrays: &[ArrayRef], + mut emit: E, + ) -> Result<()> { + let rows = self.converter.convert_columns(arrays)?; Review Comment: another note here is we could use a scartch pad in the `BatchPartitioner` to keep appending to and clearing rather thn allocating each batch ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec<DataType> = if !split_points.is_empty() { + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + + // Try single-column primitive fast path + if data_types.len() == 1 + && !sort_options.is_empty() + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(&data_types, sort_options, split_points)?; + Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Split points configured in this router. + pub(crate) fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Sort options configured in this router. + pub(crate) fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + self.split_points.len() + } + + /// Generic routing entry point that calls `emit(row_idx, partition)` for every row. + pub(crate) fn route_with<E>(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + if self.split_points.is_empty() { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + emit(row_idx, 0); + } + return Ok(()); + } + + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_with(first_col.as_ref(), emit) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_with(arrays, emit), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec<u32>], + ) -> 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<u64>, + ) -> 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<f32>), + Float64(FloatValuesRouter<f64>), + } + + impl PrimitiveRangeRouter { + fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option<Self> { + 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::<PrimitiveArray<$target_arrow_type>>()?; + 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::<Float32Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float32(FloatValuesRouter::new(vals, sort_options))) + } + DataType::Float64 => { + let arr = split_array.as_any().downcast_ref::<Float64Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float64(FloatValuesRouter::new(vals, sort_options))) + } + _ => None, + } + } + + fn route_with<E>(&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::<Float32Array>().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::<Float64Array>().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<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> PrimitiveValuesRouter<T> { + fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } + + fn route_with<A: ArrowPrimitiveType<Native = T>, E: FnMut(usize, usize)>( + &self, + array: &PrimitiveArray<A>, + 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<T: Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: Copy + Send + Sync + 'static> FloatValuesRouter<T> { + fn new(split_points: Vec<T>, 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<E: FnMut(usize, usize)>(&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<RowConverter>, + split_point_rows: Vec<OwnedRow>, Review Comment: another point is that this may not be the ebst shape for rows. Since we call `convert_columns()` we get Arrow rows which store all the rows in a buffer like this: ```text offsets: [0, 12 , 24, 39, 52] ^ ^ ^ ^ | | | | buffer: [row0 row1 row2 row3] its one contiguous allocation ``` then we immediately call `owned()` on every split point right here: ```rust row_converter .convert_columns(&split_point_arrays)? .iter() .map(|r| r.owned()) .collect() ``` So this changes the internal representation to be a vec or owner rows which you can think of as a pointer to a `Box<[row bytes]>` for each row. This will have more overhead. We could just keep the `Rows` that convert_columns` gives use then use those in binary search. I think this would help at high partition counts but can also be marked as follow up. I am mostly documenting for me to have notes. 😄 ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec<DataType> = if !split_points.is_empty() { + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + + // Try single-column primitive fast path + if data_types.len() == 1 + && !sort_options.is_empty() + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(&data_types, sort_options, split_points)?; + Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Split points configured in this router. + pub(crate) fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Sort options configured in this router. + pub(crate) fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + self.split_points.len() + } + + /// Generic routing entry point that calls `emit(row_idx, partition)` for every row. + pub(crate) fn route_with<E>(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + if self.split_points.is_empty() { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + emit(row_idx, 0); + } + return Ok(()); + } + + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_with(first_col.as_ref(), emit) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_with(arrays, emit), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec<u32>], + ) -> 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<u64>, + ) -> 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<f32>), + Float64(FloatValuesRouter<f64>), + } + + impl PrimitiveRangeRouter { + fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option<Self> { + 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::<PrimitiveArray<$target_arrow_type>>()?; + 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::<Float32Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float32(FloatValuesRouter::new(vals, sort_options))) + } + DataType::Float64 => { + let arr = split_array.as_any().downcast_ref::<Float64Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float64(FloatValuesRouter::new(vals, sort_options))) + } + _ => None, + } + } + + fn route_with<E>(&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::<Float32Array>().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::<Float64Array>().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<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> PrimitiveValuesRouter<T> { + fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } + + fn route_with<A: ArrowPrimitiveType<Native = T>, E: FnMut(usize, usize)>( + &self, + array: &PrimitiveArray<A>, + 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<T: Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: Copy + Send + Sync + 'static> FloatValuesRouter<T> { + fn new(split_points: Vec<T>, 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<E: FnMut(usize, usize)>(&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<RowConverter>, + split_point_rows: Vec<OwnedRow>, +} + +impl RowConverterRangeRouter { + fn try_new( + data_types: &[DataType], + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + let sort_fields = data_types + .iter() + .zip(sort_options) + .map(|(dt, opt)| SortField::new_with_options(dt.clone(), *opt)) + .collect::<Vec<_>>(); + + if !RowConverter::supports_fields(&sort_fields) { + return not_impl_err!( + "Range partitioning is not supported for data types: {:?}", + data_types + ); + } + + let row_converter = RowConverter::new(sort_fields)?; + let num_cols = data_types.len(); + + let split_point_rows = if split_points.is_empty() { + vec![] + } else { + let split_point_arrays = (0..num_cols) + .map(|col_idx| { + let col_scalars = + split_points.iter().map(|sp| sp.values()[col_idx].clone()); + ScalarValue::iter_to_array(col_scalars) + }) + .collect::<Result<Vec<_>>>()?; + + row_converter + .convert_columns(&split_point_arrays)? + .iter() + .map(|r| r.owned()) + .collect() + }; + + Ok(Self { + converter: Arc::new(row_converter), + split_point_rows, + }) + } + + fn route_with<E: FnMut(usize, usize)>( + &self, + arrays: &[ArrayRef], + mut emit: E, + ) -> Result<()> { + let rows = self.converter.convert_columns(arrays)?; Review Comment: then later on we will evaluate the arrays and use this `convert_columns()` from arrow: https://arrow.apache.org/rust/arrow_row/struct.RowConverter.html#method.convert_columns This documents that it will panic if the schemas dont match what was ever put in `RowConverter::new` we do not check this in this path and could panic if these dont match. These should be checked up front then if they are not use the scalar comparartor if the `ScalarValues` can be compared or error if not ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,720 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{OwnedRow, RowConverter, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, not_impl_err, validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + split_points: Vec<SplitPoint>, + sort_options: Vec<SortOptions>, + inner: RangeRouterInner, +} + +#[derive(Debug, Clone)] +enum RangeRouterInner { + /// Specialized fast path for a single primitive column with non-null split points. + Primitive(PrimitiveRangeRouter), + /// Universal fast path using Arrow's RowConverter for arbitrary types and composite keys. + Row(RowConverterRangeRouter), +} + +impl RangeRouter { + /// Constructs the best router for the given sort options and split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + validate_range_split_points(split_points, sort_options)?; + + let data_types: Vec<DataType> = if !split_points.is_empty() { + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect() + } else { + vec![] + }; + + // Try single-column primitive fast path + if data_types.len() == 1 + && !sort_options.is_empty() + && let Some(primitive_router) = + PrimitiveRangeRouter::try_new(split_points, sort_options[0]) + { + return Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Primitive(primitive_router), + }); + } + + // Try RowConverter path + let row_router = + RowConverterRangeRouter::try_new(&data_types, sort_options, split_points)?; + Ok(Self { + split_points: split_points.to_vec(), + sort_options: sort_options.to_vec(), + inner: RangeRouterInner::Row(row_router), + }) + } + + /// Split points configured in this router. + pub(crate) fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Sort options configured in this router. + pub(crate) fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } + + /// Number of split points configured in this router. + pub(crate) fn num_split_points(&self) -> usize { + self.split_points.len() + } + + /// Generic routing entry point that calls `emit(row_idx, partition)` for every row. + pub(crate) fn route_with<E>(&self, arrays: &[ArrayRef], mut emit: E) -> Result<()> + where + E: FnMut(usize, usize), + { + if self.split_points.is_empty() { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + emit(row_idx, 0); + } + return Ok(()); + } + + match &self.inner { + RangeRouterInner::Primitive(r) => { + if let Some(first_col) = arrays.first() { + r.route_with(first_col.as_ref(), emit) + } else { + Ok(()) + } + } + RangeRouterInner::Row(r) => r.route_with(arrays, emit), + } + } + + /// Groups row indices from `arrays` into partition index buckets. + pub(crate) fn route_indices( + &self, + arrays: &[ArrayRef], + indices: &mut [Vec<u32>], + ) -> 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<u64>, + ) -> 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<f32>), + Float64(FloatValuesRouter<f64>), + } + + impl PrimitiveRangeRouter { + fn try_new(split_points: &[SplitPoint], sort_options: SortOptions) -> Option<Self> { + 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::<PrimitiveArray<$target_arrow_type>>()?; + 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::<Float32Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float32(FloatValuesRouter::new(vals, sort_options))) + } + DataType::Float64 => { + let arr = split_array.as_any().downcast_ref::<Float64Array>()?; + let vals = arr.values().to_vec(); + Some(Self::Float64(FloatValuesRouter::new(vals, sort_options))) + } + _ => None, + } + } + + fn route_with<E>(&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::<Float32Array>().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::<Float64Array>().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<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: ArrowNativeTypeOp + Ord + Copy + Send + Sync + 'static> PrimitiveValuesRouter<T> { + fn new(split_points: Vec<T>, sort_options: SortOptions) -> Self { + Self { + split_points, + sort_options, + } + } + + fn route_with<A: ArrowPrimitiveType<Native = T>, E: FnMut(usize, usize)>( + &self, + array: &PrimitiveArray<A>, + 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<T: Copy + Send + Sync + 'static> { + split_points: Vec<T>, + sort_options: SortOptions, +} + +impl<T: Copy + Send + Sync + 'static> FloatValuesRouter<T> { + fn new(split_points: Vec<T>, 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<E: FnMut(usize, usize)>(&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<RowConverter>, + split_point_rows: Vec<OwnedRow>, +} + +impl RowConverterRangeRouter { + fn try_new( + data_types: &[DataType], + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + let sort_fields = data_types + .iter() + .zip(sort_options) + .map(|(dt, opt)| SortField::new_with_options(dt.clone(), *opt)) + .collect::<Vec<_>>(); + + 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)?; Review Comment: then the `RowConverter` is configured with the sort fields that used those data types -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
