Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
peter-toth wants to merge 1 commit into
base: master
Choose a base branch
from
+49 −4
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peter-toth Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


/**
* Projects a sequence of partition keys by selecting only the specified positions.
Expand Down
3 changes: 3 additions & 0 deletions ...lyst/src/main/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapper.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peter-toth Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


/** Creates a shared factory method for a given row schema to avoid excessive cache lookups. */
def getInternalRowComparableWrapperFactory(
dataTypes: Seq[DataType]): InternalRow => InternalRowComparableWrapper = {
Expand Down
39 changes: 39 additions & 0 deletions ...src/test/scala/org/apache/spark/sql/catalyst/util/InternalRowComparableWrapperSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
}
}
5 changes: 2 additions & 3 deletions ...c/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -92,8 +92,7 @@ trait DataSourceV2ScanExecBase
keyGroupedPartitioning match {

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peter-toth Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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)

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 peter-toth Sep 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Expand Down