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..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 @@ -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,22 @@ 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. + 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. 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() +}