[SPARK-59410][SQL] Derive PartitionPredicate from identity fields of a mixed partitioning

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/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateField.scala

@@ -20,14 +20,19 @@ package org.apache.spark.sql.internal.connector
 import org.apache.spark.sql.catalyst.expressions.AttributeReference
 
 /**
- * Metadata for one partition field.
+ * Metadata for one field of `Table.partitioning()`. A partition predicate is built over the
+ * fields in partitioning order, so their ordinals match the partition key a connector passes to
+ * `PartitionPredicate.eval`.
  *
  * @param fieldNames  the multi-part field name from the table's partitioning
- *                    (e.g. `Seq("s", "tz")`).
- * @param attrRef  the [[AttributeReference]] for the partition field.
- *                 Created from the resolved partition field so it carries the
- *                 flattened dotted name (e.g. `"s.tz"`) for nested fields.
+ *                    (e.g. `Seq("s", "tz")`) for an identity transform, or the transform's
+ *                    description (e.g. `Seq("bucket(4, id)")`) otherwise.
+ * @param attrRef  the [[AttributeReference]] a filter can reference, for an identity transform.
+ *                 Created from the resolved partition field so it carries the flattened dotted
+ *                 name (e.g. `"s.tz"`) for nested fields. None for any other transform: for now
+ *                 Spark does not evaluate a filter against its partition value, so no filter
+ *                 references it, but the field keeps its ordinal.
  */
 case class PartitionPredicateField(
     fieldNames: Seq[String],
-    attrRef: AttributeReference)
+    attrRef: Option[AttributeReference])

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala

@@ -19,9 +19,10 @@ package org.apache.spark.sql.internal.connector
 
 import org.apache.spark.internal.{Logging, LogKeys}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{BindReferences, Expression => CatalystExpression, ExprId, Predicate => CatalystPredicate}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Expression => CatalystExpression, ExprId, Predicate => CatalystPredicate}
 import org.apache.spark.sql.connector.expressions.NamedReference
 import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate
+import org.apache.spark.sql.types.NullType
 
 /**
  * An implementation for [[PartitionPredicate]] that wraps a Catalyst Expression representing a
@@ -32,15 +33,24 @@ class PartitionPredicateImpl private (
     private val partitionFields: Seq[PartitionPredicateField])
   extends PartitionPredicate with Logging {
 
+  /** Ordinal of each identity partition field, keyed by the attribute a filter references. */
   @transient private lazy val exprIdToIndex: Map[ExprId, Int] =
-    partitionFields.zipWithIndex.map { case (f, i) => f.attrRef.exprId -> i }.toMap
+    partitionFields.zipWithIndex.collect {
+      case (PartitionPredicateField(_, Some(attr)), i) => attr.exprId -> i
+    }.toMap
 
   /** The wrapped partition filter Catalyst Expression. */
   def expression: CatalystExpression = catalystExpr
 
   /** Bound predicate, computed once and reused for all partition rows. */
   @transient private lazy val boundPredicate: InternalRow => Boolean = {
-    val boundExpr = BindReferences.bindReference(catalystExpr, partitionFields.map(_.attrRef))
+    // One attribute per partition field, so that ordinals match the full partition key. A field
+    // of a non-identity transform has no attribute a filter can reference; a placeholder keeps
+    // its slot.
+    val input = partitionFields.map { f =>
+      f.attrRef.getOrElse(AttributeReference(f.fieldNames.mkString("."), NullType)())
+    }
+    val boundExpr = BindReferences.bindReference(catalystExpr, input)
     val predicate = CatalystPredicate.createInterpreted(boundExpr)
     predicate.eval
   }
@@ -102,16 +112,18 @@ object PartitionPredicateImpl extends Logging {
       return None
     }
 
-    val partitionExprIds = partitionFields.map(_.attrRef.exprId).toSet
+    val partitionExprIds = partitionFields.flatMap(_.attrRef).map(_.exprId).toSet
     val unmatchedRefs = catalystExpr.references.filterNot(r => partitionExprIds.contains(r.exprId))
     if (unmatchedRefs.nonEmpty) {
       logWarning(
         log"Cannot create partition predicate ${MDC(LogKeys.EXPR, catalystExpr.sql)}: " +
         log"expression references " +
         log"${MDC(LogKeys.FIELD_NAME, unmatchedRefs.map(_.name).mkString(", "))} " +
-        log"not found in partition fields " +
+        log"not found in identity partition fields " +
         log"${MDC(LogKeys.PARTITION_SPECIFICATION,
-          partitionFields.map(_.fieldNames.mkString(".")).mkString(", "))}. " +
+          partitionFields.collect {
+            case PartitionPredicateField(names, Some(_)) => names.mkString(".")
+          }.mkString(", "))}. " +
         log"Skipping pushdown for this predicate.")
       return None
     }

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/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala

