saadtajwar commented on code in PR #23854:
URL: https://github.com/apache/datafusion/pull/23854#discussion_r3642336858
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -656,38 +661,99 @@ impl SharedBuildAccumulator {
}
}
- let filter_expr = if has_canceled_unknown {
- let mut when_then_branches = empty_partition_ids
- .into_iter()
- .map(|partition_id| {
- (
- lit(ScalarValue::UInt64(Some(partition_id as
u64))),
- lit(false),
- )
- })
- .collect::<Vec<_>>();
- when_then_branches.extend(real_branches);
-
- if when_then_branches.is_empty() {
- lit(true)
- } else {
- Arc::new(CaseExpr::try_new(
- Some(modulo_expr),
- when_then_branches,
- Some(lit(true)),
- )?) as Arc<dyn PhysicalExpr>
- }
- } else if real_branches.is_empty() {
+ 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 real_branches.len() == 1
+ } else if !has_canceled_unknown
+ && real_partition_ids.len() == 1
&& empty_partition_ids.len() + 1 == num_partitions
{
- Arc::clone(&real_branches[0].1)
+ Arc::clone(&partition_filters[real_partition_ids[0]])
+ } else if let Some(range_partitioning) =
&self.probe_range_partitioning {
+ // Range partitioning
+ assert_eq!(
+ partition_filters.len(),
+ range_partitioning.partition_count()
+ );
+ assert_eq!(self.on_right.len(),
range_partitioning.ordering().len());
+ let sort_exprs = self
+ .on_right
+ .iter()
+ .zip(range_partitioning.ordering())
+ .map(|(expr, sort_expr)| {
+ PhysicalSortExpr::new(Arc::clone(expr),
sort_expr.options)
+ })
+ .collect::<Vec<_>>();
+ let else_expr = partition_filters
+ .pop()
+ .expect("Range partitioning always has at least one
partition");
+ let mut when_then_expr =
Vec::with_capacity(partition_filters.len());
+ // CASE evaluates in order
+ //
+ // CASE
+ // WHEN key <range split[0] THEN F0
+ // WHEN key <range split[1] THEN F1
+ // ...
+ // ELSE Fn
+ // END
+ for (split_point, then_expr) in range_partitioning
+ .split_points()
+ .iter()
+ .zip(partition_filters)
+ {
+ let when_expr = build_lexicographic_filter(
+ &sort_exprs,
+ split_point.values(),
+ )?;
+ when_then_expr.push((when_expr, then_expr));
+ }
+
Review Comment:
Nit: is it possible to make this more idiomatic? Ex:
```
let when_then_expr = range_partitioning
.split_points()
.iter()
.zip(partition_filters)
.map(|(split_point, then_expr)| {
let when_expr =
build_lexicographic_filter(&sort_exprs, split_point.values())?;
Ok((when_expr, then_expr))
})
.collect::<Result<Vec<_>>>()?;
```
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -398,6 +404,15 @@ impl SharedBuildAccumulator {
),
};
+ let probe_range_partitioning = if partition_mode ==
PartitionMode::Partitioned {
+ match right_child.output_partitioning() {
+ crate::Partitioning::Range(range) => Some(range.clone()),
+ _ => None,
+ }
+ } else {
+ None
+ };
+
Review Comment:
nit: could we make this a bit more idiomatic by using a single match?
Something like
```
let probe_range_partitioning = match (
partition_mode,
right_child.output_partitioning(),
) {
(PartitionMode::Partitioned, crate::Partitioning::Range(range)) => {
Some(range.clone())
}
_ => None,
};
```
##########
datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs:
##########
@@ -398,6 +404,15 @@ impl SharedBuildAccumulator {
),
};
+ let probe_range_partitioning = if partition_mode ==
PartitionMode::Partitioned {
+ match right_child.output_partitioning() {
+ crate::Partitioning::Range(range) => Some(range.clone()),
+ _ => None,
+ }
+ } else {
+ None
+ };
+
Review Comment:
Also, should we import the `Partitioning` as opposed to using
`crate:::Partitioning`? (Genuine question, I'm not sure of the standards for
this file/crate haha so if this is the norm then feel free to ignore :) )
##########
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(
+ sort_exprs: &[PhysicalSortExpr],
+ thresholds: &[ScalarValue],
+) -> Result<Arc<dyn PhysicalExpr>> {
+ assert_or_internal_err!(!sort_exprs.is_empty(), "Sort expressions must not
be empty");
+ assert_or_internal_err!(
+ sort_exprs.len() == thresholds.len(),
+ "Sort expressions and thresholds must have the same length"
+ );
+
+ // Create filter expressions for each threshold
+ let mut filters: Vec<Arc<dyn PhysicalExpr>> =
Vec::with_capacity(thresholds.len());
+
+ let mut prev_sort_expr: Option<Arc<dyn PhysicalExpr>> = None;
+ for (sort_expr, value) in sort_exprs.iter().zip(thresholds.iter()) {
Review Comment:
meganit: this is a pretty long for loop with a lot of logic in it, is it
possible to make this more concise? Specifically I think having a lot of
mutable variables can make this a bit hard to reason about - but that could
just be me & my smooth brain lol
--
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]
