From 4d4f1cf3d3e87e462f9924e3b131ec0e4e444944 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Mon, 3 Aug 2026 19:37:28 +0000 Subject: [PATCH 1/9] [SPARK-58523][SQL] Add Catalyst runtime filter pushdown interface for DSv2 scans The existing runtime filtering interfaces receive connector predicates, so a runtime filter that has no data source V2 translation (a complex expression, a UDF, RLIKE) is dropped and never reaches the source. Add an internal SupportsPushDownCatalystRuntimeFiltering that hands scans the Catalyst expressions directly, and let a scan declare which attributes it fully evaluates so Spark can stop re-evaluating those filters after the scan. --- .../read/SupportsRuntimeFiltering.java | 1 + .../read/SupportsRuntimeV2Filtering.java | 8 +- .../datasources/v2/DataSourceV2Relation.scala | 32 +- ...ortsPushDownCatalystRuntimeFiltering.scala | 77 +++++ .../InMemoryCatalystRuntimeFilterTable.scala | 124 ++++++++ ...oryTableCatalystRuntimeFilterCatalog.scala | 50 +++ .../datasources/v2/DataSourceV2Strategy.scala | 10 +- .../datasources/v2/PushDownUtils.scala | 36 ++- .../dynamicpruning/PartitionPruning.scala | 9 + ...taSourceV2CatalystRuntimeFilterSuite.scala | 289 ++++++++++++++++++ 10 files changed, 618 insertions(+), 18 deletions(-) create mode 100644 sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala create mode 100644 sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala 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 927d4a53e22fc..a1613b50b07e1 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 @@ -26,6 +26,7 @@ /** * A mix-in interface for {@link Scan}. Data sources can implement this interface if they can * filter initially planned {@link InputPartition}s using predicates Spark infers at runtime. + * Only one runtime filtering interface should be implemented by a data source. *

* Note that Spark will push runtime filters only if they are beneficial. * 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 94dbc3865958a..2939b5b574081 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 @@ -28,8 +28,8 @@ * filter initially planned {@link InputPartition}s using predicates Spark infers at runtime. * This interface is very similar to {@link SupportsRuntimeFiltering} except it uses * data source V2 {@link Predicate} instead of data source V1 {@link Filter}. - * {@link SupportsRuntimeV2Filtering} is preferred over {@link SupportsRuntimeFiltering} - * and only one of them should be implemented by the data sources. + * {@link SupportsRuntimeV2Filtering} is preferred over {@link SupportsRuntimeFiltering}. + * Only one runtime filtering interface should be implemented by a data source. *

* Iterative filtering: When {@link #supportsIterativePushdown()} returns true, * {@link #filter(Predicate[])} may be called multiple times on the same @@ -82,6 +82,10 @@ public interface SupportsRuntimeV2Filtering extends Scan { * Returns the predicates that are pushed to the data source via * {@link #filter(Predicate[])}. *

+ * These are not fully pushed predicates: Spark may still evaluate them after the scan. + * They are predicates that fully or partially help the data source prune initially planned + * {@link InputPartition}s. + *

* When iterative filtering is supported and {@link #filter(Predicate[])} was called * multiple times, this method must return predicates from all calls. *

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 52c2674adc2c6..cc6e1f33e1622 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 @@ -35,7 +35,7 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics, Histogram => V2Histogram, HistogramBin => V2HistogramBin} import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream} -import org.apache.spark.sql.internal.connector.V2StatisticsUtils +import org.apache.spark.sql.internal.connector.{SupportsPushDownCatalystRuntimeFiltering, V2StatisticsUtils} import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -196,14 +196,30 @@ case class DataSourceV2ScanRelation( /** * Resolved attributes that the scan declares for runtime filtering via - * [[SupportsRuntimeV2Filtering.filterAttributes]]. Empty when the scan - * does not implement [[SupportsRuntimeV2Filtering]] or exposes no attributes. + * [[SupportsRuntimeV2Filtering.filterAttributes]] or + * [[SupportsPushDownCatalystRuntimeFiltering.filterAttributes]]. Empty when the scan + * implements neither interface or exposes no attributes. */ - lazy val runtimeFilterAttrs: AttributeSet = scan match { - case s: SupportsRuntimeV2Filtering => - AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( - s.filterAttributes.toImmutableArraySeq, this)) - case _ => AttributeSet.empty + lazy val runtimeFilterAttrs: AttributeSet = { + val filterAttrs = scan match { + case s: SupportsRuntimeV2Filtering => s.filterAttributes + case s: SupportsPushDownCatalystRuntimeFiltering => s.filterAttributes + case _ => Array.empty[NamedReference] + } + AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( + filterAttrs.toImmutableArraySeq, this)) + } + + /** + * Resolved attributes for which a Catalyst runtime-filtering scan fully evaluates predicates. + */ + lazy val fullyPushedRuntimeFilterAttrs: AttributeSet = { + val filterAttrs = scan match { + case s: SupportsPushDownCatalystRuntimeFiltering => s.fullyPushedFilterAttributes() + case _ => Array.empty[NamedReference] + } + AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( + filterAttrs.toImmutableArraySeq, this)) } override def name: String = relation.name diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala new file mode 100644 index 0000000000000..dea28b6db7001 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.internal.connector + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.read.Scan + +/** + * A mix-in interface for [[Scan]]. Data sources can implement this interface if they can + * filter initially planned [[org.apache.spark.sql.connector.read.InputPartition]]s using + * Catalyst [[Expression]]s Spark infers at runtime. + * Only one runtime filtering interface should be implemented by a data source. + * + * Spark considers a runtime predicate fully pushed when all attributes referenced by the + * predicate are returned by [[fullyPushedFilterAttributes]]. Fully pushed predicates are not + * evaluated again after the scan. + * + * Note that Spark will push runtime filters only if they are beneficial. + */ +trait SupportsPushDownCatalystRuntimeFiltering extends Scan { + + /** + * Returns attributes this scan can be filtered by at runtime. + * + * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes. + */ + def filterAttributes: Array[NamedReference] + + /** + * Returns attributes for which this scan fully evaluates runtime predicates. + * + * Any runtime predicate that references only attributes in this set is considered fully pushed + * and will not be evaluated again after the scan. These attributes must also be returned by + * [[filterAttributes]]. + */ + def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty + + /** + * Filters this scan using runtime Catalyst expressions. + * + * The provided expressions must be interpreted as a set of predicates that are ANDed together. + * Implementations may use the expressions to prune initially planned + * [[org.apache.spark.sql.connector.read.InputPartition]]s. + * + * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime. + */ + def filter(expressions: Array[Expression]): Unit + + /** + * Returns the predicates that are pushed to the data source via [[filter]]. + * + * This method does not indicate whether a predicate is fully pushed. Spark infers that from + * [[fullyPushedFilterAttributes]]. The returned predicates may fully or partially help the data + * source prune initially planned + * [[org.apache.spark.sql.connector.read.InputPartition]]s. + * + * It's possible that there are no runtime predicates and [[filter]] is never called; + * an empty array should be returned for this case. + */ + def pushedPredicates(): Array[Expression] = Array.empty +} 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 new file mode 100644 index 0000000000000..8dff2c887eaa8 --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.catalog + +import java.util + +import scala.collection.mutable.ArrayBuffer + +import InMemoryCatalystRuntimeFilterTable._ + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} +import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystRuntimeFiltering +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.util.ArrayImplicits._ + +/** + * In-memory table whose batch scan implements + * [[SupportsPushDownCatalystRuntimeFiltering]], so runtime filters arrive as Catalyst + * [[Expression]]s rather than connector predicates. + * + * Table properties: + * - `filter-attributes` (default: all partition cols): comma-separated list of + * column names to expose from `filterAttributes`. + * - `fully-pushed-filter-attributes` (default: none): comma-separated list of + * column names to expose from `fullyPushedFilterAttributes`. + */ +class InMemoryCatalystRuntimeFilterTable( + name: String, + columns: Array[Column], + partitioning: Array[Transform], + properties: util.Map[String, String]) + extends InMemoryTableWithV2Filter(name, columns, partitioning, properties) { + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + new InMemoryCatalystRuntimeFilterScanBuilder(schema, options) + } + + class InMemoryCatalystRuntimeFilterScanBuilder( + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends InMemoryScanBuilder(tableSchema, options) { + override def build: Scan = InMemoryCatalystRuntimeFilterBatchScan( + data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, + schema, tableSchema, options) + } + + /** + * Scan that receives runtime filters as Catalyst expressions. + * Records what was pushed; pruning is left to the + * [[org.apache.spark.sql.execution.FilterExec]] above the scan, so the recorded + * expressions are the only observable effect. + */ + case class InMemoryCatalystRuntimeFilterBatchScan( + var _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + with SupportsPushDownCatalystRuntimeFiltering { + + private val _catalystPredicates = ArrayBuffer.empty[Expression] + + private val restrictedFilterAttrs: Option[Set[String]] = + Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey)) + .map(_.split(",").map(_.trim).toSet) + + override def filterAttributes: Array[NamedReference] = { + val scanFields = readSchema.fields.map(_.name).toSet + partitioning.flatMap(_.references()).filter { ref => + val name = ref.fieldNames.mkString(".") + scanFields.contains(name) && + restrictedFilterAttrs.forall(_.contains(name)) + } + } + + override def fullyPushedFilterAttributes(): Array[NamedReference] = { + val fullyPushedFilterAttrs = Option( + InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey)) + .map(_.split(",").map(_.trim).toSet) + .getOrElse(Set.empty) + filterAttributes.filter { ref => + fullyPushedFilterAttrs.contains(ref.fieldNames.mkString(".")) + } + } + + override def filter(expressions: Array[Expression]): Unit = + _catalystPredicates ++= expressions + + override def pushedPredicates(): Array[Expression] = + _catalystPredicates.toArray + } +} + +object InMemoryCatalystRuntimeFilterTable { + /** + * Table property: comma-separated column names to expose from + * filterAttributes. Default: all partition columns. + */ + private[catalog] val FilterAttributesKey = "filter-attributes" + + /** + * Table property: comma-separated column names to expose from + * fullyPushedFilterAttributes. Default: none. + */ + private[catalog] val FullyPushedFilterAttributesKey = "fully-pushed-filter-attributes" +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala new file mode 100644 index 0000000000000..b8415eae3e15b --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableCatalystRuntimeFilterCatalog.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.catalog + +import java.util + +import org.apache.spark.sql.catalyst.analysis.TableAlreadyExistsException +import org.apache.spark.sql.connector.expressions.Transform + +class InMemoryTableCatalystRuntimeFilterCatalog extends InMemoryTableCatalog { + import CatalogV2Implicits._ + + override def createTable( + ident: Identifier, + columns: Array[Column], + partitions: Array[Transform], + properties: util.Map[String, String]): Table = { + if (tables.containsKey(ident)) { + throw new TableAlreadyExistsException(ident.asMultipartIdentifier) + } + + InMemoryTableCatalog.maybeSimulateFailedTableCreation(properties) + + val tableName = s"$name.${ident.quoted}" + val table = new InMemoryCatalystRuntimeFilterTable( + tableName, columns, partitions, properties) + tables.put(ident, table) + namespaces.putIfAbsent(ident.namespace.toList, Map()) + table + } + + override def createTable(ident: Identifier, tableInfo: TableInfo): Table = { + createTable(ident, tableInfo.columns(), tableInfo.partitions(), tableInfo.properties) + } +} diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index 72807a5242c7c..bc6347a473253 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -171,7 +171,9 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat // Extract scalar subquery filters on runtime-filterable columns for runtime pushdown. // These filters stay in postScanFilters for correctness (FilterExec above scan), // but are also routed into runtimeFilters so BatchScanExec can use them for - // partition pruning via SupportsRuntimeV2Filtering.filter(). + // partition pruning via SupportsRuntimeV2Filtering.filter(). The exception is filters + // that only reference attributes the scan fully evaluates, which are dropped from + // postScanFilters below. val scalarSubqueryFilters = if (relation.runtimeFilterAttrs.nonEmpty) { postScanFilters.filter { f => f.containsPattern(SCALAR_SUBQUERY) && @@ -181,12 +183,16 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat } else { Seq.empty } + val fullyPushedRuntimeFilters = scalarSubqueryFilters.filter { f => + f.references.subsetOf(relation.fullyPushedRuntimeFilterAttrs) + } val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters, relation.ordering, relation.relation.table, relation.keyGroupedPartitioning) DataSourceV2Strategy.withProjectAndFilter( - project, postScanFilters, batchExec, !batchExec.supportsColumnar) :: Nil + project, postScanFilters.diff(fullyPushedRuntimeFilters), + batchExec, !batchExec.supportsColumnar) :: Nil case PhysicalOperation(p, f, r: StreamingDataSourceV2ScanRelation) if r.startOffset.isDefined && r.endOffset.isDefined => 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 46c4e98595827..387dd9aaa3a03 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 @@ -22,7 +22,7 @@ import scala.collection.mutable import org.apache.spark.SparkException import org.apache.spark.internal.{Logging, LogKeys} import org.apache.spark.sql.AnalysisException -import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, DynamicPruning, DynamicPruningExpression, Expression, ExpressionSet, GetStructField, Literal, NamedExpression, PythonUDF, SchemaPruning, SubqueryExpression, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.logical.SampleMethod import org.apache.spark.sql.catalyst.plans.physical.{KeyedPartitioning, Partitioning} import org.apache.spark.sql.catalyst.types.DataTypeUtils @@ -35,7 +35,7 @@ import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, Sam import org.apache.spark.sql.execution.{InSubqueryExec, ScalarSubquery => ExecScalarSubquery} import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, DataSourceUtils} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters} +import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters, SupportsPushDownCatalystRuntimeFiltering} import org.apache.spark.sql.sources import org.apache.spark.sql.types.{StructField, StructType} import org.apache.spark.util.ArrayImplicits.SparkArrayOps @@ -170,6 +170,10 @@ object PushDownUtils extends Logging { * Note: Do not call multiple times for the same `scan` instance; * [[SupportsRuntimeV2Filtering.filter]] is mutating. * + * A scan implementing [[SupportsPushDownCatalystRuntimeFiltering]] takes a separate path: all + * runtime filters are pushed as Catalyst expressions in a single call, with no translation to + * connector predicates and no `filterAttributes` gating. The two paths are mutually exclusive. + * * @return true if any filters were pushed to the data source */ def pushRuntimeFilters( @@ -213,6 +217,22 @@ object PushDownUtils extends Logging { } translatedFiltersPushed || partPredicatesPushed + + case catalystScan: SupportsPushDownCatalystRuntimeFiltering if runtimeFilters.nonEmpty => + // A DPP filter degrades to TrueLiteral when its subquery is pruned away; it carries no + // information for the source. The V2 path above drops these implicitly because + // translateRuntimeFilterV2 returns None; here we push Catalyst expressions directly, + // so filter them out explicitly. + val catalystFilters = runtimeFilters + .flatMap(unwrapRuntimeFilterExpression) + .filterNot(_ == Literal.TrueLiteral) + if (catalystFilters.nonEmpty) { + catalystScan.filter(catalystFilters.toArray) + true + } else { + false + } + case _ => false } @@ -433,16 +453,20 @@ object PushDownUtils extends Logging { private[v2] def createRuntimePartitionPredicates( runtimeFilters: Seq[Expression], partitionFields: Seq[PartitionPredicateField]): Seq[PartitionPredicateImpl] = { - val catalystExprs = runtimeFilters.flatMap { + val catalystExprs = runtimeFilters.flatMap(unwrapRuntimeFilterExpression) + val flattened = flattenNestedPartitionFilters(catalystExprs, partitionFields).keys + createPartitionPredicates(flattened.toSeq, partitionFields)._1 + } + + /** Unwraps a runtime filter to the Catalyst predicate for pushdown. */ + private def unwrapRuntimeFilterExpression(rf: Expression): Option[Expression] = + rf match { case DynamicPruningExpression(in: InSubqueryExec) if in.isResultUnavailable => None case DynamicPruningExpression(e) => Some(e) case _: DynamicPruning => None case f => Some(f.transform { case s: ExecScalarSubquery => s.toLiteral }) } - val flattened = flattenNestedPartitionFilters(catalystExprs, partitionFields).keys - createPartitionPredicates(flattened.toSeq, partitionFields)._1 - } private def isPushablePartitionFilter(f: Expression) = f.deterministic && 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 0e885b994edec..efa387360c5ed 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 @@ -28,6 +28,7 @@ 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.SupportsPushDownCatalystRuntimeFiltering /** * Dynamic partition pruning optimization is performed based on the type and * selectivity of the join operation. During query optimization, we insert a @@ -86,6 +87,14 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } else { None } + case (resExp, r @ ExtractV2Scan(scan: SupportsPushDownCatalystRuntimeFiltering)) => + val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( + scan.filterAttributes, r.output) + if (resExp.references.subsetOf(filterAttrs)) { + Some(r) + } else { + None + } case _ => None } } 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 new file mode 100644 index 0000000000000..2cd9cd9d5c3ec --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -0,0 +1,289 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.scalatest.BeforeAndAfter + +import org.apache.spark.sql.{DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GreaterThan, Literal} +import org.apache.spark.sql.connector.catalog.{InMemoryCatalystRuntimeFilterTable, InMemoryTableCatalystRuntimeFilterCatalog} +import org.apache.spark.sql.execution.{FilterExec, ScalarSubquery => ExecScalarSubquery} +import org.apache.spark.sql.execution.ExplainUtils.stripAQEPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.IntegerType + +/** + * Tests for scans that implement + * [[org.apache.spark.sql.internal.connector.SupportsPushDownCatalystRuntimeFiltering]], + * where runtime filters are pushed once as Catalyst expressions instead of connector + * predicates. + */ +class DataSourceV2CatalystRuntimeFilterSuite + extends SharedSparkSession with BeforeAndAfter { + + protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + protected val catalogName = "testcatalystruntimefilter" + + before { + spark.conf.set(s"spark.sql.catalog.$catalogName", + classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) + } + + after { + spark.sessionState.catalogManager.reset() + } + + private def withDPPConf(f: => Unit): Unit = { + 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")(f) + } + + test("scalar subquery on partition column -> pushed as Catalyst expression") { + val tbl = s"$catalogName.tbl1" + val dim = s"$catalogName.dim1" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") + checkAnswer(df, Row(3, 3)) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3))) + // `part` is not declared fully pushed, so Spark still evaluates the filter after the scan. + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + } + } + + test("predicate on fully pushed filter attributes -> not evaluated after the scan") { + val tbl = s"$catalogName.tbl_fully_pushed" + val dim = s"$catalogName.dim_fully_pushed" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, 3)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") + checkAnswer(df, (0 until 5).map(i => Row(i, 3))) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(part, Literal(3))) + assertScalarSubqueryEvaluatedAfterScan(df, expected = false) + } + } + + test("untranslatable filter -> pushed instead of dropped") { + val tbl = s"$catalogName.tbl2" + val dim = s"$catalogName.dim2" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (2)") + + // `part > sub + 1` has no data source V2 translation, so the V2 interfaces would never + // see it. The scalar subquery is literalized but the surrounding expression is kept. + val df = sql(s"SELECT * FROM $tbl WHERE part > (SELECT max(val) FROM $dim) + 1") + checkAnswer(df, Row(4, 4)) + + assertScalarSubqueryRuntimeFilters(df) + val part = AttributeReference("part", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual( + df, GreaterThan(part, Add(Literal(2), Literal(1)))) + } + } + + test("DPP filter -> pushed as InSubqueryExec expression") { + val fact = s"$catalogName.fact3" + val dim = s"$catalogName.dim3" + withTable(fact, dim) { + sql(s"CREATE TABLE $fact (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + 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) + val dppPredicate = collectBatchScan(df).runtimeFilters.collectFirst { + case DynamicPruningExpression(e) => e + }.get + assertPushedCatalystPredicatesEqual(df, dppPredicate) + } + } + } + + test("filter on column outside filterAttributes -> not pushed") { + val tbl = s"$catalogName.tbl4" + val dim = s"$catalogName.dim4" + 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, 10)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (10)") + + // p2 is a partition column but is not declared filterable, so no runtime filter is derived. + val df = sql(s"SELECT * FROM $tbl WHERE p2 = (SELECT max(val) FROM $dim)") + checkAnswer(df, (0 until 5).map(i => Row(i, i, 10))) + + assert(collectBatchScan(df).runtimeFilters.isEmpty, + "Expected no runtime filters for a column outside filterAttributes") + assertPushedCatalystPredicates(df, 0) + } + } + + test("no runtime filter -> filter() is never called") { + val tbl = s"$catalogName.tbl5" + withTable(tbl) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part)") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + + val df = sql(s"SELECT * FROM $tbl WHERE part = 3") + checkAnswer(df, Row(3, 3)) + + assert(collectBatchScan(df).runtimeFilters.isEmpty) + assertPushedCatalystPredicates(df, 0) + } + } + + // --------------------------------------------------------------------------- + // Helper methods + // --------------------------------------------------------------------------- + + private def assertDPPRuntimeFilters( + df: DataFrame, expectedCount: Int = 1): Unit = { + val batchScan = collectBatchScan(df) + val dppFilters = batchScan.runtimeFilters.collect { + case d: DynamicPruningExpression => d + } + assert(dppFilters.size === expectedCount, + s"Expected $expectedCount DynamicPruningExpression(s) " + + s"in runtimeFilters, got ${dppFilters.size}") + } + + private def assertScalarSubqueryRuntimeFilters( + df: DataFrame, expectedCount: Int = 1): Unit = { + val batchScan = collectBatchScan(df) + val scalarFilters = batchScan.runtimeFilters.collect { + case f if !f.isInstanceOf[DynamicPruning] => f + } + val dppFilters = batchScan.runtimeFilters.collect { + case d: DynamicPruning => d + } + assert(scalarFilters.size === expectedCount, + s"Expected $expectedCount scalar subquery runtime filter(s), " + + s"got ${scalarFilters.size}") + assert(dppFilters.isEmpty, + "Expected non-DPP runtime filters (scalar subquery)") + } + + /** + * Checks whether a scalar subquery runtime filter is still evaluated by a [[FilterExec]] above + * the scan. Filters that only reference `fullyPushedFilterAttributes` are dropped from it. + */ + private def assertScalarSubqueryEvaluatedAfterScan( + df: DataFrame, + expected: Boolean): Unit = { + val postScanConditions = stripAQEPlan(df.queryExecution.executedPlan).collect { + case f: FilterExec => f.condition + } + val evaluated = postScanConditions.exists(_.exists(_.isInstanceOf[ExecScalarSubquery])) + assert(evaluated === expected, + s"Expected scalar subquery evaluated after scan to be $expected, " + + s"post-scan filter conditions: $postScanConditions") + } + + private def collectBatchScan(df: DataFrame): BatchScanExec = { + stripAQEPlan(df.queryExecution.executedPlan).collectFirst { + case b: BatchScanExec => b + }.getOrElse(fail("Expected BatchScanExec in plan")) + } + + private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { + collectBatchScan(df).scan match { + case s: InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan => + s.pushedPredicates().toSeq + case _ => Seq.empty + } + } + + private def assertPushedCatalystPredicates(df: DataFrame, expected: Int): Unit = { + val preds = getPushedCatalystPredicates(df) + assert(preds.size === expected, + s"Expected $expected pushed Catalyst runtime predicate(s), got ${preds.size}: $preds") + } + + /** + * Binds [[AttributeReference]]s in `expected` to the scan output (by name) and checks that the + * pushed Catalyst runtime predicates match exactly via [[Expression.semanticEquals]]. + */ + private def assertPushedCatalystPredicatesEqual( + df: DataFrame, + expected: Expression*): Unit = { + val batchScan = collectBatchScan(df) + val actual = getPushedCatalystPredicates(df) + val normalizedExpected = expected.map(bindToScanOutput(_, batchScan.output)) + assert(actual.size === normalizedExpected.size, + s"Expected ${normalizedExpected.size} pushed Catalyst predicate(s), " + + s"got ${actual.size}: $actual") + actual.zip(normalizedExpected).foreach { case (a, e) => + assert(a.semanticEquals(e), + s"Pushed Catalyst predicate mismatch.\nExpected: $e\nActual: $a") + } + } + + private def bindToScanOutput( + expr: Expression, + output: Seq[AttributeReference]): Expression = { + val resolver = SQLConf.get.resolver + expr.transformUp { + case a: AttributeReference => + output.find(o => resolver(o.name, a.name)) + .map(_.withNullability(a.nullable).withQualifier(a.qualifier)) + .getOrElse(a) + } + } +} From 9d6dab16941ac9ab2b3d9944aedef68cebcf6dec Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Mon, 3 Aug 2026 20:23:10 +0000 Subject: [PATCH 2/9] Rename SupportsPushDownCatalystRuntimeFiltering to SupportsRuntimeCatalystFiltering In DSv2 the SupportsPushDown* prefix is used by ScanBuilder mix-ins, whose pushdown happens at query compilation. Runtime filtering interfaces are Scan mix-ins named SupportsRuntimeFiltering, so align the new trait with SupportsRuntimeV2Filtering instead. Also drop the empty parameter lists on the trait's accessors so all of them agree with filterAttributes. --- .../datasources/v2/DataSourceV2Relation.scala | 8 ++++---- ...ng.scala => SupportsRuntimeCatalystFiltering.scala} | 6 +++--- .../catalog/InMemoryCatalystRuntimeFilterTable.scala | 10 +++++----- .../sql/execution/datasources/v2/PushDownUtils.scala | 6 +++--- .../execution/dynamicpruning/PartitionPruning.scala | 4 ++-- .../DataSourceV2CatalystRuntimeFilterSuite.scala | 4 ++-- 6 files changed, 19 insertions(+), 19 deletions(-) rename sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/{SupportsPushDownCatalystRuntimeFiltering.scala => SupportsRuntimeCatalystFiltering.scala} (94%) 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 cc6e1f33e1622..f9e0c320ec724 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 @@ -35,7 +35,7 @@ import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReferenc import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.read.colstats.{ColumnStatistics, Histogram => V2Histogram, HistogramBin => V2HistogramBin} import org.apache.spark.sql.connector.read.streaming.{Offset, SparkDataStream} -import org.apache.spark.sql.internal.connector.{SupportsPushDownCatalystRuntimeFiltering, V2StatisticsUtils} +import org.apache.spark.sql.internal.connector.{SupportsRuntimeCatalystFiltering, V2StatisticsUtils} import org.apache.spark.sql.types.{DataType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ @@ -197,13 +197,13 @@ case class DataSourceV2ScanRelation( /** * Resolved attributes that the scan declares for runtime filtering via * [[SupportsRuntimeV2Filtering.filterAttributes]] or - * [[SupportsPushDownCatalystRuntimeFiltering.filterAttributes]]. Empty when the scan + * [[SupportsRuntimeCatalystFiltering.filterAttributes]]. Empty when the scan * implements neither interface or exposes no attributes. */ lazy val runtimeFilterAttrs: AttributeSet = { val filterAttrs = scan match { case s: SupportsRuntimeV2Filtering => s.filterAttributes - case s: SupportsPushDownCatalystRuntimeFiltering => s.filterAttributes + case s: SupportsRuntimeCatalystFiltering => s.filterAttributes case _ => Array.empty[NamedReference] } AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( @@ -215,7 +215,7 @@ case class DataSourceV2ScanRelation( */ lazy val fullyPushedRuntimeFilterAttrs: AttributeSet = { val filterAttrs = scan match { - case s: SupportsPushDownCatalystRuntimeFiltering => s.fullyPushedFilterAttributes() + case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes case _ => Array.empty[NamedReference] } AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala similarity index 94% rename from sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala rename to sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala index dea28b6db7001..fffbfde9ea7eb 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsPushDownCatalystRuntimeFiltering.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala @@ -33,7 +33,7 @@ import org.apache.spark.sql.connector.read.Scan * * Note that Spark will push runtime filters only if they are beneficial. */ -trait SupportsPushDownCatalystRuntimeFiltering extends Scan { +trait SupportsRuntimeCatalystFiltering extends Scan { /** * Returns attributes this scan can be filtered by at runtime. @@ -49,7 +49,7 @@ trait SupportsPushDownCatalystRuntimeFiltering extends Scan { * and will not be evaluated again after the scan. These attributes must also be returned by * [[filterAttributes]]. */ - def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty + def fullyPushedFilterAttributes: Array[NamedReference] = Array.empty /** * Filters this scan using runtime Catalyst expressions. @@ -73,5 +73,5 @@ trait SupportsPushDownCatalystRuntimeFiltering extends Scan { * It's possible that there are no runtime predicates and [[filter]] is never called; * an empty array should be returned for this case. */ - def pushedPredicates(): Array[Expression] = Array.empty + def pushedPredicates: Array[Expression] = Array.empty } 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 8dff2c887eaa8..0033371200ab3 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 @@ -26,14 +26,14 @@ import InMemoryCatalystRuntimeFilterTable._ import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} -import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystRuntimeFiltering +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.util.ArrayImplicits._ /** * In-memory table whose batch scan implements - * [[SupportsPushDownCatalystRuntimeFiltering]], so runtime filters arrive as Catalyst + * [[SupportsRuntimeCatalystFiltering]], so runtime filters arrive as Catalyst * [[Expression]]s rather than connector predicates. * * Table properties: @@ -74,7 +74,7 @@ class InMemoryCatalystRuntimeFilterTable( tableSchema: StructType, options: CaseInsensitiveStringMap) extends BatchScanBaseClass(_data, readSchema, tableSchema) - with SupportsPushDownCatalystRuntimeFiltering { + with SupportsRuntimeCatalystFiltering { private val _catalystPredicates = ArrayBuffer.empty[Expression] @@ -91,7 +91,7 @@ class InMemoryCatalystRuntimeFilterTable( } } - override def fullyPushedFilterAttributes(): Array[NamedReference] = { + override def fullyPushedFilterAttributes: Array[NamedReference] = { val fullyPushedFilterAttrs = Option( InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey)) .map(_.split(",").map(_.trim).toSet) @@ -104,7 +104,7 @@ class InMemoryCatalystRuntimeFilterTable( override def filter(expressions: Array[Expression]): Unit = _catalystPredicates ++= expressions - override def pushedPredicates(): Array[Expression] = + override def pushedPredicates: Array[Expression] = _catalystPredicates.toArray } } 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 387dd9aaa3a03..ac6ca3f97ce3a 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 @@ -35,7 +35,7 @@ import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, Sam import org.apache.spark.sql.execution.{InSubqueryExec, ScalarSubquery => ExecScalarSubquery} import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, DataSourceUtils} import org.apache.spark.sql.internal.SQLConf -import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters, SupportsPushDownCatalystRuntimeFiltering} +import org.apache.spark.sql.internal.connector.{PartitionPredicateField, PartitionPredicateImpl, SupportsPushDownCatalystFilters, SupportsRuntimeCatalystFiltering} import org.apache.spark.sql.sources import org.apache.spark.sql.types.{StructField, StructType} import org.apache.spark.util.ArrayImplicits.SparkArrayOps @@ -170,7 +170,7 @@ object PushDownUtils extends Logging { * Note: Do not call multiple times for the same `scan` instance; * [[SupportsRuntimeV2Filtering.filter]] is mutating. * - * A scan implementing [[SupportsPushDownCatalystRuntimeFiltering]] takes a separate path: all + * A scan implementing [[SupportsRuntimeCatalystFiltering]] takes a separate path: all * runtime filters are pushed as Catalyst expressions in a single call, with no translation to * connector predicates and no `filterAttributes` gating. The two paths are mutually exclusive. * @@ -218,7 +218,7 @@ object PushDownUtils extends Logging { translatedFiltersPushed || partPredicatesPushed - case catalystScan: SupportsPushDownCatalystRuntimeFiltering if runtimeFilters.nonEmpty => + case catalystScan: SupportsRuntimeCatalystFiltering if runtimeFilters.nonEmpty => // A DPP filter degrades to TrueLiteral when its subquery is pruned away; it carries no // information for the source. The V2 path above drops these implicitly because // translateRuntimeFilterV2 returns None; here we push Catalyst expressions directly, 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 efa387360c5ed..056fccba92c48 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 @@ -28,7 +28,7 @@ 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.SupportsPushDownCatalystRuntimeFiltering +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering /** * Dynamic partition pruning optimization is performed based on the type and * selectivity of the join operation. During query optimization, we insert a @@ -87,7 +87,7 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } else { None } - case (resExp, r @ ExtractV2Scan(scan: SupportsPushDownCatalystRuntimeFiltering)) => + case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) => val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( scan.filterAttributes, r.output) if (resExp.references.subsetOf(filterAttrs)) { 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 2cd9cd9d5c3ec..9aac5a25d2d29 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 @@ -31,7 +31,7 @@ import org.apache.spark.sql.types.IntegerType /** * Tests for scans that implement - * [[org.apache.spark.sql.internal.connector.SupportsPushDownCatalystRuntimeFiltering]], + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], * where runtime filters are pushed once as Catalyst expressions instead of connector * predicates. */ @@ -245,7 +245,7 @@ class DataSourceV2CatalystRuntimeFilterSuite private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { collectBatchScan(df).scan match { case s: InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan => - s.pushedPredicates().toSeq + s.pushedPredicates.toSeq case _ => Seq.empty } } From 51daf4aa74067ce7d682384cf7819caec0d6691e Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Mon, 3 Aug 2026 22:26:20 +0000 Subject: [PATCH 3/9] Use empty parameter lists on SupportsRuntimeCatalystFiltering accessors Matches the Java runtime filtering interfaces, where filterAttributes() and pushedPredicates() carry empty parameter lists, and keeps the trait's own members consistent with each other. --- .../execution/datasources/v2/DataSourceV2Relation.scala | 4 ++-- .../connector/SupportsRuntimeCatalystFiltering.scala | 6 +++--- .../catalog/InMemoryCatalystRuntimeFilterTable.scala | 8 ++++---- .../sql/execution/dynamicpruning/PartitionPruning.scala | 2 +- .../DataSourceV2CatalystRuntimeFilterSuite.scala | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) 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 f9e0c320ec724..bcc2b6041be25 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 @@ -203,7 +203,7 @@ case class DataSourceV2ScanRelation( lazy val runtimeFilterAttrs: AttributeSet = { val filterAttrs = scan match { case s: SupportsRuntimeV2Filtering => s.filterAttributes - case s: SupportsRuntimeCatalystFiltering => s.filterAttributes + case s: SupportsRuntimeCatalystFiltering => s.filterAttributes() case _ => Array.empty[NamedReference] } AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( @@ -215,7 +215,7 @@ case class DataSourceV2ScanRelation( */ lazy val fullyPushedRuntimeFilterAttrs: AttributeSet = { val filterAttrs = scan match { - case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes + case s: SupportsRuntimeCatalystFiltering => s.fullyPushedFilterAttributes() case _ => Array.empty[NamedReference] } AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( 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 fffbfde9ea7eb..406b888c2478d 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 @@ -40,7 +40,7 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes. */ - def filterAttributes: Array[NamedReference] + def filterAttributes(): Array[NamedReference] /** * Returns attributes for which this scan fully evaluates runtime predicates. @@ -49,7 +49,7 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * and will not be evaluated again after the scan. These attributes must also be returned by * [[filterAttributes]]. */ - def fullyPushedFilterAttributes: Array[NamedReference] = Array.empty + def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty /** * Filters this scan using runtime Catalyst expressions. @@ -73,5 +73,5 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * It's possible that there are no runtime predicates and [[filter]] is never called; * an empty array should be returned for this case. */ - def pushedPredicates: Array[Expression] = Array.empty + def pushedPredicates(): Array[Expression] = Array.empty } 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 0033371200ab3..3a728c10b034d 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 @@ -82,7 +82,7 @@ class InMemoryCatalystRuntimeFilterTable( Option(InMemoryCatalystRuntimeFilterTable.this.properties.get(FilterAttributesKey)) .map(_.split(",").map(_.trim).toSet) - override def filterAttributes: Array[NamedReference] = { + override def filterAttributes(): Array[NamedReference] = { val scanFields = readSchema.fields.map(_.name).toSet partitioning.flatMap(_.references()).filter { ref => val name = ref.fieldNames.mkString(".") @@ -91,12 +91,12 @@ class InMemoryCatalystRuntimeFilterTable( } } - override def fullyPushedFilterAttributes: Array[NamedReference] = { + override def fullyPushedFilterAttributes(): Array[NamedReference] = { val fullyPushedFilterAttrs = Option( InMemoryCatalystRuntimeFilterTable.this.properties.get(FullyPushedFilterAttributesKey)) .map(_.split(",").map(_.trim).toSet) .getOrElse(Set.empty) - filterAttributes.filter { ref => + filterAttributes().filter { ref => fullyPushedFilterAttrs.contains(ref.fieldNames.mkString(".")) } } @@ -104,7 +104,7 @@ class InMemoryCatalystRuntimeFilterTable( override def filter(expressions: Array[Expression]): Unit = _catalystPredicates ++= expressions - override def pushedPredicates: Array[Expression] = + override def pushedPredicates(): Array[Expression] = _catalystPredicates.toArray } } 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 056fccba92c48..1bf12a695bc4c 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 @@ -89,7 +89,7 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) => val filterAttrs = V2ExpressionUtils.resolveAttributeRefs( - scan.filterAttributes, r.output) + scan.filterAttributes(), r.output) if (resExp.references.subsetOf(filterAttrs)) { Some(r) } else { 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 9aac5a25d2d29..3a0a83c8bb588 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 @@ -245,7 +245,7 @@ class DataSourceV2CatalystRuntimeFilterSuite private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { collectBatchScan(df).scan match { case s: InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan => - s.pushedPredicates.toSeq + s.pushedPredicates().toSeq case _ => Seq.empty } } From c181f4ce052cc852f2525bce1e5df9483a097edc Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 00:02:40 +0000 Subject: [PATCH 4/9] Register the test catalog via sparkConf instead of before/after The catalog only needs to be registered once for the suite's session, and every test already scopes its tables with withTable, so the per-test registration and catalogManager.reset() are not needed. --- .../DataSourceV2CatalystRuntimeFilterSuite.scala | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) 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 3a0a83c8bb588..4b08a89f8cbae 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 @@ -17,8 +17,7 @@ package org.apache.spark.sql.connector -import org.scalatest.BeforeAndAfter - +import org.apache.spark.SparkConf import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GreaterThan, Literal} import org.apache.spark.sql.connector.catalog.{InMemoryCatalystRuntimeFilterTable, InMemoryTableCatalystRuntimeFilterCatalog} @@ -35,20 +34,14 @@ import org.apache.spark.sql.types.IntegerType * where runtime filters are pushed once as Catalyst expressions instead of connector * predicates. */ -class DataSourceV2CatalystRuntimeFilterSuite - extends SharedSparkSession with BeforeAndAfter { +class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName protected val catalogName = "testcatalystruntimefilter" - before { - spark.conf.set(s"spark.sql.catalog.$catalogName", + override def sparkConf: SparkConf = super.sparkConf + .set(s"spark.sql.catalog.$catalogName", classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) - } - - after { - spark.sessionState.catalogManager.reset() - } private def withDPPConf(f: => Unit): Unit = { withSQLConf( From 6563ae2e8aba97edc3e4d11e4c771186db611458 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 00:09:33 +0000 Subject: [PATCH 5/9] Add test for predicate referencing partly fully pushed attributes Pins down that a runtime predicate is only dropped from the post-scan filter when every attribute it references is declared fully pushed. --- ...taSourceV2CatalystRuntimeFilterSuite.scala | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) 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 4b08a89f8cbae..7c349077c622a 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 @@ -94,6 +94,32 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { } } + test("predicate on partly fully pushed filter attributes -> evaluated after the scan") { + val tbl = s"$catalogName.tbl_partly_pushed" + val dim = s"$catalogName.dim_partly_pushed" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, p1 INT, p2 INT) USING $v2Source " + + "PARTITIONED BY (p1, p2) " + + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'p1')") + for (i <- 0 until 5) { + sql(s"INSERT INTO $tbl VALUES ($i, 1, 2)") + } + sql(s"CREATE TABLE $dim (val INT) USING $v2Source") + sql(s"INSERT INTO $dim VALUES (3)") + + // The predicate also references p2, which is not declared fully pushed, so it is not + // considered fully pushed and Spark keeps evaluating it after the scan. + val df = sql(s"SELECT * FROM $tbl WHERE p1 + p2 = (SELECT max(val) FROM $dim)") + checkAnswer(df, (0 until 5).map(i => Row(i, 1, 2))) + + assertScalarSubqueryRuntimeFilters(df) + val p1 = AttributeReference("p1", IntegerType, nullable = false)() + val p2 = AttributeReference("p2", IntegerType, nullable = false)() + assertPushedCatalystPredicatesEqual(df, EqualTo(Add(p1, p2), Literal(3))) + assertScalarSubqueryEvaluatedAfterScan(df, expected = true) + } + } + test("untranslatable filter -> pushed instead of dropped") { val tbl = s"$catalogName.tbl2" val dim = s"$catalogName.dim2" From c3fc23e72fd29b7342affe1b4072501faa5b8421 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 18:23:46 -0700 Subject: [PATCH 6/9] Evaluate pushed predicates in the Catalyst runtime filtering test scan The fixture only recorded what was pushed, so the fully-pushed test could not catch an incorrect post-scan filter removal. Bind and interpret predicates over partition columns to prune partitions, and test with both matching and nonmatching partitions. Fail instead of returning empty on an unexpected scan type. --- .../InMemoryCatalystRuntimeFilterTable.scala | 50 ++++++++++++++++--- ...taSourceV2CatalystRuntimeFilterSuite.scala | 11 ++-- 2 files changed, 50 insertions(+), 11 deletions(-) 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 3a728c10b034d..b696c09fc07bf 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 @@ -23,9 +23,10 @@ import scala.collection.mutable.ArrayBuffer import InMemoryCatalystRuntimeFilterTable._ -import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Expression, Predicate => CatalystPredicate} import org.apache.spark.sql.connector.expressions.{NamedReference, Transform} import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap @@ -64,9 +65,9 @@ class InMemoryCatalystRuntimeFilterTable( /** * Scan that receives runtime filters as Catalyst expressions. - * Records what was pushed; pruning is left to the - * [[org.apache.spark.sql.execution.FilterExec]] above the scan, so the recorded - * expressions are the only observable effect. + * Records what was pushed and evaluates expressions that reference only + * partition columns against each partition key, so fully-pushed predicates + * are enforced when Spark drops the post-scan [[org.apache.spark.sql.execution.FilterExec]]. */ case class InMemoryCatalystRuntimeFilterBatchScan( var _data: Seq[InputPartition], @@ -101,11 +102,46 @@ class InMemoryCatalystRuntimeFilterTable( } } - override def filter(expressions: Array[Expression]): Unit = + override def filter(expressions: Array[Expression]): Unit = { _catalystPredicates ++= expressions + val partAttrs = partitionAttributes + if (partAttrs.isEmpty) return + + val resolver = SQLConf.get.resolver + expressions.foreach { expr => + val remapped = expr.transform { + case a: AttributeReference => + partAttrs.find(p => resolver(p.name, a.name)).getOrElse(a) + } + // 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 => partAttrs.exists(_.exprId == r.exprId))) { + val bound = BindReferences.bindReference(remapped, partAttrs) + val pred = CatalystPredicate.createInterpreted(bound) + data = data.filter { p => + try { + pred.eval(p.asInstanceOf[BufferedRows].partitionKey()) + } catch { + // Keep the partition on eval failure, matching PartitionPredicateImpl. + case _: Exception => true + } + } + } + } + } + + /** Predicates recorded by [[filter]], for test assertions only. */ + def pushedCatalystPredicates: Seq[Expression] = _catalystPredicates.toSeq - override def pushedPredicates(): Array[Expression] = - _catalystPredicates.toArray + /** AttributeReferences matching the partition-key InternalRow field order. */ + private def partitionAttributes: Seq[AttributeReference] = { + partitioning.flatMap(_.references()).flatMap { ref => + val name = ref.fieldNames.mkString(".") + readSchema.find(_.name == name).orElse(tableSchema.find(_.name == name)).map { f => + AttributeReference(f.name, f.dataType, f.nullable)() + } + }.toSeq + } } } 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 7c349077c622a..e07367e538a9a 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 @@ -78,14 +78,16 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { withTable(tbl, dim) { sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Source PARTITIONED BY (part) " + "TBLPROPERTIES('fully-pushed-filter-attributes' = 'part')") + // Matching and nonmatching partitions: the scan must prune nonmatching ones itself + // because Spark drops the post-scan FilterExec for fully pushed attributes. for (i <- 0 until 5) { - sql(s"INSERT INTO $tbl VALUES ($i, 3)") + sql(s"INSERT INTO $tbl VALUES ($i, $i)") } sql(s"CREATE TABLE $dim (val INT) USING $v2Source") sql(s"INSERT INTO $dim VALUES (3)") val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") - checkAnswer(df, (0 until 5).map(i => Row(i, 3))) + checkAnswer(df, Row(3, 3)) assertScalarSubqueryRuntimeFilters(df) val part = AttributeReference("part", IntegerType, nullable = false)() @@ -264,8 +266,9 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { private def getPushedCatalystPredicates(df: DataFrame): Seq[Expression] = { collectBatchScan(df).scan match { case s: InMemoryCatalystRuntimeFilterTable#InMemoryCatalystRuntimeFilterBatchScan => - s.pushedPredicates().toSeq - case _ => Seq.empty + s.pushedCatalystPredicates + case other => + fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other") } } From 5623e3024ac9d83ebb5ca378b6dad388e0e5ac22 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 18:23:56 -0700 Subject: [PATCH 7/9] Document runtime filter attribute contracts and drop pushedPredicates Filter attributes must be top-level fields present in the read schema, otherwise they fail to resolve when Spark builds the scan relation. The same requirement is undocumented on SupportsRuntimeFiltering and SupportsRuntimeV2Filtering, so note it there too. Carry over the partitioning preservation contract from SupportsRuntimeV2Filtering.filter, which replanWithRuntimeFilters enforces for this interface as well. Nothing in Spark reads pushedPredicates on this trait, so drop it and let the test fixture expose its own accessor. --- .../read/SupportsRuntimeFiltering.java | 4 +++ .../read/SupportsRuntimeV2Filtering.java | 4 +++ .../SupportsRuntimeCatalystFiltering.scala | 32 +++++++++++-------- .../datasources/v2/DataSourceV2Strategy.scala | 2 +- 4 files changed, 28 insertions(+), 14 deletions(-) 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 a1613b50b07e1..88107a6ae7d08 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 @@ -39,6 +39,10 @@ public interface SupportsRuntimeFiltering extends SupportsRuntimeV2Filtering { *

* Spark will call {@link #filter(Filter[])} if it can derive a runtime * predicate for any of the filter attributes. + *

+ * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); 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 2939b5b574081..f23de66fdc0af 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 @@ -50,6 +50,10 @@ public interface SupportsRuntimeV2Filtering extends Scan { *

* Spark will call {@link #filter(Predicate[])} if it can derive a runtime * predicate for any of the filter attributes. + *

+ * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); 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 406b888c2478d..6823757595b54 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 @@ -39,6 +39,9 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * Returns attributes this scan can be filtered by at runtime. * * Spark will call [[filter]] if it can derive a runtime filter for any of these attributes. + * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested + * references and attributes pruned out of the read schema fail to resolve when Spark builds + * the scan relation. */ def filterAttributes(): Array[NamedReference] @@ -48,6 +51,10 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * Any runtime predicate that references only attributes in this set is considered fully pushed * and will not be evaluated again after the scan. These attributes must also be returned by * [[filterAttributes]]. + * + * Each reference must be a top-level attribute present in [[Scan.readSchema]]. Nested + * references and attributes pruned out of the read schema fail to resolve when Spark builds + * the scan relation. */ def fullyPushedFilterAttributes(): Array[NamedReference] = Array.empty @@ -58,20 +65,19 @@ trait SupportsRuntimeCatalystFiltering extends Scan { * Implementations may use the expressions to prune initially planned * [[org.apache.spark.sql.connector.read.InputPartition]]s. * + * If the scan also implements + * [[org.apache.spark.sql.connector.read.SupportsReportPartitioning]], it must preserve + * the originally reported partitioning during runtime filtering. While applying runtime + * predicates, the scan may detect that some + * [[org.apache.spark.sql.connector.read.InputPartition]]s have no matching data, in which + * case it can either replace the initially planned + * [[org.apache.spark.sql.connector.read.InputPartition]]s that have no matching data with + * empty [[org.apache.spark.sql.connector.read.InputPartition]]s, or report only a subset of + * the original partition values (omitting those with no data) via + * [[org.apache.spark.sql.connector.read.Batch#planInputPartitions]]. The scan must not + * report new partition values that were not present in the original partitioning. + * * Note that Spark will call [[Scan.toBatch]] again after filtering the scan at runtime. */ def filter(expressions: Array[Expression]): Unit - - /** - * Returns the predicates that are pushed to the data source via [[filter]]. - * - * This method does not indicate whether a predicate is fully pushed. Spark infers that from - * [[fullyPushedFilterAttributes]]. The returned predicates may fully or partially help the data - * source prune initially planned - * [[org.apache.spark.sql.connector.read.InputPartition]]s. - * - * It's possible that there are no runtime predicates and [[filter]] is never called; - * an empty array should be returned for this case. - */ - def pushedPredicates(): Array[Expression] = Array.empty } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala index bc6347a473253..bc6630b3c6d61 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala @@ -171,7 +171,7 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat // Extract scalar subquery filters on runtime-filterable columns for runtime pushdown. // These filters stay in postScanFilters for correctness (FilterExec above scan), // but are also routed into runtimeFilters so BatchScanExec can use them for - // partition pruning via SupportsRuntimeV2Filtering.filter(). The exception is filters + // partition pruning via SupportsRuntimeV2Filtering.filter(). The exceptions are filters // that only reference attributes the scan fully evaluates, which are dropped from // postScanFilters below. val scalarSubqueryFilters = if (relation.runtimeFilterAttrs.nonEmpty) { From d4141e417250170e1d09b2ebd4a99de651befad3 Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 18:23:58 -0700 Subject: [PATCH 8/9] Screen non-deterministic filters in the Catalyst runtime filtering branch Apply the same pushability guard as the V2 PartitionPredicate path so a non-deterministic filter cannot become the sole evaluator when fullyPushedFilterAttributes drops the post-scan FilterExec. --- .../spark/sql/execution/datasources/v2/PushDownUtils.scala | 4 ++++ 1 file changed, 4 insertions(+) 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 ac6ca3f97ce3a..d0c105ac59283 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 @@ -223,9 +223,13 @@ object PushDownUtils extends Logging { // information for the source. The V2 path above drops these implicitly because // translateRuntimeFilterV2 returns None; here we push Catalyst expressions directly, // so filter them out explicitly. + // Screen with the same pushability guard as the V2 PartitionPredicate path + // (deterministic, no subquery, no Python UDF). Keeps non-deterministic filters + // from being the sole evaluator when fullyPushedFilterAttributes drops FilterExec. val catalystFilters = runtimeFilters .flatMap(unwrapRuntimeFilterExpression) .filterNot(_ == Literal.TrueLiteral) + .filter(isPushablePartitionFilter) if (catalystFilters.nonEmpty) { catalystScan.filter(catalystFilters.toArray) true From d04741492cbc38ab5193ef54cb2fc33f8afed69e Mon Sep 17 00:00:00 2001 From: Szehon Ho Date: Tue, 4 Aug 2026 18:24:09 -0700 Subject: [PATCH 9/9] Inject row-level group filters for Catalyst runtime filtering scans RowLevelOperationRuntimeGroupFiltering only matched SupportsRuntimeV2Filtering, so a source adopting the Catalyst interface silently lost runtime group filtering for DELETE, UPDATE and MERGE. The rule needs nothing but filterAttributes and injects a plain DynamicPruningExpression, so add the two cases and type the helpers on the attributes rather than the interface. Cover both group-based and delta-based operations. A delta-based DELETE does not scan the group key, so it asserts that no filter is injected. --- .../connector/catalog/InMemoryBaseTable.scala | 36 ++- .../InMemoryRowLevelOperationTable.scala | 113 ++++++++- ...wLevelOperationRuntimeGroupFiltering.scala | 37 ++- ...lOperationCatalystRuntimeFilterSuite.scala | 50 ++++ ...lOperationCatalystRuntimeFilterSuite.scala | 60 +++++ ...rationCatalystRuntimeFilterSuiteBase.scala | 229 ++++++++++++++++++ 6 files changed, 491 insertions(+), 34 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala 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 fba80eb3d4cbe..662fe826993b2 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 @@ -521,23 +521,35 @@ abstract class InMemoryBaseTable( private var _pushedFilters: Array[Filter] = Array.empty override def build: Scan = { - val scan = if (InMemoryBaseTable.this.ordering.nonEmpty) { - new InMemoryBatchScanWithOrdering( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } else { - InMemoryBatchScan( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } - if (evaluableFilters.nonEmpty) { - scan.filter(evaluableFilters) + val scan = createScan( + data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, options) + scan match { + case s: InMemoryBatchScan => + if (evaluableFilters.nonEmpty) { + s.filter(evaluableFilters) + } + s.pushedFilters = _pushedFilters + case _ => } - scan.pushedFilters = _pushedFilters recordScanEvent(_pushedFilters) scan } + /** + * Creates the batch scan for [[build]]. + */ + protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (InMemoryBaseTable.this.ordering.nonEmpty) { + new InMemoryBatchScanWithOrdering(partitions, readSchema, tableSchema, options) + } else { + InMemoryBatchScan(partitions, readSchema, tableSchema, options) + } + } + override def pruneColumns(requiredSchema: StructType): Unit = { // The required schema could contain conflict-renamed metadata columns, so we need to match // them by their logical (original) names, not their current names. 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 fd02f926e8141..6de76647a0859 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 @@ -19,15 +19,19 @@ package org.apache.spark.sql.connector.catalog import java.util +import scala.collection.mutable.ArrayBuffer + import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Expression, GenericInternalRow, Predicate => CatalystPredicate} import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} import org.apache.spark.sql.connector.expressions.{FieldReference, LogicalExpressions, NamedReference, SortDirection, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder} +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.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.CaseInsensitiveStringMap import org.apache.spark.unsafe.types.UTF8String @@ -78,7 +82,10 @@ class InMemoryRowLevelOperationTable private ( private final val SUPPORTS_DELTAS = "supports-deltas" private final val SPLIT_UPDATES = "split-updates" private final val NO_METADATA = "no-metadata" + private final val USE_CATALYST_RUNTIME_FILTERING = "use-catalyst-runtime-filtering" private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" + private final val useCatalystRuntimeFiltering = + properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty @@ -133,7 +140,7 @@ class InMemoryRowLevelOperationTable private ( case class PartitionBasedOperation(command: Command, options: CaseInsensitiveStringMap) extends RowLevelOperation with RowLevelOperationWithOptions { - var configuredScan: InMemoryBatchScan = _ + var configuredScan: BatchScanBaseClass = _ override def requiredMetadataAttributes(): Array[NamedReference] = { if (noMetadata) { @@ -144,12 +151,8 @@ class InMemoryRowLevelOperationTable private ( } override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) { - override def build: Scan = { - val scan = super.build() - configuredScan = scan.asInstanceOf[InMemoryBatchScan] - scan - } + newRowLevelScanBuilder(options) { scan => + configuredScan = scan } } @@ -186,7 +189,7 @@ class InMemoryRowLevelOperationTable private ( override def description(): String = "InMemoryPartitionReplaceOperation" } - private case class PartitionBasedReplaceData(scan: InMemoryBatchScan) + private case class PartitionBasedReplaceData(scan: BatchScanBaseClass) extends TestBatchWrite { override protected def doCommit( @@ -216,7 +219,7 @@ class InMemoryRowLevelOperationTable private ( override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) + newRowLevelScanBuilder(options)(_ => ()) } override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = { @@ -267,6 +270,94 @@ class InMemoryRowLevelOperationTable private ( override def abort(messages: Array[WriterCommitMessage]): Unit = {} } + + /** + * Builds a scan for row-level operations. When + * `use-catalyst-runtime-filtering` is set, the scan implements + * [[SupportsRuntimeCatalystFiltering]] so group filtering goes through the Catalyst path. + */ + private def newRowLevelScanBuilder( + options: CaseInsensitiveStringMap)( + onBuild: BatchScanBaseClass => Unit): ScanBuilder = { + new InMemoryScanBuilder(schema, options) { + override protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (useCatalystRuntimeFiltering) { + InMemoryCatalystRowLevelBatchScan(partitions, readSchema, tableSchema, options) + } else { + super.createScan(partitions, readSchema, tableSchema, options) + } + } + + override def build: Scan = { + val scan = super.build().asInstanceOf[BatchScanBaseClass] + onBuild(scan) + scan + } + } + } + + /** + * Row-level batch scan that receives runtime filters as Catalyst expressions. + * Evaluates partition-column predicates against each partition key so group filtering + * actually prunes partitions (needed for `replacedPartitions` assertions). + */ + case class InMemoryCatalystRowLevelBatchScan( + var _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + with SupportsRuntimeCatalystFiltering { + + private val _catalystPredicates = ArrayBuffer.empty[Expression] + + override def filterAttributes(): Array[NamedReference] = { + val scanFields = readSchema.fields.map(_.name).toSet + partitioning.flatMap(_.references()) + .filter(ref => scanFields.contains(ref.fieldNames.mkString("."))) + } + + override def filter(expressions: Array[Expression]): Unit = { + _catalystPredicates ++= expressions + val partAttrs = partitionAttributes + if (partAttrs.isEmpty) return + + val resolver = SQLConf.get.resolver + expressions.foreach { expr => + val remapped = expr.transform { + case a: AttributeReference => + partAttrs.find(p => resolver(p.name, a.name)).getOrElse(a) + } + if (remapped.references.forall(r => partAttrs.exists(_.exprId == r.exprId))) { + val bound = BindReferences.bindReference(remapped, partAttrs) + val pred = CatalystPredicate.createInterpreted(bound) + data = data.filter { p => + try { + pred.eval(p.asInstanceOf[BufferedRows].partitionKey()) + } catch { + case _: Exception => true + } + } + } + } + } + + /** Predicates recorded by [[filter]], for test assertions only. */ + def pushedCatalystPredicates: Seq[Expression] = _catalystPredicates.toSeq + + private def partitionAttributes: Seq[AttributeReference] = { + partitioning.flatMap(_.references()).flatMap { ref => + val name = ref.fieldNames.mkString(".") + readSchema.find(_.name == name).orElse(tableSchema.find(_.name == name)).map { f => + AttributeReference(f.name, f.dataType, f.nullable)() + } + }.toSeq + } + } } private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory { 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 9f8409efa360e..87139a3a20e15 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 @@ -24,9 +24,11 @@ import org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation, GroupBasedRowLevelOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, RowLevelWrite} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering} 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.util.ArrayImplicits._ /** @@ -51,26 +53,39 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(replaceData, cond, scan) + 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) case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(writeDelta, cond, scan) + ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) + if canInjectGroupFilters(cond, scan.filterAttributes()) => + injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes()) } private def canInjectGroupFilters( cond: Expression, - scan: SupportsRuntimeV2Filtering): Boolean = { + filterAttrs: Array[NamedReference]): Boolean = { conf.runtimeRowLevelOperationGroupFilterEnabled && cond != TrueLiteral && - scan.filterAttributes.nonEmpty + filterAttrs.nonEmpty } private def injectGroupFilters( write: RowLevelWrite, cond: Expression, - scan: SupportsRuntimeV2Filtering): LogicalPlan = { + scan: Scan, + filterAttrs: Array[NamedReference]): LogicalPlan = { // use reference equality on scan to find required scan relations val newQuery = write.query transformUp { case r: DataSourceV2ScanRelation if r.scan eq scan => @@ -79,9 +94,9 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla val originalTable = r.relation.table.asRowLevelOperationTable.table val relation = r.relation.copy(table = originalTable) val matchingRowsPlan = buildMatchingRowsPlan(write, relation, cond) - val filterAttrs = scan.filterAttributes.toImmutableArraySeq - val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, matchingRowsPlan) - val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r) + val filterAttrsSeq = filterAttrs.toImmutableArraySeq + val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, matchingRowsPlan) + val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, r) Filter(buildDynamicPruningCond(matchingRowsPlan, buildKeys, pruningKeys), r) } // optimize subqueries to rewrite them as joins and trigger job planning diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..ce104f5631f11 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.sql.Row + +class DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = super.extraTableProps + props.put("supports-deltas", "true") + props + } + + test("delete does not use group filtering when the group key is not scanned") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + // a delta-based delete scans the row ID, the condition columns and the metadata columns, so + // `dep` is not read and the scan cannot declare it as a filter attribute + assertNoCatalystGroupFilter(executedPlan) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } +} 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 new file mode 100644 index 0000000000000..34b43e2354abc --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.connector.catalog.InMemoryTable +import org.apache.spark.sql.connector.write.DeleteSummary + +class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + test("delete runtime group filtering with SupportsRuntimeCatalystFiltering") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // only pk 1 matches, so hr is rewritten and its other row (pk 3) is copied over + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "salary INT, dep STRING", groups = Seq("hr"))) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + + checkReplacedPartitions(Seq("hr")) + checkDeleteMetrics(numDeletedRows = 1, numCopiedRows = 1) + } + + private def checkDeleteMetrics(numDeletedRows: Long, numCopiedRows: Long): Unit = { + val t = catalog.loadTable(ident).asInstanceOf[InMemoryTable] + val summary = t.commits.last.writeSummary.get.asInstanceOf[DeleteSummary] + assert(summary.numDeletedRows() === numDeletedRows, + s"Expected numDeletedRows=$numDeletedRows, got ${summary.numDeletedRows()}") + assert(summary.numCopiedRows() === numCopiedRows, + s"Expected numCopiedRows=$numCopiedRows, got ${summary.numCopiedRows()}") + } +} 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 new file mode 100644 index 0000000000000..810131eb54cef --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala @@ -0,0 +1,229 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.{BufferedRows, InMemoryRowLevelOperationTable} +import org.apache.spark.sql.execution.InSubqueryExec +import org.apache.spark.sql.execution.ReusedSubqueryExec +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.types.StructType +import org.apache.spark.unsafe.types.UTF8String + +/** + * Verifies that row-level runtime group filtering injects filters for scans that implement + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], where the filter + * reaches the connector as a Catalyst expression instead of a connector predicate. + * + * The tests here apply to both group-based and delta-based row-level operations. DELETE is left + * to the concrete suites because delta-based deletes only scan the row ID and the condition + * columns, so the group key is not available to filter on. + */ +abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase + extends RowLevelOperationSuiteBase { + + import testImplicits._ + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("use-catalyst-runtime-filtering", "true") + props + } + + test("update runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("updated_id") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // the subquery blocks planning-time pushdown, leaving group filtering to do the pruning; + // only id 1 matches, and it lives in hr + val updatedIdDF = Seq(Some(1), None).toDF() + updatedIdDF.createOrReplaceTempView("updated_id") + + val executedPlan = executeAndKeepPlan { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM updated_id)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "id INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 1, -1, "hr") :: Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } + } + + test("merge runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 200, "dep": "hr" } + |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" } + |{ "pk": 4, "id": 4, "salary": 400, "dep": "software" } + |{ "pk": 5, "id": 5, "salary": 500, "dep": "software" } + |""".stripMargin) + + // pk 1 to 3 match rows in hr, pk 6 matches nothing and becomes an insert, so hr is the + // only group that has to be rewritten + val sourceDF = Seq(1, 2, 3, 6).toDF("pk") + sourceDF.createOrReplaceTempView("source") + + val executedPlan = executeAndKeepPlan { + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = t.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, id, salary, dep) VALUES (s.pk, 0, 0, 'hr') + |""".stripMargin) + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "pk INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq( + Row(1, 1, 101, "hr"), + Row(2, 2, 201, "hr"), + Row(3, 3, 301, "hr"), + Row(4, 4, 400, "software"), + Row(5, 5, 500, "software"), + Row(6, 0, 0, "hr"))) + } + } + + /** + * Asserts the injected group filter down to its contents: the scan declares + * `expectedFilterAttrs` 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`. + * + * A group-based UPDATE is rewritten as a union of two branches sharing one scan, so the plan + * can hold more than one scan node. Each carries its own copy of the filter, keyed on that + * branch's own attributes and with its own expr IDs, and pushes it separately (see the UPDATE + * case in RowLevelOperationRuntimeGroupFiltering.buildMatchingRowsPlan). Every copy is checked + * against `expectedFilter` rather than against one another, so the expectation stays explicit + * and does not depend on how many copies the rewrite happens to produce. + */ + protected def assertCatalystGroupFilter( + executedPlan: SparkPlan, + expectedFilterAttrs: 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") + + batchScans.foreach { batchScan => + batchScan.runtimeFilters match { + case Seq(DynamicPruningExpression(inSubquery: InSubqueryExec)) => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => fail(s"expected a single dynamic pruning group filter, got $other") + } + } + + // the scan must receive the Catalyst subquery Spark planned, not a translated connector + // predicate, once per scan node + val pushed = scan.pushedCatalystPredicates + assert(pushed.size === batchScans.size, + s"expected each of the ${batchScans.size} scan node(s) to push the filter once, got $pushed") + pushed.foreach { + case inSubquery: InSubqueryExec => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => + fail(s"expected the group filter pushed as an InSubqueryExec, got $other") + } + + val scannedGroups = scan.data.map(_.asInstanceOf[BufferedRows].keyString()).distinct + assert(scannedGroups.sorted === expectedFilter.groups.sorted, + s"scan must read only the filtered groups, got ${scannedGroups.mkString(", ")}") + } + + /** + * The expected shape of a group filter subquery: the columns it reads, which must be only those + * needed to evaluate the row-level condition, and the groups it resolves to at runtime. + */ + protected case class GroupFilter(scanSchema: String, groups: Seq[String]) + + private def assertGroupFilter( + filter: InSubqueryExec, + expectedFilterAttrs: Seq[String], + expectedFilter: GroupFilter): Unit = { + assert(filter.child.references.toSeq.map(_.name) === expectedFilterAttrs, + s"expected the group filter keyed on $expectedFilterAttrs, got ${filter.child}") + + // the second branch of a group-based UPDATE reuses the first branch's subquery, and + // ReusedSubqueryExec is a leaf node, so unwrap it to reach the plan underneath + val subqueryPlan = filter.plan match { + case reused: ReusedSubqueryExec => reused.child + case plan => plan + } + val subqueryScan = find(subqueryPlan) { case _: BatchScanExec => true; case _ => false } + .getOrElse(fail(s"could not find the scan of group filter subquery ${filter.plan.name}")) + assert( + DataTypeUtils.sameType(subqueryScan.schema, StructType.fromDDL(expectedFilter.scanSchema)), + s"unexpected group filter subquery scan schema ${subqueryScan.schema.sql}") + + val groups = filter.values() + .getOrElse(fail("group filter subquery produced no values")) + .map(_.asInstanceOf[UTF8String].toString) + assert(groups.toSeq.sorted === expectedFilter.groups.sorted, + s"group filter must select the groups holding matching rows, got ${groups.mkString(", ")}") + } + + /** Asserts no group filter was injected, e.g. because the scan does not read the group key. */ + protected def assertNoCatalystGroupFilter(executedPlan: SparkPlan): Unit = { + val batchScan = collect(executedPlan) { case s: BatchScanExec => s }.head + val scan = catalystScan(batchScan) + assert(scan.filterAttributes().isEmpty, + s"expected no filter attributes, got ${scan.filterAttributes().mkString(", ")}") + assert(batchScan.runtimeFilters.isEmpty, + s"expected no runtime filters, got ${batchScan.runtimeFilters}") + assert(scan.pushedCatalystPredicates.isEmpty, + s"expected no pushed predicates, got ${scan.pushedCatalystPredicates}") + } + + private type CatalystRowLevelScan = + InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan + + private def catalystScan(batchScan: BatchScanExec): CatalystRowLevelScan = { + batchScan.scan match { + case s: InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan => s + case other => fail(s"expected InMemoryCatalystRowLevelBatchScan, got ${other.getClass}") + } + } +}