From d4bfb9452b6d3a1061e422ff88975e38e21a8704 Mon Sep 17 00:00:00 2001 From: Cheng Pan Date: Thu, 10 Sep 2026 20:34:23 +0800 Subject: [PATCH] [SPARK-59410][SQL] Derive PartitionPredicate from identity fields of a mixed partitioning PushDownUtils.getPartitionPredicateSchema returned a schema only when every transform in Table.partitioning() is an identity transform, so a table partitioned by e.g. dt (identity) and bucket(16, user_id) never received a PartitionPredicate, and a Catalyst-only filter on dt such as the cast(dt AS DATE) = DATE'...' produced by type coercion could not prune partitions in the static pass, via DPP, or in a metadata-only DELETE. The schema now has one field per transform, in partitioning order. Identity fields carry an attribute a filter can reference; other fields have none but keep their ordinal, so a predicate still binds against the full partition key and the connector contract is unchanged. A filter on the source column of a non-identity transform stays a data filter. A partitioning with no identity transform still yields no schema. The in-memory V2 filter test table now accepts only column-vs-literal predicates and returns anything else, e.g. a predicate over a cast, as a real connector would. Assisted-by: Claude Fable 5.1 --- .../connector/PartitionPredicateField.scala | 17 ++-- .../connector/PartitionPredicateImpl.scala | 18 +++- .../catalog/InMemoryTableWithV2Filter.scala | 23 +++-- .../PartitionPredicateImplSuite.scala | 31 +++++- .../datasources/v2/PushDownUtils.scala | 31 +++--- ...ataSourceV2EnhancedDeleteFilterSuite.scala | 44 +++++++++ ...SourceV2EnhancedPartitionFilterSuite.scala | 96 +++++++++++++++++++ ...2EnhancedRuntimePartitionFilterSuite.scala | 55 +++++++++++ 8 files changed, 285 insertions(+), 30 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateField.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateField.scala index 993c8706feeca..37581a3329ab6 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateField.scala +++ b/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]) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala index 0550344fe4e8a..fb6ff6d8bced7 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImpl.scala +++ b/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( diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala index 6922c874cf73a..605bdd5a40b3c 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala +++ b/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 + } } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala index 7dd6ae8717f0f..0cef9e0f160fd 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/internal/connector/PartitionPredicateImplSuite.scala +++ b/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]( 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 d2b6c0bfa13c4..eb8ee6d6c6d4c 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 @@ -367,8 +367,8 @@ 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]] = { @@ -376,8 +376,8 @@ 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 `table.partitioning`, if the + * partitioning supports [[PartitionPredicate]] push down. None if not supported. */ def getPartitionPredicateSchema(table: Table, output: Seq[AttributeReference]) : Option[Seq[PartitionPredicateField]] = { @@ -385,8 +385,15 @@ 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, 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 } 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..11cc2d7592e0b 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 @@ -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") { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala index 3fe928daf9e03..e308fca937904 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedPartitionFilterSuite.scala +++ b/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. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala index 07a214c24fbc3..c6215d39dd0b7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2EnhancedRuntimePartitionFilterSuite.scala @@ -219,6 +219,61 @@ class DataSourceV2EnhancedRuntimePartitionFilterSuite } } + test("mixed partitioning: DPP on the identity col next to a bucket transform -> " + + "PartitionPredicate") { + val fact = s"$catalogName.fact_mixed" + val dim = s"$catalogName.dim_mixed" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source " + + "PARTITIONED BY (part, bucket(4, id))") + for (i <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (dim_id INT, dim_val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2, 'two')") + + withDPPConf { + val df = sql( + s"""SELECT f.id, f.part FROM $fact f JOIN $dim d + |ON f.part = d.dim_id WHERE d.dim_val = 'two'""".stripMargin) + checkAnswer(df, Row(2, 2)) + + assertDPPRuntimeFilters(df) + assertPushedPartitionPredicates(df, 1) + assertScanReturnsPartitionKeys(df, Set("2/2")) + assertReferencedPartitionFieldOrdinals(df, Array(0), Array("part", "bucket(4, id)")) + } + } + } + + test("mixed partitioning: DPP on the source col of a bucket transform -> " + + "no PartitionPredicate") { + val fact = s"$catalogName.fact_mixed2" + val dim = s"$catalogName.dim_mixed2" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source " + + "PARTITIONED BY (part, bucket(4, id))") + for (i <- 0 until 5) { + sql(s"INSERT INTO $fact VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (dim_id INT, dim_val STRING) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2, 'two')") + + withDPPConf { + val df = sql( + s"""SELECT f.id, f.part FROM $fact f JOIN $dim d + |ON f.id = d.dim_id WHERE d.dim_val = 'two'""".stripMargin) + checkAnswer(df, Row(2, 2)) + + // Spark cannot evaluate `id IN (...)` against the bucket value, so no PartitionPredicate + // is derived and the scan keeps every partition. + assertDPPRuntimeFilters(df) + assertPushedPartitionPredicates(df, 0) + assertScanReturnsPartitionKeys(df, Set("0/0", "1/1", "2/2", "3/3", "4/0")) + } + } + } + test("case 3: scalar subquery translatable, rejected in 1st pass -> PartitionPredicate") { val tbl = s"$catalogName.tbl" val dim = s"$catalogName.dim"