[SPARK-58120][SQL] doCanonicalize preserves keyGroupedPartitioning expression order

[wangyum]

What changes were proposed in this pull request?

This PR fixes BatchScanExec.doCanonicalize to use QueryPlan.normalizeExpressions instead of QueryPlan.normalizePredicates when canonicalizing keyGroupedPartitioning expressions. normalizePredicates combines expressions with And, canonicalizes, then splits back — which can reorder the expressions. normalizeExpressions canonicalizes each expression individually, preserving the original order.

Why are the changes needed?

BatchScanExec.doCanonicalize previously used QueryPlan.normalizePredicates on keyGroupedPartitioning, which combines the partition expressions with And, canonicalizes, then splits back. This reordering causes a mismatch between expression data types and partition key row values, leading to a ClassCastException at runtime. For example, if keyGroupedPartitioning is [id (IntegerType), data (StringType)], canonicalization could return [data (StringType), id (IntegerType)], and partition values no longer align with their corresponding expressions.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Added a regression test in KeyGroupedPartitioningSuite: SPARK-58120: doCanonicalize preserves keyGroupedPartitioning expression order, which verifies that the canonicalized keyGroupedPartitioning expressions preserve the same order and data types as the original.

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

Authored with assistance by GLM 5.2.

[peter-toth]

Took a quick look and the fix makes sense to me, thanks @wangyum. Pending CI.

[dongjoon-hyun]

Oh, thank you, @wangyum and @peter-toth !

+1, LGTM.

[dongjoon-hyun]

Since this is a regression at Apache Spark 4.2.0, cc @huaxingao , although I guess CCE unlikely happens in general cases, too.

• [SPARK-55535][SPARK-55092][SQL] Refactor KeyGroupedPartitioning and Storage Partition Join #54330

[dongjoon-hyun]

cc @viirya, @cloud-fan, @pan3793 too , FYI.

[cloud-fan]

0 blocking, 1 non-blocking, 0 nits.
Correct, minimally scoped fix. LGTM.

Verification

Traced the normalizePredicates -> per-element normalizeExpressions change. normalizePredicates reduces the seq to a single And, canonicalizes, and splits back; And is a CommutativeExpression, so its canonicalization reorders operands by hashCode (Expression.orderCommutative), permuting an order-significant sequence. Each expression's per-element canonical form is identical between the two forms (the old one applies the same normalizeExpressions to each And operand) — only ordering differs, so the fix is equivalent per-element and correct on ordering. Empty/single-element/None cases are all order-preserving no-ops. This also converges keyGroupedPartitioning onto the pattern already used for output one line above and by FileScan for its filter sequences; runtimeFilters correctly keeps normalizePredicates since a filter conjunction is order-insensitive.

Suggestions (1)

• Non-blocking, informational: the regression test discriminates the fix well — it asserts on canonicalized dataType metadata (Seq(IntegerType, StringType)), not just checkAnswer rows, which is the right signal for a canonicalization-ordering bug (a rows-only comparison could mask it). No change requested.

Release note

This is a 4.2.0-internal regression (the normalizePredicates call was introduced in #54330, in the same 4.2 line), not a regression from any released version — so no GA user is losing previously-working behavior yet. The failure is also loud (a ClassCastException, not a silent wrong result) and confined to storage-partitioned-join / key-grouped-partitioning paths whose expressions happen to reorder under hashCode. No need to fail the current RC

[huaxingao]

Thanks @dongjoon-hyun for the heads up. Agreed this is a real regression introduced in 4.2.0 by the KGP/SPJ refactor in #54330.

Before I announce the RC6 result, I want to sync on release impact. I don't think this rises to a blocker, for a few reasons:

• Results are never silently wrong. The reorder only affects the canonical copy of the plan, the one Spark uses to compare plans for reuse/subquery/AQE, not the copy it executes. Result-producing execution always runs the original, correctly-ordered plan.
• The SPJ planning decision is unaffected. Join-key matching compares partition expressions individually (order-preserving), so the decision to apply SPJ, and the actual join, never depends on the reordered list.
• Worst case is a loud, contained failure. Building the canonical copy only rearranges a column list in memory, so it can't fail on its own. Even if the reordered metadata were ever lined up against mixed-type partition values, it would surface as a visible ClassCastException confined to the SPJ / key-grouped-partitioning paths, not silent corruption.
• Narrow trigger and a safe fallback. It needs multi-column key-grouped partitioning whose expressions reorder under hashCode. Anyone who hits it can disable SPJ with spark.sql.sources.v2.bucketing.enabled=false and fall back to a normal shuffle join (correct results, just without the SPJ optimization) until 4.2.1.

So I'm leaning toward not failing RC6. If anyone feels it should block 4.2.0, please -1 on the vote thread. I'm happy to hold the announcement to discuss first.

Also cc @sunchao @szehon-ho

[viirya]

The fix looks right, and asserting on the canonicalized dataType order is the correct observable for this bug. One question on the test's discriminating power: whether the old code actually reorders this particular pair depends on the hashCode ordering of the two canonicalized attributes (orderCommutative sorts And operands by hash), which is deterministic but essentially arbitrary per (name, type, exprId) combination. If [id, data] happened to sort in the original order, this test would pass even without the fix and wouldn't guard against a reintroduction. Could you confirm the test fails on master without the fix? If it does (which the JIRA repro suggests), all good - might be worth a brief note in the test comment that the id/data pair is known to reorder under hashCode, so a future refactor of the test columns doesn't silently lose the coverage.

[wangyum]

Thanks @viirya. I confirmed the test fails on master without the fix — the [int, string] pair does reorder under normalizePredicates's hashCode sorting.

I also added a catalyst-level unit test in QueryPlanSuite — "SPARK-58120: normalizePredicates reorders expressions by hashCode via orderCommutative" — that directly demonstrates the reorder: normalizePredicates produces [StringType, IntegerType] while normalizeExpressions preserves [IntegerType, StringType].

[viirya]

Thanks, this fully resolves it - and more thoroughly than I suggested. The QueryPlanSuite test pins the reorder down at the right level: it exercises normalizePredicates vs normalizeExpressions directly with fixed ExprIds, so it doesn't depend on any end-to-end path happening to trigger the reorder, and it can't go stale under a test-column refactor. I confirmed it locally - the test passes on this branch, and since orderCommutative sorts by hashCode over attributes with stable exprIds, the asserted [StringType, IntegerType] order is deterministic rather than flaky. The cross-reference comment you added in KeyGroupedPartitioningSuite is exactly the note I had in mind. LGTM on this thread.

[sunchao]

To me, I'm fine to make this not blocking RC6. Spark uses canonicalized plans only for AQE and reuse comparisons, not execution. A plugin or fallback path could reconstruct a scan with reordered mixed-type partition keys, causing a deterministic ClassCastException that visibly fails the query, not silent incorrect results. Users can still disable SPJ with spark.sql.sources.v2.bucketing.enabled=false. Therefore, I'm OK to proceed with RC6 and promptly back-port the fix unless broader impact is demonstrated.

[szehon-ho]

Agree, it is only runtime exception in this case and not wrong result, sounds like ok to fix in minor release

[szehon-ho]

if we are just testing the canonicalization, we can just move to SparkPlanSuite by creating BatchScanExec without the round trips, wdyt? Unless we want to add more assert as @viirya point out

[wangyum]

+1 to move to SparkPlanSuite

[viirya]

LGTM. Minimally scoped, correct fix; the catalyst-level regression test pins the reorder down deterministically and I verified it passes locally. Thanks @wangyum!