gene-bordegaray commented on code in PR #24598:
URL: https://github.com/apache/datafusion/pull/24598#discussion_r3858596383
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -999,25 +1009,31 @@ 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<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<SortOptions> = ordering.iter().map(|e|
e.options).collect();
+ let data_types: Vec<DataType> = if !split_points.is_empty() {
+ (0..ordering.len())
+ .map(|col_idx| split_points[0].values()[col_idx].data_type())
Review Comment:
could we call validate_range_split_points before indexing here?
##########
datafusion/physical-plan/src/repartition/range.rs:
##########
Review Comment:
I am looking at the overall design here and I think a good amount of things
are getting spread around / responsibilities kinda jumbled up and making this
hard to follow.
What I imaging is the `RangeRouter` owning the range metadata with the enum
to either use the `Primitive` or the `Rows` routes.
Botht he `RangeExpr` and the `BatchPartitioner` are doing all the
`RangeRouter` construction separtely. I think this should just be consilidated
into `RangeRouter::try_new`. Thus `RangeRouter` becomes the single place we
ensure that split point invariants, guarantees, type matching, etc.
##########
datafusion/physical-plan/src/repartition/range.rs:
##########
@@ -0,0 +1,830 @@
+// 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};
+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),
+}
+
+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<Self> {
+ // 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 {
+ inner: RangeRouterInner::Primitive(primitive_router),
+ });
+ }
+
+ // 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.
+ 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(),
+ }
+ }
+
+ /// Groups row indices from `arrays` into partition index buckets.
+ pub(crate) fn route_indices(
+ &self,
+ arrays: &[ArrayRef],
+ indices: &mut [Vec<u32>],
+ ) -> 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),
+ }
+ }
+
+ /// Appends output partition IDs to `partition_ids`.
+ pub(crate) fn route_partition_ids(
+ &self,
+ arrays: &[ArrayRef],
+ partition_ids: &mut Vec<u64>,
+ ) -> 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),
+ }
+ }
+}
+
+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 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<u32>]) -> Result<()> {
Review Comment:
route_indices and route_partition_ids duplicate a lot
Could we have one generic routing loop that emits (row_idx, partition) like:
```rust
fn route_with<E>(&self, arrays: &[ArrayRef], emit: E) -> Result<()>
where
E: FnMut(usize, usize);
```
--
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]
