From c484bddf655104862c16462285db00c70af2a9cd 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` 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) --- .../plans/physical/partitioning.scala | 6 ++- .../util/InternalRowComparableWrapper.scala | 3 ++ .../InternalRowComparableWrapperSuite.scala | 39 +++++++++++++++++++ .../v2/DataSourceV2ScanExecBase.scala | 5 +-- 4 files changed, 49 insertions(+), 4 deletions(-) create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala 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 2354cf69205a0..2f07bc213e7b5 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 @@ -880,9 +880,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.orderingFor(dataTypes) /** * 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 3566319c39ad9..d33e32453f140 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 @@ -115,6 +115,9 @@ object InternalRowComparableWrapper { new InternalRowComparableWrapper(partitionRow, partitionExpression.map(_.dataType)) } + /** The cached ordering a wrapper of these `dataTypes` compares its rows with in `equals`. */ + def orderingFor(dataTypes: Seq[DataType]): BaseOrdering = orderingCache.get(dataTypes) + /** Creates a shared factory method for a given row schema to avoid excessive cache lookups. */ def getInternalRowComparableWrapperFactory( dataTypes: Seq[DataType]): InternalRow => InternalRowComparableWrapper = { 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 new file mode 100644 index 0000000000000..b720174239c09 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +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.{IntegerType, LongType} + +class InternalRowComparableWrapperSuite extends SparkFunSuite { + + 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) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala index f00d8b9b82cb4..56b9219e3039c 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Expression, RowOrdering, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Ascending, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.physical import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning import org.apache.spark.sql.catalyst.util.truncatedString @@ -92,8 +92,7 @@ trait DataSourceV2ScanExecBase keyGroupedPartitioning match { case Some(exprs) if conf.v2BucketingEnabled && KeyedPartitioning.supportsExpressions(exprs) && inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) => - val dataTypes = exprs.map(_.dataType) - val rowOrdering = RowOrdering.createNaturalAscendingOrdering(dataTypes) + val rowOrdering = KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType)) val partitionKeys = inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering) KeyedPartitioning(exprs, partitionKeys)