* Note that Spark will push runtime filters only if they are beneficial. * @@ -38,6 +39,10 @@ public interface SupportsRuntimeFiltering extends SupportsRuntimeV2Filtering { *
* Spark will call {@link #filter(Filter[])} if it can derive a runtime * predicate for any of the filter attributes. + *
+ * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); 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..f23de66fdc0af 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 @@ -50,6 +50,10 @@ public interface SupportsRuntimeV2Filtering extends Scan { *
* Spark will call {@link #filter(Predicate[])} if it can derive a runtime * predicate for any of the filter attributes. + *
+ * Each reference must be a top-level attribute present in {@link Scan#readSchema()}. + * Nested references and attributes pruned out of the read schema fail to resolve when + * Spark builds the scan relation. */ NamedReference[] filterAttributes(); @@ -82,6 +86,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..6823757595b54 --- /dev/null +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/SupportsRuntimeCatalystFiltering.scala @@ -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 + + /** + * 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 +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala index fba80eb3d4cbe..662fe826993b2 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala @@ -521,23 +521,35 @@ abstract class InMemoryBaseTable( private var _pushedFilters: Array[Filter] = Array.empty override def build: Scan = { - val scan = if (InMemoryBaseTable.this.ordering.nonEmpty) { - new InMemoryBatchScanWithOrdering( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } else { - InMemoryBatchScan( - data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, - options) - } - if (evaluableFilters.nonEmpty) { - scan.filter(evaluableFilters) + val scan = createScan( + data.map(_.asInstanceOf[InputPartition]).toImmutableArraySeq, schema, tableSchema, options) + scan match { + case s: InMemoryBatchScan => + if (evaluableFilters.nonEmpty) { + s.filter(evaluableFilters) + } + s.pushedFilters = _pushedFilters + case _ => } - scan.pushedFilters = _pushedFilters recordScanEvent(_pushedFilters) scan } + /** + * Creates the batch scan for [[build]]. + */ + protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (InMemoryBaseTable.this.ordering.nonEmpty) { + new InMemoryBatchScanWithOrdering(partitions, readSchema, tableSchema, options) + } else { + InMemoryBatchScan(partitions, readSchema, tableSchema, options) + } + } + override def pruneColumns(requiredSchema: StructType): Unit = { // The required schema could contain conflict-renamed metadata columns, so we need to match // them by their logical (original) names, not their current names. 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..b696c09fc07bf --- /dev/null +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryCatalystRuntimeFilterTable.scala @@ -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 + 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" +} diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala index fd02f926e8141..6de76647a0859 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala @@ -19,15 +19,19 @@ package org.apache.spark.sql.connector.catalog import java.util +import scala.collection.mutable.ArrayBuffer + import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BindReferences, Expression, GenericInternalRow, Predicate => CatalystPredicate} import org.apache.spark.sql.connector.catalog.constraints.Constraint import org.apache.spark.sql.connector.distributions.{Distribution, Distributions} import org.apache.spark.sql.connector.expressions.{FieldReference, LogicalExpressions, NamedReference, SortDirection, SortOrder, Transform} import org.apache.spark.sql.connector.expressions.filter.Predicate -import org.apache.spark.sql.connector.read.{Scan, ScanBuilder} +import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder} import org.apache.spark.sql.connector.write.{BatchWrite, DeltaBatchWrite, DeltaWrite, DeltaWriteBuilder, DeltaWriter, DeltaWriterFactory, LogicalWriteInfo, PhysicalWriteInfo, RequiresDistributionAndOrdering, RowLevelOperation, RowLevelOperationBuilder, RowLevelOperationInfo, SupportsDelta, Write, WriteBuilder, WriterCommitMessage} import org.apache.spark.sql.connector.write.RowLevelOperation.Command +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.unsafe.types.UTF8String @@ -78,7 +82,10 @@ class InMemoryRowLevelOperationTable private ( private final val SUPPORTS_DELTAS = "supports-deltas" private final val SPLIT_UPDATES = "split-updates" private final val NO_METADATA = "no-metadata" + private final val USE_CATALYST_RUNTIME_FILTERING = "use-catalyst-runtime-filtering" private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" + private final val useCatalystRuntimeFiltering = + properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty @@ -133,7 +140,7 @@ class InMemoryRowLevelOperationTable private ( case class PartitionBasedOperation(command: Command, options: CaseInsensitiveStringMap) extends RowLevelOperation with RowLevelOperationWithOptions { - var configuredScan: InMemoryBatchScan = _ + var configuredScan: BatchScanBaseClass = _ override def requiredMetadataAttributes(): Array[NamedReference] = { if (noMetadata) { @@ -144,12 +151,8 @@ class InMemoryRowLevelOperationTable private ( } override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) { - override def build: Scan = { - val scan = super.build() - configuredScan = scan.asInstanceOf[InMemoryBatchScan] - scan - } + newRowLevelScanBuilder(options) { scan => + configuredScan = scan } } @@ -186,7 +189,7 @@ class InMemoryRowLevelOperationTable private ( override def description(): String = "InMemoryPartitionReplaceOperation" } - private case class PartitionBasedReplaceData(scan: InMemoryBatchScan) + private case class PartitionBasedReplaceData(scan: BatchScanBaseClass) extends TestBatchWrite { override protected def doCommit( @@ -216,7 +219,7 @@ class InMemoryRowLevelOperationTable private ( override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { - new InMemoryScanBuilder(schema, options) + newRowLevelScanBuilder(options)(_ => ()) } override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = { @@ -267,6 +270,94 @@ class InMemoryRowLevelOperationTable private ( override def abort(messages: Array[WriterCommitMessage]): Unit = {} } + + /** + * Builds a scan for row-level operations. When + * `use-catalyst-runtime-filtering` is set, the scan implements + * [[SupportsRuntimeCatalystFiltering]] so group filtering goes through the Catalyst path. + */ + private def newRowLevelScanBuilder( + options: CaseInsensitiveStringMap)( + onBuild: BatchScanBaseClass => Unit): ScanBuilder = { + new InMemoryScanBuilder(schema, options) { + override protected def createScan( + partitions: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap): BatchScanBaseClass = { + if (useCatalystRuntimeFiltering) { + InMemoryCatalystRowLevelBatchScan(partitions, readSchema, tableSchema, options) + } else { + super.createScan(partitions, readSchema, tableSchema, options) + } + } + + override def build: Scan = { + val scan = super.build().asInstanceOf[BatchScanBaseClass] + onBuild(scan) + scan + } + } + } + + /** + * Row-level batch scan that receives runtime filters as Catalyst expressions. + * Evaluates partition-column predicates against each partition key so group filtering + * actually prunes partitions (needed for `replacedPartitions` assertions). + */ + case class InMemoryCatalystRowLevelBatchScan( + var _data: Seq[InputPartition], + readSchema: StructType, + tableSchema: StructType, + options: CaseInsensitiveStringMap) + extends BatchScanBaseClass(_data, readSchema, tableSchema) + with SupportsRuntimeCatalystFiltering { + + private val _catalystPredicates = ArrayBuffer.empty[Expression] + + override def filterAttributes(): Array[NamedReference] = { + val scanFields = readSchema.fields.map(_.name).toSet + partitioning.flatMap(_.references()) + .filter(ref => scanFields.contains(ref.fieldNames.mkString("."))) + } + + override def filter(expressions: Array[Expression]): Unit = { + _catalystPredicates ++= expressions + 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) + } + 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 { + case _: Exception => true + } + } + } + } + } + + /** Predicates recorded by [[filter]], for test assertions only. */ + def pushedCatalystPredicates: Seq[Expression] = _catalystPredicates.toSeq + + 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 + } + } } private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory { 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..bc6630b3c6d61 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 exceptions are 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..d0c105ac59283 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,26 @@ 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. + // Screen with the same pushability guard as the V2 PartitionPredicate path + // (deterministic, no subquery, no Python UDF). Keeps non-deterministic filters + // from being the sole evaluator when fullyPushedFilterAttributes drops FilterExec. + val catalystFilters = runtimeFilters + .flatMap(unwrapRuntimeFilterExpression) + .filterNot(_ == Literal.TrueLiteral) + .filter(isPushablePartitionFilter) + if (catalystFilters.nonEmpty) { + catalystScan.filter(catalystFilters.toArray) + true + } else { + false + } + case _ => false } @@ -433,16 +457,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/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala index 9f8409efa360e..87139a3a20e15 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala @@ -24,9 +24,11 @@ import org.apache.spark.sql.catalyst.optimizer.RewritePredicateSubquery import org.apache.spark.sql.catalyst.planning.{DeltaBasedRowLevelOperation, GroupBasedRowLevelOperation} import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, LogicalPlan, RowLevelWrite} import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering +import org.apache.spark.sql.connector.expressions.NamedReference +import org.apache.spark.sql.connector.read.{Scan, SupportsRuntimeV2Filtering} import org.apache.spark.sql.connector.write.RowLevelOperation.Command.{DELETE, MERGE, UPDATE} import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Implicits, DataSourceV2Relation, DataSourceV2ScanRelation, ExtractV2Scan} +import org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering import org.apache.spark.util.ArrayImplicits._ /** @@ -51,26 +53,39 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown { case GroupBasedRowLevelOperation(replaceData, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(replaceData, cond, scan) + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) + if canInjectGroupFilters(cond, scan.filterAttributes) => + injectGroupFilters(replaceData, cond, scan, scan.filterAttributes) + + case GroupBasedRowLevelOperation(replaceData, _, Some(cond), + ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) + if canInjectGroupFilters(cond, scan.filterAttributes()) => + injectGroupFilters(replaceData, cond, scan, scan.filterAttributes()) + + case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), + ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) + if canInjectGroupFilters(cond, scan.filterAttributes) => + injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes) case DeltaBasedRowLevelOperation(writeDelta, _, Some(cond), - ExtractV2Scan(scan: SupportsRuntimeV2Filtering)) if canInjectGroupFilters(cond, scan) => - injectGroupFilters(writeDelta, cond, scan) + ExtractV2Scan(scan: SupportsRuntimeCatalystFiltering)) + if canInjectGroupFilters(cond, scan.filterAttributes()) => + injectGroupFilters(writeDelta, cond, scan, scan.filterAttributes()) } private def canInjectGroupFilters( cond: Expression, - scan: SupportsRuntimeV2Filtering): Boolean = { + filterAttrs: Array[NamedReference]): Boolean = { conf.runtimeRowLevelOperationGroupFilterEnabled && cond != TrueLiteral && - scan.filterAttributes.nonEmpty + filterAttrs.nonEmpty } private def injectGroupFilters( write: RowLevelWrite, cond: Expression, - scan: SupportsRuntimeV2Filtering): LogicalPlan = { + scan: Scan, + filterAttrs: Array[NamedReference]): LogicalPlan = { // use reference equality on scan to find required scan relations val newQuery = write.query transformUp { case r: DataSourceV2ScanRelation if r.scan eq scan => @@ -79,9 +94,9 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla val originalTable = r.relation.table.asRowLevelOperationTable.table val relation = r.relation.copy(table = originalTable) val matchingRowsPlan = buildMatchingRowsPlan(write, relation, cond) - val filterAttrs = scan.filterAttributes.toImmutableArraySeq - val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, matchingRowsPlan) - val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrs, r) + val filterAttrsSeq = filterAttrs.toImmutableArraySeq + val buildKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, matchingRowsPlan) + val pruningKeys = V2ExpressionUtils.resolveRefs[Attribute](filterAttrsSeq, r) Filter(buildDynamicPruningCond(matchingRowsPlan, buildKeys, pruningKeys), r) } // optimize subqueries to rewrite them as joins and trigger job planning 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..e07367e538a9a --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -0,0 +1,311 @@ +/* + * 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')") + // Matching and nonmatching partitions: the scan must prune nonmatching ones itself + // because Spark drops the post-scan FilterExec for fully pushed attributes. + 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))) + 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.pushedCatalystPredicates + case other => + fail(s"Expected InMemoryCatalystRuntimeFilterBatchScan, got $other") + } + } + + 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) + } + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..ce104f5631f11 --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite.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 + +import org.apache.spark.sql.Row + +class DeltaBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = super.extraTableProps + props.put("supports-deltas", "true") + props + } + + test("delete does not use group filtering when the group key is not scanned") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + // a delta-based delete scans the row ID, the condition columns and the metadata columns, so + // `dep` is not read and the scan cannot declare it as a filter attribute + assertNoCatalystGroupFilter(executedPlan) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala new file mode 100644 index 0000000000000..34b43e2354abc --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedRowLevelOperationCatalystRuntimeFilterSuite.scala @@ -0,0 +1,60 @@ +/* + * 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.sql.Row +import org.apache.spark.sql.connector.catalog.InMemoryTable +import org.apache.spark.sql.connector.write.DeleteSummary + +class GroupBasedRowLevelOperationCatalystRuntimeFilterSuite + extends RowLevelOperationCatalystRuntimeFilterSuiteBase { + + test("delete runtime group filtering with SupportsRuntimeCatalystFiltering") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // only pk 1 matches, so hr is rewritten and its other row (pk 3) is copied over + val executedPlan = executeAndKeepPlan { + sql(s"DELETE FROM $tableNameAsString WHERE salary IN (300, 400, 500)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "salary INT, dep STRING", groups = Seq("hr"))) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + + checkReplacedPartitions(Seq("hr")) + checkDeleteMetrics(numDeletedRows = 1, numCopiedRows = 1) + } + + private def checkDeleteMetrics(numDeletedRows: Long, numCopiedRows: Long): Unit = { + val t = catalog.loadTable(ident).asInstanceOf[InMemoryTable] + val summary = t.commits.last.writeSummary.get.asInstanceOf[DeleteSummary] + assert(summary.numDeletedRows() === numDeletedRows, + s"Expected numDeletedRows=$numDeletedRows, got ${summary.numDeletedRows()}") + assert(summary.numCopiedRows() === numCopiedRows, + s"Expected numCopiedRows=$numCopiedRows, got ${summary.numCopiedRows()}") + } +} diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala new file mode 100644 index 0000000000000..810131eb54cef --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationCatalystRuntimeFilterSuiteBase.scala @@ -0,0 +1,229 @@ +/* + * 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.sql.Row +import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.{BufferedRows, InMemoryRowLevelOperationTable} +import org.apache.spark.sql.execution.InSubqueryExec +import org.apache.spark.sql.execution.ReusedSubqueryExec +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.types.StructType +import org.apache.spark.unsafe.types.UTF8String + +/** + * Verifies that row-level runtime group filtering injects filters for scans that implement + * [[org.apache.spark.sql.internal.connector.SupportsRuntimeCatalystFiltering]], where the filter + * reaches the connector as a Catalyst expression instead of a connector predicate. + * + * The tests here apply to both group-based and delta-based row-level operations. DELETE is left + * to the concrete suites because delta-based deletes only scan the row ID and the condition + * columns, so the group key is not available to filter on. + */ +abstract class RowLevelOperationCatalystRuntimeFilterSuiteBase + extends RowLevelOperationSuiteBase { + + import testImplicits._ + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("use-catalyst-runtime-filtering", "true") + props + } + + test("update runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("updated_id") { + // the table is partitioned by dep, so hr and software are the two groups + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 300, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 150, "dep": "software" } + |{ "pk": 3, "id": 3, "salary": 120, "dep": "hr" } + |""".stripMargin) + + // the subquery blocks planning-time pushdown, leaving group filtering to do the pruning; + // only id 1 matches, and it lives in hr + val updatedIdDF = Seq(Some(1), None).toDF() + updatedIdDF.createOrReplaceTempView("updated_id") + + val executedPlan = executeAndKeepPlan { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM updated_id)") + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "id INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 1, -1, "hr") :: Row(2, 2, 150, "software") :: Row(3, 3, 120, "hr") :: Nil) + } + } + + test("merge runtime group filtering with SupportsRuntimeCatalystFiltering") { + withTempView("source") { + createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING", + """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "id": 2, "salary": 200, "dep": "hr" } + |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" } + |{ "pk": 4, "id": 4, "salary": 400, "dep": "software" } + |{ "pk": 5, "id": 5, "salary": 500, "dep": "software" } + |""".stripMargin) + + // pk 1 to 3 match rows in hr, pk 6 matches nothing and becomes an insert, so hr is the + // only group that has to be rewritten + val sourceDF = Seq(1, 2, 3, 6).toDF("pk") + sourceDF.createOrReplaceTempView("source") + + val executedPlan = executeAndKeepPlan { + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN + | UPDATE SET t.salary = t.salary + 1 + |WHEN NOT MATCHED THEN + | INSERT (pk, id, salary, dep) VALUES (s.pk, 0, 0, 'hr') + |""".stripMargin) + } + assertCatalystGroupFilter( + executedPlan, + expectedFilterAttrs = Seq("dep"), + expectedFilter = GroupFilter(scanSchema = "pk INT, dep STRING", groups = Seq("hr"))) + + // software was never read, so its rows must come back untouched + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Seq( + Row(1, 1, 101, "hr"), + Row(2, 2, 201, "hr"), + Row(3, 3, 301, "hr"), + Row(4, 4, 400, "software"), + Row(5, 5, 500, "software"), + Row(6, 0, 0, "hr"))) + } + } + + /** + * Asserts the injected group filter down to its contents: the scan declares + * `expectedFilterAttrs` in `filterAttributes`, every scan node carries one dynamic pruning + * filter matching `expectedFilter`, the connector received that same filter as a Catalyst + * expression, and the scan then read only `expectedFilter.groups`. + * + * A group-based UPDATE is rewritten as a union of two branches sharing one scan, so the plan + * can hold more than one scan node. Each carries its own copy of the filter, keyed on that + * branch's own attributes and with its own expr IDs, and pushes it separately (see the UPDATE + * case in RowLevelOperationRuntimeGroupFiltering.buildMatchingRowsPlan). Every copy is checked + * against `expectedFilter` rather than against one another, so the expectation stays explicit + * and does not depend on how many copies the rewrite happens to produce. + */ + protected def assertCatalystGroupFilter( + executedPlan: SparkPlan, + expectedFilterAttrs: Seq[String], + expectedFilter: GroupFilter): Unit = { + val batchScans = collect(executedPlan) { case s: BatchScanExec => s } + assert(batchScans.nonEmpty, "expected a batch scan for the row-level operation") + val scan = catalystScan(batchScans.head) + assert(batchScans.forall(_.scan eq scan), + s"expected all ${batchScans.size} scan nodes to share one scan") + + val filterAttrs = scan.filterAttributes().map(_.fieldNames.mkString(".")).toSeq + assert(filterAttrs === expectedFilterAttrs, + s"expected the scan to declare $expectedFilterAttrs as filter attributes, got $filterAttrs") + + batchScans.foreach { batchScan => + batchScan.runtimeFilters match { + case Seq(DynamicPruningExpression(inSubquery: InSubqueryExec)) => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => fail(s"expected a single dynamic pruning group filter, got $other") + } + } + + // the scan must receive the Catalyst subquery Spark planned, not a translated connector + // predicate, once per scan node + val pushed = scan.pushedCatalystPredicates + assert(pushed.size === batchScans.size, + s"expected each of the ${batchScans.size} scan node(s) to push the filter once, got $pushed") + pushed.foreach { + case inSubquery: InSubqueryExec => + assertGroupFilter(inSubquery, expectedFilterAttrs, expectedFilter) + case other => + fail(s"expected the group filter pushed as an InSubqueryExec, got $other") + } + + val scannedGroups = scan.data.map(_.asInstanceOf[BufferedRows].keyString()).distinct + assert(scannedGroups.sorted === expectedFilter.groups.sorted, + s"scan must read only the filtered groups, got ${scannedGroups.mkString(", ")}") + } + + /** + * The expected shape of a group filter subquery: the columns it reads, which must be only those + * needed to evaluate the row-level condition, and the groups it resolves to at runtime. + */ + protected case class GroupFilter(scanSchema: String, groups: Seq[String]) + + private def assertGroupFilter( + filter: InSubqueryExec, + expectedFilterAttrs: Seq[String], + expectedFilter: GroupFilter): Unit = { + assert(filter.child.references.toSeq.map(_.name) === expectedFilterAttrs, + s"expected the group filter keyed on $expectedFilterAttrs, got ${filter.child}") + + // the second branch of a group-based UPDATE reuses the first branch's subquery, and + // ReusedSubqueryExec is a leaf node, so unwrap it to reach the plan underneath + val subqueryPlan = filter.plan match { + case reused: ReusedSubqueryExec => reused.child + case plan => plan + } + val subqueryScan = find(subqueryPlan) { case _: BatchScanExec => true; case _ => false } + .getOrElse(fail(s"could not find the scan of group filter subquery ${filter.plan.name}")) + assert( + DataTypeUtils.sameType(subqueryScan.schema, StructType.fromDDL(expectedFilter.scanSchema)), + s"unexpected group filter subquery scan schema ${subqueryScan.schema.sql}") + + val groups = filter.values() + .getOrElse(fail("group filter subquery produced no values")) + .map(_.asInstanceOf[UTF8String].toString) + assert(groups.toSeq.sorted === expectedFilter.groups.sorted, + s"group filter must select the groups holding matching rows, got ${groups.mkString(", ")}") + } + + /** Asserts no group filter was injected, e.g. because the scan does not read the group key. */ + protected def assertNoCatalystGroupFilter(executedPlan: SparkPlan): Unit = { + val batchScan = collect(executedPlan) { case s: BatchScanExec => s }.head + val scan = catalystScan(batchScan) + assert(scan.filterAttributes().isEmpty, + s"expected no filter attributes, got ${scan.filterAttributes().mkString(", ")}") + assert(batchScan.runtimeFilters.isEmpty, + s"expected no runtime filters, got ${batchScan.runtimeFilters}") + assert(scan.pushedCatalystPredicates.isEmpty, + s"expected no pushed predicates, got ${scan.pushedCatalystPredicates}") + } + + private type CatalystRowLevelScan = + InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan + + private def catalystScan(batchScan: BatchScanExec): CatalystRowLevelScan = { + batchScan.scan match { + case s: InMemoryRowLevelOperationTable#InMemoryCatalystRowLevelBatchScan => s + case other => fail(s"expected InMemoryCatalystRowLevelBatchScan, got ${other.getClass}") + } + } +}