This is an automated email from the ASF dual-hosted git repository.
gengliangwang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 7ce9afec1992 [SPARK-60558][SQL] Expose subquery planning context to
DataSource V2 scan builders
7ce9afec1992 is described below
commit 7ce9afec19929e902c123ca821a99caec2ada45e
Author: yyanyy <[email protected]>
AuthorDate: Mon Oct 26 19:37:40 2026 -0700
[SPARK-60558][SQL] Expose subquery planning context to DataSource V2 scan
builders
### What changes were proposed in this pull request?
Add the evolving `ScanBuilderInfo` API with an `isInSubquery()` flag and a
default `SupportsRead.newScanBuilder(ScanBuilderInfo,
CaseInsensitiveStringMap)` overload. Spark supplies this context when creating
batch scan builders, including the destination context when rebuilding merged
scans.
The default overload delegates to the existing options-only method,
preserving compatibility with existing connectors and passing the original read
options unchanged. Each builder receives a fresh context object. The context is
advisory and must not change the rows returned by a scan.
### Why are the changes needed?
Different stages of query planning may call for different connector
behavior. Providing subquery-planning context allows connectors to adjust their
behavior and avoid unnecessary work. This also provides a foundation for Spark
to expose additional planning context to connectors in the future.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Passed 199 tests across `V2ScanRelationPushDownSubquerySuite`,
`MergeSubplansSuite`, and `DataSourceV2Suite`:
```text
build/sbt 'sql/testOnly *V2ScanRelationPushDownSubquerySuite
*MergeSubplansSuite org.apache.spark.sql.connector.DataSourceV2Suite'
```
The new coverage checks context availability before pushdown and statistics
selection, compatibility with options-only connectors, unchanged option-map
identity and contents, fresh context across repeated table use,
scalar/IN/EXISTS subqueries, local/global temporary-view reuse and replacement,
and the destination context of rebuilt merged scans.
Java Checkstyle, Scala style checks for SQL production and test sources,
and `git diff --check` also passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex 0.163.0 (GPT-6)
Closes #59662 from yyanyy/spark-subquery-stats-20261020.
Authored-by: yyanyy <[email protected]>
Signed-off-by: Gengliang Wang <[email protected]>
(cherry picked from commit 4c9713ee755a7c814693e4f53843b513dfdd4103)
Signed-off-by: Gengliang Wang <[email protected]>
---
.../spark/sql/connector/catalog/SupportsRead.java | 18 ++
.../spark/sql/connector/read/ScanContext.java | 64 +++++
.../datasources/v2/V2ScanRelationPushDown.scala | 26 +-
.../sql/execution/planmerging/PlanMerger.scala | 5 +-
.../v2/V2ScanRelationPushDownSubquerySuite.scala | 303 +++++++++++++++++++++
.../execution/planmerging/MergeSubplansSuite.scala | 36 ++-
6 files changed, 442 insertions(+), 10 deletions(-)
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsRead.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsRead.java
index e0ed9772dced..83684e4242ab 100644
---
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsRead.java
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsRead.java
@@ -20,6 +20,7 @@ package org.apache.spark.sql.connector.catalog;
import org.apache.spark.annotation.Evolving;
import org.apache.spark.sql.connector.read.Scan;
import org.apache.spark.sql.connector.read.ScanBuilder;
+import org.apache.spark.sql.connector.read.ScanContext;
import org.apache.spark.sql.util.CaseInsensitiveStringMap;
/**
@@ -40,4 +41,21 @@ public interface SupportsRead extends Table {
* string-to-string map.
*/
ScanBuilder newScanBuilder(CaseInsensitiveStringMap options);
+
+ /**
+ * Returns a {@link ScanBuilder} configured with read options and planning
context.
+ *
+ * <p>The default implementation delegates to
+ * {@link #newScanBuilder(CaseInsensitiveStringMap)}. Data sources can
override this method to use
+ * the planning context during builder creation. Calls to the options-only
overload do not provide
+ * subquery context.
+ *
+ * @param context the planning context for this scan
+ * @param options the options for reading, which is an immutable
case-insensitive
+ * string-to-string map
+ * @since 4.4.0
+ */
+ default ScanBuilder newScanBuilder(ScanContext context,
CaseInsensitiveStringMap options) {
+ return newScanBuilder(options);
+ }
}
diff --git
a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/ScanContext.java
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/ScanContext.java
new file mode 100644
index 000000000000..b9ae61263a90
--- /dev/null
+++
b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/ScanContext.java
@@ -0,0 +1,64 @@
+/*
+ * 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.read;
+
+import org.apache.spark.annotation.Evolving;
+
+/**
+ * Planning context provided to a data source when creating a {@link
ScanBuilder}.
+ *
+ * @since 4.4.0
+ */
+@Evolving
+public class ScanContext {
+ private final boolean isInSubquery;
+
+ public ScanContext(boolean isInSubquery) {
+ this.isInSubquery = isInSubquery;
+ }
+
+ /**
+ * Returns whether Spark is constructing this scan within a subquery.
+ *
+ * <p>The context describes the planning scope when the scan builder is
created. Subsequent
+ * optimizer rewrites may move or reuse the scan in another scope.
+ *
+ * <p>Data sources can use this context to adjust planning work, such as
statistics computation.
+ * It must not change the rows returned by the scan.
+ */
+ public boolean isInSubquery() {
+ return isInSubquery;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (!(o instanceof ScanContext that)) return false;
+ return isInSubquery == that.isInSubquery;
+ }
+
+ @Override
+ public int hashCode() {
+ return Boolean.hashCode(isInSubquery);
+ }
+
+ @Override
+ public String toString() {
+ return "ScanContext{isInSubquery=" + isInSubquery + "}";
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
index 768ed370c4cc..7140fc44adaa 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala
@@ -28,14 +28,14 @@ import
org.apache.spark.sql.catalyst.expressions.{aggregate, Alias, And, Attribu
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.catalyst.optimizer.{CollapseGroupedSumOfCount,
CollapseProject}
import org.apache.spark.sql.catalyst.planning.{PhysicalOperation,
ScanOperation}
-import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join,
LeafNode, Limit, LimitAndOffset, LocalLimit, LocalRelation, LogicalPlan,
Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join,
LeafNode, Limit, LimitAndOffset, LocalLimit, LocalRelation, LogicalPlan,
Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort, Subquery}
import org.apache.spark.sql.catalyst.plans.logical.PlanHelper
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder}
import org.apache.spark.sql.connector.expressions.aggregate.{Aggregation, Avg,
Count, CountStar, Max, Min, Sum}
import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, Statistics =>
V2Statistics, SupportsPushDownAggregates, SupportsPushDownFilters,
SupportsPushDownJoin, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters, SupportsPushDownVariantExtractions,
SupportsReportStatistics, V1Scan, VariantExtraction}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, ScanContext,
Statistics => V2Statistics, SupportsPushDownAggregates,
SupportsPushDownFilters, SupportsPushDownJoin, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters, SupportsPushDownVariantExtractions,
SupportsReportStatistics, V1Scan, VariantExtraction}
import org.apache.spark.sql.execution.datasources.{DataSourceStrategy,
VariantInRelation, VariantMetadata}
import org.apache.spark.sql.internal.SQLConf
import
org.apache.spark.sql.internal.connector.{SupportsPushDownCatalystFilters,
VariantExtractionImpl}
@@ -48,9 +48,12 @@ import org.apache.spark.util.Utils
object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
import DataSourceV2Implicits._
- def apply(plan: LogicalPlan): LogicalPlan = {
+ def apply(plan: LogicalPlan): LogicalPlan =
+ applyWithContext(plan, isSubquery = plan.isInstanceOf[Subquery])
+
+ private def applyWithContext(plan: LogicalPlan, isSubquery: Boolean):
LogicalPlan = {
val pushdownRules = Seq[LogicalPlan => LogicalPlan] (
- createScanBuilder,
+ createScanBuilder(_, isSubquery),
pushDownSample,
pushDownFilters,
pushDownJoin,
@@ -79,13 +82,17 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
* translation, column pruning, the iterative PartitionPredicate second
pass) stays owned here.
* `conditions` are pushed as a single `Filter` directly above the relation;
the caller decides
* which of them must come back fully enforced.
+ * `isSubquery` describes the destination plan's scope, absent from this
synthetic input.
*/
def rebuildScan(
relation: DataSourceV2Relation,
projectList: Seq[NamedExpression],
- conditions: Seq[Expression]): Option[DataSourceV2ScanRelation] = {
+ conditions: Seq[Expression],
+ isSubquery: Boolean): Option[DataSourceV2ScanRelation] = {
val child = conditions.reduceOption(And).map(Filter(_,
relation)).getOrElse(relation)
- apply(Project(projectList, child)).collectFirst { case s:
DataSourceV2ScanRelation => s }
+ applyWithContext(Project(projectList, child), isSubquery).collectFirst {
+ case scan: DataSourceV2ScanRelation => scan
+ }
}
private def collapseGroupedSumOfCount(plan: LogicalPlan): LogicalPlan = {
@@ -97,9 +104,12 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan]
with PredicateHelper {
}
}
- private def createScanBuilder(plan: LogicalPlan) = plan.transform {
+ private def createScanBuilder(
+ plan: LogicalPlan,
+ isSubquery: Boolean): LogicalPlan = plan.transform {
case r: DataSourceV2Relation =>
- ScanBuilderHolder(r.output, r,
r.table.asReadable.newScanBuilder(r.options))
+ val context = new ScanContext(isSubquery)
+ ScanBuilderHolder(r.output, r,
r.table.asReadable.newScanBuilder(context, r.options))
}
private def pushDownFilters(plan: LogicalPlan) = plan.transform {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
index bc57176e98dc..3bbf52a646be 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/planmerging/PlanMerger.scala
@@ -903,7 +903,10 @@ class PlanMerger(
// here: SPARK-58207 keeps non-deterministic filters from being pushed to
a V2 source, so an
// offered one is simply not pushed (and dropped when the scan is
extracted).
val conds = strictFilters ++
bestEffortFilter.filter(_.references.subsetOf(relationOut))
- V2ScanRelationPushDown.rebuildScan(relation, unionAttrs, conds).filter {
scan =>
+ // Spark reads the merged result through scalar subqueries, even when the
original
+ // scans came from the main query. Rebuild the scan with the destination's
context.
+ V2ScanRelationPushDown.rebuildScan(
+ relation, unionAttrs, conds, isSubquery = true).filter { scan =>
// The rebuilt scan must itself be mergeable. rebuildScan re-runs the
full pushdown, so any
// non-reproducible pushdown it introduces would make the merged scan
unsound. Today's
// Project-over-Filter input only triggers the plain path (always
mergeable), so this is
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDownSubquerySuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDownSubquerySuite.scala
new file mode 100644
index 000000000000..54744be9f70f
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDownSubquerySuite.scala
@@ -0,0 +1,303 @@
+/*
+ * 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.execution.datasources.v2
+
+import java.util
+import java.util.OptionalLong
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.expressions.{GreaterThan, Literal}
+import org.apache.spark.sql.catalyst.plans.logical.{Filter, LogicalPlan,
Project, Subquery}
+import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryTable,
InMemoryTableCatalog, SupportsRead, Table, TableCapability}
+import org.apache.spark.sql.connector.expressions.aggregate.Aggregation
+import org.apache.spark.sql.connector.read.{Batch, Scan, ScanBuilder,
ScanContext, Statistics, SupportsPushDownAggregates, SupportsPushDownFilters,
SupportsPushDownRequiredColumns, SupportsReportStatistics}
+import org.apache.spark.sql.sources
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class V2ScanRelationPushDownSubquerySuite extends QueryTest with
SharedSparkSession {
+ import SubqueryScanTestCatalog.calls
+
+ override protected def sparkConf: SparkConf = super.sparkConf.set(
+ "spark.sql.catalog.subquery_stats",
classOf[SubqueryScanTestCatalog].getName)
+
+ override protected def beforeEach(): Unit = {
+ super.beforeEach()
+ calls.clear()
+ }
+
+ private def newRelation(
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty):
DataSourceV2Relation = {
+ val schema = new StructType().add("id", IntegerType).add("value",
IntegerType)
+ val table = new InMemoryTable("test", schema, Array.empty,
util.Collections.emptyMap())
+ DataSourceV2Relation.create(new SubqueryScanTestTable(table), None, None,
options)
+ }
+
+ private def optimize(query: String): LogicalPlan =
sql(query).queryExecution.optimizedPlan
+
+ private def assertContexts(expected: Set[Boolean]): Unit = {
+ assert(calls.nonEmpty)
+ assert(calls.forall(_.context.isDefined))
+ assert(calls.map(_.context.get.isInSubquery()).toSet == expected)
+ calls.indices.foreach { i =>
+ (0 until i).foreach { j => assert(calls(i).context.get ne
calls(j).context.get) }
+ }
+ }
+
+ test("scan context equality, hashing and string representation") {
+ for (inSubquery <- Seq(false, true)) {
+ val context = new ScanContext(inSubquery)
+ val sameContext = new ScanContext(inSubquery)
+ assert(context.equals(context))
+ assert(context == sameContext)
+ assert(context.hashCode() == sameContext.hashCode())
+ assert(context != new ScanContext(!inSubquery))
+ assert(!context.equals(null))
+ assert(!context.equals(new Object))
+ assert(context.toString == s"ScanContext{isInSubquery=$inSubquery}")
+ }
+ }
+
+ test("the context overload delegates to an existing options-only
implementation") {
+ val options = new CaseInsensitiveStringMap(Map("MiXeD" -> "value").asJava)
+ var receivedOptions: CaseInsensitiveStringMap = null
+ val table: SupportsRead = new InMemoryTable(
+ "legacy", new StructType().add("id", IntegerType), Array.empty,
+ util.Collections.emptyMap()) {
+ override def newScanBuilder(options: CaseInsensitiveStringMap):
ScanBuilder = {
+ receivedOptions = options
+ super.newScanBuilder(options)
+ }
+ }
+ table.newScanBuilder(new ScanContext(true), options).build()
+ assert(receivedOptions eq options)
+ assert(receivedOptions.get("mixed") == "value")
+ }
+
+ test("direct options-only calls leave subquery context unspecified") {
+ val relation = newRelation()
+ val table = relation.table.asInstanceOf[SupportsRead]
+ table.newScanBuilder(relation.options).build()
+ assert(calls.size == 1)
+ assert(calls.head.context.isEmpty)
+ assert(calls.head.options eq relation.options)
+ }
+
+ for (inSubquery <- Seq(false, true)) {
+ test(s"builder context is available before pushdown:
inSubquery=$inSubquery") {
+ val input = Map("MiXeD" -> "value", "isInSubquery" -> "user option")
+ val options = new CaseInsensitiveStringMap(input.asJava)
+ val relation = newRelation(options)
+ val plan = Project(Seq(relation.output.head),
+ Filter(GreaterThan(relation.output.head, Literal(0)), relation))
+ val result = V2ScanRelationPushDown(
+ if (inSubquery) Subquery(plan, correlated = false) else plan)
+
+ assertContexts(Set(inSubquery))
+ val call = calls.head
+ assert(call.options eq options)
+ assert(options.asCaseSensitiveMap().asScala.toMap == input)
+ assert(call.events.head == "select statistics")
+ assert(call.events.contains("filters"))
+ assert(call.events.contains("prune"))
+ assert(call.requiredSchema.fieldNames.toSeq == Seq("id"))
+ val scan = result.collectFirst { case r: DataSourceV2ScanRelation =>
r.scan }.get
+ val stats =
scan.asInstanceOf[SupportsReportStatistics].estimateStatistics()
+ assert(stats.sizeInBytes().getAsLong == (if (inSubquery) 10L else 100L))
+ }
+ }
+
+ test("each planning pass receives a fresh context even when it reuses a
relation") {
+ val relation = newRelation()
+ V2ScanRelationPushDown(relation)
+ V2ScanRelationPushDown(Subquery(relation, correlated = false))
+ V2ScanRelationPushDown(relation)
+ assert(calls.map(_.context.get.isInSubquery()).toSeq == Seq(false, true,
false))
+ assertContexts(Set(false, true))
+ assert(calls.forall(_.options eq relation.options))
+ }
+
+ test("rebuilding a scan uses its destination context") {
+ val relation = newRelation()
+ V2ScanRelationPushDown(Subquery(relation, correlated = false))
+ for (inSubquery <- Seq(false, true)) {
+ val rebuilt = V2ScanRelationPushDown.rebuildScan(
+ relation, Seq(relation.output.head), Nil, isSubquery = inSubquery)
+ assert(rebuilt.isDefined)
+ assert(calls.last.context.get.isInSubquery() == inSubquery)
+ assert(calls.last.options eq relation.options)
+ }
+ assertContexts(Set(false, true))
+ }
+
+ test("scalar, IN and EXISTS subqueries provide context during SQL
optimization") {
+ withTable("subquery_stats.t") {
+ sql("CREATE TABLE subquery_stats.t (id INT, value INT) USING foo")
+ Seq(
+ "SELECT (SELECT max(id) FROM subquery_stats.t)",
+ "SELECT 1 WHERE 1 IN (SELECT id FROM subquery_stats.t)",
+ "SELECT 1 WHERE EXISTS (SELECT id FROM subquery_stats.t)").foreach {
query =>
+ calls.clear()
+ optimize(query)
+ assertContexts(Set(true))
+ }
+
+ calls.clear()
+ optimize("SELECT max(id) FROM subquery_stats.t")
+ assertContexts(Set(false))
+ assert(calls.exists(_.events.contains("aggregate")))
+
+ calls.clear()
+ optimize("SELECT (SELECT max(id) FROM subquery_stats.t)")
+ assertContexts(Set(true))
+ assert(calls.exists(_.events.contains("aggregate")))
+ }
+ }
+
+ test("the same table can be planned in the main query and a scalar
subquery") {
+ withTable("subquery_stats.t") {
+ sql("CREATE TABLE subquery_stats.t (id INT, value INT) USING foo")
+ optimize("SELECT id, (SELECT max(id) FROM subquery_stats.t) FROM
subquery_stats.t")
+ assertContexts(Set(false, true))
+ }
+ }
+
+ for (global <- Seq(false, true)) {
+ test(s"temporary view reuse preserves the planning context:
global=$global") {
+ val viewName = "subquery_stats_view"
+ val qualifiedView = if (global) s"global_temp.$viewName" else viewName
+ def withView(f: => Unit): Unit = {
+ if (global) withGlobalTempView(viewName)(f) else
withTempView(viewName)(f)
+ }
+ withTable("subquery_stats.t") {
+ withView {
+ sql("CREATE TABLE subquery_stats.t (id INT, value INT) USING foo")
+ val df = sql("SELECT id, (SELECT max(id) FROM subquery_stats.t) AS
maximum " +
+ "FROM subquery_stats.t")
+ df.queryExecution.optimizedPlan
+ assertContexts(Set(false, true))
+ if (global) df.createOrReplaceGlobalTempView(viewName)
+ else df.createOrReplaceTempView(viewName)
+
+ calls.clear()
+ optimize(s"SELECT * FROM $qualifiedView")
+ assertContexts(Set(false, true))
+
+ calls.clear()
+ optimize(s"SELECT (SELECT max(id) FROM $qualifiedView)")
+ assertContexts(Set(true))
+
+ calls.clear()
+ optimize(s"SELECT * FROM $qualifiedView")
+ assertContexts(Set(false, true))
+
+ val replacement = sql("SELECT id FROM subquery_stats.t")
+ if (global) replacement.createOrReplaceGlobalTempView(viewName)
+ else replacement.createOrReplaceTempView(viewName)
+ calls.clear()
+ optimize(s"SELECT * FROM $qualifiedView")
+ assertContexts(Set(false))
+ }
+ }
+ }
+ }
+}
+
+class SubqueryScanTestCatalog extends InMemoryTableCatalog {
+ override def loadTable(ident: Identifier): Table = {
+ new
SubqueryScanTestTable(super.loadTable(ident).asInstanceOf[InMemoryTable])
+ }
+}
+
+object SubqueryScanTestCatalog {
+ val calls: ArrayBuffer[SubqueryScanTestCall] = ArrayBuffer.empty
+}
+
+class SubqueryScanTestCall(
+ val context: Option[ScanContext],
+ val options: CaseInsensitiveStringMap) {
+ val events: ArrayBuffer[String] = ArrayBuffer("select statistics")
+ val sizeInBytes: Long = if (context.exists(_.isInSubquery())) 10L else 100L
+ var requiredSchema: StructType = new StructType()
+}
+
+class SubqueryScanTestTable(table: InMemoryTable) extends Table with
SupportsRead {
+ override def name(): String = table.name
+ override def schema(): StructType = table.schema()
+ override def capabilities(): util.Set[TableCapability] = {
+ util.EnumSet.of(TableCapability.BATCH_READ)
+ }
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder
= {
+ newBuilder(None, options)
+ }
+
+ override def newScanBuilder(
+ context: ScanContext,
+ options: CaseInsensitiveStringMap): ScanBuilder = {
+ newBuilder(Some(context), options)
+ }
+
+ private def newBuilder(
+ context: Option[ScanContext],
+ options: CaseInsensitiveStringMap): ScanBuilder = {
+ val call = new SubqueryScanTestCall(context, options)
+ SubqueryScanTestCatalog.calls += call
+ val delegate = table.newScanBuilder(options)
+ new ScanBuilder with SupportsPushDownRequiredColumns with
SupportsPushDownFilters
+ with SupportsPushDownAggregates {
+ override def pruneColumns(requiredSchema: StructType): Unit = {
+ call.events += "prune"
+ call.requiredSchema = requiredSchema
+
delegate.asInstanceOf[SupportsPushDownRequiredColumns].pruneColumns(requiredSchema)
+ }
+
+ override def pushFilters(filters: Array[sources.Filter]):
Array[sources.Filter] = {
+ call.events += "filters"
+ delegate.asInstanceOf[SupportsPushDownFilters].pushFilters(filters)
+ }
+
+ override def pushedFilters(): Array[sources.Filter] = {
+ delegate.asInstanceOf[SupportsPushDownFilters].pushedFilters()
+ }
+
+ override def pushAggregation(aggregation: Aggregation): Boolean = {
+ call.events += "aggregate"
+ false
+ }
+
+ override def build(): Scan = {
+ call.events += "build"
+ val scan = delegate.build()
+ new Scan with SupportsReportStatistics {
+ override def readSchema(): StructType = scan.readSchema()
+ override def toBatch: Batch = scan.toBatch
+ override def estimateStatistics(): Statistics = new Statistics {
+ override def sizeInBytes(): OptionalLong =
OptionalLong.of(call.sizeInBytes)
+ override def numRows(): OptionalLong = OptionalLong.empty()
+ }
+ }
+ }
+ }
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
index 564fd617d6d1..2f2a8082978d 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/planmerging/MergeSubplansSuite.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.planmerging
import java.util.OptionalLong
+import scala.collection.mutable.ArrayBuffer
+
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.dsl.plans._
import org.apache.spark.sql.catalyst.expressions.{Alias, And, Ascending,
Attribute, AttributeReference, CreateNamedStruct, Expression, ExpressionSet,
ExprId, GetStructField, GreaterThan, If, LessThan, Literal, Or, ScalarSubquery,
SortOrder, TransformExpression}
@@ -31,7 +33,7 @@ import
org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier, Supp
import org.apache.spark.sql.connector.catalog.functions.{BoundFunction,
ScalarFunction, UnboundFunction}
import org.apache.spark.sql.connector.expressions.{Expressions,
FieldReference, SortDirection => V2SortDirection, SortOrder => V2SortOrder}
import org.apache.spark.sql.connector.expressions.filter.Predicate
-import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, Statistics =>
V2Statistics, SupportsPushDownLimit, SupportsPushDownOffset,
SupportsPushDownRequiredColumns, SupportsPushDownTableSample,
SupportsPushDownTopN, SupportsPushDownV2Filters, SupportsReportOrdering,
SupportsReportPartitioning, SupportsReportStatistics}
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, ScanContext,
Statistics => V2Statistics, SupportsPushDownLimit, SupportsPushDownOffset,
SupportsPushDownRequiredColumns, SupportsPushDownTableSample,
SupportsPushDownTopN, SupportsPushDownV2Filters, SupportsReportOrdering,
SupportsReportPartitioning, SupportsReportStatistics}
import
org.apache.spark.sql.connector.read.partitioning.{KeyGroupedPartitioning =>
V2KeyGroupedPartitioning, Partitioning => V2Partitioning, UnknownPartitioning
=> V2UnknownPartitioning}
import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
DataSourceV2ScanRelation, V2ScanPartitioningAndOrdering, V2ScanRelationPushDown}
import org.apache.spark.sql.internal.SQLConf
@@ -2527,6 +2529,38 @@ class MergeSubplansSuite extends PlanTest {
comparePlans(Optimize.execute(originalQuery.analyze),
correctAnswer.analyze)
}
+ test("merged DSv2 scans receive subquery context for main-query aggregates")
{
+ val scanContexts = ArrayBuffer.empty[Boolean]
+ val table = new TestV2Table(StructType(Seq(
+ StructField("a", IntegerType), StructField("b", IntegerType)))) {
+ override def newScanBuilder(
+ context: ScanContext,
+ options: CaseInsensitiveStringMap): ScanBuilder = {
+ scanContexts += context.isInSubquery()
+ super.newScanBuilder(options)
+ }
+ }
+
+ def aggregate(column: String): LogicalPlan = {
+ val relation = DataSourceV2Relation.create(table, None, None)
+ val attribute = relation.output.find(_.name == column).get
+ val pushed = V2ScanRelationPushDown(Project(Seq(attribute), relation))
+ pushed.groupBy()(sum(attribute).as(s"sum_$column"))
+ }
+
+ val originalQuery = aggregate("a").join(aggregate("b")).analyze
+ assert(scanContexts.toSeq == Seq(false, false))
+ scanContexts.clear()
+
+ // Merging moves these aggregates into scalar subqueries and rebuilds
their shared scan.
+ val optimized = Optimize.execute(originalQuery)
+ assert(optimized.isInstanceOf[WithCTE])
+ val scans = v2Scans(optimized)
+ assert(scans.size == 1, s"expected the scans to merge:\n$optimized")
+ assert(scans.head.output.map(_.name).toSet == Set("a", "b"))
+ assert(scanContexts.nonEmpty && scanContexts.forall(identity))
+ }
+
gridTest("SPARK-58892: inferred filters survive rebuilding a merged scan")(
Seq(false, true)) { estimateInferred =>
val table = new TestInferredFilterTable(estimateInferred)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
