Repository navigation
refactor: move EnforceDistribution per-child flags into ChildPolicy - #26716
Conversation
|
@andygrove gentle ping, could you take a look when you get a chance? CI also needs approval to run. |
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
@Aaryan-Dadu, thanks for this refactor. Extracting the interacting distribution-policy decisions makes them easier to understand and test independently. The approach is appropriate for the issue's scope, and the focused tests cover the important threshold and preservation interactions.
The extraction stays at the existing decision boundary without introducing a competing implementation or changing public contracts. Static comparison confirms that the predicates, lazy aggregate checks, and repartitioning decisions retain their behavior. Source repartitioning, range scaling, co-partitioning enforcement, and ordering restoration remain consistent with the previous implementation.
I'm approving because the change improves local maintainability without altering optimizer behavior. The targeted unit and integration suites passed. I've left two optional suggestions concerning legacy partitioning coverage and documentation of the partition-count snapshot boundary; neither is blocking.
| !policy(&config, TARGET, hash_on(&["a"])).allow_subset_satisfy_partitioning | ||
| ); | ||
| } | ||
|
|
There was a problem hiding this comment.
Optional coverage suggestion: could we test the grouping-ID policy with both KeyPartitioned and deprecated HashPartitioned? ChildPolicy::new supports both, but hash_on only constructs KeyPartitioned. A small table-driven test would help protect the legacy case during migration.
| maintains, | ||
| RepartitionRequirementStatus { | ||
| |(child_idx, mut child, required_input_ordering, maintains, status)| { | ||
| let policy = ChildPolicy::new( |
There was a problem hiding this comment.
Optional documentation suggestion: could we clarify that the policy intentionally captures the partition count before source repartitioning, while hash eligibility uses the updated count afterward? A short comment would make this snapshot boundary easier to understand without changing the existing behavior.
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Looks nice -- thank you @Aaryan-Dadu and @kosiew
FYI @jayzhan211 and @zhuqi-lucas
|
🚀 |
Which issue does this PR close?
Rationale for this change
Step 1 from the issue. The per-child closure in
ensure_distribution_with_statsworks out several flags inline (allow_subset_satisfy_partitioning,add_roundrobin,needs_repartition, ...). They depend on each other and can only be tested through full plan tests right now. This moves them into one place so they can be unit tested. No behaviour change.What changes are included in this PR?
ChildPolicyConfig: the settings that are the same for every child.ChildPolicy::new: works out the flags for one child. The partial aggregate check (including thepartial_aggregate_input_keeps_groups_localcase from feat: reuse range partitioning for date_trunc aggregation #26449) is passed as a closure, so it still only runs when it did before.ChildPolicy::needs_hash_repartition: the repartition decision from the hash arm.What is the testing strategy for this PR?
New unit tests in
enforce_distribution.rsfor the flags. The existingphysical_optimizertests and sqllogictests pass without changes.Are there any user-facing changes?
No