* 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..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 @@ -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.{SupportsRuntimeCatalystFiltering, 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 + * [[SupportsRuntimeCatalystFiltering.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: SupportsRuntimeCatalystFiltering => 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: SupportsRuntimeCatalystFiltering => 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/SupportsRuntimeCatalystFiltering.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala new file mode 100644 index 0000000000000..406b888c2478d --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.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 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. + */ + 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..3a728c10b034d --- /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.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 + * [[SupportsRuntimeCatalystFiltering]], 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 SupportsRuntimeCatalystFiltering { + + 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..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 @@ -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, SupportsRuntimeCatalystFiltering} 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 [[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. + * * @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: 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, + // 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..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 @@ -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.SupportsRuntimeCatalystFiltering /** * 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: SupportsRuntimeCatalystFiltering)) => + 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..7c349077c622a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -0,0 +1,308 @@ +/* + * 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.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} +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.SupportsRuntimeCatalystFiltering]], + * where runtime filters are pushed once as Catalyst expressions instead of connector + * predicates. + */ +class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { + + protected val v2Source = classOf[FakeV2ProviderWithCustomSchema].getName + protected val catalogName = "testcatalystruntimefilter" + + override def sparkConf: SparkConf = super.sparkConf + .set(s"spark.sql.catalog.$catalogName", + classOf[InMemoryTableCatalystRuntimeFilterCatalog].getName) + + 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("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" + 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) + } + } +}