@@ -203,14 +203,19 @@ object InMemoryTableWithV2Filter {
     }
   }
 
+  /**
+   * Whether every predicate has a shape [[evalPredicate]] can evaluate: a plain column, or a
+   * column and a literal. A predicate over an expression, e.g. a cast, is not supported and
+   * returned to Spark, as a real connector without expression support would do.
+   */
   def supportsPredicates(predicates: Array[Predicate]): Boolean = {
-    predicates.flatMap(splitAnd).forall {
-      case p: Predicate if p.name().equals("=") => true
-      case p: Predicate if p.name().equals("<=>") => true
-      case p: Predicate if p.name().equals("IS_NULL") => true
-      case p: Predicate if p.name().equals("IS_NOT_NULL") => true
-      case p: Predicate if p.name().equals("ALWAYS_TRUE") => true
-      case _ => false
+    predicates.flatMap(splitAnd).forall { p =>
+      (p.name(), p.children().toSeq) match {
+        case ("=" | "<=>", Seq(_: NamedReference, _: LiteralValue[_])) => true
+        case ("IS_NULL" | "IS_NOT_NULL", Seq(_: NamedReference)) => true
+        case ("ALWAYS_TRUE", _) => true
+        case _ => false
+      }
     }
   }
 

sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala

@@ -52,11 +52,38 @@ class PartitionPredicateImplSuite extends SparkFunSuite {
     checkNestedPartitionPathReferencesAfterSerialization(serializer)
   }
 
+  test("non-identity partition field: predicate binds by ordinal and never references it") {
+    val ref = DataTypeUtils.toAttribute(StructField("p", StringType, nullable = true))
+    val fields = Seq(
+      PartitionPredicateField(Seq("bucket(4, id)"), None),
+      PartitionPredicateField(Seq("p"), Some(ref)))
+    val predicate = PartitionPredicateImpl(GreaterThan(ref, Literal("m")), fields).get
+
+    // The partition key carries one value per field; the bucket value at ordinal 0 is skipped.
+    assert(predicate.eval(InternalRow(3, UTF8String.fromString("z"))) === true)
+    assert(predicate.eval(InternalRow(3, UTF8String.fromString("a"))) === false)
+    assert(refsWithOrdinals(predicate.references.toSeq) === Seq(("p", 1)))
+
+    // A filter on the source column of the bucket transform has no field to bind to.
+    val id = DataTypeUtils.toAttribute(StructField("id", IntegerType, nullable = true))
+    assert(PartitionPredicateImpl(GreaterThan(id, Literal(1)), fields).isEmpty)
+
+    Seq(new JavaSerializer(new SparkConf()), new KryoSerializer(new SparkConf())).foreach { s =>
+      val serializer = s.newInstance()
+      val deserialized = serializer.deserialize[PartitionPredicateImpl](
+        serializer.serialize(predicate))
+      assert(deserialized.eval(InternalRow(3, UTF8String.fromString("z"))) === true)
+      assert(deserialized.eval(InternalRow(3, UTF8String.fromString("a"))) === false)
+      assert(refsWithOrdinals(deserialized.references.toSeq) === Seq(("p", 1)))
+      assert(deserialized.equals(predicate))
+    }
+  }
+
   private def checkPartitionPredicateImplAfterSerialization(
       serializer: SerializerInstance): Unit = {
     val ref = DataTypeUtils.toAttribute(StructField("p", IntegerType, nullable = true))
     val expr = GreaterThan(ref, Literal(5))
-    val fields = Seq(PartitionPredicateField(Seq("p"), ref))
+    val fields = Seq(PartitionPredicateField(Seq("p"), Some(ref)))
     val predicate = PartitionPredicateImpl(expr, fields).get
 
     val deserialized = serializer.deserialize[PartitionPredicateImpl](
@@ -77,7 +104,7 @@ class PartitionPredicateImplSuite extends SparkFunSuite {
       serializer: SerializerInstance): Unit = {
     val ref = DataTypeUtils.toAttribute(StructField("ts.timezone", StringType, nullable = false))
     val expr = GreaterThan(ref, Literal("x"))
-    val fields = Seq(PartitionPredicateField(Seq("ts", "timezone"), ref))
+    val fields = Seq(PartitionPredicateField(Seq("ts", "timezone"), Some(ref)))
     val predicate = PartitionPredicateImpl(expr, fields).get
 
     val deserialized = serializer.deserialize[PartitionPredicateImpl](

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala

@@ -367,26 +367,33 @@ object PushDownUtils extends Logging {
   }
 
   /**
-   * Returns a Seq of [[PartitionPredicateField]] representing partition transform expression types,
-   * if schema is supported for [[PartitionPredicate]] push down. None if not supported.
+   * Returns one [[PartitionPredicateField]] per transform of `relation.table.partitioning`, if
+   * the partitioning supports [[PartitionPredicate]] push down. None if not supported.
    */
   def getPartitionPredicateSchema(relation: DataSourceV2Relation)
   : Option[Seq[PartitionPredicateField]] = {
     getPartitionPredicateSchema(relation.table, relation.output)
   }
 
   /**
-   * Returns a Seq of [[PartitionPredicateField]] representing partition transform expression types,
-   * if schema is supported for [[PartitionPredicate]] push down. None if not supported.
+   * Returns one [[PartitionPredicateField]] per transform of `table.partitioning`, if the
+   * partitioning supports [[PartitionPredicate]] push down. None if not supported.
    */
   def getPartitionPredicateSchema(table: Table, output: Seq[AttributeReference])
   : Option[Seq[PartitionPredicateField]] = {
     getPartitionPredicateSchema(table.partitioning, output)
   }
 
   /**
-   * Returns a Seq of [[PartitionPredicateField]] representing partition transform expression types,
-   * if schema is supported for [[PartitionPredicate]] push down. None if not supported.
+   * Returns one [[PartitionPredicateField]] per transform, in partitioning order, if the
+   * partitioning supports [[PartitionPredicate]] push down. None if not supported.
+   *
+   * Only an identity transform yields a field with an attribute, so only filters over identity
+   * partition columns become partition predicates. Any other transform is kept as a field without
+   * an attribute: Spark cannot evaluate a filter against its partition value, but the field must
+   * keep its ordinal since a predicate is evaluated against the full partition key. The
+   * partitioning is not supported when it is empty, has no identity transform, or has an identity
+   * transform that does not resolve against `output`.
    *
    * Use this overload when the caller has access to the partition transforms but not the
    * full [[Table]].
@@ -402,11 +409,11 @@ object PushDownUtils extends Logging {
       val fields = transforms.flatMap {
         case t: IdentityTransform =>
           resolveIdentityPartitionField(t, rootStruct).map { sf =>
-            PartitionPredicateField(t.ref.fieldNames().toSeq, DataTypeUtils.toAttribute(sf))
+            PartitionPredicateField(t.ref.fieldNames().toSeq, Some(DataTypeUtils.toAttribute(sf)))
           }
-        case _ => None
+        case t => Some(PartitionPredicateField(Seq(t.describe()), None))
       }
-      if (fields.length == transforms.length) {
+      if (fields.length == transforms.length && fields.exists(_.attrRef.isDefined)) {
         Some(fields.toSeq)
       } else {
         None
@@ -451,7 +458,7 @@ object PushDownUtils extends Logging {
       flattenedFilters: Seq[Expression],
       partitionFields: Seq[PartitionPredicateField])
   : (Seq[PartitionPredicateImpl], Seq[Expression]) = {
-    val partitionAttributes = partitionFields.map(_.attrRef)
+    val partitionAttributes = partitionFields.flatMap(_.attrRef)
     val (partFilters, nonPartitionFilters) =
       DataSourceUtils.getPartitionFiltersAndDataFilters(partitionAttributes, flattenedFilters)
     val (pushable, nonPushable) = partFilters.partition(isPushablePartitionFilter(_))
@@ -539,7 +546,9 @@ object PushDownUtils extends Logging {
       filters: Seq[Expression],
       partitionFields: Seq[PartitionPredicateField])
   : Map[Expression, Expression] = {
-    val pathToAttr = partitionFields.map(f => f.fieldNames -> f.attrRef).toMap
+    val pathToAttr = partitionFields.collect {
+      case PartitionPredicateField(names, Some(attr)) => names -> attr
+    }.toMap
     filters.map(f => doNormalizePartitionFilters(f, pathToAttr) -> f).toMap
   }