[SPARK-59249][SQL] Take the grouped key-row ordering from the shared InternalRowComparableWrapper cache

[peter-toth]

What changes were proposed in this pull request?

InternalRowComparableWrapper gains a public orderingFor(dataTypes) that returns the ordering from its existing orderingCache, and the two storage-partitioned join sites that built the same ordering by hand now go through it.

KeyedPartitioning.groupedKeyRowOrdering called RowOrdering.createNaturalAscendingOrdering directly, which is byte-for-byte the loadFunc of that cache. DataSourceV2ScanExecBase.outputPartitioning did the same to sort the partition keys it then hands to KeyedPartitioning.

Why are the changes needed?

Two reasons, and the second is worth more than it looks.

One definition instead of two. InternalRowComparableWrapper.equals is ordering.compare(row, other.row) == 0 over the cached ordering, while groupedKeyRowOrdering is what lays grouped partition keys out. So "two partition keys are equal" and "two partition keys sort together" already had to be the same relation, and they were, only by both call sites happening to name the same function.

The direct call is expensive. Only the Janino step inside GenerateOrdering is cached, keyed on the generated source. Everything upstream re-runs per call: ExpressionCanonicalizer over each SortOrder, building the whole Java source, CodeFormatter.stripOverlappingComments rebuilding it line by line, then hashing the multi-KB body to probe the compile cache. Measured in this worktree, for a two-column key:

RowOrdering.createNaturalAscendingOrdering    93-105 us per call
InternalRowComparableWrapper.orderingFor      0.16-0.6 us per call

DataSourceV2ScanExecBase.outputPartitioning is the site where that repeats. It is a def, and an instrumented run of KeyGroupedPartitioningSuite counted 31,922 calls across 1,167 scan instances, a mean of 27 per scan. The other callers are once-per-something and get the same ~100 us back as noise: GroupPartitionsExec.groupAndSortByKeys runs from a lazy val, EnsureRequirements runs once per join per pass, and KeyedPartitioning.keyRowOrdering is itself a lazy val.

Nothing gets slower. A NonFateSharingCache hit measured 0.163 us single-threaded and 0.334 us with 16 threads on one key. The cached value is a stateless comparator, so eviction costs a rebuild and nothing else, and almost every type list the new callers look up is one the wrapper path already loads for the same rows. The exception is KeyedPartitioning.keyRowOrdering, which keys on keyDataTypes, and that falls back to the expression types when there are no partition keys, so it can add an entry of its own.

SortMergeJoinEvaluatorFactory also calls createNaturalAscendingOrdering directly and is deliberately left alone. Those are join-key rows from an UnsafeProjection, never wrapped, and their counterpart is the child Sort's ordering rather than wrapper equality. They also run per partition on executors, so feeding their schemas into this cache would evict partition-key entries for no gain.

Does this PR introduce any user-facing change?

No. The cache's load function is the call that was being made directly, so it produces the same ordering.

One caveat, pre-existing on the wrapper path and now extended to these two sites. createNaturalAscendingOrdering goes through CodeGeneratorWithInterpretedFallback, which reads spark.sql.codegen.factoryMode on every call, while the cache reads it once per type list for the JVM's lifetime. So a test that flips that internal conf after an entry is loaded now gets the ordering built under the earlier mode. The two orderings agree semantically, so this costs coverage rather than correctness.

How was this patch tested?

A new InternalRowComparableWrapperSuite, which the class did not have, only a benchmark. It asserts that groupedKeyRowOrdering and a wrapper hold one ordering instance, and fails if the delegation is reverted.

Identity is what the test asserts because identity is what changes. Both sides were already built by the same function, so no comparison of rows can tell the two apart. What one instance buys is that neither side can later be given a definition the other does not have.

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

Generated-by: Claude Code (Opus 5)

[peter-toth]

cc @dongjoon-hyun

[dongjoon-hyun]

Went through the diff and the callers of groupedKeyRowOrdering (GroupPartitionsExec, EnsureRequirements, KeyedPartitioning.keyRowOrdering, DataSourceV2ScanExecBase). No correctness issue in the change itself: the shared BaseOrdering instance is stateless (natural ascending over BoundReferences, no mutable state, no references), it never leaves the driver, and keyDataTypes always resolves to the same list the wrappers were built with, so the identity the new test pins holds. Inline comments below are non-blocking things to weigh; two follow-up notes that fall outside the diff:

