kosiew commented on code in PR #25554:
URL: https://github.com/apache/datafusion/pull/25554#discussion_r4266915626
##########
datafusion/physical-plan/src/joins/utils.rs:
##########
@@ -1984,14 +1990,67 @@ pub(crate) fn symmetric_join_output_partitioning(
JoinType::Inner | JoinType::Right => {
adjust_right_output_partitioning(right_partitioning,
left_columns_len)?
}
- JoinType::Full => {
- // We could also use left partition count as they are necessarily
equal.
-
Partitioning::UnknownPartitioning(right_partitioning.partition_count())
- }
+ JoinType::Full => full_join_output_partitioning(left, right, on,
eq_properties)?,
};
Ok(result)
}
+fn full_join_output_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ on: JoinOnRef,
+ eq_properties: &mut EquivalenceProperties,
+) -> Result<Partitioning> {
+ let left_columns_len = left.schema().fields.len();
+ let left_partitioning = left.output_partitioning();
+ let unknown =
+
Partitioning::UnknownPartitioning(right.output_partitioning().partition_count());
+ if on.is_empty() || left_partitioning.partition_count() < 2 {
+ return Ok(unknown);
+ }
+ let (left_keys, right_keys) = on
+ .iter()
+ .map(|(l, r)| (Arc::clone(l), Arc::clone(r)))
+ .unzip();
+ let requirements = InputDistributionRequirements::co_partitioned(vec![
+ Distribution::KeyPartitioned(left_keys),
+ Distribution::KeyPartitioned(right_keys),
+ ]);
+ let children = [left.as_ref(), right.as_ref()];
+ if !requirements
+ .unsatisfied_co_partitioned_children("Full join", &children)?
+ .is_empty()
+ {
+ return Ok(unknown);
+ }
+ // Either side of a row may be null, so each key becomes `coalesce(left,
right)`.
+ let mut coalesced = Vec::with_capacity(on.len());
+ for (l, r) in on {
+ let r = add_offset_to_expr(Arc::clone(r), left_columns_len as _)?;
+ let left_first = coalesce_keys(Arc::clone(l), Arc::clone(&r))?;
Review Comment:
Computed join keys can invalidate the reported partitioning and
coalesce-order equivalence. For example, with a left key `CASE WHEN a IS NOT
NULL THEN a ELSE 0 END` and right key `c`, right-only rows produce 0 for the
left-first coalesce but `c` for the right-first one, potentially causing
incorrect downstream aggregation. Please restrict this optimization to column
keys or validate that computed keys preserve their values and become NULL on
null-extended rows, falling back to `UnknownPartitioning` otherwise. Please
also add a regression test covering unmatched rows and a downstream aggregate.
##########
datafusion/physical-plan/src/joins/utils.rs:
##########
@@ -4648,6 +4713,332 @@ mod tests {
Ok(())
}
+ fn int32_schema(first: &str, second: &str) -> Arc<Schema> {
+ Arc::new(Schema::new(vec![
+ Field::new(first, DataType::Int32, true),
+ Field::new(second, DataType::Int32, true),
+ ]))
+ }
+
+ fn unknown_partitioned(
+ schema: Arc<Schema>,
+ partitions: usize,
+ ) -> Arc<dyn ExecutionPlan> {
+ TestMemoryExec::try_new_exec(&vec![vec![]; partitions], schema, None)
+ .expect("memory source")
+ }
+
+ fn repartitioned(
+ schema: Arc<Schema>,
+ partitioning: Partitioning,
+ ) -> Arc<dyn ExecutionPlan> {
+ let source = unknown_partitioned(schema, 1);
+ Arc::new(RepartitionExec::try_new(source,
partitioning).expect("repartition"))
+ }
+
+ fn col_at(name: &str, index: usize) -> PhysicalExprRef {
+ Arc::new(Column::new(name, index))
+ }
+
+ fn coalesced(first: PhysicalExprRef, second: PhysicalExprRef) ->
PhysicalExprRef {
+ coalesce_keys(first, second).expect("coalesce")
+ }
+
+ fn range_on(
+ expr: PhysicalExprRef,
+ options: SortOptions,
+ split_points: &[i32],
+ ) -> Partitioning {
+ let ordering =
+ LexOrdering::new([PhysicalSortExpr::new(expr,
options)]).expect("one key");
+ let split_points = split_points
+ .iter()
+ .map(|value|
SplitPoint::new(vec![ScalarValue::Int32(Some(*value))]))
+ .collect();
+ let range = RangePartitioning::try_new(ordering,
split_points).expect("range");
+ Partitioning::Range(range)
+ }
+
+ fn a_equals_c() -> JoinOn {
+ vec![(col_at("a", 0), col_at("c", 0))]
+ }
+
+ fn join_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ join_type: JoinType,
+ on: JoinOnRef,
+ ) -> Partitioning {
+ let (schema, _) = build_join_schema(&left.schema(), &right.schema(),
&join_type);
+ let mut eq_properties = EquivalenceProperties::new(Arc::new(schema));
+ symmetric_join_output_partitioning(
+ left,
+ right,
+ &join_type,
+ on,
+ &mut eq_properties,
+ )
+ .expect("join partitioning")
+ }
+
+ fn full_join_partitioning(
+ left: &Arc<dyn ExecutionPlan>,
+ right: &Arc<dyn ExecutionPlan>,
+ on: JoinOnRef,
+ ) -> Partitioning {
+ join_partitioning(left, right, JoinType::Full, on)
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_hash_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+ let on = a_equals_c();
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &on),
+ Partitioning::Hash(vec![coalesced(col_at("a", 0), col_at("c",
2))], 4)
+ );
+ assert_eq!(
+ join_partitioning(&left, &right, JoinType::Inner, &on),
+ Partitioning::Hash(vec![col_at("c", 2)], 4)
+ );
+ assert_eq!(
+ join_partitioning(&left, &right, JoinType::Left, &on),
+ Partitioning::Hash(vec![col_at("a", 0)], 4)
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_range_keys() {
+ let options = SortOptions::new(false, true);
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ range_on(col_at("a", 0), options, &[10, 20, 30]),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ range_on(col_at("c", 0), options, &[10, 20, 30]),
+ );
+
+ let expected = range_on(
+ coalesced(col_at("a", 0), col_at("c", 2)),
+ options,
+ &[10, 20, 30],
+ );
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()),
+ expected
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_coalesces_every_key_pair() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0), col_at("b", 1)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0), col_at("d", 1)], 4),
+ );
+ let on = vec![
+ (col_at("a", 0), col_at("c", 0)),
+ (col_at("b", 1), col_at("d", 1)),
+ ];
+
+ let expected = Partitioning::Hash(
+ vec![
+ coalesced(col_at("a", 0), col_at("c", 2)),
+ coalesced(col_at("b", 1), col_at("d", 3)),
+ ],
+ 4,
+ );
+ assert_eq!(full_join_partitioning(&left, &right, &on), expected);
+ }
+
+ #[test]
+ fn
full_join_output_partitioning_is_unknown_for_different_range_split_points() {
+ let options = SortOptions::new(false, true);
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ range_on(col_at("a", 0), options, &[10, 20, 30]),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ range_on(col_at("c", 0), options, &[10, 20, 40]),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_is_unknown_for_single_partition_inputs() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 1),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 1),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(1)"
+ );
+ }
+
+ #[test]
+ fn
full_join_output_partitioning_is_unknown_when_keys_are_not_the_join_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("b", 1)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &a_equals_c()).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_join_output_partitioning_is_unknown_without_join_keys() {
+ let left = repartitioned(
+ int32_schema("a", "b"),
+ Partitioning::Hash(vec![col_at("a", 0)], 4),
+ );
+ let right = repartitioned(
+ int32_schema("c", "d"),
+ Partitioning::Hash(vec![col_at("c", 0)], 4),
+ );
+
+ assert_eq!(
+ full_join_partitioning(&left, &right, &[]).to_string(),
+ "UnknownPartitioning(4)"
+ );
+ }
+
+ #[test]
+ fn full_sort_merge_join_reports_coalesced_partitioning() {
Review Comment:
Could we also test sort-merge join projections? `SortMergeJoinExec` handles
property projection separately, so it would be useful to verify that retaining
or reordering both keys preserves the remapped coalesced partitioning, while
dropping one key invalidates it. This is optional coverage and can be deferred.
--
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]
