gengliangwang commented on code in PR #59662:
URL: https://github.com/apache/spark/pull/59662#discussion_r4351939311
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -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])
Review Comment:
`isSubquery` comes from whether the plan root is a `Subquery` when this rule
runs, but IN/EXISTS and correlated scalar/lateral subqueries are optimized in
the "Subquery" batch (`OptimizeSubqueries`) and only later rewritten into
semi/left joins in the main query by `RewritePredicateSubquery` /
`RewriteCorrelatedScalarSubquery`. Their scans are then built with
`isInSubquery=true` even though they end up as main-query join inputs.
If a connector follows the javadoc and skips or cheapens statistics when the
flag is true, that join side falls back to default size estimates, which can
lead to a bad broadcast-vs-shuffle choice or join order. The test "scalar, IN
and EXISTS subqueries provide context" asserts `Set(true)` for IN/EXISTS, which
locks this in. Should the flag be restricted to subqueries that actually stay
subqueries (e.g. uncorrelated scalar subqueries), or documented as a
best-effort hint?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsRead.java:
##########
@@ -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(ScanBuilderContext context,
CaseInsensitiveStringMap options) {
Review Comment:
Several batch read paths still call the options-only overload, so connectors
get no context there: `GroupBasedRowLevelOperationScanPlanning` (which runs
right before `V2ScanRelationPushDown` in `earlyScanPushDownRules`),
`RowLevelOperationTable.newScanBuilder`, `ChangelogTable`, and
`DataSourceV2Relation.computeStats`. A connector overriding this overload would
behave differently on those paths without any signal. Should those callers be
migrated too (or the description's "Spark supplies this context when creating
batch scan builders" narrowed)?
##########
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 =>
Review Comment:
Hard-coding `isSubquery = true` covers `MergeSubplans` only.
`PartitionPruning` and `InjectRuntimeFilter` also put existing
`DataSourceV2ScanRelation`s (built with `isInSubquery=false`) into subquery
plans, so equivalent subquery scans report different values depending on which
rule produced them. Together with the comments in `V2ScanRelationPushDown`, it
may be worth either deriving the context where the scan's consumer is actually
decided, or documenting that the flag is a best-effort, stage-specific hint.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDownSubquerySuite.scala:
##########
@@ -0,0 +1,289 @@
+/*
+ * 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,
ScanBuilderContext, 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("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 ScanBuilderContext(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
Review Comment:
`calls` is a global, unsynchronized `ArrayBuffer` shared by every session in
the JVM that loads `subquery_stats`. Concurrent planning or another suite using
the catalog could append extra entries and make `assertContexts` flaky. Could
this be scoped per test (or at least synchronized)?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -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] = {
Review Comment:
`rebuildScan` is on a public object; making `isSubquery` required breaks
existing 3-arg callers (extensions / forks). Would a default (`isSubquery:
Boolean = false`) or an overload keeping the old behavior work here?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/ScanBuilderContext.java:
##########
@@ -0,0 +1,44 @@
+/*
+ * 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 ScanBuilderContext {
Review Comment:
Since this is meant as a foundation for more planning context, a concrete
class with a public constructor will need a new public constructor per field,
and connectors/tests can instantiate it directly so old constructors can never
be removed. How about following the `LogicalWriteInfo` / `LogicalWriteInfoImpl`
pattern: a public interface plus a Spark-internal implementation?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala:
##########
@@ -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 = {
Review Comment:
Related: a non-inlined CTE consumed only from subqueries, e.g. `WITH c AS
(SELECT ... FROM v2t) SELECT (SELECT max(x) FROM c), (SELECT min(x) FROM c)`,
is pushed down as part of the main plan (root is `WithCTE`), so its scan gets
`isInSubquery=false` although it is only read through scalar subqueries. The
flag ends up depending on where the tree was rooted during optimization rather
than on how the scan is consumed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
