From cc432ae762eaa5838d5c9369a16105ba422e9ae9 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Fri, 4 Sep 2026 14:37:12 +0200 Subject: [PATCH] [SPARK-59249][SQL] Take the grouped key-row ordering from the shared InternalRowComparableWrapper cache ### What changes were proposed in this pull request? `InternalRowComparableWrapper.Factory` exposes the `ordering` it already holds, and `KeyedPartitioning.groupedKeyRowOrdering` reads it instead of building one. `groupedKeyRowOrdering` called `RowOrdering.createNaturalAscendingOrdering` directly, which is byte-for-byte the `loadFunc` of the wrapper's `orderingCache`. Its four callers therefore all get the cached instance: `KeyedPartitioning.keyRowOrdering`, `EnsureRequirements`' merge ordering, `GroupPartitionsExec.groupAndSortByKeys` and `DataSourceV2ScanExecBase`. Going through the factory rather than a separate `orderingFor` accessor also settles the types. SPARK-59187 made the factory erase the naming of what it is given, so a caller that reached the cache directly would key on a raw type list and get a different entry than the wrappers over the same rows. ### 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 a cached lookup 0.16-0.6 us per call Every caller is once-per-something, so each gets that ~100 us back rather than a count of them: `EnsureRequirements` runs once per join per pass, `GroupPartitionsExec.groupAndSortByKeys` and `KeyedPartitioning.keyRowOrdering` run from `lazy val`s, and `DataSourceV2ScanExecBase.outputPartitioning` became one in SPARK-59252. 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. `DataSourceV2ScanExecBase` keeps its own direct call in this PR. SPARK-59285 ([#58552](https://github.com/apache/spark/pull/58552)) rewrites that site to take the ordering off the wrapper factory it already builds, which is one cache lookup rather than two, so changing it here would only be undone there. `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) --- .../sql/catalyst/plans/physical/partitioning.scala | 6 +++++- .../util/InternalRowComparableWrapper.scala | 9 +++++---- .../util/InternalRowComparableWrapperSuite.scala | 14 ++++++++++++++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index 3168f552f4c33..1d41624112a90 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -936,9 +936,13 @@ object KeyedPartitioning { * `PartitioningCollection`, whose invariant requires equal partition keys -- but join types that * expose only one side's partitioning (e.g. LEFT OUTER) run nothing that compares the two * orders, and silently return wrong results. + * + * It is the keys' own ordering, the one `InternalRowComparableWrapper.equals` compares with, so + * one definition answers both. `EnsureRequirements`' `OrderedDistribution` arm is the one place + * that lays grouped keys out in another order, the distribution's own. */ def groupedKeyRowOrdering(dataTypes: Seq[DataType]): BaseOrdering = - RowOrdering.createNaturalAscendingOrdering(dataTypes) + InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(dataTypes).ordering /** * Projects a sequence of partition keys by selecting only the specified positions. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala index 5eaac05ade156..d6b841f1dafcb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala @@ -152,15 +152,16 @@ object InternalRowComparableWrapper { * Builds wrappers over one row schema, holding the cache lookups that schema needs so a caller * does not repeat them per row. * - * `dataTypes` is what the rows it builds compare at, which is `comparableTypes` of what it was - * given. A caller reporting a type list beside those rows takes it from here rather than erasing - * on its own, so the two cannot answer differently. + * It also answers for the schema it settled on. `dataTypes` is what the rows it builds compare + * at, which is `comparableTypes` of what it was given, and `ordering` sorts them. A caller + * reporting a type list beside those rows, or sorting them, takes it from here rather than + * deriving it again, so the two cannot answer differently. */ final class Factory private[InternalRowComparableWrapper] (val dataTypes: Seq[DataType]) extends (InternalRow => InternalRowComparableWrapper) { + val ordering: BaseOrdering = orderingCache.get(dataTypes) private[this] val structType = structTypeCache.get(dataTypes) - private[this] val ordering = orderingCache.get(dataTypes) override def apply(row: InternalRow): InternalRowComparableWrapper = new InternalRowComparableWrapper(row, dataTypes, structType, ordering) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala index 2723eeecfeb3b..26692ca6530be 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala @@ -19,6 +19,7 @@ package org.apache.spark.sql.catalyst.util import org.apache.spark.SparkFunSuite import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning import org.apache.spark.sql.types._ class InternalRowComparableWrapperSuite extends SparkFunSuite { @@ -81,6 +82,19 @@ class InternalRowComparableWrapperSuite extends SparkFunSuite { assert(InternalRowComparableWrapper.comparableTypes(alreadyErased) eq alreadyErased) } + test("SPARK-59249: the grouped key layout and wrapper equality hold one ordering instance") { + // Identity is the property to assert, because behaviour is not what changes here: both sides + // were already built by the same function, so they already compared the same way. What one + // instance buys is that neither side can later be given a definition the other does not have. + // `InternalRowComparableWrapper.equals` compares its rows with the instance below, and + // `KeyedPartitioning` sorts and groups partition keys with it. The two type lists are built + // separately, so this pins the shared cache as well. + val wrapper = InternalRowComparableWrapper + .getInternalRowComparableWrapperFactory(Seq(IntegerType, LongType))(InternalRow(1, 2L)) + + assert(KeyedPartitioning.groupedKeyRowOrdering(Seq(IntegerType, LongType)) eq wrapper.ordering) + } + test("SPARK-59187: a factory answers for the types it settled on") { // A caller that reports a type list beside the rows a factory built takes it from here, so the // two cannot answer differently.