[SPARK-59457][SQL] RowLevelOperationTable should delegate partitioning() to the wrapped table

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 = {

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],

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) {

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 {