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

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: Spark
+ *                 cannot 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,7 +112,7 @@ 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(

sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala

@@ -203,14 +203,23 @@ 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 =>
+      def column = p.children().length == 1 && p.children()(0).isInstanceOf[NamedReference]
+      def columnAndLiteral = p.children().length == 2 &&
+        p.children()(0).isInstanceOf[NamedReference] &&
+        p.children()(1).isInstanceOf[LiteralValue[_]]
+      p.name() match {
+        case "=" | "<=>" => columnAndLiteral
+        case "IS_NULL" | "IS_NOT_NULL" => column
+        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
   }
 

sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedDeleteFilterSuite.scala

@@ -184,6 +184,50 @@ class DataSourceV2EnhancedDeleteFilterSuite extends SharedSparkSession {
     }
   }
 
+  // Mixed partitioning: the bucket field keeps its ordinal but is never referenced, so the
+  // IN on the identity column still becomes a PartitionPredicate over the full partition key.
+  test("second pass accepted: identity column next to a bucket transform") {
+    withTable(deleteTableName) {
+      sql(s"CREATE TABLE $deleteTableName (pk INT, dep STRING, salary INT) " +
+        s"USING $v2Source PARTITIONED BY (dep, bucket(4, pk))")
+      sql(s"INSERT INTO $deleteTableName VALUES " +
+        "(1, 'hr', 100), (2, 'software', 200), (3, 'marketing', 300)")
+
+      assertDeleteWithFilters(
+        s"DELETE FROM $deleteTableName WHERE dep IN ('hr', 'software')",
+        expectedNumConditions = 1,
+        expectedNumPartitionPredicates = 1,
+        expectedOrdinalsPerPredicate = Seq(Array(0)),
+        expectedPartitionFieldNames = Array("dep", "bucket(4, pk)"))
+
+      checkAnswer(
+        sql(s"SELECT * FROM $deleteTableName"),
+        Row(3, "marketing", 300) :: Nil)
+    }
+  }
+
+  // `dt = DATE'...'` is analyzed as `cast(dt AS DATE) = DATE'...'`, which the table cannot
+  // evaluate in the first pass; the second pass turns it into a PartitionPredicate.
+  test("second pass accepted: cast on a string identity column next to a bucket transform") {
+    withTable(deleteTableName) {
+      sql(s"CREATE TABLE $deleteTableName (pk INT, dt STRING, salary INT) " +
+        s"USING $v2Source PARTITIONED BY (dt, bucket(4, pk))")
+      sql(s"INSERT INTO $deleteTableName VALUES " +
+        "(1, '2026-09-01', 100), (2, '2026-09-02', 200), (3, '2026-09-03', 300)")
+
+      assertDeleteWithFilters(
+        s"DELETE FROM $deleteTableName WHERE dt = DATE'2026-09-02'",
+        expectedNumConditions = 1,
+        expectedNumPartitionPredicates = 1,
+        expectedOrdinalsPerPredicate = Seq(Array(0)),
+        expectedPartitionFieldNames = Array("dt", "bucket(4, pk)"))
+
+      checkAnswer(
+        sql(s"SELECT * FROM $deleteTableName"),
+        Seq(Row(1, "2026-09-01", 100), Row(3, "2026-09-03", 300)))
+    }
+  }
+
   // Table property disables PartitionPredicate acceptance;
   // both passes rejected, falls back to row-level operation.
   test("first and second pass rejected: table rejects all") {

sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala

@@ -403,6 +403,102 @@ class DataSourceV2EnhancedPartitionFilterSuite
     }
   }
 
