cloud-fan commented on code in PR #57727: URL: https://github.com/apache/spark/pull/57727#discussion_r3709846457
########## 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: 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 `checkAnswer` cannot catch an incorrect post-scan-filter removal. Please make this fixture filter its partitions (or use a dedicated fully-evaluating fixture) and test with both matching and nonmatching partitions. ########## 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: The plural subject needs agreement here: `The exceptions are filters`. -- 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]