• The cache is placed in the wrapper (a catalyst.util equality helper) while the cost sits in GenerateOrdering.create upstream of CodeGenerator.compile's cache. That works for the planning-time sites this PR targets, but SortMergeJoinEvaluatorFactory and any future caller keep paying the full construction cost, and the reason for excluding SMJ (evicting the wrapper's entries) only exists because the cache is the wrapper's. A RowOrdering-level memo keyed on (dataTypes, factoryMode, classLoader) would serve everyone, but that is a bigger change than this PR should carry. Just noting it for a possible follow-up.
• NonFateSharingLoadingCache.get takes the per-key KeyLock even on a hit (allocates a lock object, putIfAbsent, remove, notifyAll). A getIfPresent-first fast path would keep the non-fate-sharing guarantee and would help the factory's two lookups too. Again a core follow-up, not this PR.

[dongjoon-hyun]

One pre-existing property of this cache that the PR extends to the sort/group sites: the key is Seq[DataType], but PythonUserDefinedType.equals/hashCode compare only pyUDT and ignore sqlType (sql/api/.../UserDefinedType.scala). supportsExpressions does not gate UDTs out and OrderUtils.isOrderable accepts them via sqlType, so two Python UDTs with the same class name but different sqlType (e.g. two Connect clients on different versions of the same class) share one entry, and the second gets an ordering built for the first's sqlType. Before this PR the sort sites built a fresh ordering from the exact types passed; now they can get the mismatched one. Very much an edge case and the root cause is UDT equality, not the cache, but worth knowing that groupedKeyRowOrdering now inherits it.

[peter-toth]

Good catch, and I had not traced that. PythonUserDefinedType.equals/hashCode comparing only pyUDT collapses two types the sort sites used to keep apart. The root cause is UDT equality rather than the cache, and it already governs every wrapper-based comparison of such keys, so I would rather not paper over it here. Worth a follow-up on UserDefinedType itself.

[dongjoon-hyun]

The caveat in the description is accurate, and I checked that nothing in the current suites hits it: every SharedSparkSession suite runs under CODEGEN_ONLY with fallback off, NO_CODEGEN appears only inside withSQLConf blocks that never build a KeyedPartitioning, and no other construction-time conf affects comparison results (collation is already part of StringType). Two things it does change for these four sites, which used to re-read the mode per call: the first loader in the JVM decides codegen vs interpreted for the rest of the JVM's life, and in FALLBACK mode a transient codegen failure pins the interpreted instance permanently. If you'd rather keep per-mode fidelity, keying the cache on (SQLConf.get.codegenFactoryMode, dataTypes) in the load path is cheap and keeps the identity the new test asserts within any one mode. Accepting the caveat as written is also defensible.

[peter-toth]

Thank you for checking the suites. I am taking the caveat as written rather than keying on the mode. The two orderings agree semantically, so what is at stake is coverage of the interpreted path, not a result.

Your FALLBACK point is the sharper half: a transient codegen failure would pin the interpreted instance for the JVM's life. If that ever bites, (codegenFactoryMode, dataTypes) is the change and it stays cheap.

[dongjoon-hyun]

The 27-calls-per-scan count that motivates this PR comes from outputPartitioning being a def. After this change each call still maps every input partition to its key, sorts them, and then KeyedPartitioning.apply wraps each key and runs a Murmur3-hashed .distinct, so for scans with hundreds or thousands of partitions the per-call work is still well above the ~100us removed here. Making it @transient override lazy val outputPartitioning looks safe: all four concrete scans define inputPartitions as a lazy val, keyGroupedPartitioning is a constructor parameter, doCanonicalize goes through copy, and BatchScanExec.filteredPartitions consumes the pre-filter partitioning, which is exactly the cached value. FileSourceScanExec already does this (lazy val (outputPartitioning, outputOrdering)). Fine as a follow-up if you'd rather keep this PR minimal, but it would subsume most of the benefit claimed here.

[peter-toth]

Agreed, and we arrived at the same place from the other direction while measuring this PR, so it is already filed as SPARK-59252 and in progress.

Measured by instrumenting the body and running KeyGroupedPartitioningSuite: 33,051 executions as a def against 1,262 as a lazy val. The same run shows 19,105 of those reads arriving through PartitioningPreservingUnaryExecNode.outputPartitioning, itself a def doing more work per call, so the ticket covers the three nodes together rather than the scan alone.

[dongjoon-hyun]

Minor: KeyedPartitioning.apply recomputes expressions.map(_.dataType) and its factory does a second orderingCache.get for the same key, so per call this path builds the type list twice and takes the key lock three times. Since apply does not need pre-sorted input (isGrouped is order-independent), building the partitioning first and sorting with its own keyOrdering, the way PushDownUtils already does with keyRowOrdering, would drop the duplicate. Micro-level next to the sort itself, so only if you are touching this anyway.

[peter-toth]

Right. The clean form needs KeyedPartitioning.apply to take the types instead of recomputing them, which is what SPARK-59187 is about, so I would rather not add an overload here and remove it there. Leaving the line as it is.