[SPARK-46367][SQL][FOLLOWUP] Assert same arity in projectKeyedPartitionings

[cloud-fan]

What changes were proposed in this pull request?

Follow-up to #55519.

PartitioningPreservingUnaryExecNode.projectKeyedPartitionings assumes all input KPs share the same partitionKeys, which implies the same expression arity. This invariant is asserted by GroupPartitionsExec and is established by every upstream constructor of PartitioningCollection that feeds this method (a join's PartitioningCollection(left.outputPartitioning, right.outputPartitioning) combines KPs that EnsureRequirements has aligned to the same join keys).

If the invariant is ever violated upstream, indexing kp.expressions(i) for i >= kp.expressions.length throws an opaque IndexOutOfBoundsException that points at this method rather than at the producer.

This PR adds an assert that surfaces the real cause with a clear message. The invariant is unchanged; this just turns silent misuse into a debuggable failure, so a producer-side bug can be fixed at its source.

Why are the changes needed?

Improve error message when an upstream node violates the same-arity invariant. No behavior change on valid plans.

Does this PR introduce any user-facing change?

No (assertion-only; planner internal).

How was this patch tested?

New unit test SPARK-46367: mixed-arity KeyedPartitionings in input fail with a clear assertion in ProjectedOrderingAndPartitioningSuite that constructs a mixed-arity PartitioningCollection child and verifies the assert fires with the expected message instead of throwing IndexOutOfBoundsException.

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

Generated-by: Claude Opus 4.7

[cloud-fan]

cc @peter-toth

[peter-toth]

@cloud-fan , I don't think this is valid.

There simply can't be mixed-arity KeyedPartitionings in a PartitioningCollection. This is an invariant required by GroupPartitions as well:

spark/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

Lines 70 to 72 in 5949ab3

	// There can be multiple `KeyedPartitioning` in an output partitioning of a join, but they
	// can only differ in `expressions`. `partitionKeys` must match so we can calculate it only
	// once via `groupedPartitions`.

and actualy, it doesn't make sense to have/keep lower arity KPs in a collection if there are higher arity ones in it because the expressions of the lower must be the subset of the expressions of the higher.
If we know more details of the keys of partitions, like KP([a, b], [(1, 1), (2, 2)]) so we know that a = 1 and b = 1 in the first; and a = 2 and b = 2 in the second partition, why would we have/keep the less granular KP([a], [(1), (2)]) or KP([b], [(1), (2)]) in the collection?
Maybe the wording of the comment is not the best but the "partitionKeys must match" means that arity of KPs in a collection must be the same.

projectKeyedPartitionings() assumes that its input holds that invariant and because it produces only the maximum achievable granularity (arity) KPs it also maintains that invariant.

[cloud-fan]

@peter-toth you're right — the invariant should stay strict. I've reworked this PR to just add an assert so an upstream violation surfaces with a clear message instead of an opaque IndexOutOfBoundsException. The producer-side bug I was working around lives in our internal fork; I'll fix it there. PTAL.

[peter-toth]

Maybe you could use kps.tail.forall(...).

Please note that in GroupPartitionsExec we use .partitionKeys comparison, which is even more stricter, but more costly.

[peter-toth]

I wonder if we should move these invariant checks into PartitioningCollection.
Maybe expressions.length = + partitionKeys eq checks on KPs for maximum efficiency.

[dongjoon-hyun]

+1, LGTM. Thank you, @cloud-fan and @peter-toth .

[cloud-fan]

The parquet test failure is unrelated, thanks for review, merging to master/4.x!