jayzhan211 commented on code in PR #24598: URL: https://github.com/apache/datafusion/pull/24598#discussion_r3904106399
########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,907 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{Row, RowConverter, Rows, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_err, + validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + data_types: Vec<DataType>, + 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, + /// inferring key data types from the split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + let data_types: Option<Vec<DataType>> = if !split_points.is_empty() { + Some( + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect(), + ) + } else { + None + }; + Self::try_new_with_optional_data_types( + sort_options, + split_points, + data_types.as_deref(), + ) Review Comment: ```suggestion Self::try_new_with_optional_data_types(sort_options, split_points, None) ``` ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -667,34 +703,26 @@ impl RangeExpr { on_columns: Vec<PhysicalExprRef>, range_partitioning: &RangePartitioning, Review Comment: ```suggestion range_partitioning: &RangePartitioning, schema: &Schema, ``` ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -667,34 +703,26 @@ impl RangeExpr { on_columns: Vec<PhysicalExprRef>, range_partitioning: &RangePartitioning, ) -> Result<Self> { - let sort_options = range_partitioning + let sort_options: Vec<SortOptions> = 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) Review Comment: ```suggestion Self::try_new_parts( on_columns, range_partitioning.split_points(), &sort_options, schema, ) ``` ########## datafusion/physical-plan/src/repartition/range.rs: ########## @@ -0,0 +1,907 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Routers for range partitioning. + +use std::cmp::Ordering; +use std::sync::Arc; + +use arrow::array::*; +use arrow::compute::SortOptions; +use arrow::datatypes::*; +use arrow::row::{Row, RowConverter, Rows, SortField}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_err, not_impl_err, plan_err, + validate_range_split_points, +}; +use datafusion_physical_expr::SplitPoint; + +/// A router for assigning rows to range partitions. +#[derive(Debug, Clone)] +pub(crate) struct RangeRouter { + data_types: Vec<DataType>, + 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, + /// inferring key data types from the split points. + pub(crate) fn try_new( + sort_options: &[SortOptions], + split_points: &[SplitPoint], + ) -> Result<Self> { + let data_types: Option<Vec<DataType>> = if !split_points.is_empty() { + Some( + (0..sort_options.len()) + .map(|col_idx| split_points[0].values()[col_idx].data_type()) + .collect(), + ) + } else { + None + }; + Self::try_new_with_optional_data_types( + sort_options, + split_points, + data_types.as_deref(), + ) Review Comment: **Split-point width mismatch panics instead of returning a plan error** `RangeRouter::try_new` indexes into `values()` before `validate_range_split_points` gets a chance to run (validation happens inside `try_new_with_optional_data_types`): ```rust // range.rs:60 .map(|col_idx| split_points[0].values()[col_idx].data_type()) // thread panicked: index out of bounds: the len is 1 but the index is 1 ``` A width-1 split point with a width-2 ordering panics. Before this PR, RangeExpr::try_new_parts called validate_range_split_points first and returned plan_err!("Range partitioning split point 0 has width 1, but ordering has width 2"). This is reachable from RangeExpr::try_from_proto, so a malformed serialized plan panics the process rather than erroring. The pre-computation is redundant anyway — try_new_with_optional_data_types infers the identical types after validating That also makes the else if !split_points.is_empty() arm in try_new_with_optional_data_types live again (it's currently unreachable, since try_new always passes Some for non-empty split points). Worth a test asserting the plan error for a mismatched width. ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -667,34 +703,26 @@ impl RangeExpr { on_columns: Vec<PhysicalExprRef>, range_partitioning: &RangePartitioning, ) -> Result<Self> { - let sort_options = range_partitioning + let sort_options: Vec<SortOptions> = 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<PhysicalExprRef>, - split_points: Vec<SplitPoint>, - sort_options: Vec<SortOptions>, + split_points: &[SplitPoint], + sort_options: &[SortOptions], Review Comment: ```suggestion fn try_new_parts( on_columns: Vec<PhysicalExprRef>, split_points: &[SplitPoint], sort_options: &[SortOptions], schema: &Schema, ``` ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -667,34 +703,26 @@ impl RangeExpr { on_columns: Vec<PhysicalExprRef>, range_partitioning: &RangePartitioning, ) -> Result<Self> { - let sort_options = range_partitioning + let sort_options: Vec<SortOptions> = 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<PhysicalExprRef>, - split_points: Vec<SplitPoint>, - sort_options: Vec<SortOptions>, + split_points: &[SplitPoint], + sort_options: &[SortOptions], ) -> Result<Self> { assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key"); assert_or_internal_err!( on_columns.len() == sort_options.len(), "RangeExpr key count must match sort options" ); - validate_range_split_points(&split_points, &sort_options)?; - Ok(Self { - on_columns, - split_points, - sort_options, - }) + let router = RangeRouter::try_new(sort_options, split_points)?; Review Comment: **`RangeExpr` now rejects key types that `compare_rows` accepted** `try_new_parts` builds the router with `RangeRouter::try_new`, which infers `data_types` from the split-point scalars, and `RangeRouter::route_with` then hard-rejects any array whose `DataType` isn't exactly equal. `RepartitionExecState` sidesteps this by calling `try_new_with_data_types` with the schema types, but the `RangeExpr` path (`shared_bounds.rs:764`) does no coercion — so the dynamic filter fails on exactly the mismatches your own new repartition tests assert must work: ```rust // split point Decimal128(10, 2), probe column Decimal128(20, 2) let expr = RangeExpr::try_new(vec![col("k", &schema)?], &range_part)?; expr.evaluate(&batch) // Err(Execution("Range partitioning expected column 0 to be of type Decimal128(10, 2), but got Decimal128(20, 2)")) ``` Before this PR compare_rows routed that fine (per the comment on range_repartition_routes_decimal_with_wider_column_precision). A range-partitioned probe side whose split points come from stats/metadata — tz-less timestamps, narrower decimals — goes from working to a hard query failure. BatchPartitioner::try_new_range_partitioner has the same gap for external callers. A schema is available at every construction site, so resolve the key types there Callers: shared_bounds.rs:764 has self.probe_schema; try_from_proto has ctx.schema(). with_new_children can keep the router as-is (Self { on_columns: children, router: self.router.clone() }), which also drops a rebuild+revalidate. Please add a regression test through RangeExpr::evaluate mirroring the two repartition tests you added — the current ones only cover the RepartitionExec path, which is the one that already coerces ########## datafusion/physical-plan/src/repartition/mod.rs: ########## @@ -667,34 +703,26 @@ impl RangeExpr { on_columns: Vec<PhysicalExprRef>, range_partitioning: &RangePartitioning, ) -> Result<Self> { - let sort_options = range_partitioning + let sort_options: Vec<SortOptions> = 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<PhysicalExprRef>, - split_points: Vec<SplitPoint>, - sort_options: Vec<SortOptions>, + split_points: &[SplitPoint], + sort_options: &[SortOptions], ) -> Result<Self> { assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key"); assert_or_internal_err!( on_columns.len() == sort_options.len(), "RangeExpr key count must match sort options" ); - validate_range_split_points(&split_points, &sort_options)?; - Ok(Self { - on_columns, - split_points, - sort_options, - }) + let router = RangeRouter::try_new(sort_options, split_points)?; Review Comment: ```suggestion // Route on the key's actual types, not the split points' types: split // points are not required to carry the key's exact type (`Decimal128` // precision, timestamp timezone), and `RangeRouter` rejects any array // whose type differs from the one it was built for. let data_types = on_columns .iter() .map(|expr| expr.data_type(schema)) .collect::<Result<Vec<_>>>()?; let router = RangeRouter::try_new_with_data_types( sort_options, split_points, &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]
