gene-bordegaray commented on code in PR #23854:
URL: https://github.com/apache/datafusion/pull/23854#discussion_r3751472331
##########
datafusion/physical-plan/src/ordering.rs:
##########
@@ -52,3 +59,95 @@ pub enum InputOrderMode {
/// existing ordering.
Sorted,
}
+
+/// Build the filter expression with the given thresholds.
+/// This is now called outside of any locks to reduce critical section time.
+pub(crate) fn build_lexicographic_filter(
Review Comment:
this is now only used in topK and the diff can be smaller if we revertd this
👍
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -638,6 +645,218 @@ enum BatchPartitionerState {
/// executions and runs.
pub const REPARTITION_RANDOM_STATE: SeededRandomState =
SeededRandomState::with_seed(0);
+/// Physical expression that returns the Range partition for each input row.
+///
+/// This uses the same routing function as [`BatchPartitioner`], so dynamic
+/// filtering and repartitioning agree for every [`ScalarValue`] comparison.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct RangeExpr {
Review Comment:
this is a grat implementation. Thank you!
this file is growing quite large. maybe we make a follow up to break some of
this up a bit 👍
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -2263,6 +2469,38 @@ mod tests {
use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning,
SplitPoint};
use insta::assert_snapshot;
+ #[test]
+ fn range_expr_preserves_duplicate_remapped_children() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::UInt32, false),
+ Field::new("b", DataType::UInt32, false),
+ ]));
+ let range_partitioning = RangePartitioning::try_new(
+ [
+ PhysicalSortExpr::new(col("a", &schema)?,
SortOptions::default()),
+ PhysicalSortExpr::new(col("b", &schema)?,
SortOptions::default()),
+ ]
+ .into(),
+ vec![SplitPoint::new(vec![
+ ScalarValue::UInt32(Some(10)),
+ ScalarValue::UInt32(Some(20)),
+ ])],
+ )?;
+ let expr = Arc::new(RangeExpr::try_new(
+ vec![col("a", &schema)?, col("b", &schema)?],
+ &range_partitioning,
+ )?);
+ let remapped = col("a", &schema)?;
+ let rewritten =
+ expr.with_new_children(vec![Arc::clone(&remapped),
Arc::clone(&remapped)])?;
+
+ let children = rewritten.children();
+ assert_eq!(children.len(), 2);
Review Comment:
we are only checking that the children exist here, could we check the range
and sort properties as well
like the sort options, split points
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -638,6 +645,218 @@ enum BatchPartitionerState {
/// executions and runs.
pub const REPARTITION_RANDOM_STATE: SeededRandomState =
SeededRandomState::with_seed(0);
+/// Physical expression that returns the Range partition for each input row.
+///
+/// This uses the same routing function as [`BatchPartitioner`], so dynamic
+/// filtering and repartitioning agree for every [`ScalarValue`] comparison.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct RangeExpr {
+ on_columns: Vec<PhysicalExprRef>,
+ split_points: Vec<SplitPoint>,
+ sort_options: Vec<SortOptions>,
+}
+
+impl RangeExpr {
+ /// Creates a Range expression for `on_columns` using the supplied routing
+ /// metadata.
+ pub fn try_new(
+ on_columns: Vec<PhysicalExprRef>,
+ range_partitioning: &RangePartitioning,
+ ) -> Result<Self> {
+ let sort_options = range_partitioning
+ .ordering()
+ .iter()
+ .map(|expr| expr.options)
+ .collect();
+ Self::try_new_parts(
+ on_columns,
+ range_partitioning.split_points().to_vec(),
+ sort_options,
+ )
+ }
+
+ fn try_new_parts(
+ on_columns: Vec<PhysicalExprRef>,
+ split_points: Vec<SplitPoint>,
+ sort_options: Vec<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,
+ })
+ }
+
+ /// Returns the Range split points used for routing.
+ pub fn split_points(&self) -> &[SplitPoint] {
+ &self.split_points
+ }
+
+ /// Returns the per-key sort options used for routing.
+ pub fn sort_options(&self) -> &[SortOptions] {
+ &self.sort_options
+ }
+}
+
+impl Display for RangeExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "range_partition")
+ }
+}
+
+impl PhysicalExpr for RangeExpr {
+ fn children(&self) -> Vec<&PhysicalExprRef> {
+ self.on_columns.iter().collect()
Review Comment:
since we access this field we shoudl add a public accessor like `on_columns`
like `HashExpr` has 👍
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -638,6 +645,218 @@ enum BatchPartitionerState {
/// executions and runs.
pub const REPARTITION_RANDOM_STATE: SeededRandomState =
SeededRandomState::with_seed(0);
+/// Physical expression that returns the Range partition for each input row.
+///
+/// This uses the same routing function as [`BatchPartitioner`], so dynamic
+/// filtering and repartitioning agree for every [`ScalarValue`] comparison.
+#[derive(Debug, Hash, PartialEq, Eq)]
+pub struct RangeExpr {
+ on_columns: Vec<PhysicalExprRef>,
+ split_points: Vec<SplitPoint>,
+ sort_options: Vec<SortOptions>,
+}
+
+impl RangeExpr {
+ /// Creates a Range expression for `on_columns` using the supplied routing
+ /// metadata.
+ pub fn try_new(
+ on_columns: Vec<PhysicalExprRef>,
+ range_partitioning: &RangePartitioning,
+ ) -> Result<Self> {
+ let sort_options = range_partitioning
+ .ordering()
+ .iter()
+ .map(|expr| expr.options)
+ .collect();
+ Self::try_new_parts(
+ on_columns,
+ range_partitioning.split_points().to_vec(),
+ sort_options,
+ )
+ }
+
+ fn try_new_parts(
+ on_columns: Vec<PhysicalExprRef>,
+ split_points: Vec<SplitPoint>,
+ sort_options: Vec<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,
+ })
+ }
+
+ /// Returns the Range split points used for routing.
+ pub fn split_points(&self) -> &[SplitPoint] {
+ &self.split_points
+ }
+
+ /// Returns the per-key sort options used for routing.
+ pub fn sort_options(&self) -> &[SortOptions] {
+ &self.sort_options
+ }
+}
+
+impl Display for RangeExpr {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "range_partition")
+ }
+}
+
+impl PhysicalExpr for RangeExpr {
+ fn children(&self) -> Vec<&PhysicalExprRef> {
+ self.on_columns.iter().collect()
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<PhysicalExprRef>,
+ ) -> Result<PhysicalExprRef> {
+ assert_or_internal_err!(
+ children.len() == self.on_columns.len(),
+ "RangeExpr expected {} children, got {}",
+ self.on_columns.len(),
+ children.len()
+ );
+ Ok(Arc::new(Self::try_new_parts(
+ children,
+ self.split_points.clone(),
+ self.sort_options.clone(),
+ )?))
+ }
+
+ fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
+ Ok(DataType::UInt64)
+ }
+
+ fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
+ Ok(false)
+ }
+
+ fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+ 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);
+ }
+ Ok(ColumnarValue::Array(Arc::new(UInt64Array::from(
+ partition_ids,
+ ))))
+ }
+
+ fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "range_partition")
+ }
+
+ #[cfg(feature = "proto")]
+ fn try_to_proto(
+ &self,
+ ctx:
&datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+ ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
+ // Encode the raw ordered children: rebuilding a `LexOrdering` would
+ // deduplicate equivalent children after dynamic-filter remapping.
+ use datafusion_proto_models::protobuf;
+
+ let sort_exprs = self
+ .on_columns
+ .iter()
+ .zip(&self.sort_options)
+ .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr),
*options))
+ .collect::<Vec<_>>();
+ let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?;
+ let split_point = self
+ .split_points
+ .iter()
+ .map(|split_point| {
+ let value = split_point
+ .values()
+ .iter()
+ .map(|value| value.try_into().map_err(Into::into))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(protobuf::PhysicalRangeSplitPoint { value })
+ })
+ .collect::<Result<Vec<_>>>()?;
+ Ok(Some(protobuf::PhysicalExprNode {
+ expr_id: None,
+ expr_type: Some(protobuf::physical_expr_node::ExprType::RangeExpr(
+ protobuf::PhysicalRangeExprNode {
+ sort_expr,
+ split_point,
+ },
+ )),
+ }))
+ }
+}
+
+#[cfg(feature = "proto")]
+impl RangeExpr {
+ /// Reconstructs a [`RangeExpr`] from its protobuf representation.
+ pub fn try_from_proto(
+ node: &datafusion_proto_models::protobuf::PhysicalExprNode,
+ ctx:
&datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+ ) -> Result<PhysicalExprRef> {
+ // Decode the raw ordered children for the same reason as
`try_to_proto`.
+ use datafusion_proto_models::protobuf;
Review Comment:
lets move shared imports to modeule level 👍
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -688,38 +692,95 @@ impl SharedBuildAccumulator {
}
}
- let filter_expr = if has_canceled_unknown {
- let mut when_then_branches = empty_partition_ids
+ let filter_expr = if has_canceled_unknown
+ && real_partition_ids.is_empty()
+ && empty_partition_ids.is_empty()
+ {
+ lit(true)
+ } else if !has_canceled_unknown &&
real_partition_ids.is_empty() {
+ lit(false)
+ } else if !has_canceled_unknown
+ && real_partition_ids.len() == 1
+ && empty_partition_ids.len() + 1 == num_partitions
+ {
+ Arc::clone(&partition_filters[real_partition_ids[0]])
+ } else if let Some(range_partitioning) =
&self.probe_range_partitioning {
+ // Range partitioning
+ assert_eq!(
Review Comment:
this could also be a `assert_or_internal_err` I think
--
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]
