[SPARK-53074][SQL] Avoid partial clustering in SPJ to meet a child's required distribution

[chirag-s-db]

What changes were proposed in this pull request?

Currently, SPJ logic can apply partial clustering (when enabled) to either side of an inner JOIN as long as the nodes between the scan and JOIN preserve partitioning. This doesn't work if one of these nodes is using the scan's key-grouped partitioning to satisfy its required distribution (for example, a grouping agg or window function).

This PR avoids this issue by avoiding applying a partially clustered distribution to a JOIN's child if any node in that child relies on the KeyGroupedPartitioning to satisfy its required distribution (since it's not safe to do so with a partially clustered distribution).

Why are the changes needed?

Without this fix, using a partially-clustered distribution with SPJ may cause correctness issues.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

See test changes.

Was this patch authored or co-authored using generative AI tooling?

No.

[chirag-s-db]

@szehon-ho @sunchao Could you take a look at this PR? Thanks!

[violetnspct]

@chirag-s-db Should you be adding the tests to cover the following scenarios?

1. Aggregate operation requiring key-grouped partitioning. Need test that verifies behavior with GROUP BY operation. Aggregates also require specific partitioning that must be preserved.

2. Edge cases:

   • Multiple unary nodes in plan chain. Important because need to verify behavior when multiple nodes have distribution requirements
   • Mixed nodes with and without distribution requirements. Important because need to verify partial clustering decision is correct when only some nodes require specific distribution

[szehon-ho]

should we check number of partitions to make sure that partial cluster distribution replication did not kick in?

[chirag-s-db]

Good point, done.

[szehon-ho]

nit: can we add a comment about side-effects, like select row_number()

[chirag-s-db]

Added expansion to this comment.

[szehon-ho]

sorry just some question here:

1. unary node is just to make sure we exclude window functions, and not joins? Should we check window function node specifically, if that's the goal? Just trying to understand the reason why unary node in particular here.
2. Is it a better check to do something like:
   u.child.outputPartitioning.satisfies(u.requiredChildDistribution.head)?

[chirag-s-db]

unary node is just to make sure we exclude window functions, and not joins? Should we check window function node specifically, if that's the goal? Just trying to understand the reason why unary node in particular here.

We don't want to restrict this only to window functions - any exec node that has a specified required distribution applies (including windows, grouping aggregates, etc.). Checking for unary nodes just generally makes this easier to reason about (since it guarantees that there's only a single required child distribution to check).

Note that we also exclude any non-unary nodes (see the line below) like JOINs. The reason that it's not safe to apply to a JOIN is that when applying the partially clustered distribution, all scans on the partially clustered side are marked as partially clustered, which would create incorrect results on the lower JOIN. For example, suppose we had 3 partitioned tables A, B, and C and joined like A JOIN (b JOIN C) (assuming all keys lined up with the partitioning etc). It would be safe to apply a partially clustered distribution to A (since the tasks on b JOIN c can simply be replicated), but it would not be safe to apply a partially clustered distribution to the right side (since the results of b JOIN c would be incorrect.

Is it a better check to do something like:
u.child.outputPartitioning.satisfies(u.requiredChildDistribution.head)?

I don't think this would work - the KeyGroupedPartitioning would still (on paper) return that it satisfies the required distribution, even if the distribution actually will be partially clustered. In addition, we haven't yet actually applied the partially clustered distribution here, so the output partitioning is still the original KeyGroupedPartitioning (without any of the pushed down spjParams applied).