diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala index 50179824e255f..0ed78b667db3e 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala @@ -21,6 +21,7 @@ import java.util import org.apache.spark.sql.connector.catalog.{Column, SupportsRead, SupportsRowLevelOperations, SupportsWrite, Table, TableCapability} import org.apache.spark.sql.connector.catalog.constraints.Constraint +import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.read.ScanBuilder import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -40,6 +41,7 @@ private[sql] case class RowLevelOperationTable( override def columns: Array[Column] = table.columns() override def capabilities: util.Set[TableCapability] = table.capabilities override def constraints(): Array[Constraint] = table.constraints() + override def partitioning(): Array[Transform] = table.partitioning() override def toString: String = table.toString override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryPartitionPredicateDeleteTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryPartitionPredicateDeleteTable.scala index 497f715776e56..11b5b966aa04c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryPartitionPredicateDeleteTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryPartitionPredicateDeleteTable.scala @@ -19,17 +19,23 @@ package org.apache.spark.sql.connector.catalog import java.util +import scala.collection.mutable.ArrayBuffer + import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException +import org.apache.spark.sql.catalyst.expressions.MetadataStructFieldWithLogicalName import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.MultipartIdentifierHelper import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.connector.expressions.filter.{PartitionPredicate, Predicate} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ /** * In-memory table that supports row-level operations and accepts [[PartitionPredicate]]s - * in V2 [[canDeleteWhere]]/[[deleteWhere]] for metadata-only deletes. + * in V2 [[canDeleteWhere]]/[[deleteWhere]] for metadata-only deletes, and in the scan of a + * group-based UPDATE, MERGE or DELETE, which pushes V2 predicates iteratively. * * Contains some knobs to control acceptance of various partition and data predicates. */ @@ -107,6 +113,72 @@ class InMemoryPartitionPredicateDeleteTable( } } + /** + * Row-level scans push V2 predicates iteratively, so a group-based operation receives a + * second-pass [[PartitionPredicate]] the same way a metadata-only DELETE does. Only partition + * predicates prune, by partition key; a data predicate is always returned since the scan + * cannot filter rows. + */ + override protected def newRowLevelScanBuilder( + options: CaseInsensitiveStringMap)( + onBuild: BatchScanBaseClass => Unit): ScanBuilder = { + new PartitionPredicateRowLevelScanBuilder(onBuild) + } + + class PartitionPredicateRowLevelScanBuilder(onBuild: BatchScanBaseClass => Unit) + extends ScanBuilder with SupportsPushDownV2Filters with SupportsPushDownRequiredColumns { + + private var readSchema: StructType = schema + private val pushed = ArrayBuffer.empty[Predicate] + + override def supportsIterativePushdown(): Boolean = true + + override def pushPredicates(predicates: Array[Predicate]): Array[Predicate] = { + val (accepted, returned) = predicates.partition { + case _: PartitionPredicate => acceptPartitionPredicates + case p => refsOnlyPartCols(p) && InMemoryTableWithV2Filter.supportsPredicates(Array(p)) + } + pushed ++= accepted + returned + } + + override def pushedPredicates(): Array[Predicate] = pushed.toArray + + override def pruneColumns(requiredSchema: StructType): Unit = { + val metadataNames = metadataColumns.map(_.name).toSet + val schemaNames = schema.map(_.name).toSet + readSchema = StructType(requiredSchema.filter { + case MetadataStructFieldWithLogicalName(_, name) => metadataNames.contains(name) + case f => schemaNames.contains(f.name) + }) + } + + override def build(): Scan = { + val (partPreds, stdPreds) = pushed.toArray.partition(_.isInstanceOf[PartitionPredicate]) + val partitionPredicates = partPreds.map(_.asInstanceOf[PartitionPredicate]) + val keys = InMemoryTableWithV2Filter.filtersToKeys( + data.map(_.key).toImmutableArraySeq, + partCols.map(_.toSeq.quoted).toImmutableArraySeq, + stdPreds).toSet + val partitions = data.filter { p => + keys.contains(p.key) && partitionPredicates.forall(_.eval(p.partitionKey())) + } + val scan = PartitionPredicateRowLevelBatchScan( + partitions.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, + readSchema, schema, partitionPredicates.toImmutableArraySeq) + onBuild(scan) + scan + } + } + + /** Row-level batch scan that records the [[PartitionPredicate]]s it was pruned by. */ + case class PartitionPredicateRowLevelBatchScan( + _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + pushedPartitionPredicates: Seq[PartitionPredicate]) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + private def rowMatchesAll( row: InternalRow, preds: Array[Predicate], diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala index 7556bc96912f6..c08dfef04a5b4 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala @@ -272,7 +272,7 @@ class InMemoryRowLevelOperationTable private ( * scan mixes in [[CatalystRuntimeFilteringScan]] so group filtering goes through the Catalyst * path. */ - private def newRowLevelScanBuilder( + protected def newRowLevelScanBuilder( options: CaseInsensitiveStringMap)( onBuild: BatchScanBaseClass => Unit): ScanBuilder = { new InMemoryScanBuilder(schema, options) { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala index ee2a95febeb8e..975ab11f31bc4 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala @@ -19,20 +19,24 @@ package org.apache.spark.sql.connector import org.apache.spark.SparkConf import org.apache.spark.sql.Row -import org.apache.spark.sql.connector.catalog.InMemoryPartitionPredicateDeleteCatalog +import org.apache.spark.sql.connector.catalog.{InMemoryPartitionPredicateDeleteCatalog, InMemoryPartitionPredicateDeleteTable} import org.apache.spark.sql.connector.expressions.PartitionFieldReference import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate +import org.apache.spark.sql.connector.write.RowLevelOperationTable import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} -import org.apache.spark.sql.execution.datasources.v2.{DeleteFromTableExec, ReplaceDataExec, WriteDeltaExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DeleteFromTableExec, ReplaceDataExec, WriteDeltaExec} import org.apache.spark.sql.functions.udf import org.apache.spark.sql.test.SharedSparkSession import org.apache.spark.sql.util.QueryExecutionListener /** * Tests for metadata-only delete optimization using second-pass - * PartitionPredicate (see SPARK-55596). + * PartitionPredicate (see SPARK-55596), and for the second pass in the scan of a + * group-based UPDATE or MERGE (GroupBasedRowLevelOperationScanPlanning). */ -class DataSourceV2EnhancedDeleteFilterSuite extends SharedSparkSession { +class DataSourceV2EnhancedDeleteFilterSuite extends SharedSparkSession + with AdaptiveSparkPlanHelper { private val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName private val catalogName = "ppd_cat" @@ -241,6 +245,31 @@ class DataSourceV2EnhancedDeleteFilterSuite extends SharedSparkSession { } } + // A group-based UPDATE reads the table through RowLevelOperationTable. The wrapper reports + // the table's partitioning, so the row-level scan, which pushes V2 predicates iteratively, + // receives the IN on the partition column as a second-pass PartitionPredicate, and only the + // matching partitions are read and replaced. + test("SPARK-59457: group-based UPDATE receives a second-pass PartitionPredicate") { + withTable(deleteTableName) { + sql(s"CREATE TABLE $deleteTableName (pk INT, dep STRING, salary INT) " + + s"USING $v2Source PARTITIONED BY (dep)") + sql(s"INSERT INTO $deleteTableName VALUES " + + "(1, 'hr', 100), (2, 'software', 200), (3, 'marketing', 300)") + + val plan = executeAndKeepPlan { + sql(s"UPDATE $deleteTableName SET salary = salary + 1 WHERE dep IN ('hr', 'software')") + } + assertRowLevelScanPrunedByPartitionPredicate(plan, + expectedOrdinals = Array(0), + expectedPartitionFieldNames = Array("dep"), + expectedReplacedDeps = Set("hr", "software")) + + checkAnswer( + sql(s"SELECT * FROM $deleteTableName"), + Seq(Row(1, "hr", 101), Row(2, "software", 201), Row(3, "marketing", 300))) + } + } + private def executeAndKeepPlan(func: => Unit): SparkPlan = { var executedPlan: SparkPlan = null @@ -357,6 +386,35 @@ class DataSourceV2EnhancedDeleteFilterSuite extends SharedSparkSession { } } + /** + * Asserts that the group-based plan's row-level scan was pruned by one PartitionPredicate with + * the given references, and that only the partitions with the given `dep` values, the first + * partition field, were replaced. + */ + private def assertRowLevelScanPrunedByPartitionPredicate( + plan: SparkPlan, + expectedOrdinals: Array[Int], + expectedPartitionFieldNames: Array[String], + expectedReplacedDeps: Set[String]): Unit = { + assert(plan.isInstanceOf[ReplaceDataExec], + s"Expected ReplaceDataExec but got: ${plan.getClass.getSimpleName}") + val scans = collect(plan) { case s: BatchScanExec => s } + val scan = scans.map(_.scan).collectFirst { + case s: InMemoryPartitionPredicateDeleteTable#PartitionPredicateRowLevelBatchScan => s + }.getOrElse(fail("Expected the row-level scan of the in-memory table")) + assertPartitionFieldReferences( + scan.pushedPartitionPredicates.toArray, Seq(expectedOrdinals), expectedPartitionFieldNames) + + val table = scans.map(_.table).collectFirst { + case RowLevelOperationTable(t: InMemoryPartitionPredicateDeleteTable, _) => t + }.getOrElse(fail("Expected the row-level operation table")) + val replacedDeps = table.replacedPartitions.map(_.head.toString) + assert( + replacedDeps.toSet === expectedReplacedDeps && + replacedDeps.size === expectedReplacedDeps.size, + s"Expected replaced partitions for $expectedReplacedDeps, got ${table.replacedPartitions}") + } + private def assertDeleteWithRowLevel(query: String): Unit = { val plan = executeAndKeepPlan { sql(query) } plan match {