+  test("mixed partitioning: second-pass PartitionPredicate on the identity field") {
+    withTable(partFilterTableName) {
+      sql(s"CREATE TABLE $partFilterTableName (part_col string, id int, data string) " +
+        s"USING $v2Source PARTITIONED BY (part_col, bucket(4, id))")
+      sql(s"INSERT INTO $partFilterTableName VALUES ('a', 1, 'x'), ('A', 2, 'y'), ('b', 3, 'z')")
+
+      spark.udf.register("my_upper", (s: String) =>
+        if (s == null) null else s.toUpperCase(Locale.ROOT))
+
+      // Untranslatable, Partition Filter on the identity field; 2nd Pass Accepted.
+      // The bucket field keeps ordinal 1 but is never referenced.
+      val df = sql(s"SELECT * FROM $partFilterTableName WHERE my_upper(part_col) = 'A'")
+      checkAnswer(df, Seq(Row("a", 1, "x"), Row("A", 2, "y")))
+      assertPushedPartitionPredicates(df, 1)
+      assertScanReturnsPartitionKeys(df, Set("a/1", "A/2"))
+      assertReferencedPartitionFieldOrdinals(df, Array(0), Array("part_col", "bucket(4, id)"))
+    }
+  }
+
+  test("mixed partitioning: filter on the source column of a bucket transform stays post-scan") {
+    withTable(partFilterTableName) {
+      sql(s"CREATE TABLE $partFilterTableName (part_col string, id int, data string) " +
+        s"USING $v2Source PARTITIONED BY (part_col, bucket(4, id))")
+      sql(s"INSERT INTO $partFilterTableName VALUES ('a', 1, 'x'), ('A', 2, 'y'), ('b', 3, 'z')")
+
+      spark.udf.register("my_upper", (s: String) =>
+        if (s == null) null else s.toUpperCase(Locale.ROOT))
+      spark.udf.register("my_plus1", (i: Int) => i + 1)
+
+      // Only the identity conjunct becomes a PartitionPredicate. Spark cannot evaluate the
+      // bucket conjunct against the partition key, so it is a data filter applied after the scan.
+      val df = sql(s"SELECT * FROM $partFilterTableName " +
+        "WHERE my_upper(part_col) = 'A' AND my_plus1(id) = 3")
+      checkAnswer(df, Seq(Row("A", 2, "y")))
+      assertPushedPartitionPredicates(df, 1)
+      assertScanReturnsPartitionKeys(df, Set("a/1", "A/2"))
+      assertReferencedPartitionFieldOrdinals(df, Array(0), Array("part_col", "bucket(4, id)"))
+      assert(df.queryExecution.executedPlan.exists(_.isInstanceOf[FilterExec]),
+        "Filter on the bucket source column should remain as a post-scan Filter")
+    }
+  }
+
+  test("mixed partitioning: no identity transform -> no PartitionPredicate") {
+    withTable(partFilterTableName) {
+      sql(s"CREATE TABLE $partFilterTableName (id int, data string) " +
+        s"USING $v2Source PARTITIONED BY (bucket(4, id))")
+      sql(s"INSERT INTO $partFilterTableName VALUES (1, 'x'), (2, 'y'), (3, 'z')")
+
+      spark.udf.register("my_plus1", (i: Int) => i + 1)
+
+      val df = sql(s"SELECT * FROM $partFilterTableName WHERE my_plus1(id) = 3")
+      checkAnswer(df, Seq(Row(2, "y")))
+      assertPushedPartitionPredicates(df, 0)
+      assertScanReturnsPartitionKeys(df, Set("1", "2", "3"))
+    }
+  }
+
+  test("mixed partitioning: identity field after a bucket transform -> ordinal 1") {
+    withTable(partFilterTableName) {
+      sql(s"CREATE TABLE $partFilterTableName (part_col string, id int, data string) " +
+        s"USING $v2Source PARTITIONED BY (bucket(4, id), part_col)")
+      sql(s"INSERT INTO $partFilterTableName VALUES ('a', 1, 'x'), ('A', 2, 'y'), ('b', 3, 'z')")
+
+      spark.udf.register("my_upper_second", (s: String) =>
+        if (s == null) null else s.toUpperCase(Locale.ROOT))
+
+      // The bucket field at ordinal 0 keeps its slot, so `part_col` binds to the second
+      // partition-key value and the reference reports ordinal 1.
+      val df = sql(s"SELECT * FROM $partFilterTableName WHERE my_upper_second(part_col) = 'A'")
+      checkAnswer(df, Seq(Row("a", 1, "x"), Row("A", 2, "y")))
+      assertPushedPartitionPredicates(df, 1)
+      assertScanReturnsPartitionKeys(df, Set("1/a", "2/A"))
+      assertReferencedPartitionFieldOrdinals(
+        df, Array(1), Array("bucket(4, id)", "part_col"))
+    }
+  }
+
+  test("mixed partitioning: cast from type coercion on the identity field is pruned by " +
+    "the second pass") {
+    withTable(partFilterTableName) {
+      sql(s"CREATE TABLE $partFilterTableName (dt string, id int, data string) " +
+        s"USING $v2Source PARTITIONED BY (dt, bucket(4, id))")
+      sql(s"INSERT INTO $partFilterTableName VALUES " +
+        "('2026-09-01', 1, 'x'), ('2026-09-02', 2, 'y'), ('2026-09-03', 3, 'z')")
+
+      // `dt = DATE'...'` is analyzed as `cast(dt AS DATE) = DATE'...'`. The source cannot
+      // evaluate a predicate over a cast and returns it in the first pass; the second pass
+      // evaluates it against the `dt` value of the full partition key.
+      val df = sql(s"SELECT * FROM $partFilterTableName WHERE dt = DATE'2026-09-02'")
+      checkAnswer(df, Seq(Row("2026-09-02", 2, "y")))
+      assertPushedPartitionPredicates(df, 1)
+      assertScanReturnsPartitionKeys(df, Set("2026-09-02/2"))
+      assertReferencedPartitionFieldOrdinals(df, Array(0), Array("dt", "bucket(4, id)"))
+    }
+  }
+
   test("non-deterministic partition filter not pushed as PartitionPredicate") {
     // Same checks as FileSourceStrategy/PruneFileSourcePartitions: non-deterministic
     // partition filters must not be pushed as PartitionPredicate; they are applied after scan.