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 d3c9cf3b402e1..19371dcb94dec 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 @@ -22,7 +22,7 @@ import java.util.{Optional, OptionalLong} import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.analysis.{MultiInstanceRelation, NamedRelation, TimeTravelSpec} import org.apache.spark.sql.catalyst.catalog.{CatalogColumnStat, CatalogStatistics} -import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, Expression, SortOrder} +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeMap, AttributeReference, AttributeSet, Expression, SortOrder, V2ExpressionUtils} import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.logical.{ColumnStat, ExposesMetadataColumns, Histogram, HistogramBin, LeafNode, LogicalPlan, Statistics} import org.apache.spark.sql.catalyst.streaming.{StreamingSourceIdentifyingName, Unassigned} @@ -31,11 +31,12 @@ import org.apache.spark.sql.catalyst.util.{truncatedString, CharVarcharUtils} import org.apache.spark.sql.connector.catalog.{CatalogPlugin, FunctionCatalog, Identifier, SupportsMetadataColumns, Table, TableCapability, TableCatalog, V2TableUtil} import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} -import org.apache.spark.sql.connector.read.{Scan, Statistics => V2Statistics, SupportsReportStatistics} +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.types.{DataType, StructType} import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils /** @@ -174,6 +175,18 @@ case class DataSourceV2ScanRelation( // skip adding IsNotNull when the scan already implies it, or infer new filters across // joins), so plan stability testing is needed first. + /** + * Resolved attributes that the scan declares for runtime filtering via + * [[SupportsRuntimeV2Filtering.filterAttributes]]. Empty when the scan + * does not implement [[SupportsRuntimeV2Filtering]] or exposes no attributes. + */ + lazy val runtimeFilterAttrs: AttributeSet = scan match { + case s: SupportsRuntimeV2Filtering => + AttributeSet(V2ExpressionUtils.resolveRefs[Attribute]( + s.filterAttributes.toImmutableArraySeq, this)) + case _ => AttributeSet.empty + } + override def name: String = relation.name override def simpleString(maxFields: Int): String = { diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala index 4f509f688ba26..f2827faf59435 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryTableWithV2Filter.scala @@ -87,6 +87,22 @@ class InMemoryTableWithV2Filter( }) } } + case p : Predicate if p.name().equals("=") => + if (p.children().length == 2) { + val filterRef = p.children()(0).asInstanceOf[FieldReference].references.head + if (filterRef.toString.equals(ref.toString)) { + val matchingKey = p.children()(1).asInstanceOf[LiteralValue[_]].value + if (matchingKey != null) { + data = data.filter(partition => { + val key = partition.asInstanceOf[BufferedRows].keyString() + key == matchingKey.toString + }) + } else { + data = Seq.empty // NULL = anything is always false + } + } + } + case _ => // Ignore unsupported predicate types } } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala index 4551f3512018a..dde294c6019b0 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala @@ -63,7 +63,7 @@ case class BatchScanExec( @transient private[sql] lazy val filteredPartitions: Seq[Option[InputPartition]] = { val dataSourceFilters = runtimeFilters.flatMap { case DynamicPruningExpression(e) => DataSourceV2Strategy.translateRuntimeFilterV2(e) - case _ => None + case f => DataSourceV2Strategy.translateScalarSubqueryFilterV2(f) } val originalPartitioning = outputPartitioning 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 3d3b4d1cae11c..5934cfd587623 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 @@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.{And, Attribute, DynamicPruning import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.planning.PhysicalOperation import org.apache.spark.sql.catalyst.plans.logical._ +import org.apache.spark.sql.catalyst.trees.TreePattern.SCALAR_SUBQUERY import org.apache.spark.sql.catalyst.util.{toPrettySQL, GeneratedColumn, IdentityColumn, ResolveDefaultColumns, ResolveTableConstraints, V2ExpressionBuilder} import org.apache.spark.sql.classic.SparkSession import org.apache.spark.sql.connector.catalog.{Identifier, StagingTableCatalog, SupportsDeleteV2, SupportsNamespaces, SupportsPartitionManagement, SupportsWrite, TableCapability, TableCatalog, TruncatableTable, V1Table} @@ -42,7 +43,7 @@ import org.apache.spark.sql.connector.read.LocalScan import org.apache.spark.sql.connector.read.streaming.{ContinuousStream, MicroBatchStream, SupportsRealTimeMode} import org.apache.spark.sql.connector.write.V1Write import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors} -import org.apache.spark.sql.execution.{FilterExec, InSubqueryExec, LeafExecNode, LocalTableScanExec, ProjectExec, RowDataSourceScanExec, SparkPlan, SparkStrategy => Strategy} +import org.apache.spark.sql.execution.{FilterExec, InSubqueryExec, LeafExecNode, LocalTableScanExec, ProjectExec, RowDataSourceScanExec, ScalarSubquery => ExecScalarSubquery, SparkPlan, SparkStrategy => Strategy} import org.apache.spark.sql.execution.command.CommandUtils import org.apache.spark.sql.execution.datasources.{DataSourceStrategy, LogicalRelationWithTable, PushableColumnAndNestedColumn} import org.apache.spark.sql.execution.streaming.continuous.{WriteToContinuousDataSource, WriteToContinuousDataSourceExec} @@ -155,10 +156,26 @@ class DataSourceV2Strategy(session: SparkSession) extends Strategy with Predicat // projection and filters were already pushed down in the optimizer. // this uses PhysicalOperation to get the projection and ensure that if the batch scan does // not support columnar, a projection is added to convert the rows to UnsafeRow. - val (runtimeFilters, postScanFilters) = filters.partition { + val (dynamicFilters, postScanFilters) = filters.partition { case _: DynamicPruning => true case _ => false } + + // 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(). + val scalarSubqueryFilters = if (relation.runtimeFilterAttrs.nonEmpty) { + postScanFilters.filter { f => + f.containsPattern(SCALAR_SUBQUERY) && + f.references.nonEmpty && + f.references.subsetOf(relation.runtimeFilterAttrs) + } + } else { + Seq.empty + } + val runtimeFilters = dynamicFilters ++ scalarSubqueryFilters + val batchExec = BatchScanExec(relation.output, relation.scan, runtimeFilters, relation.ordering, relation.relation.table, relation.keyGroupedPartitioning) DataSourceV2Strategy.withProjectAndFilter( @@ -746,6 +763,19 @@ private[sql] object DataSourceV2Strategy extends Logging { None } + /** + * Literalizes scalar subqueries in the given expression and translates the result to a V2 + * [[Predicate]]. Used at runtime in [[BatchScanExec]] after scalar subqueries have been + * evaluated. + */ + protected[sql] def translateScalarSubqueryFilterV2( + expr: Expression): Option[Predicate] = { + val literalized = expr.transform { + case s: ExecScalarSubquery => s.toLiteral + } + translateFilterV2(literalized) + } + /** * Creates new spark plan that should apply given filters and projections to given scan node * @param project Projection list that should be output of returned spark plan diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala index b10eeebb2b84b..d2cc342f48112 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2SQLSuite.scala @@ -4315,7 +4315,44 @@ class DataSourceV2SQLSuiteV1Filter } class DataSourceV2SQLSuiteV2Filter extends DataSourceV2SQLSuite { + import org.apache.spark.sql.catalyst.expressions.DynamicPruning + import org.apache.spark.sql.execution.datasources.v2.BatchScanExec + override protected val catalogAndNamespace = "testv2filter.ns1.ns2." + + test("SPARK-56467: scalar subquery filters on partition columns are pushed into runtimeFilters") { + val tbl = s"${catalogAndNamespace}tbl" + val dim = s"${catalogAndNamespace}dim" + withTable(tbl, dim) { + sql(s"CREATE TABLE $tbl (id INT, part INT) USING $v2Format PARTITIONED BY (part)") + for (i <- 0 until 10) { + sql(s"INSERT INTO $tbl VALUES ($i, $i)") + } + + sql(s"CREATE TABLE $dim (val INT) USING $v2Format") + sql(s"INSERT INTO $dim VALUES (3)") + + val df = sql(s"SELECT * FROM $tbl WHERE part = (SELECT max(val) FROM $dim)") + + // Verify query correctness + checkAnswer(df, Row(3, 3)) + + // Verify runtime filters contain the scalar subquery filter + val batchScan = collect(df.queryExecution.executedPlan) { + case b: BatchScanExec => b + }.head + assert(batchScan.runtimeFilters.nonEmpty, + "Expected runtimeFilters to contain scalar subquery filter") + assert(!batchScan.runtimeFilters.exists( + _.isInstanceOf[DynamicPruning]), + "Expected non-DPP runtime filter (scalar subquery)") + + // Verify partition pruning: only 1 of 10 partitions should remain + val numPartitions = batchScan.filteredPartitions.count(_.isDefined) + assert(numPartitions == 1, + s"Expected 1 partition after scalar subquery pruning, got $numPartitions") + } + } } class ReserveSchemaNullabilityCatalog extends InMemoryCatalog {