[SPARK-58523][SQL] Add Catalyst runtime filtering interface for DSv2 scans #57727
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| /* | ||
| * 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. | ||
| * 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] | ||
|
|
||
| /** | ||
| * 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]]. | ||
| * | ||
| * 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 | ||
|
Contributor
There was a problem hiding this comment. Finding 2. No objection to the attribute-level shape — v2 file sources already work this way.
Something along these lines: /**
* 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]].
*
* Only declare an attribute here if this scan evaluates an arbitrary deterministic Catalyst
* predicate over it exactly, for every row it returns -- e.g. an identity partition column,
* whose value is known for every row of a surviving partition. Do not declare an attribute
* whose predicates only guide approximate pruning, such as file or row-group statistics.
* Spark may push any expression that references only these attributes, so do not assume a
* fixed set of operators: bind and evaluate the expression (see
* [[PartitionPredicateImpl]]) instead of pattern matching it.
*/While you're here, it would help to name the intended implementor in the PR description — it makes the contract judgeable and tells a reader why the interface is internal.
Contributor
There was a problem hiding this comment. One correction here in light of #57760: that PR adds
Member
Author
There was a problem hiding this comment. Thanks -- both failure modes are real, and your 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 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. |
||
|
|
||
| /** | ||
| * 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. | ||
| * | ||
| * 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 | ||
|
Contributor
There was a problem hiding this comment. Finding 4.
Member
Author
There was a problem hiding this comment. Done, carried the paragraph over. An SPJ-active adopter reading only this Javadoc should not have to learn the contract from a |
||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,160 @@ | ||
| /* | ||
| * 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.{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 | ||
| 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 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], | ||
| 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 | ||
|
Contributor
There was a problem hiding this comment. A scan that declares a filter fully pushed must ensure its returned partitions satisfy that predicate. This implementation only records the expression, while the fully-pushed test uses rows that all happen to match, so
Member
Author
There was a problem hiding this comment. Good catch, thanks. The fixture now prunes rather than just recording. The fully-pushed test inserts |
||
| 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 | ||
|
|
||
| /** 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 | ||
| } | ||
| } | ||
| } | ||
|
|
||
| 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" | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Finding 6.
resolveRefs→V2ExpressionUtils.resolveRefthrowscannotResolveAttributeErrorwhen a reference doesn't resolve against the plan's output, and casts the result toAttribute— so a nested reference, whichLogicalPlan.resolvehands back as anAlias(GetStructField(...)), throws aClassCastException.fullyPushedFilterAttributes()therefore has an unwritten requirement: top-level attributes only, and only ones that survived column pruning into the scan'sreadSchema. Break it and the query fails at planning time.filterAttributescarries the same requirement and is equally undocumented, but it's forced for every scan relation, so an adopter trips it on the first query. This one is only forced when a scalar-subquery runtime filter is present (scalarSubqueryFilters.filterdoesn't evaluate its closure on an empty Seq), which makes it a query-shape-dependent failure. Worth a line on the trait alongside finding 2; the new fixture quietly depends on it via thescanFields.contains(name)guard atInMemoryCatalystRuntimeFilterTable.scala:267.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Documented on both
filterAttributes()andfullyPushedFilterAttributes(): each reference must be a top-level attribute present inreadSchema, 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()andSupportsRuntimeV2Filtering.filterAttributes(). Documentation only, no behaviour change there.