diff --git a/common/utils/src/main/resources/error/error-conditions.json b/common/utils/src/main/resources/error/error-conditions.json index 98b5f9f7560e6..03815febd2523 100644 --- a/common/utils/src/main/resources/error/error-conditions.json +++ b/common/utils/src/main/resources/error/error-conditions.json @@ -2292,6 +2292,11 @@ "The attribute cannot be resolved." ] }, + "NOT_IN_FILTER_ATTRIBUTES" : { + "message" : [ + "The attribute must also be returned by `filterAttributes()`." + ] + }, "NOT_TOP_LEVEL" : { "message" : [ "The attribute must be top-level, but it is a nested reference." diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java index 51531234a6a64..cdb1e4c842ae8 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeFiltering.java @@ -50,7 +50,8 @@ public interface SupportsRuntimeFiltering extends SupportsRuntimeV2Filtering { * The provided expressions must be interpreted as a set of filters that are ANDed together. * Implementations may use the filters to prune initially planned {@link InputPartition}s. *

- * Spark tracks runtime-filter eligibility by root attribute. If {@link #filterAttributes()} + * Spark currently tracks runtime-filter eligibility by root attribute. If + * {@link #filterAttributes()} * returns a nested reference, this method may receive a filter on another nested field under * the same root. Implementations must inspect each filter and use only filters they can apply. * Nested paths are encoded in a V1 {@link Filter} as dot-separated names, with each path part diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java index 671a822096bbe..d43b3eab545a2 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/SupportsRuntimeV2Filtering.java @@ -63,7 +63,8 @@ public interface SupportsRuntimeV2Filtering extends Scan { * The provided expressions must be interpreted as a set of predicates that are ANDed together. * Implementations may use the predicates to prune initially planned {@link InputPartition}s. *

- * Spark tracks runtime-filter eligibility by root attribute. If {@link #filterAttributes()} + * Spark currently tracks runtime-filter eligibility by root attribute. If + * {@link #filterAttributes()} * returns a nested reference, this method may receive a predicate on another nested field under * the same root. Implementations must inspect each predicate and use only predicates they can * apply. diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala index 349a59a1aaf8f..af84df8fcc068 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala @@ -27,6 +27,7 @@ import org.apache.spark.sql.catalyst.analysis.{NoSuchFunctionException, Unresolv import org.apache.spark.sql.catalyst.encoders.EncoderUtils import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke} import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, SampleMethod} +import org.apache.spark.sql.catalyst.types.DataTypeUtils.fromAttributes import org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier} import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.connector.catalog.functions.ScalarFunction.MAGIC_METHOD_NAME @@ -70,6 +71,34 @@ object V2ExpressionUtils extends SQLConfHelper with Logging { AttributeSet(resolveRefs[NamedExpression](refs.toImmutableArraySeq, plan)) } + /** + * Resolves data source runtime-filter attributes and wraps resolution failures with connector + * context. + */ + private[sql] def resolveDataSourceRuntimeFilterRefs( + refs: Array[NamedReference], + output: Seq[Attribute], + method: String, + scanClass: String): AttributeSet = { + if (refs.isEmpty) return AttributeSet.empty + + val plan = LocalRelation(output) + val resolvedAttrs = refs.map { ref => + try { + resolveRef[NamedExpression](ref, plan) + } catch { + case e: AnalysisException => + throw QueryCompilationErrors.cannotResolveDataSourceRuntimeFilterAttributeError( + attribute = ref.fieldNames, + method = method, + scanClass = scanClass, + relationOutput = fromAttributes(output), + cause = e) + } + } + AttributeSet(resolvedAttrs) + } + /** * Converts the array of input V2 [[V2SortOrder]] into their counterparts in catalyst. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala index 3b5d7accd65b0..45efdbe372347 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala @@ -4680,6 +4680,19 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat None) } + def fullyPushedDataSourceRuntimeFilterAttributeNotFilterableError( + attribute: Array[String], + scanClass: String, + relationOutput: StructType): AnalysisException = { + invalidDataSourceRuntimeFilterAttributeError( + attribute, + "fullyPushedFilterAttributes()", + scanClass, + relationOutput, + "NOT_IN_FILTER_ATTRIBUTES", + None) + } + private def invalidDataSourceRuntimeFilterAttributeError( attribute: Array[String], method: String, diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala index 26af15574a27a..d74cffb14f41b 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala @@ -20,10 +20,9 @@ package org.apache.spark.sql.execution.datasources.v2 import java.util.{Collections, Optional, OptionalLong} import org.apache.spark.SparkException -import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.analysis.{MultiInstanceRelation, NamedRelation, TimeTravelSpec} import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics} -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, NamedExpression, SortOrder, V2ExpressionUtils} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics} import org.apache.spark.sql.catalyst.plans.logical.statsEstimation.EstimationUtils @@ -208,12 +207,13 @@ case class DataSourceV2ScanRelation( lazy val runtimeFilterAttrs: AttributeSet = { checkRuntimeFilteringInterfaces() resolvedFullyPushedRuntimeFilterAttrs - val filterAttrs = scan match { - case s: SupportsRuntimeV2Filtering => s.filterAttributes - case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() - case _ => Array.empty[NamedReference] - } - resolveFilterAttrs(filterAttrs, "filterAttributes()") + resolvedRuntimeFilterAttrs + } + + private[sql] lazy val declaredRuntimeFilterAttrs: Array[NamedReference] = scan match { + case s: SupportsRuntimeV2Filtering => s.filterAttributes + case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() + case _ => Array.empty } private lazy val declaredFullyPushedRuntimeFilterAttrs: Array[NamedReference] = scan match { @@ -221,10 +221,17 @@ case class DataSourceV2ScanRelation( case _ => Array.empty } + private lazy val resolvedRuntimeFilterAttrs: AttributeSet = { + resolveFilterAttrs(declaredRuntimeFilterAttrs, "filterAttributes()") + } + private lazy val resolvedFullyPushedRuntimeFilterAttrs: AttributeSet = { - checkFullyPushedFilterAttrs() - resolveFilterAttrs( + checkFullyPushedFilterAttrsAreTopLevel() + val resolvedAttrs = resolveFilterAttrs( declaredFullyPushedRuntimeFilterAttrs, "fullyPushedFilterAttributes()") + resolvedRuntimeFilterAttrs + checkFullyPushedFilterAttrsAreFilterable() + resolvedAttrs } /** @@ -245,20 +252,8 @@ case class DataSourceV2ScanRelation( private def resolveFilterAttrs( filterAttrs: Array[NamedReference], method: String): AttributeSet = { - val resolvedAttrs = filterAttrs.map { ref => - try { - V2ExpressionUtils.resolveRef[NamedExpression](ref, this) - } catch { - case e: AnalysisException => - throw QueryCompilationErrors.cannotResolveDataSourceRuntimeFilterAttributeError( - attribute = ref.fieldNames, - method = method, - scanClass = scan.getClass.getName, - relationOutput = fromAttributes(output), - cause = e) - } - } - AttributeSet(resolvedAttrs) + V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs( + filterAttrs, output, method, scan.getClass.getName) } override val nodePatterns: Seq[TreePattern] = Seq(DATA_SOURCE_V2_SCAN_RELATION) @@ -319,7 +314,7 @@ case class DataSourceV2ScanRelation( } } - private def checkFullyPushedFilterAttrs(): Unit = { + private def checkFullyPushedFilterAttrsAreTopLevel(): Unit = { declaredFullyPushedRuntimeFilterAttrs.find(_.fieldNames.length > 1).foreach { ref => throw QueryCompilationErrors.nestedDataSourceFullyPushedRuntimeFilterAttributeError( attribute = ref.fieldNames, @@ -328,6 +323,20 @@ case class DataSourceV2ScanRelation( } } + private def checkFullyPushedFilterAttrsAreFilterable(): Unit = { + declaredFullyPushedRuntimeFilterAttrs.find { fullyPushedRef => + !declaredRuntimeFilterAttrs.exists { filterRef => + fullyPushedRef.fieldNames.length == filterRef.fieldNames.length && + fullyPushedRef.fieldNames.lazyZip(filterRef.fieldNames).forall(conf.resolver) + } + }.foreach { ref => + throw QueryCompilationErrors.fullyPushedDataSourceRuntimeFilterAttributeNotFilterableError( + attribute = ref.fieldNames, + scanClass = scan.getClass.getName, + relationOutput = fromAttributes(output)) + } + } + override def doCanonicalize(): DataSourceV2ScanRelation = { this.copy( relation = this.relation.copy( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala index ec34efed26252..bdba27f20e8b5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala @@ -78,9 +78,10 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * Implementations may use the expressions to prune initially planned * [[org.apache.spark.sql.connector.read.InputPartition]]s. * - * Spark tracks runtime-filter eligibility by root attribute. If [[filterAttributes]] returns a - * nested reference, an expression may access another nested field under the same root. The scan - * must match each access against its own partition layout and use only expressions it can apply. + * Spark currently tracks runtime-filter eligibility by root attribute. If [[filterAttributes]] + * returns a nested reference, an expression may access another nested field under the same root. + * The scan must match each access against its own partition layout and use only expressions it + * can apply. * * Spark may call this method more than once for the same scan instance: a plan can hold several * scan nodes sharing one scan (e.g. the two branches of a group-based UPDATE), and each pushes diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala index fde2149caac15..d0f7a19583874 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala @@ -29,7 +29,7 @@ import scala.collection.mutable.{ArrayBuffer, ListBuffer} import scala.jdk.CollectionConverters._ import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BoundReference, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate} import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, CaseInsensitiveMap, CharVarcharUtils, DateTimeUtils, GenericArrayData, MapData, ResolveDefaultColumns} import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} @@ -231,6 +231,10 @@ abstract class InMemoryBaseTable( } } + protected def identityPartitionReferences: Array[NamedReference] = { + partitioning.collect { case IdentityTransform(ref) => ref } + } + private val UTC = ZoneId.of("UTC") private val EPOCH_LOCAL_DATE = Instant.EPOCH.atZone(UTC).toLocalDate @@ -513,9 +517,10 @@ abstract class InMemoryBaseTable( } private def canEvaluate(filter: Filter): Boolean = { - if (partitioning.length == 1 && partitioning.head.references.length == 1) { + val identityRefs = identityPartitionReferences + if (partitioning.length == 1 && identityRefs.length == 1) { filter match { - case In(attrName, _) if attrName == partitioning.head.references.head.toString => true + case In(attrName, _) if attrName == identityRefs.head.toString => true case _ => false } } else { @@ -633,6 +638,12 @@ abstract class InMemoryBaseTable( override def toBatch: Batch = this + protected def identityPartitionAttributes: Array[NamedReference] = { + identityPartitionReferences.distinct + .filter(ref => readSchema.findNestedField( + ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined) + } + override def estimateStatistics(): Statistics = { if (data.isEmpty) { return InMemoryStats(OptionalLong.of(0L), OptionalLong.of(0L), new util.HashMap()) @@ -717,11 +728,10 @@ abstract class InMemoryBaseTable( /** * Reference implementation of [[SupportsRuntimeCatalystFiltering.filter]] for the in-memory - * fixtures: records what was pushed, and for expressions referencing only partition columns - * binds them against the partition key and drops partitions that do not match. Binding and - * interpreting rather than pattern matching a fixed set of operators is what lets the fixture - * honor an arbitrary pushed expression, the same way `PartitionPredicateImpl` does. Mixing - * classes supply their own `filterAttributes()`. + * fixtures: records what was pushed, and binds expressions referencing only identity partition + * columns against the partition key to drop partitions that do not match. Interpreting the + * bound expression lets the fixture honor arbitrary pushed expressions. Mixing classes supply + * their own `filterAttributes()`. */ trait CatalystRuntimeFilteringScan extends SupportsRuntimeCatalystFiltering { self: BatchScanBaseClass => @@ -737,18 +747,15 @@ abstract class InMemoryBaseTable( filterCalls += 1 val partAttrs = partitionAttributes if (partAttrs.isEmpty) return - val partAttrRefs = partAttrs.map(_._2) expressions.foreach { expr => // Top down, so `s.part` is rewritten before its `s` child is considered. val remapped = expr.transformDown { case e => partitionAttrFor(e, partAttrs).getOrElse(e) } - // Only evaluate expressions whose refs are all partition columns, so we can bind - // against the partition key InternalRow (same approach as PartitionPredicateImpl). - if (remapped.references.forall(r => partAttrRefs.exists(_.exprId == r.exprId))) { - val bound = BindReferences.bindReference(remapped, partAttrRefs) - val pred = CatalystPredicate.createInterpreted(bound) + // Evaluate expressions only when every reference maps to an identity partition-key slot. + if (remapped.references.isEmpty) { + val pred = CatalystPredicate.createInterpreted(remapped) self.data = self.data.filter { p => try { pred.eval(p.asInstanceOf[BufferedRows].partitionKey()) @@ -772,37 +779,33 @@ abstract class InMemoryBaseTable( def filterCallCount: Int = filterCalls /** - * The `AttributeReference`s standing for the partition key InternalRow fields, in its field - * order, each paired with the name-part sequence of its partition column. The parts are kept - * unflattened so a quoted top-level column `a.b` (parts `Seq("a.b")`) stays distinct from a - * nested column `a`.`b` (parts `Seq("a", "b")`). Example: - * - `PARTITIONED BY (part, s.nested)` -> `(Seq("part"), AttributeReference(part))`, then - * `(Seq("s", "nested"), AttributeReference(s.nested))` + * Identity partition columns paired with their bound partition-key slots. + * + * Only identity transforms expose a source path because their partition-key slot retains the + * source value. Name parts stay separate so a quoted top-level column `a.b` remains distinct + * from a nested column `a`.`b`. */ - private def partitionAttributes: Seq[(Seq[String], AttributeReference)] = { - partitioning.flatMap(_.references()).flatMap { ref => - val path = ref.fieldNames.toImmutableArraySeq - val resolver = SQLConf.get.resolver - readSchema.findNestedField(path, resolver = resolver) - .orElse(tableSchema.findNestedField(path, resolver = resolver)).map { - case (_, f) => - path -> AttributeReference(ref.fieldNames.mkString("."), f.dataType, f.nullable)() - } + private def partitionAttributes: Seq[(Seq[String], BoundReference)] = { + partitioning.zipWithIndex.flatMap { + case (IdentityTransform(ref), ordinal) => + val path = ref.fieldNames.toImmutableArraySeq + val resolver = SQLConf.get.resolver + readSchema.findNestedField(path, resolver = resolver) + .orElse(tableSchema.findNestedField(path, resolver = resolver)).map { + case (_, f) => path -> BoundReference(ordinal, f.dataType, f.nullable) + } + case _ => None }.toSeq } /** - * The partition key `AttributeReference` that `e` reads, or None if `e` reads no partition - * column. The path `e` reads is compared to each partition column's name parts component-wise - * with the resolver, so a quoted top-level column `a.b` cannot collide with a nested column - * `a`.`b`. Examples, under `PARTITIONED BY (part, s.nested)` where `nested` is field 0 of `s`: - * - `AttributeReference(part)` -> `AttributeReference(part)` - * - `GetStructField(AttributeReference(s), 0)` -> `AttributeReference(s.nested)` - * - `AttributeReference(s)` -> None if `s` itself is not a partition column, only `s.nested` + * The partition-key slot that `e` reads, or None if `e` reads no identity partition column. + * The path `e` reads is compared to each partition column's name parts component-wise with the + * resolver, so a quoted top-level column `a.b` cannot collide with a nested column `a`.`b`. */ private def partitionAttrFor( e: CatalystExpression, - partAttrs: Seq[(Seq[String], AttributeReference)]): Option[AttributeReference] = { + partAttrs: Seq[(Seq[String], BoundReference)]): Option[BoundReference] = { val resolver = SQLConf.get.resolver partitionKeyPath(e).flatMap { path => partAttrs.collectFirst { @@ -841,14 +844,12 @@ abstract class InMemoryBaseTable( var pushedFilters: Array[Filter] = Array.empty override def filterAttributes(): Array[NamedReference] = { - partitioning.flatMap(_.references) - .filter(ref => readSchema.findNestedField( - ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined) + identityPartitionAttributes } override def filter(filters: Array[Filter]): Unit = { - if (partitioning.length == 1 && partitioning.head.references().length == 1) { - val ref = partitioning.head.references().head + if (partitioning.length == 1 && identityPartitionReferences.length == 1) { + val ref = identityPartitionReferences.head filters.foreach { case In(attrName, values) if attrName == ref.toString => val matchingKeys = values.map { value => diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala index 58f07307dcbe1..e0f300ac125ce 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala @@ -107,7 +107,9 @@ class InMemoryCatalystRuntimeFilterTable( // Not intersected with `filterAttributes()`, so a table can declare a fully pushed attribute // that is not a filter attribute, a combination the interface forbids. override def fullyPushedFilterAttributes(): Array[NamedReference] = { - partitionAttrs.filter(ref => fullyPushedFilterAttrs.contains(ref.fieldNames.mkString("."))) + identityPartitionAttributes.filter { ref => + fullyPushedFilterAttrs.contains(ref.fieldNames.mkString(".")) + } } } } 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 428a9215dcce0..7556bc96912f6 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 @@ -28,7 +28,6 @@ import org.apache.spark.sql.connector.expressions.filter.Predicate import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} import org.apache.spark.sql.connector.write.{BatchWrite, DeltaBatchWrite, DeltaWrite, DeltaWriteBuilder, DeltaWriter, DeltaWriterFactory, LogicalWriteInfo, PhysicalWriteInfo, RequiresDistributionAndOrdering, RowLevelOperation, RowLevelOperationBuilder, RowLevelOperationInfo, SupportsDelta, Write, WriteBuilder, WriterCommitMessage} import org.apache.spark.sql.connector.write.RowLevelOperation.Command -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.unsafe.types.UTF8String @@ -311,9 +310,7 @@ class InMemoryRowLevelOperationTable private ( with CatalystRuntimeFilteringScan { override def filterAttributes(): Array[NamedReference] = { - partitioning.flatMap(_.references()) - .filter(ref => readSchema.findNestedField( - ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined) + identityPartitionAttributes } } } 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 fa50da6732bf2..1228b70b103d0 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 @@ -27,7 +27,6 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, LiteralValue, import org.apache.spark.sql.connector.expressions.filter.{And, Predicate} import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.write.{LogicalWriteInfo, SupportsOverwriteV2, WriteBuilder, WriterCommitMessage} -import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -78,14 +77,12 @@ class InMemoryTableWithV2Filter( extends BatchScanBaseClass(_data, readSchema, tableSchema) with SupportsRuntimeV2Filtering { override def filterAttributes(): Array[NamedReference] = { - partitioning.flatMap(_.references) - .filter(ref => readSchema.findNestedField( - ref.fieldNames.toImmutableArraySeq, resolver = SQLConf.get.resolver).isDefined) + identityPartitionAttributes } override def filter(filters: Array[Predicate]): Unit = { - if (partitioning.length == 1 && partitioning.head.references().length == 1) { - val ref = partitioning.head.references().head + if (partitioning.length == 1 && identityPartitionReferences.length == 1) { + val ref = identityPartitionReferences.head filters.foreach { case p : Predicate if p.name().equals("IN") => if (p.children().length > 1) { 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 9a1d00a96a685..3a5796d3448fc 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 @@ -212,8 +212,11 @@ object PushDownUtils extends Logging { // filters whose translation was not already accepted in the first pass. (See SPARK-55596) // Only candidates whose referenced columns are declared in filterAttributes() are eligible. val partPredicatesPushed = filterableScan.supportsIterativePushdown() && { - val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( - filterableScan.filterAttributes(), output) + val filterAttrs = V2ExpressionUtils.resolveDataSourceRuntimeFilterRefs( + filterableScan.filterAttributes(), + output, + "filterAttributes()", + filterableScan.getClass.getName) val pushed = filterableScan.pushedPredicates().toSet val candidates = runtimeFilters.filter { f => !filtersToTranslated.get(f).exists(pushed.contains) && diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala index 1bf12a695bc4c..c0747cbfb1130 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala @@ -23,12 +23,11 @@ import org.apache.spark.sql.catalyst.optimizer.{JoinSelectionHelper, ReusableBro import org.apache.spark.sql.catalyst.planning.ExtractEquiJoinKeys import org.apache.spark.sql.catalyst.plans.logical._ import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering import org.apache.spark.sql.execution.LogicalRDD import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRelation} -import org.apache.spark.sql.execution.datasources.v2.ExtractV2Scan -import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation + /** * Dynamic partition pruning optimization is performed based on the type and * selectivity of the join operation. During query optimization, we insert a @@ -79,18 +78,8 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } else { None } - case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) => - val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( - scan.filterAttributes, r.output) - if (resExp.references.subsetOf(filterAttrs)) { - Some(r) - } else { - None - } - case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) => - val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( - scan.filterAttributes(), r.output) - if (resExp.references.subsetOf(filterAttrs)) { + case (resExp, r: DataSourceV2ScanRelation) => + if (resExp.references.subsetOf(r.runtimeFilterAttrs)) { Some(r) } else { None diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala index b17e224c4a5ad..a26a0cea74f1e 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala @@ -26,10 +26,9 @@ import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPl import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreePattern.{REPLACE_DATA, WRITE_DELTA} import org.apache.spark.sql.connector.expressions.NamedReference -import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering} +import org.apache.spark.sql.connector.read.Scan import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} -import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan} -import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering +import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation} import org.apache.spark.util.ArrayImplicits._ /** @@ -55,32 +54,22 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla override def apply(plan: LogicalPlan): LogicalPlan = plan.transformDownWithPruning( _.containsAnyPattern(REPLACE_DATA, WRITE_DELTA)) { case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) - if canInjectGroupFilters(cond, scan.filterAttributes) => - injectGroupFilters(replaceData, cond, scan, scan.filterAttributes) - - case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) - if canInjectGroupFilters(cond, scan.filterAttributes()) => - injectGroupFilters(replaceData, cond, scan, scan.filterAttributes()) - - case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) - if canInjectGroupFilters(cond, scan.filterAttributes) => - injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes) + r: DataSourceV2ScanRelation) + if canInjectGroupFilters(cond, r) => + injectGroupFilters(replaceData, cond, r.scan, r.declaredRuntimeFilterAttrs) case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) - if canInjectGroupFilters(cond, scan.filterAttributes()) => - injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes()) + r: DataSourceV2ScanRelation) + if canInjectGroupFilters(cond, r) => + injectGroupFilters(writeDelta, cond, r.scan, r.declaredRuntimeFilterAttrs) } private def canInjectGroupFilters( cond: Expression, - filterAttrs: Array[NamedReference]): Boolean = { + scanRelation: DataSourceV2ScanRelation): Boolean = { conf.runtimeRowLevelOperationGroupFilterEnabled && cond != TrueLiteral && - filterAttrs.nonEmpty + scanRelation.runtimeFilterAttrs.nonEmpty } private def injectGroupFilters( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala index defa376dac50c..45cca0c40a68d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -116,6 +116,30 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } + test("Catalyst runtime filtering fixture fully pushes only identity partition sources") { + val tbl = s"$catalogName.tbl_transformed_fully_pushed" + val dim = s"$catalogName.dim_transformed_fully_pushed" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part DATE) USING $v2Source " + + "PARTITIONED BY (days(part)) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')") + sql(s"INSERT INTO $tbl VALUES " + + "(1, DATE '2026-08-01'), (2, DATE '2026-08-02')") + sql(s"CREATE TABLE $dim (value DATE) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (DATE '2026-08-02')") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(value) FROM $dim)") + checkAnswer(df, Row(2, java.sql.Date.valueOf("2026-08-02"))) + + assertScalarSubqueryRuntimeFilters(df) + assertPushedCatalystPredicates(df, expected = 1) + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + val scan = collectBatchScan(df) + assert(scan.inputPartitions.size === 2) + assert(scan.filteredPartitions.flatten.size === 2) + } + } + test("nested fully pushed filter attribute -> rejected without a runtime filter") { val tbl = s"$catalogName.tbl_nested_fully_pushed" withTable(tbl) { @@ -142,6 +166,29 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } + test("fully pushed root attribute outside filterAttributes -> rejected") { + val tbl = s"$catalogName.tbl_fully_pushed_root" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, s STRUCT) USING $v2Source") + val scanRelation = sql(s"SELECT * FROM $tbl").queryExecution.optimizedPlan.collectFirst { + case r: DataSourceV2ScanRelation => r + }.getOrElse(fail("Expected a DataSourceV2ScanRelation")) + val scan = new FullyPushedRootAttributeScan + val e = intercept[AnalysisException] { + scanRelation.copy(scan = scan).runtimeFilterAttrs + } + checkError( + exception = e, + condition = "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.NOT_IN_FILTER_ATTRIBUTES", + parameters = Map( + "attribute" -> "`s`", + "method" -> "fullyPushedFilterAttributes()", + "scanClass" -> scan.getClass.getName, + "relationOutput" -> "\"STRUCT>\""), + sqlState = "KD000") + } + } + test("nested filter attribute under a non-struct column -> rejected during resolution") { val tbl = s"$catalogName.tbl_malformed_nested_filter_attr" withTable(tbl) { @@ -432,17 +479,38 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } - test("filter on column outside filterAttributes -> not pushed, even if declared fully pushed") { + test("fully pushed attribute outside filterAttributes -> rejected") { val tbl = s"$catalogName.tbl4" - val dim = s"$catalogName.dim4" - withTable(tbl, dim) { - // p2 is a partition column but is not declared filterable, so no runtime filter is derived - // for it. Declaring it fully pushed as well, which the interface forbids for an attribute - // that is not filterable, must not cost it the post-scan filter: nothing was pushed, so the - // scan prunes nothing and the nonmatching rows would come back. + withTable(tbl) { sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + "PARTITIONED BY (p1, p2) " + "TBLPROPERTIES('filter-attributes' = 'p1', 'fully-pushed-filter-attributes' = 'p2')") + val df = sql(s"SELECT * FROM $tbl") + val scanClass = df.queryExecution.optimizedPlan.collectFirst { + case r: DataSourceV2ScanRelation => r.scan.getClass.getName + }.getOrElse(fail("Expected a DataSourceV2ScanRelation")) + val e = intercept[AnalysisException] { + df.queryExecution.executedPlan + } + checkError( + exception = e, + condition = "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.NOT_IN_FILTER_ATTRIBUTES", + parameters = Map( + "attribute" -> "`p2`", + "method" -> "fullyPushedFilterAttributes()", + "scanClass" -> scanClass, + "relationOutput" -> "\"STRUCT\""), + sqlState = "KD000") + } + } + + test("filter on column outside filterAttributes -> not pushed") { + val tbl = s"$catalogName.tbl_restricted_filter_attrs" + val dim = s"$catalogName.dim_restricted_filter_attrs" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + + "PARTITIONED BY (p1, p2) " + + "TBLPROPERTIES('filter-attributes' = 'p1')") for (i <- 0 until 5) { sql(s"INSERT INTO $tbl VALUES ($i, $i, $i)") } @@ -452,10 +520,13 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM $dim)") checkAnswer(df, Row(3, 3, 3)) - assert(collectBatchScan(df).runtimeFilters.isEmpty, + val scan = collectBatchScan(df) + assert(scan.runtimeFilters.isEmpty, "Expected no runtime filters for a column outside filterAttributes") - assertPushedCatalystPredicates(df, 0) + assertPushedCatalystPredicates(df, expected = 0) assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + assert(scan.inputPartitions.size === 5) + assert(scan.filteredPartitions.flatten.size === 5) } } @@ -755,6 +826,21 @@ private class MissingFullyPushedFilterAttributeScan extends SupportsRuntimeCatal override def filter(expressions: Array[Expression]): Unit = {} } +/** A scan declaring a root struct fully pushed without declaring it filterable. */ +private class FullyPushedRootAttributeScan extends SupportsRuntimeCatalystFiltering { + + override def readSchema(): StructType = + StructType.fromDDL("id INT, s STRUCT") + + override def filterAttributes(): Array[NamedReference] = + Array(FieldReference(Seq("s", "part"))) + + override def fullyPushedFilterAttributes(): Array[NamedReference] = + Array(FieldReference("s")) + + override def filter(expressions: Array[Expression]): Unit = {} +} + /** A scan declaring a nested runtime-filter attribute beneath an integer column. */ private class NestedFilterAttributeScan extends SupportsRuntimeCatalystFiltering { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index cc003d440eb57..bd2eb571da700 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -178,6 +178,60 @@ abstract class DataSourceV2SQLSuite } } + test("transformed nested partition source does not receive a DPP runtime filter") { + val fact = s"${catalogAndNamespace}fact_transformed_nested_runtime_filter" + val dim = s"${catalogAndNamespace}dim_transformed_nested_runtime_filter" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact " + + "(id INT, derives STRUCT) USING " + + s"$v2Format PARTITIONED BY (truncate(derives.toStr, 1))") + sql(s"INSERT INTO $fact VALUES " + + "(1, named_struct('toStr', 'AA', 'other', 'a')), " + + "(2, named_struct('toStr', 'BB', 'other', 'b')), " + + "(3, named_struct('toStr', 'CC', 'other', 'c'))") + sql(s"CREATE TABLE $dim (value STRING, selected INT) USING $v2Format") + sql(s"INSERT INTO $dim VALUES ('AA', 0), ('BB', 1)") + + withSQLConf( + SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "false", + SQLConf.DYNAMIC_PARTITION_PRUNING_FALLBACK_FILTER_RATIO.key -> "10") { + val df = sql( + s"""SELECT f.id FROM $fact f JOIN $dim d + |ON f.derives.toStr = d.value WHERE d.selected = 1""".stripMargin) + checkAnswer(df, Row(2)) + + val factScan = collect(df.queryExecution.executedPlan) { + case b: BatchScanExec if b.output.exists(_.name == "derives") => b + }.head + assert(factScan.runtimeFilters.isEmpty, + s"expected no runtime filters on a transform source, got ${factScan.runtimeFilters}") + assert(factScan.partitions.size === 3) + assert(factScan.filteredPartitions.flatten.size === 3) + } + } + } + + test("filter on transformed partition source remains post-scan") { + val table = s"${catalogAndNamespace}transformed_source_filter" + withTable(table) { + sql(s"CREATE TABLE $table (id INT, part DATE) USING $v2Format " + + "PARTITIONED BY (days(part))") + sql(s"INSERT INTO $table VALUES " + + "(1, DATE '2026-08-01'), (2, DATE '2026-08-02'), (3, DATE '2026-08-03')") + + val df = sql(s"SELECT * FROM $table WHERE " + + "part IN (DATE '2026-08-01', DATE '2026-08-03')") + checkAnswer(df, Seq( + Row(1, java.sql.Date.valueOf("2026-08-01")), + Row(3, java.sql.Date.valueOf("2026-08-03")))) + + val scan = collect(df.queryExecution.executedPlan) { case b: BatchScanExec => b }.head + assert(scan.partitions.size === 3) + assert(scan.filteredPartitions.flatten.size === 3) + } + } + private def checkExplain(query: String, relationPattern: Regex): Unit = { val explain = spark.sql(s"EXPLAIN EXTENDED $query").head().getString(0) val relations = explain.split("\n").filter(_.contains("RelationV2")) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala index becdde5aa2f4a..39138ceb0d39d 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -40,7 +40,7 @@ class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite } assertCatalystGroupFilter( executedPlan, - expectedFilterAttrs = Seq("dep"), + expectedFilterPaths = Seq(Seq("dep")), expectedFilter = GroupFilter(scanSchema = "salary INT, dep STRING", groups = Seq("hr"))) checkAnswer( @@ -66,10 +66,9 @@ class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite } assertCatalystGroupFilter( executedPlan, - expectedFilterAttrs = Seq("dep.name"), + expectedFilterPaths = Seq(Seq("dep", "name")), expectedFilter = GroupFilter( - scanSchema = "salary INT, dep STRUCT", groups = Seq("hr")), - expectedFilterPaths = Some(Seq(Seq("dep", "name")))) + scanSchema = "salary INT, dep STRUCT", groups = Seq("hr"))) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala index 43ed7797fb5bc..01fd6591ffdaf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala @@ -21,6 +21,7 @@ import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.expressions.{Attribute, CreateNamedStruct, DynamicPruningExpression, Expression, GetStructFieldObject} import org.apache.spark.sql.catalyst.types.DataTypeUtils import org.apache.spark.sql.connector.catalog.{BufferedRows, InMemoryRowLevelOperationTable} +import org.apache.spark.sql.connector.expressions.Expressions.bucket import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference} import org.apache.spark.sql.connector.expressions.Transform import org.apache.spark.sql.execution.InSubqueryExec @@ -69,7 +70,7 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase } assertCatalystGroupFilter( executedPlan, - expectedFilterAttrs = Seq("dep"), + expectedFilterPaths = Seq(Seq("dep")), expectedFilter = GroupFilter(scanSchema = "id INT, dep STRING", groups = Seq("hr"))) // software was never read, so its rows must come back untouched @@ -107,7 +108,7 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase } assertCatalystGroupFilter( executedPlan, - expectedFilterAttrs = Seq("dep"), + expectedFilterPaths = Seq(Seq("dep")), expectedFilter = GroupFilter(scanSchema = "pk INT, dep STRING", groups = Seq("hr"))) // software was never read, so its rows must come back untouched @@ -147,10 +148,9 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase } assertCatalystGroupFilter( executedPlan, - expectedFilterAttrs = Seq("dep.name"), + expectedFilterPaths = Seq(Seq("dep", "name")), expectedFilter = GroupFilter( - scanSchema = "pk INT, dep STRUCT", groups = Seq("hr")), - expectedFilterPaths = Some(Seq(Seq("dep", "name")))) + scanSchema = "pk INT, dep STRUCT", groups = Seq("hr"))) checkAnswer( sql(s"SELECT * FROM $tableNameAsString"), @@ -160,9 +160,26 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase } } + test("non-identity partition transform does not enable runtime group filtering") { + val schema = "pk INT NOT NULL, id INT, salary INT, dep STRING" + createTable(schema, Array[Transform](bucket(4, "dep"))) + append(schema, + """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 200, "dep": "software" } + |""".stripMargin) + + val executedPlan = executeAndKeepPlan { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id = 1") + } + val batchScans = collect(executedPlan) { case s: BatchScanExec => s } + assert(batchScans.nonEmpty, "expected a batch scan for the row-level operation") + assert(batchScans.forall(_.runtimeFilters.isEmpty), + s"expected no runtime group filters, got ${batchScans.flatMap(_.runtimeFilters)}") + } + /** * Asserts the injected group filter down to its contents: the scan declares - * `expectedFilterAttrs` in `filterAttributes`, every scan node carries one dynamic pruning + * `expectedFilterPaths` in `filterAttributes`, every scan node carries one dynamic pruning * filter matching `expectedFilter`, the connector received that same filter as a Catalyst * expression, and the scan then read only `expectedFilter.groups`. * @@ -175,24 +192,22 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase */ protected def assertCatalystGroupFilter( executedPlan: SparkPlan, - expectedFilterAttrs: Seq[String], - expectedFilter: GroupFilter, - expectedFilterPaths: Option[Seq[Seq[String]]] = None): Unit = { + expectedFilterPaths: Seq[Seq[String]], + expectedFilter: GroupFilter): Unit = { val batchScans = collect(executedPlan) { case s: BatchScanExec => s } assert(batchScans.nonEmpty, "expected a batch scan for the row-level operation") val scan = catalystScan(batchScans.head) assert(batchScans.forall(_.scan eq scan), s"expected all ${batchScans.size} scan nodes to share one scan") - val filterAttrs = scan.filterAttributes().map(_.fieldNames.mkString(".")).toSeq - assert(filterAttrs === expectedFilterAttrs, - s"expected the scan to declare $expectedFilterAttrs as filter attributes, got $filterAttrs") - val filterPaths = expectedFilterPaths.getOrElse(expectedFilterAttrs.map(Seq(_))) + val filterPaths = scan.filterAttributes().map(_.fieldNames.toSeq).toSeq + assert(filterPaths === expectedFilterPaths, + s"expected the scan to declare $expectedFilterPaths as filter attributes, got $filterPaths") batchScans.foreach { batchScan => batchScan.runtimeFilters match { case Seq(DynamicPruningExpression(inSubquery: InSubqueryExec)) => - assertGroupFilter(inSubquery, filterPaths, expectedFilter) + assertGroupFilter(inSubquery, expectedFilterPaths, expectedFilter) case other => fail(s"expected a single dynamic pruning group filter, got $other") } } @@ -204,7 +219,7 @@ abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase s"expected each of the ${batchScans.size} scan node(s) to push the filter once, got $pushed") pushed.foreach { case inSubquery: InSubqueryExec => - assertGroupFilter(inSubquery, filterPaths, expectedFilter) + assertGroupFilter(inSubquery, expectedFilterPaths, expectedFilter) case other => fail(s"expected the group filter pushed as an InSubqueryExec, got $other") }