From 20182e178de4eb3a94220cb4b6cf68f758f69c5e Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Thu, 17 Sep 2026 18:50:31 +0000 Subject: [PATCH 1/5] [SPARK-59642][SQL] Validate the arity of connector-reported partition keys Reject a DSv2 partition key row whose field count differs from the number of reported partition expressions, at the KeyedPartitioning construction boundary and on the runtime-filter re-ingestion path. Co-authored-by: Isaac --- .../plans/physical/partitioning.scala | 16 ++++++++++++++ .../spark/sql/catalyst/ShuffleSpecSuite.scala | 22 ++++++++++++++++++- .../v2/DataSourceV2ScanExecBase.scala | 6 +++-- .../datasources/v2/PushDownUtils.scala | 5 +++++ ...taSourceV2CatalystRuntimeFilterSuite.scala | 14 +++++++++++- 5 files changed, 59 insertions(+), 4 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 9b178e5047b26..f100f8183b75d 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 @@ -962,6 +962,7 @@ object KeyedPartitioning { def apply( expressions: Seq[Expression], partitionKeys: Seq[InternalRow]): KeyedPartitioning = { + checkPartitionKeyArity(expressions, partitionKeys) val factory = InternalRowComparableWrapper .getInternalRowComparableWrapperFactory(expressions.map(_.dataType)) val comparablePartitionKeys = partitionKeys.map(factory) @@ -972,6 +973,21 @@ object KeyedPartitioning { KeyLayout(comparablePartitionKeys, factory.dataTypes, isGrouped, isCollapsed = false)) } + // The key's arity is implicit in `HasPartitionKey` and read positionally downstream, so an + // inconsistent key would otherwise fail far away as an opaque `ArrayIndexOutOfBoundsException`. + def checkPartitionKeyArity( + expressions: Seq[Expression], + partitionKeys: Seq[InternalRow]): Unit = { + partitionKeys.foreach { key => + if (key.numFields != expressions.length) { + throw new SparkException("Data source reported a partition key with " + + s"${key.numFields} field(s) from HasPartitionKey.partitionKey(), but reported " + + s"${expressions.length} partition expression(s). Every reported partition key " + + "must have one field per reported partition expression.") + } + } + } + /** * Concatenates partitionings that agree on their expressions, which is what a `UnionExec` does to * its children's partitions. The result reports one key per output partition, so its keys are the diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala index 1ee8002eb4b1f..fe02951605ca1 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala @@ -17,7 +17,7 @@ package org.apache.spark.sql.catalyst -import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException} +import org.apache.spark.{SparkException, SparkFunSuite, SparkUnsupportedOperationException} import org.apache.spark.sql.catalyst.dsl.expressions._ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression} import org.apache.spark.sql.catalyst.plans.SQLHelper @@ -980,6 +980,26 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper { assert(e.getMessage.contains("expected specs to be non-empty")) } + test("SPARK-59642: a reported partition key of a different arity is rejected at construction") { + val expressions = Seq($"a".int, $"b".int) + val wellFormed = InternalRow(1, 1) + + val tooWide = intercept[SparkException] { + KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2, 2, 99))) + } + assert(tooWide.getMessage.contains("partition key with 3 field(s)")) + assert(tooWide.getMessage.contains("2 partition expression(s)")) + + val tooNarrow = intercept[SparkException] { + KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2))) + } + assert(tooNarrow.getMessage.contains("partition key with 1 field(s)")) + assert(tooNarrow.getMessage.contains("2 partition expression(s)")) + + val partitioning = KeyedPartitioning(expressions, Seq(wellFormed, InternalRow(2, 2))) + assert(partitioning.numPartitions === 2) + } + test("SPARK-59256: flattening reaches the members of a nested collection") { val distribution = ClusteredDistribution(Seq($"a", $"b")) val buried = HashShuffleSpec(HashPartitioning(Seq($"a"), 10), distribution) 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 a42a8751a92dd..7f4a58d04c390 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 @@ -109,8 +109,10 @@ trait DataSourceV2ScanExecBase // `PartitioningCollection.fromPartitionings` refuses them. See `KeyedPartitioning.apply` // for why this is the ordering to sort with. val keys = inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()) - .sorted(KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType))) - Some(KeyedPartitioning(exprs, keys)) + // Ahead of the sort, which reads every key at the declared positions. + KeyedPartitioning.checkPartitionKeyArity(exprs, keys) + val ordering = KeyedPartitioning.groupedKeyRowOrdering(exprs.map(_.dataType)) + Some(KeyedPartitioning(exprs, keys.sorted(ordering))) case _ => None } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala index f1cf9452f4f21..3b8a7cbffcf17 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala @@ -306,6 +306,11 @@ object PushDownUtils extends Logging { "filtering") } + // Re-reported rows from `filter()` that never pass through `KeyedPartitioning.apply`. + KeyedPartitioning.checkPartitionKeyArity( + k.expressions, + newPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).toSeq) + val inputMap = k.partitionKeys.groupBy(identity).view.mapValues(_.size) val comparableKeyWrapperFactory = InternalRowComparableWrapper .getInternalRowComparableWrapperFactory(k.keyDataTypes) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala index 520d05e5c48ba..cd0dde1fdf02d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -655,7 +655,8 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { /** * While SPJ is active the scan's partitioning has to survive runtime filtering, so the * post-filter partitions still line up with the other side of the join: splits may be pruned, - * but the source may not drop a partition key, invent one, or grow a key's split count. + * but the source may not drop a partition key, invent one, grow a key's split count, or report a + * key of a different shape than the partitioning was built from. */ test("data source that breaks the partitioning it reported -> rejected") { val partAttr = AttributeReference("part", IntegerType)() @@ -680,6 +681,12 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { replanAfterFiltering(Seq(KeyedInputPartition(1), KeyedInputPartition(1))) } assert(splitsGrown.getMessage.contains("must not report new partitions for a given key")) + + val wrongArity = intercept[SparkException] { + replanAfterFiltering(Seq(WideKeyedInputPartition(1, 99))) + } + assert(wrongArity.getMessage.contains("partition key with 2 field(s)")) + assert(wrongArity.getMessage.contains("1 partition expression(s)")) } // --------------------------------------------------------------------------- @@ -853,6 +860,11 @@ private case class KeyedInputPartition(key: Int) extends InputPartition with Has override def partitionKey(): InternalRow = InternalRow(key) } +private case class WideKeyedInputPartition(key: Int, extra: Int) + extends InputPartition with HasPartitionKey { + override def partitionKey(): InternalRow = InternalRow(key, extra) +} + /** * A scan reporting one set of partitions before filtering and another after, so it can break the * requirement to preserve the partitioning it originally reported. From 51948434ae1f6a3f8f9eb1a254b5c5fe00cccde9 Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Fri, 18 Sep 2026 18:40:49 +0000 Subject: [PATCH 2/5] [SPARK-59642][SQL] Name both arity failure modes in the guard comment The rationale comment above `checkPartitionKeyArity` named only the opaque `ArrayIndexOutOfBoundsException`, which is what a too-narrow key produces. A too-wide key never crashed: the grouped-key ordering and the comparable-wrapper grouping are built over `expressions.length`, so its trailing fields were silently dropped and keys grouped too loosely. Name both directions so the symmetric exact-arity check reads as contract enforcement, not crash cosmetics. Co-authored-by: Isaac --- .../spark/sql/catalyst/plans/physical/partitioning.scala | 5 +++-- 1 file changed, 3 insertions(+), 2 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 f100f8183b75d..c616140266603 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 @@ -973,8 +973,9 @@ object KeyedPartitioning { KeyLayout(comparablePartitionKeys, factory.dataTypes, isGrouped, isCollapsed = false)) } - // The key's arity is implicit in `HasPartitionKey` and read positionally downstream, so an - // inconsistent key would otherwise fail far away as an opaque `ArrayIndexOutOfBoundsException`. + // The key's arity is implicit in `HasPartitionKey` and read positionally downstream: a short key + // runs off the end as an opaque `ArrayIndexOutOfBoundsException`, while a long key's trailing + // fields are silently dropped by readers built over `expressions.length`, grouping too loosely. def checkPartitionKeyArity( expressions: Seq[Expression], partitionKeys: Seq[InternalRow]): Unit = { From b28584046c66379b9891ad289b5197894b084919 Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Fri, 18 Sep 2026 20:14:42 +0000 Subject: [PATCH 3/5] [SPARK-59642][SQL] Test the arity guard ahead of the reported-key sort `DataSourceV2ScanExecBase` is the one caller whose `checkPartitionKeyArity` call is load-bearing rather than redundant with `KeyedPartitioning.apply`: it runs before `keys.sorted(groupedKeyRowOrdering(...))` reads every key at the declared positions, so a too-narrow key becomes a `SparkException` instead of an `ArrayIndexOutOfBoundsException` from inside the ordering. The existing tests drive `KeyedPartitioning.apply` and `PushDownUtils.replanWithRuntimeFilters`, so neither fails if that call is removed or moved after the sort. The new test reports two keys that tie on the leading field, which is what forces the ordering to read the short key's missing field. Removing the pre-sort call makes it fail with `ArrayIndexOutOfBoundsException`. Co-authored-by: Isaac --- .../v2/DataSourceV2ScanExecBaseSuite.scala | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBaseSuite.scala diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBaseSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBaseSuite.scala new file mode 100644 index 0000000000000..31b69b034b4c4 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBaseSuite.scala @@ -0,0 +1,66 @@ +/* + * 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.execution.datasources.v2 + +import org.apache.spark.SparkException +import org.apache.spark.sql.QueryTest +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.connector.read.{Batch, HasPartitionKey, InputPartition, PartitionReaderFactory, Scan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{IntegerType, StructType} + +class DataSourceV2ScanExecBaseSuite extends QueryTest with SharedSparkSession { + + private val exprA = AttributeReference("a", IntegerType)() + private val exprB = AttributeReference("b", IntegerType)() + + test("SPARK-59642: a reported key of the wrong arity is rejected before the keys are sorted") { + // The keys tie on `a`, so the sort that follows has to read `b` off the one-field key: were the + // check to run after it, this would be an ArrayIndexOutOfBoundsException from the ordering. + val exec = BatchScanExec( + output = Seq(exprA, exprB), + scan = new KeyedPartitionsScan(Seq(InternalRow(1, 5), InternalRow(1))), + runtimeFilters = Seq.empty, + table = null, + keyGroupedPartitioning = Some(Seq(exprA, exprB))) + + withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "true") { + val e = intercept[SparkException](exec.outputPartitioning) + assert(e.getMessage.contains("partition key with 1 field(s)")) + assert(e.getMessage.contains("2 partition expression(s)")) + } + } +} + +private case class KeyedPartition(key: InternalRow) extends InputPartition with HasPartitionKey { + override def partitionKey(): InternalRow = key +} + +private class KeyedPartitionsScan(keys: Seq[InternalRow]) extends Scan with Batch { + override def readSchema(): StructType = + new StructType().add("a", IntegerType).add("b", IntegerType) + + override def toBatch: Batch = this + + override def planInputPartitions(): Array[InputPartition] = keys.map(KeyedPartition(_)).toArray + + override def createReaderFactory(): PartitionReaderFactory = + throw new UnsupportedOperationException() +} From 78530c7bf6165af7c132207311bd9468b3ad439f Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Mon, 21 Sep 2026 09:58:23 +0000 Subject: [PATCH 4/5] [SPARK-59642][SQL] Document the enforced partition-key arity The check added here states the contract only in its own exception text, inside catalyst, which `project/SparkBuild.scala` strips from the published unidoc. State the obligation on the connector-facing interface an implementor reads before writing the method, and record the rejection in the upgrade notes an operator reads before upgrading a connector that already ships. Two shapes stop working rather than merely failing differently, so the notes name both and give `spark.sql.sources.v2.bucketing.enabled=false` as the only lever back, together with its cost. Co-authored-by: Isaac --- docs/sql-migration-guide.md | 1 + .../org/apache/spark/sql/connector/read/HasPartitionKey.java | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/sql-migration-guide.md b/docs/sql-migration-guide.md index dfe63fc83e4fa..38f311dbfe2b9 100644 --- a/docs/sql-migration-guide.md +++ b/docs/sql-migration-guide.md @@ -33,6 +33,7 @@ license: | - Since Spark 4.4, the CSV `extension` option must be non-empty and contain only letters. In Spark 4.3 and earlier, due to a bug in the extension validation check, empty strings and three-character suffixes containing non-letters (for example `ab1`) were accepted. This affects both writes and reads of CSV tables whose persisted `OPTIONS` contain such values; update the option to a non-empty letters-only suffix to restore access. - Since Spark 4.4, when `array_repeat` or `array_insert` is asked to build an array larger than the maximum supported array length, generated code raises the same error as interpreted evaluation. `array_repeat` now fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER` instead of the internal error `_LEGACY_ERROR_TEMP_2176`, and `array_insert` fails with `COLLECTION_SIZE_LIMIT_EXCEEDED.FUNCTION` instead of `COLLECTION_SIZE_LIMIT_EXCEEDED.PARAMETER`, which named a `count` parameter that `array_insert` does not have. Both functions raise an error under exactly the same conditions as before; only the reported error condition changes. - Since Spark 4.4, `/*+ ... */` inside a bracketed comment is parsed as a nested comment, so its closing `*/` no longer closes the outer comment. For example, `/* note /*+ x */ SELECT 1` previously returned `1`, but now raises `UNCLOSED_BRACKETED_COMMENT`. Close the outer comment explicitly, for example `/* note /*+ x */ */ SELECT 1`. +- Since Spark 4.4, a Data Source V2 scan that reports a partitioning through `SupportsReportPartitioning` must return a partition key row from `HasPartitionKey.partitionKey()` holding one field per reported partition expression. A key of any other width is rejected with an error naming both counts. Most mismatched widths already failed the query, as an `ArrayIndexOutOfBoundsException` or a bare `AssertionError`, and for those only the diagnostic changes. Two shapes are newly rejected and previously let a query complete: a scan reporting exactly one partition whose key is wider than its partition expressions, and a key narrower than them whose leading fields never tie with another key's, which could also group partitions incorrectly. To restore the previous behavior, set `spark.sql.sources.v2.bucketing.enabled` to `false`, which also disables storage-partitioned joins. ## Upgrading from Spark SQL 4.2 to 4.3 diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/HasPartitionKey.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/HasPartitionKey.java index a4421aad80fff..ae504ee2c5833 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/HasPartitionKey.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/HasPartitionKey.java @@ -36,7 +36,9 @@ *

* It is implementor's responsibility to ensure that when an input partition implements this * interface, its records all have the same value for the partition keys. Spark doesn't check - * this property. + * that the records agree. It does check the key's width: the row returned by + * {@link #partitionKey()} must hold one field per partition expression the scan reports through + * {@link SupportsReportPartitioning}, and Spark rejects a key of any other width. * * @see org.apache.spark.sql.connector.read.SupportsReportPartitioning * @see org.apache.spark.sql.connector.read.partitioning.Partitioning From 93ed6117189ce1e921637a9671211c17f41fd71f Mon Sep 17 00:00:00 2001 From: Tim Meehan Date: Mon, 21 Sep 2026 18:05:43 +0000 Subject: [PATCH 5/5] [SPARK-59642][SQL] Correct the partition-key arity rationale comment The rationale above `KeyedPartitioning.checkPartitionKeyArity` gave the `ArrayIndexOutOfBoundsException` to a key narrower than the reported partition expressions and the silent, too-loose grouping to a wider one. It is the other way round: the interpreted struct hash bounds its loop on the key row's own field count while indexing a type array built from the declared expressions, so a wider key indexes past that array and throws, and a narrower one hashes its short prefix without error. Both silent outcomes are conditional as well -- the wide key completes only when the scan reports a single partition, and the narrow one only until two keys tie on their leading fields -- where the comment stated them unconditionally. Comment only. The guard, its exception, its message, and its three call sites are unchanged, as are the migration-guide entry and the `HasPartitionKey` javadoc, which already state the contract in the correct direction. Co-authored-by: Isaac --- .../spark/sql/catalyst/plans/physical/partitioning.scala | 7 ++++--- 1 file changed, 4 insertions(+), 3 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 c616140266603..b033f476243ae 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 @@ -973,9 +973,10 @@ object KeyedPartitioning { KeyLayout(comparablePartitionKeys, factory.dataTypes, isGrouped, isCollapsed = false)) } - // The key's arity is implicit in `HasPartitionKey` and read positionally downstream: a short key - // runs off the end as an opaque `ArrayIndexOutOfBoundsException`, while a long key's trailing - // fields are silently dropped by readers built over `expressions.length`, grouping too loosely. + // The key's arity is implicit in `HasPartitionKey` and read positionally downstream: a key wider + // than `expressions` raises an `ArrayIndexOutOfBoundsException` in the struct hash unless the + // scan reports one partition, and a narrower one silently groups too loosely until two keys tie + // on their leading fields and the ordering reads past its end. def checkPartitionKeyArity( expressions: Seq[Expression], partitionKeys: Seq[InternalRow]): Unit = {