szehon-ho commented on code in PR #57727: URL: https://github.com/apache/spark/pull/57727#discussion_r3717301072
########## 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 Review Comment: Good catch, thanks. The fixture now prunes rather than just recording. `filter()` remaps the expression's references onto the partition attributes, and when every reference is a partition column it binds the expression and evaluates it against the partition key, dropping partitions that don't match -- the same bind-and-interpret approach as `PartitionPredicateImpl`, rather than pattern matching operators. An incorrect post-scan filter removal now shows up as a wrong answer. The fully-pushed test inserts `($i, $i)` for `i` in 0..4, so four of the five partitions don't match, and it expects a single row back. ########## 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 Review Comment: Fixed. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala: ########## @@ -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) Review Comment: Great catch, and the most valuable one in the review -- a filter that is pushed, evaluated by the source on its own roll of `rand()`, and then re-rolled by Spark is exactly the kind of bug that never shows up in a test with a deterministic fixture. Thank you as well for not just reporting it but fixing it at the root in apache/spark#57760, and for working out that the `fullyPushedFilterAttributes` case makes it a correctness issue rather than a cosmetic one. I applied your suggestion, so the Catalyst branch now screens with `isPushablePartitionFilter`. I kept it even though #57760 is still open: until that lands this branch needs the guard on its own, and afterwards it is redundant but harmless and keeps the two passes symmetric. Happy to drop it once #57760 is in if you would rather have the single gate in `DataSourceV2Strategy`. Agreed on the DPP half. `PlanExpression.deterministic` folds in the filtering plan, so a non-deterministic `DynamicPruningSubquery` fails `NodeWithOnlyDeterministicProjectAndFilter` and `CleanupDynamicPruningFilters` rewrites it to `TrueLiteral` before planning -- it cannot reach the source either way. ########## 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 Review Comment: Thanks -- both failure modes are real, and your `RLIKE` example is a good demonstration that the source cannot refuse an individual predicate once `FilterExec` is gone. On the wording, I landed on keeping the doc shorter. "Returns attributes for which this scan fully evaluates runtime predicates", together with "will not be evaluated again after the scan", already carries the obligation: fully evaluating a predicate means evaluating it exactly, for every row returned, whatever the predicate looks like. That rules out approximate statistics pruning and hand-matching a fixed set of operators without spelling either out. The longer text reads as implementation guidance for one particular way of satisfying the contract, and it is fairly technical for a trait Javadoc, so I would rather leave it out. I did take the top-level attribute requirement from your finding 6, since nothing in the existing wording implies it and getting it wrong fails at planning time. On naming the implementor, the description explains the class of source this targets -- Spark-integrated sources that already bind and interpret Catalyst expressions for partition pruning, the same ones that go through `PartitionPredicateImpl` -- which is also why the interface is internal. Noted on your follow-up: after rebasing onto #57760 the set of expressions is bounded to deterministic ones, so the remaining concern is shape, which is what the current wording covers. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/PartitionPruning.scala: ########## @@ -86,6 +87,14 @@ object PartitionPruning extends Rule[LogicalPlan] with PredicateHelper with Join } else { None } + case (resExp, r @ ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) => Review Comment: Good catch -- silently losing group filtering for MERGE/UPDATE/DELETE is much worse than an error, and I had missed that this rule dispatches on the interface too. Added, and it was mechanical as you expected. `canInjectGroupFilters` and `injectGroupFilters` are now typed on `Array[NamedReference]` and `Scan` rather than on `SupportsRuntimeV2Filtering`, and there are four cases: group-based and delta-based, each for both interfaces. Also added tests, since the delta side of this had no Catalyst coverage at all: `RowLevelOperationCatalystRuntimeFilterSuiteBase` with a group-based and a delta-based subclass. They assert the injected filter is keyed on the group key, that its subquery projects only the columns the row-level condition needs and resolves to the expected groups, that the same Catalyst expression reaches the connector once per scan node, and that the scan then reads only those groups. A delta-based DELETE scans the row ID, the condition columns and the metadata columns, so the group key is not in the read schema and no filter is injected -- that case is asserted explicitly. -- 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]
