szehon-ho commented on code in PR #57727: URL: https://github.com/apache/spark/pull/57727#discussion_r3717301697
########## 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 Review Comment: Done, carried the paragraph over. An SPJ-active adopter reading only this Javadoc should not have to learn the contract from a `SparkException` thrown by `replanWithRuntimeFilters`. ########## 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 Review Comment: You're right that nothing reads it -- dropped from the trait. The fixture exposes its own `pushedCatalystPredicates` accessor instead, following the `InMemoryEnhancedRuntimePartitionFilterTable.pushedPartitionPredicates` precedent, so the assertions still work without implying a contract Spark relies on. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala: ########## @@ -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]( Review Comment: Documented on both `filterAttributes()` and `fullyPushedFilterAttributes()`: each reference must be a top-level attribute present in `readSchema`, since nested references and attributes pruned out of the read schema fail to resolve when Spark builds the scan relation. Since you noted the requirement is equally undocumented on the existing interfaces, I added the same note to `SupportsRuntimeFiltering.filterAttributes()` and `SupportsRuntimeV2Filtering.filterAttributes()`. Documentation only, no behaviour change there. ########## 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 Review Comment: Applied -- a test that passes because it couldn't find the scan is worse than no test. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
