dongjoon-hyun commented on code in PR #59943:
URL: https://github.com/apache/spark/pull/59943#discussion_r4463353250
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -261,6 +260,19 @@ class MicroBatchExecution(
}
import
org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
+ // The columns and the filters that the query reads from each Data Source
V2 source, which the
+ // scan of the source gets if it supports them. They come from the plan of
the query optimized
+ // once as its micro-batches are, with their current time and their
configuration, which a
+ // restarted query takes from the offset log. The plan is only optimized
for a scan builder
+ // that supports a pushdown.
+ lazy val sourceReads = {
+ val batchConf = streamConf.clone()
+ initialLatestOffsetSeq.flatMap(_._2.metadataOpt).foreach { metadata =>
+ OffsetSeqMetadata.setSessionConf(metadata, batchConf, logChanges =
false)
+ }
+
StreamingV2ScanPushDown.reads(sparkSessionForStream.sessionState.optimizer,
batchConf,
+ withBatchTimestamp(planWithDedupKeys, batchTimestampMs = 0L))
Review Comment:
`withBatchTimestamp` only replaces `CurrentTimestamp`,
`CurrentTimestampNanos`, `LocalTimestamp`, `LocalTimestampNanos`, and
`CurrentDate`. So, `ComputeCurrentTime` still folds `current_time()` (and the
current date which `CAST(time AS TIMESTAMP*)` uses) into a literal during this
startup optimization, and the resulting deterministic filter is pushed down to
the scan.
For example, `readStream.table(t).where("tod < current_time()")` on a TIME
column pushes `tod < TIME'09:00:00'` at startup. The later batches fold
`current_time()` to their own time (e.g., `15:00`), but the scan keeps dropping
the rows with `tod` in `[09:00, 15:00)`, which the batch plan would keep. This
is a silent data loss.
The continuous mode has the same issue because the guard of `runContinuous`
doesn't reject `CurrentTime`. After a reconfiguration, the query is planned
again with a new time while the scan keeps the old filter.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Strategy.scala:
##########
@@ -1101,6 +1103,30 @@ private[sql] object DataSourceV2Strategy extends Logging
{
withFilter
}
}
+
+ /**
+ * Returns the output of the scan node of a streaming relation: the columns
that its scan reads.
+ * A scan built to read fewer columns than the relation outputs only gets
them when the
+ * projections and the filters over the relation, which the scan node is
planned with, read no
+ * other column of the relation (see `StreamingV2ScanPushDown`).
+ */
+ private def streamingScanOutput(
Review Comment:
This validates that the scan reads every column which the micro-batch plan
needs, but nothing validates that the micro-batch plan still applies the
filters which were pushed down at startup. Any divergence between the startup
plan and a micro-batch plan silently drops rows instead of failing, e.g., the
CTE and `current_time()` cases in the other comments, a conf changed later via
the `foreachBatch` session (like excluded optimizer rules on
`sparkSessionForStream`), or user extra strategies.
What about recording the pushed filters on
`StreamingDataSourceV2ScanRelation` and requiring here that
`ExpressionSet(filters)` contains them, like `prunedOutput` for the columns? It
would turn these cases into a loud failure.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.streaming.runtime
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, ExpressionSet, SubqueryExpression}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.v2.PushDownUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pushes the top-level columns and the filters that the optimized plan of a
streaming query reads
+ * from each Data Source V2 streaming source down to the scan of the source, as
+ * `V2ScanRelationPushDown` does for the scans of batch queries.
+ *
+ * The stream of a source is created from its scan, and has to exist before
any micro-batch is
+ * planned, so the scan is built when the query starts. The query is optimized
once at that point,
+ * by the optimizer of its micro-batches and with their configuration, and the
scan of each source
+ * gets what the planner takes from the optimized plan for the relation of the
source (see
+ * `DataSourceV2Strategy`): the projections and the filters over each
occurrence of the relation.
+ *
+ * The plan of the query does not change, so each micro-batch is optimized as
without these
+ * pushdowns: the optimizer keeps deciding which columns every operator gets,
and in which order,
+ * including the stateful ones, and Spark still evaluates all the filters over
the relation. Only
+ * the scan node of the physical plan outputs the columns that the scan reads,
and the planner
+ * checks that the micro-batch reads no other column of the relation.
+ *
+ * The other operations that the scans of batch queries can apply, such as
aggregates and limits,
+ * apply to the data of all the micro-batches of a streaming query, which the
scan of a micro-batch
+ * does not read, so they are not pushed down.
+ */
+private[sql] object StreamingV2ScanPushDown extends Logging {
+
+ /**
+ * What the optimized plan of a streaming query reads from a streaming
relation.
+ *
+ * @param requiredColumns the columns of the relation that the plan reads
+ * @param filters the deterministic filters without subqueries that the plan
applies over every
+ * occurrence of the relation, which its scan may apply too
+ */
+ case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression])
+
+ /**
+ * Returns what the optimized `plan` reads from each of its streaming
relations of Data Source
+ * V2: what the planner takes from the projections and the filters over each
occurrence of the
+ * relation, or all its columns and no filter where nothing is projected or
filtered over it. A
+ * relation that the optimized plan does not have, or any relation if the
optimizer fails, has no
+ * entry, so that its scan reads all its columns and gets no filter.
+ *
+ * @param optimizer the optimizer of the micro-batches of the query
+ * @param conf the configuration that the micro-batches of the query run with
+ * @param plan the plan of the query, with the current time of a micro-batch
as the micro-batches
+ * have it, which the optimizer does not fold, so that no filter
on it is pushed down
+ */
+ def reads(
+ optimizer: RuleExecutor[LogicalPlan],
+ conf: SQLConf,
+ plan: LogicalPlan): Map[StreamingRelationV2, Reads] = {
+ SQLConf.withExistingConf(conf) {
+ try {
+ val occurrences = mutable.HashMap.empty[StreamingRelationV2,
mutable.Buffer[Reads]]
+ // Visits the nodes as the planner plans them: the projections and the
filters over a
+ // relation are planned with its scan, and the other nodes on their
own.
+ def visit(node: LogicalPlan): Unit = node match {
+ case PhysicalOperation(project, filters, relation:
StreamingRelationV2) =>
+ occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) +=
Reads(
+ AttributeSet((project ++ filters).flatMap(_.references)),
+ filters.filter(f => f.deterministic &&
!SubqueryExpression.hasSubquery(f)))
+ case _ =>
+ node.children.foreach(visit)
+ }
+ visit(optimizer.execute(plan))
+ occurrences.map { case (relation, reads) =>
+ // The scan of a relation reads the rows of all its occurrences, so
it can only apply the
+ // filters that all of them apply.
+ val filters = reads.map(r => ExpressionSet(r.filters)).reduce {
(left, right) =>
+ ExpressionSet(left.filter(right.contains))
+ }
+ relation -> Reads(
+ AttributeSet.fromAttributeSets(reads.map(_.requiredColumns)),
filters.toSeq)
+ }.toMap
+ } catch {
+ case NonFatal(e) =>
+ logWarning(log"Could not optimize the streaming query to push down
what it reads to " +
+ log"the scans of its sources, which read all their columns with no
filter.", e)
+ Map.empty
+ }
+ }
+ }
+
+ /**
+ * Builds the scan of a streaming source, and creates its stream with
`createStream`. If its
+ * scan builder supports it, and `conf` enables it, the scan only reads the
columns of `output`
+ * that the query reads, and gets the filters that the query applies to all
the rows that it reads
+ * from the source (see `reads`). `reads` is only evaluated for such a scan
builder. Returns the
+ * scan, its stream, and the columns of `output` that the scan reads, in the
order of its read
+ * schema, unless these are `output` itself.
+ *
+ * A pruned scan must read the required columns and only columns of
`output`, each once, as
+ * `SupportsPushDownRequiredColumns` requires. A scan that does not is
replaced by a scan of all
+ * the columns, as without pruning, and its stream is stopped. Spark asks a
scan for its read
+ * schema once the stream is created, since a scan may only know then what
its stream reads.
+ *
+ * @param output the output of the relation of the source, with the metadata
columns that the
+ * query reads, which the analyzer added to it, so that
connectors implementing
+ * `SupportsMetadataColumns` include them in `readSchema()`
+ * @param conf the configuration of the query
+ * @param reads what the query reads from the relation, or None to read all
its columns
+ */
+ def buildScan[S <: SparkDataStream](
+ table: SupportsRead,
+ options: CaseInsensitiveStringMap,
+ output: Seq[AttributeReference],
+ conf: SQLConf,
+ reads: => Option[Reads])(
+ createStream: Scan => S): (Scan, S, Option[Seq[AttributeReference]]) = {
+ val builder = table.newScanBuilder(options)
+ val prunesColumns = conf.streamingV2ColumnPruningEnabled &&
+ builder.isInstanceOf[SupportsPushDownRequiredColumns]
+ val pushesFilters = conf.streamingV2FilterPushdownEnabled &&
supportsFilters(builder)
+ val known = if (prunesColumns || pushesFilters) reads else None
+ val requiredOutput = known.filter(_ => prunesColumns)
+ .map(r => output.filter(r.requiredColumns.contains))
+ .getOrElse(output)
+ val filters = known.filter(_ => pushesFilters)
+ .map(r => DataSourceStrategy.normalizeExprs(r.filters, output))
+ .getOrElse(Nil)
+
+ def build(builder: ScanBuilder, columns: Seq[AttributeReference]): Scan = {
+ if (filters.nonEmpty) {
+ // The plan of the query keeps the filters, which Spark evaluates
after the scan, so the
+ // scan may apply any of them.
+ PushDownUtils.pushFilters(builder, filters, partitionFields = None)
Review Comment:
Pushing the filters down to the streaming scans changes
`MicroBatchScanExec.numOutputRows`, which `ProgressReporter` reports as
`numInputRows`, `inputRowsPerSecond`, and `processedRowsPerSecond`.
For a source which honors the pushed filters, e.g., with `where("extra >
10")`, `numInputRows` used to count all the rows in the offset range, but now
it counts only the rows passing the filters. The dashboards and the alerts on
the input rates can drop without any change of the data volume. Toggling
`spark.sql.streaming.v2FilterPushdown.enabled` across restarts also changes the
metrics. Could you mention this user-facing change in the migration guide?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala:
##########
@@ -78,6 +77,13 @@ class ContinuousExecution(
val v2ToRelationMap = MutableMap[StreamingRelationV2,
StreamingDataSourceV2ScanRelation]()
var nextSourceId = 0
import
org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
+ // The columns and the filters that the query reads from each Data Source
V2 source, which the
+ // scan of the source gets if it supports them. They come from the plan of
the query optimized
+ // once, as it is optimized when it runs. The plan is only optimized for a
scan builder that
+ // supports a pushdown.
+ val streamConf = sparkSessionForStream.sessionState.conf
+ lazy val sourceReads = StreamingV2ScanPushDown.reads(
Review Comment:
The `sourceReads` setup is duplicated between `MicroBatchExecution` and
`ContinuousExecution`, and the two copies have already diverged. This one runs
in the constructor, on the thread calling `start()`, before `runStream` applies
the stream conf overrides (e.g., CBO and AQE off), and without any timestamp
substitution (see the `current_time()` comment).
What about moving the conf cloning and the time substitution into
`StreamingV2ScanPushDown.reads` to keep a single copy?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.streaming.runtime
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, ExpressionSet, SubqueryExpression}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.v2.PushDownUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pushes the top-level columns and the filters that the optimized plan of a
streaming query reads
+ * from each Data Source V2 streaming source down to the scan of the source, as
+ * `V2ScanRelationPushDown` does for the scans of batch queries.
+ *
+ * The stream of a source is created from its scan, and has to exist before
any micro-batch is
+ * planned, so the scan is built when the query starts. The query is optimized
once at that point,
+ * by the optimizer of its micro-batches and with their configuration, and the
scan of each source
+ * gets what the planner takes from the optimized plan for the relation of the
source (see
+ * `DataSourceV2Strategy`): the projections and the filters over each
occurrence of the relation.
+ *
+ * The plan of the query does not change, so each micro-batch is optimized as
without these
+ * pushdowns: the optimizer keeps deciding which columns every operator gets,
and in which order,
+ * including the stateful ones, and Spark still evaluates all the filters over
the relation. Only
+ * the scan node of the physical plan outputs the columns that the scan reads,
and the planner
+ * checks that the micro-batch reads no other column of the relation.
+ *
+ * The other operations that the scans of batch queries can apply, such as
aggregates and limits,
+ * apply to the data of all the micro-batches of a streaming query, which the
scan of a micro-batch
+ * does not read, so they are not pushed down.
+ */
+private[sql] object StreamingV2ScanPushDown extends Logging {
+
+ /**
+ * What the optimized plan of a streaming query reads from a streaming
relation.
+ *
+ * @param requiredColumns the columns of the relation that the plan reads
+ * @param filters the deterministic filters without subqueries that the plan
applies over every
+ * occurrence of the relation, which its scan may apply too
+ */
+ case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression])
+
+ /**
+ * Returns what the optimized `plan` reads from each of its streaming
relations of Data Source
+ * V2: what the planner takes from the projections and the filters over each
occurrence of the
+ * relation, or all its columns and no filter where nothing is projected or
filtered over it. A
+ * relation that the optimized plan does not have, or any relation if the
optimizer fails, has no
+ * entry, so that its scan reads all its columns and gets no filter.
+ *
+ * @param optimizer the optimizer of the micro-batches of the query
+ * @param conf the configuration that the micro-batches of the query run with
+ * @param plan the plan of the query, with the current time of a micro-batch
as the micro-batches
+ * have it, which the optimizer does not fold, so that no filter
on it is pushed down
+ */
+ def reads(
+ optimizer: RuleExecutor[LogicalPlan],
+ conf: SQLConf,
+ plan: LogicalPlan): Map[StreamingRelationV2, Reads] = {
+ SQLConf.withExistingConf(conf) {
+ try {
+ val occurrences = mutable.HashMap.empty[StreamingRelationV2,
mutable.Buffer[Reads]]
+ // Visits the nodes as the planner plans them: the projections and the
filters over a
+ // relation are planned with its scan, and the other nodes on their
own.
+ def visit(node: LogicalPlan): Unit = node match {
+ case PhysicalOperation(project, filters, relation:
StreamingRelationV2) =>
+ occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) +=
Reads(
Review Comment:
`Reads` is keyed by the exact `StreamingRelationV2` instance, but the
optimizer can copy the relation with `newInstance()` (e.g., `InlineCTE` /
`ReplaceCTERefWithRepartition` via `DeduplicateRelations.deduplicateRight`).
The reads of the copy are recorded under a key which nobody looks up, while the
copies share the same scan in every micro-batch.
For example,
```scala
streamDf.createOrReplaceTempView("s")
spark.sql("WITH c AS (SELECT * FROM s) " +
"SELECT a.id, b.extra FROM c a JOIN c b ON a.id = b.id WHERE a.data = 'x'")
```
At startup, `InlineCTE` keeps the original relation for `a` and uses
`s.newInstance()` for `b`, so only the reads of `a` (`id`, `data`, and `data =
'x'`) reach `buildScan`. In each micro-batch, `InlineCTE` clones the scan
relation with `StreamingDataSourceV2ScanRelation.newInstance()`, which shares
the scan and the stream.
- Since `b` needs `extra`, which is not in `prunedOutput`,
`streamingScanOutput` throws `INTERNAL_ERROR` at the first batch. This query
worked before this PR.
- If `b` reads only a subset of the columns of `a`, the pushed filter of `a`
(`data = 'x'`) is applied to the rows of `b` too. It gives wrong join results
for the connectors which honor the pushed filters.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingV2ScanPushDownSuite.scala:
##########
@@ -0,0 +1,664 @@
+/*
+ * 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.streaming
+
+import java.util
+
+import org.apache.logging.log4j.Level
+import org.scalatest.BeforeAndAfter
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.Project
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.classic.Dataset
+import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryBaseTable,
InMemoryTableCatalog, SupportsRead, Table, TableCapability}
+import org.apache.spark.sql.connector.read.{InputPartition, Scan, ScanBuilder,
SupportsPushDownRequiredColumns}
+import
org.apache.spark.sql.connector.read.streaming.{ContinuousPartitionReader,
ContinuousPartitionReaderFactory, ContinuousStream, Offset, PartitionOffset}
+import org.apache.spark.sql.execution.datasources.v2.{MicroBatchScanExec,
StreamingDataSourceV2ScanRelation}
+import org.apache.spark.sql.execution.streaming.runtime.{StreamExecution,
StreamingQueryWrapper, StreamingV2ScanPushDown}
+import org.apache.spark.sql.execution.streaming.sources.ContinuousMemoryStream
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataType, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class StreamingV2ScanPushDownSuite extends StreamTest with BeforeAndAfter {
+ import testImplicits._
+
+ private val tbl = "testcat.t"
+
+ before {
+ spark.conf.set("spark.sql.catalog.testcat",
classOf[InMemoryTableCatalog].getName)
+ }
+
+ after {
+ spark.sessionState.catalogManager.reset()
+ spark.conf.unset("spark.sql.catalog.testcat")
+ }
+
+ private def withSourceTable(properties: String = "")(f: => Unit): Unit = {
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id bigint, data string, extra int) $properties")
+ sql(s"INSERT INTO $tbl VALUES (1, 'a', 10), (2, 'b', 20)")
+ f
+ }
+ }
+
+ private def insert(values: String, table: String = tbl): StreamAction =
Execute { _ =>
+ sql(s"INSERT INTO $table VALUES $values")
+ }
+
+ private def pruningEnabled(enabled: Boolean): Map[String, String] =
+ Map(SQLConf.STREAMING_V2_COLUMN_PRUNING_ENABLED.key -> enabled.toString)
+
+ /**
+ * Runs `f`, which restarts streams with other values of the confs of these
pushdowns, and then
+ * restores their values: `testStream` restores the value that a conf had
before the last
+ * `StartStream` that set it, rather than before the first.
+ */
+ private def withPushDownConfsRestored(f: => Unit): Unit = withSQLConf(
+ SQLConf.STREAMING_V2_COLUMN_PRUNING_ENABLED.key -> "true",
+ SQLConf.STREAMING_V2_FILTER_PUSHDOWN_ENABLED.key -> "true")(f)
+
+ private def filterPushdownEnabled(enabled: Boolean): Map[String, String] =
+ Map(SQLConf.STREAMING_V2_FILTER_PUSHDOWN_ENABLED.key -> enabled.toString)
+
+ /** Returns the filters pushed down to the scan of each streaming source of
the query. */
+ private def pushedFilters(query: StreamExecution): Seq[Seq[String]] = {
+ query.logicalPlan.collect {
+ case relation: StreamingDataSourceV2ScanRelation =>
+ relation.scan match {
+ case scan: InMemoryBaseTable#InMemoryBatchScan =>
+ scan.pushedFilters.map(_.toString).toSeq.sorted
+ }
+ }
+ }
+
+ /** Checks that the filters pushed down to the scan of each streaming source
are `filters`. */
+ private def assertPushedFilters(filters: String*): StreamAction =
AssertOnQuery { query =>
+ val pushed = pushedFilters(query)
+ assert(pushed.nonEmpty)
+ pushed.foreach(p => assert(p === filters.sorted))
+ true
+ }
+
+ /** Checks the read schemas of the scans of the streaming sources of the
query. */
+ private def assertReadSchemas(columns: Seq[String]*): StreamAction =
AssertOnQuery { query =>
+ val readSchemas = query.logicalPlan.collect {
+ case relation: StreamingDataSourceV2ScanRelation =>
+ relation.scan.readSchema().fieldNames.toSeq
+ }
+ assert(readSchemas.sortBy(_.mkString(",")) ===
columns.sortBy(_.mkString(",")))
+ true
+ }
+
+ /** Checks that the scan of each streaming source of the query reads
`columns`. */
+ private def assertReadSchema(columns: String*): StreamAction = AssertOnQuery
{ query =>
+ val readSchemas = query.logicalPlan.collect {
+ case relation: StreamingDataSourceV2ScanRelation =>
+ relation.scan.readSchema().fieldNames.toSeq
+ }
+ assert(readSchemas.nonEmpty)
+ readSchemas.foreach(readSchema => assert(readSchema === columns))
+ true
+ }
+
+ /** Checks the output of the scan nodes of the last micro-batch. */
+ private def assertScanOutputs(columns: Seq[String]*): StreamAction =
AssertOnQuery { query =>
+ val outputs = query.lastExecution.executedPlan.collect {
+ case scan: MicroBatchScanExec => scan.output.map(_.name)
+ }
+ assert(outputs.sortBy(_.mkString(",")) === columns.sortBy(_.mkString(",")))
+ true
+ }
+
+ test("read only the columns that a streaming query needs, also after a
restart") {
+ withSourceTable() {
+ val df = spark.readStream.table(tbl).where("extra > 10").select("data")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("data", "extra"),
+ assertScanOutputs(Seq("data", "extra")),
+ CheckAnswer(Row("b")),
+ StopStream,
+ insert("(3, 'c', 30), (4, 'd', 5)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchema("data", "extra"),
+ CheckAnswer(Row("b"), Row("c")))
+ }
+ }
+
+ // Each case is a query over the source, and the columns that the optimizer
reads for it, which
+ // fewer than those that the operators of its analyzed plan reference.
+ Seq[(String, DataFrame => DataFrame, Seq[String])](
+ ("a column that a later projection drops",
+ _.withColumn("x", lit(1)).select("data"), Seq("data")),
+ ("chained projections", _.select("id", "data").select("data"),
Seq("data")),
+ ("expressions and functions that a later projection drops",
+ _.select(upper($"data").as("u"), ($"id" * 2).as("i")).select("u"),
Seq("data")),
+ ("several aliases of a column",
+ _.select($"id".as("x"), $"data".as("y")).select($"x".as("z")).where("z >
1"), Seq("id")),
+ ("a field of a struct of columns",
+ _.select(struct($"id", $"data").as("s")).select("s.data"), Seq("data")),
+ ("the fields of a struct of columns",
+ _.select(struct($"id", $"extra").as("s")).select("s.*"), Seq("id",
"extra")),
+ ("a filter between projections",
+ _.select("id", "data").where("id > 1").select("data"), Seq("id",
"data")),
+ ("a projection of all the columns", _.select("*").select("data"),
Seq("data"))
+ ).foreach { case (name, query, columns) =>
+ test(s"read only the columns that the optimized query reads: $name") {
+ withSourceTable() {
+ val expected = query(spark.read.table(tbl)).collect().toSeq
+ testStream(query(spark.readStream.table(tbl)))(
+ ProcessAllAvailable(),
+ assertReadSchema(columns: _*),
+ assertScanOutputs(columns),
+ CheckAnswer(expected: _*))
+ }
+ }
+ }
+
+ Seq[(String, DataFrame => DataFrame)](
+ ("from_json", _.select(from_json($"json", new StructType().add("a",
"int")).as("p"))
+ .select("p.a")),
+ ("get_json_object", _.select(get_json_object($"json", "$.a").as("a"))),
+ ("parse_json", _.select(parse_json($"json").as("v"))
+ .select(variant_get($"v", "$.a", "int").as("a")))
+ ).foreach { case (name, query) =>
+ test(s"read only the parsed column of $name") {
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id bigint, json string, extra int)")
+ sql(s"""INSERT INTO $tbl VALUES (1, '{"a": 1}', 10), (2, '{"a": 2}',
20)""")
+ val expected = query(spark.read.table(tbl)).collect().toSeq
+ testStream(query(spark.readStream.table(tbl)))(
+ ProcessAllAvailable(),
+ assertReadSchema("json"),
+ CheckAnswer(expected: _*))
+ }
+ }
+ }
+
+ test("read only the columns of a streaming aggregation, also after a
restart") {
+ withSourceTable() {
+ val df = spark.readStream.table(tbl).groupBy("data").count()
+ testStream(df, OutputMode.Complete())(
+ ProcessAllAvailable(),
+ assertReadSchema("data"),
+ CheckAnswer(Row("a", 1L), Row("b", 1L)),
+ StopStream,
+ insert("(3, 'a', 30)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchema("data"),
+ CheckAnswer(Row("a", 2L), Row("b", 1L)))
+ }
+ }
+
+ test("read only the event-time column of a time window") {
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id bigint, data string, ts timestamp)")
+ sql(s"INSERT INTO $tbl VALUES " +
+ "(1, 'a', timestamp_seconds(0)), (2, 'b', timestamp_seconds(15))")
+ val df = spark.readStream.table(tbl).groupBy(window($"ts", "10
seconds")).count()
+ .select($"window.start".cast("long"), $"count")
+ testStream(df, OutputMode.Complete())(
+ ProcessAllAvailable(),
+ assertReadSchema("ts"),
+ CheckAnswer(Row(0L, 1L), Row(10L, 1L)))
+ }
+ }
+
+ test("read no column of the source when the query needs none, also after a
restart") {
+ withSourceTable() {
+ testStream(spark.readStream.table(tbl).groupBy().count(),
OutputMode.Complete())(
+ ProcessAllAvailable(),
+ assertReadSchema(),
+ assertScanOutputs(Nil),
+ CheckAnswer(Row(2L)),
+ StopStream,
+ insert("(3, 'c', 30)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchema(),
+ CheckAnswer(Row(3L)))
+ testStream(spark.readStream.table(tbl).select(lit(1).as("one")))(
+ ProcessAllAvailable(),
+ assertReadSchema(),
+ CheckAnswer(Row(1), Row(1), Row(1)))
+ }
+ }
+
+ test("read only the columns of a streaming deduplication, also after a
restart") {
+ withSourceTable() {
+ val df = spark.readStream.table(tbl).dropDuplicates("data").select("id")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("id", "data"),
+ CheckAnswer(Row(1L), Row(2L)),
+ StopStream,
+ insert("(3, 'a', 30), (4, 'c', 40)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchema("id", "data"),
+ CheckAnswer(Row(1L), Row(2L), Row(4L)))
+ }
+ }
+
+ test("read the metadata columns that a streaming query needs") {
+ withSourceTable() {
+ val df = spark.readStream.table(tbl).select("data", "index")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("data", "index"),
+ CheckAnswerRowsByFunc(rows => {
+ assert(rows.map(_.getString(0)).sorted === Seq("a", "b"))
+ assert(rows.forall(!_.isNullAt(1)))
+ }, lastOnly = false))
+ }
+ }
+
+ test("read only the columns of each child of a union") {
+ withSourceTable() {
+ // Each read of the table has its own relation and scan, and the
optimizer prunes the
+ // columns of each child of the union by position.
+ val twoReads =
spark.readStream.table(tbl).union(spark.readStream.table(tbl)).select("data")
+ testStream(twoReads)(
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("data"), Seq("data")),
+ CheckAnswer(Row("a"), Row("b"), Row("a"), Row("b")))
+
+ val differentColumns = spark.readStream.table(tbl).select("id")
+ .union(spark.readStream.table(tbl).select($"extra".cast("long")))
+ testStream(differentColumns)(
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("id"), Seq("extra")),
+ CheckAnswer(Row(1L), Row(2L), Row(10L), Row(20L)))
+
+ // The children of a self-union read the same relation, whose scan reads
the columns that
+ // both of them read.
+ val source = spark.readStream.table(tbl)
+ testStream(source.select("id", "data").union(source.select("id",
"data")).select("data"))(
+ ProcessAllAvailable(),
+ AssertOnQuery { query =>
+ val relations = query.logicalPlan.collect {
+ case relation: StreamingDataSourceV2ScanRelation => relation
+ }
+ assert(relations.length == 2 &&
relations.map(_.stream).distinct.length == 1)
+ assert(relations.head.scan.readSchema().fieldNames.toSeq ===
Seq("data"))
+ true
+ },
+ CheckAnswer(Row("a"), Row("b"), Row("a"), Row("b")))
+ }
+ }
+
+ test("read the columns that a scan returns when it reads more than the
required ones") {
+ withSourceTable(
+ s"TBLPROPERTIES
('${InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING}' = 'true')") {
+ val df = spark.readStream.table(tbl).select("data", "index")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("id", "data", "extra", "index"),
+ CheckAnswerRowsByFunc(rows => {
+ assert(rows.map(_.getString(0)).sorted === Seq("a", "b"))
+ assert(rows.forall(!_.isNullAt(1)))
+ }, lastOnly = false))
+ }
+ }
+
+ test("read all the columns when a scan does not read the required ones") {
+ withSourceTable(
+ s"TBLPROPERTIES
('${InMemoryBaseTable.SIMULATE_INCOMPLETE_COLUMN_PRUNING}' = 'true')") {
+ val table = spark.sessionState.catalogManager.catalog("testcat")
+
.asInstanceOf[InMemoryTableCatalog].liveTable(Identifier.of(Array.empty, "t"))
+ .asInstanceOf[InMemoryBaseTable]
+ val logAppender = new LogAppender("the fallback of streaming column
pruning")
+ withLogAppender(logAppender,
+ loggerNames =
Seq(StreamingV2ScanPushDown.getClass.getName.stripSuffix("$")),
+ level = Some(Level.WARN)) {
+ testStream(spark.readStream.table(tbl).select("data"))(
+ ProcessAllAvailable(),
+ // The scan that does not read "data" is replaced by one of all the
columns, and the
+ // stream of the replaced scan is stopped.
+ assertReadSchema("id", "data", "extra"),
+ AssertOnQuery(_ => table.stoppedMicroBatchStreams.get() == 1),
+ CheckAnswer(Row("a"), Row("b")))
+ }
+ val warnings =
logAppender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ assert(warnings.exists { message =>
+ message.contains("was asked to read the columns data") &&
+ message.contains("does not have the columns data") &&
+ message.contains("It reads all the columns instead")
+ }, warnings)
+ }
+ }
+
+ test("read the event-time column of a watermark that the query does not
output") {
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (id bigint, data string, ts timestamp)")
+ sql(s"INSERT INTO $tbl VALUES " +
+ "(1, 'a', timestamp_seconds(0)), (2, 'b', timestamp_seconds(10))")
+ val df = spark.readStream.table(tbl).withWatermark("ts", "5
seconds").select("data")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("data", "ts"),
+ CheckAnswer(Row("a"), Row("b")),
+ // The watermark gets the event times of the rows.
+ AssertOnQuery { query =>
+ query.recentProgress.filter(_.numInputRows >
0).map(_.eventTime.get("max")).toSeq ==
+ Seq("1970-01-01T00:00:10.000Z")
+ })
+ }
+ }
+
+ test("read only the columns of a stream-static join") {
+ withSourceTable() {
+ val static = spark.range(2, 4).toDF("sid")
+ val source = spark.readStream.table(tbl)
+ val df = source.join(static, source("id") ===
static("sid")).select("data")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("id", "data"),
+ CheckAnswer(Row("b")),
+ insert("(3, 'c', 30), (4, 'd', 40)"),
+ ProcessAllAvailable(),
+ CheckAnswer(Row("b"), Row("c")))
+ }
+ }
+
+ test("read only the columns of each side of a stream-stream join, also after
a restart") {
+ withSourceTable() {
+ val left = spark.readStream.table(tbl).select("id", "data")
+ val right = spark.readStream.table(tbl).selectExpr("id AS rid", "extra")
+ val df = left.join(right, left("id") === right("rid")).select("data",
"extra")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("id", "data"), Seq("id", "extra")),
+ CheckAnswer(Row("a", 10), Row("b", 20)),
+ StopStream,
+ insert("(3, 'c', 30)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("id", "data"), Seq("id", "extra")),
+ CheckAnswer(Row("a", 10), Row("b", 20), Row("c", 30)))
+ }
+ }
+
+ test("keep the layout of the state of a stream-stream join across the values
of the conf " +
+ "when a scan returns the columns in another order") {
+ withPushDownConfsRestored {
+ val (leftTable, rightTable) = ("testcat.l", "testcat.r")
+ withTable(leftTable, rightTable) {
+ sql(s"CREATE TABLE $leftTable (k bigint, v bigint, unused int)
TBLPROPERTIES " +
+ s"('${InMemoryBaseTable.SIMULATE_REORDERED_COLUMN_PRUNING}' =
'true')")
+ sql(s"CREATE TABLE $rightTable (rk bigint, w int)")
+ // The left row is kept in the state of the join, which is written
without pruning first.
+ sql(s"INSERT INTO $leftTable VALUES (1, 100, 0)")
+ val left = spark.readStream.table(leftTable)
+ val right = spark.readStream.table(rightTable)
+ val df = left.join(right, $"k" === $"rk").select("v", "w")
+ testStream(df)(
+ StartStream(additionalConfs = pruningEnabled(false)),
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("k", "v", "unused"), Seq("rk", "w")),
+ CheckAnswer(),
+ StopStream,
+ // The pruned scan returns (v, k), and the join still gets (k, v).
+ insert("(1, 7)", rightTable),
+ StartStream(additionalConfs = pruningEnabled(true)),
+ ProcessAllAvailable(),
+ assertReadSchemas(Seq("v", "k"), Seq("rk", "w")),
+ CheckAnswer(Row(100L, 7)),
+ // And the other way around.
+ insert("(2, 200, 0)", leftTable),
+ ProcessAllAvailable(),
+ StopStream,
+ insert("(2, 8)", rightTable),
+ StartStream(additionalConfs = pruningEnabled(false)),
+ ProcessAllAvailable(),
+ CheckAnswer(Row(100L, 7), Row(200L, 8)))
+ }
+ }
+ }
+
+ test("keep the state of a streaming aggregation across the values of the
conf when a scan " +
+ "returns the columns in another order") {
+ withPushDownConfsRestored {
+ withTable(tbl) {
+ sql(s"CREATE TABLE $tbl (k bigint, v bigint, unused int) TBLPROPERTIES
" +
+ s"('${InMemoryBaseTable.SIMULATE_REORDERED_COLUMN_PRUNING}' =
'true')")
+ sql(s"INSERT INTO $tbl VALUES (1, 10, 0), (2, 20, 0)")
+ val df = spark.readStream.table(tbl).groupBy("k").agg(sum("v"))
+ testStream(df, OutputMode.Complete())(
+ StartStream(additionalConfs = pruningEnabled(false)),
+ ProcessAllAvailable(),
+ assertReadSchema("k", "v", "unused"),
+ CheckAnswer(Row(1L, 10L), Row(2L, 20L)),
+ StopStream,
+ insert("(1, 1, 0)"),
+ StartStream(additionalConfs = pruningEnabled(true)),
+ ProcessAllAvailable(),
+ assertReadSchema("v", "k"),
+ CheckAnswer(Row(1L, 11L), Row(2L, 20L)),
+ StopStream,
+ insert("(2, 2, 0)"),
+ StartStream(additionalConfs = pruningEnabled(false)),
+ ProcessAllAvailable(),
+ CheckAnswer(Row(1L, 11L), Row(2L, 22L)))
+ }
+ }
+ }
+
+ test("fail to plan a micro-batch that reads a column that its scan does not
read") {
+ withSourceTable() {
+ testStream(spark.readStream.table(tbl).select("data"))(
+ ProcessAllAvailable(),
+ AssertOnQuery { query =>
+ val relation = query.logicalPlan.collectFirst {
+ case relation: StreamingDataSourceV2ScanRelation => relation
+ }.get
+ assert(relation.output.map(_.name) === Seq("id", "data", "extra"))
+ assert(relation.prunedOutput.map(_.map(_.name)) ===
Some(Seq("data")))
+ val offset = relation.stream.initialOffset()
+ val batch = relation.copy(startOffset = Some(offset), endOffset =
Some(offset))
+ // Unlike the query, this plan reads all the columns of the relation.
+ val e = intercept[SparkException] {
+ spark.sessionState.planner.plan(Project(relation.output,
batch)).next()
+ }
+ assert(e.getCondition === "INTERNAL_ERROR")
+ assert(e.getMessage.contains("does not read the columns"))
+
assert(e.getMessage.contains(SQLConf.STREAMING_V2_COLUMN_PRUNING_ENABLED.key))
+ true
+ })
+ }
+ }
+
+ test("read only the columns that a continuous query needs") {
+ val input = ContinuousMemoryStream.singlePartition[(Long, String, Int)]
+ val table = new PruningContinuousTable(input)
+ val df = Dataset.ofRows(spark, StreamingRelationV2(None,
"pruning_continuous", table,
+ CaseInsensitiveStringMap.empty(),
DataTypeUtils.toAttributes(table.schema()), None, None,
+ None)).toDF("id", "data", "extra").where("extra > 10").select("data")
+ val query = df.writeStream.format("memory").queryName("continuous_pruning")
Review Comment:
Test cleanup:
- This test (and the new `ColumnarDataSourceSuite` test with
`pruned_stream`) registers a memory sink temp view which is never dropped.
`continuous_pruning` and `pruned_stream` stay in the shared `SparkSession` and
can collide with the later suites using the same names. Shall we wrap them with
`withTempView`?
- `pushedFilters` uses a `match` which fails with `MatchError` for any scan
other than `InMemoryBatchScan`, instead of a clear assertion message.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.streaming.runtime
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, ExpressionSet, SubqueryExpression}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.v2.PushDownUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pushes the top-level columns and the filters that the optimized plan of a
streaming query reads
+ * from each Data Source V2 streaming source down to the scan of the source, as
+ * `V2ScanRelationPushDown` does for the scans of batch queries.
+ *
+ * The stream of a source is created from its scan, and has to exist before
any micro-batch is
+ * planned, so the scan is built when the query starts. The query is optimized
once at that point,
+ * by the optimizer of its micro-batches and with their configuration, and the
scan of each source
+ * gets what the planner takes from the optimized plan for the relation of the
source (see
+ * `DataSourceV2Strategy`): the projections and the filters over each
occurrence of the relation.
+ *
+ * The plan of the query does not change, so each micro-batch is optimized as
without these
+ * pushdowns: the optimizer keeps deciding which columns every operator gets,
and in which order,
+ * including the stateful ones, and Spark still evaluates all the filters over
the relation. Only
+ * the scan node of the physical plan outputs the columns that the scan reads,
and the planner
+ * checks that the micro-batch reads no other column of the relation.
+ *
+ * The other operations that the scans of batch queries can apply, such as
aggregates and limits,
+ * apply to the data of all the micro-batches of a streaming query, which the
scan of a micro-batch
+ * does not read, so they are not pushed down.
+ */
+private[sql] object StreamingV2ScanPushDown extends Logging {
+
+ /**
+ * What the optimized plan of a streaming query reads from a streaming
relation.
+ *
+ * @param requiredColumns the columns of the relation that the plan reads
+ * @param filters the deterministic filters without subqueries that the plan
applies over every
+ * occurrence of the relation, which its scan may apply too
+ */
+ case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression])
+
+ /**
+ * Returns what the optimized `plan` reads from each of its streaming
relations of Data Source
+ * V2: what the planner takes from the projections and the filters over each
occurrence of the
+ * relation, or all its columns and no filter where nothing is projected or
filtered over it. A
+ * relation that the optimized plan does not have, or any relation if the
optimizer fails, has no
+ * entry, so that its scan reads all its columns and gets no filter.
+ *
+ * @param optimizer the optimizer of the micro-batches of the query
+ * @param conf the configuration that the micro-batches of the query run with
+ * @param plan the plan of the query, with the current time of a micro-batch
as the micro-batches
+ * have it, which the optimizer does not fold, so that no filter
on it is pushed down
+ */
+ def reads(
+ optimizer: RuleExecutor[LogicalPlan],
+ conf: SQLConf,
+ plan: LogicalPlan): Map[StreamingRelationV2, Reads] = {
+ SQLConf.withExistingConf(conf) {
+ try {
+ val occurrences = mutable.HashMap.empty[StreamingRelationV2,
mutable.Buffer[Reads]]
+ // Visits the nodes as the planner plans them: the projections and the
filters over a
+ // relation are planned with its scan, and the other nodes on their
own.
+ def visit(node: LogicalPlan): Unit = node match {
+ case PhysicalOperation(project, filters, relation:
StreamingRelationV2) =>
+ occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) +=
Reads(
+ AttributeSet((project ++ filters).flatMap(_.references)),
+ filters.filter(f => f.deterministic &&
!SubqueryExpression.hasSubquery(f)))
+ case _ =>
+ node.children.foreach(visit)
+ }
+ visit(optimizer.execute(plan))
+ occurrences.map { case (relation, reads) =>
+ // The scan of a relation reads the rows of all its occurrences, so
it can only apply the
+ // filters that all of them apply.
+ val filters = reads.map(r => ExpressionSet(r.filters)).reduce {
(left, right) =>
+ ExpressionSet(left.filter(right.contains))
+ }
+ relation -> Reads(
+ AttributeSet.fromAttributeSets(reads.map(_.requiredColumns)),
filters.toSeq)
+ }.toMap
+ } catch {
+ case NonFatal(e) =>
+ logWarning(log"Could not optimize the streaming query to push down
what it reads to " +
+ log"the scans of its sources, which read all their columns with no
filter.", e)
+ Map.empty
+ }
+ }
+ }
+
+ /**
+ * Builds the scan of a streaming source, and creates its stream with
`createStream`. If its
+ * scan builder supports it, and `conf` enables it, the scan only reads the
columns of `output`
+ * that the query reads, and gets the filters that the query applies to all
the rows that it reads
+ * from the source (see `reads`). `reads` is only evaluated for such a scan
builder. Returns the
+ * scan, its stream, and the columns of `output` that the scan reads, in the
order of its read
+ * schema, unless these are `output` itself.
+ *
+ * A pruned scan must read the required columns and only columns of
`output`, each once, as
+ * `SupportsPushDownRequiredColumns` requires. A scan that does not is
replaced by a scan of all
+ * the columns, as without pruning, and its stream is stopped. Spark asks a
scan for its read
+ * schema once the stream is created, since a scan may only know then what
its stream reads.
+ *
+ * @param output the output of the relation of the source, with the metadata
columns that the
+ * query reads, which the analyzer added to it, so that
connectors implementing
+ * `SupportsMetadataColumns` include them in `readSchema()`
+ * @param conf the configuration of the query
+ * @param reads what the query reads from the relation, or None to read all
its columns
+ */
+ def buildScan[S <: SparkDataStream](
+ table: SupportsRead,
+ options: CaseInsensitiveStringMap,
+ output: Seq[AttributeReference],
+ conf: SQLConf,
+ reads: => Option[Reads])(
+ createStream: Scan => S): (Scan, S, Option[Seq[AttributeReference]]) = {
+ val builder = table.newScanBuilder(options)
+ val prunesColumns = conf.streamingV2ColumnPruningEnabled &&
+ builder.isInstanceOf[SupportsPushDownRequiredColumns]
+ val pushesFilters = conf.streamingV2FilterPushdownEnabled &&
supportsFilters(builder)
+ val known = if (prunesColumns || pushesFilters) reads else None
+ val requiredOutput = known.filter(_ => prunesColumns)
+ .map(r => output.filter(r.requiredColumns.contains))
+ .getOrElse(output)
+ val filters = known.filter(_ => pushesFilters)
+ .map(r => DataSourceStrategy.normalizeExprs(r.filters, output))
+ .getOrElse(Nil)
+
+ def build(builder: ScanBuilder, columns: Seq[AttributeReference]): Scan = {
+ if (filters.nonEmpty) {
+ // The plan of the query keeps the filters, which Spark evaluates
after the scan, so the
+ // scan may apply any of them.
+ PushDownUtils.pushFilters(builder, filters, partitionFields = None)
+ }
+ builder match {
+ case pruning: SupportsPushDownRequiredColumns =>
+ pruning.pruneColumns(DataTypeUtils.fromAttributes(columns))
+ case _ =>
+ }
+ builder.build()
+ }
+ def buildFullScan(builder: ScanBuilder): (Scan, S,
Option[Seq[AttributeReference]]) = {
+ val scan = build(builder, output)
+ (scan, createStream(scan), None)
+ }
+
+ if (requiredOutput.length == output.length) {
+ buildFullScan(builder)
+ } else {
+ val scan = build(builder, requiredOutput)
+ val stream = createStream(scan)
+ val (readSchema, scanOutput) = try {
+ val readSchema = scan.readSchema()
+ (readSchema, prunedScanOutput(readSchema, output, requiredOutput))
+ } catch {
+ case e: Throwable =>
+ try stream.stop() catch {
+ case NonFatal(suppressed) => e.addSuppressed(suppressed)
+ }
+ throw e
+ }
+ scanOutput match {
+ case Right(scanOutput) =>
+ val readsOutput = scanOutput.map(_.exprId) == output.map(_.exprId)
+ (scan, stream, if (readsOutput) None else Some(scanOutput))
+ case Left(reason) =>
+ val columns = requiredOutput.map(_.name).mkString(", ")
+ logWarning(log"The scan of the streaming table " +
+ log"${MDC(LogKeys.STREAMING_TABLE, table)} was asked to read the
columns " +
+ log"${MDC(LogKeys.COLUMN_NAMES, columns)}, but its read schema " +
+ log"${MDC(LogKeys.SCHEMA, readSchema.simpleString)} " +
+ log"${MDC(LogKeys.REASON, reason)}. It reads all the columns
instead.")
+ stream.stop()
+ buildFullScan(table.newScanBuilder(options))
+ }
+ }
+ }
+
+ private def supportsFilters(builder: ScanBuilder): Boolean = builder match {
Review Comment:
This is `true` for `PythonScanBuilder` because it implements
`SupportsPushDownFilters`. Contrary to the PR description, the streaming scans
of Python data sources are affected: the filter pushdown of the batch
`DataSourceReader` now runs in a Python worker at query start.
With `spark.sql.python.filterPushdown.enabled=true`,
`readStream.format(<Python source>).load().where("x > 1")` calls
`PythonScanBuilder.pushFilters`, whose worker calls
`data_source.reader(schema)`.
- A streaming-only source, which implements only `streamReader`, raises
`PySparkNotImplementedError`. So, a query which worked before fails to start.
- A source with a batch reader runs `reader.pushFilters` (and possibly
`partitions()`) at every stream start, while `PythonMicroBatchStream` ignores
the pushed filters.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.streaming.runtime
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, ExpressionSet, SubqueryExpression}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.v2.PushDownUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pushes the top-level columns and the filters that the optimized plan of a
streaming query reads
+ * from each Data Source V2 streaming source down to the scan of the source, as
+ * `V2ScanRelationPushDown` does for the scans of batch queries.
+ *
+ * The stream of a source is created from its scan, and has to exist before
any micro-batch is
+ * planned, so the scan is built when the query starts. The query is optimized
once at that point,
+ * by the optimizer of its micro-batches and with their configuration, and the
scan of each source
+ * gets what the planner takes from the optimized plan for the relation of the
source (see
+ * `DataSourceV2Strategy`): the projections and the filters over each
occurrence of the relation.
+ *
+ * The plan of the query does not change, so each micro-batch is optimized as
without these
+ * pushdowns: the optimizer keeps deciding which columns every operator gets,
and in which order,
+ * including the stateful ones, and Spark still evaluates all the filters over
the relation. Only
+ * the scan node of the physical plan outputs the columns that the scan reads,
and the planner
+ * checks that the micro-batch reads no other column of the relation.
+ *
+ * The other operations that the scans of batch queries can apply, such as
aggregates and limits,
+ * apply to the data of all the micro-batches of a streaming query, which the
scan of a micro-batch
+ * does not read, so they are not pushed down.
+ */
+private[sql] object StreamingV2ScanPushDown extends Logging {
+
+ /**
+ * What the optimized plan of a streaming query reads from a streaming
relation.
+ *
+ * @param requiredColumns the columns of the relation that the plan reads
+ * @param filters the deterministic filters without subqueries that the plan
applies over every
+ * occurrence of the relation, which its scan may apply too
+ */
+ case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression])
+
+ /**
+ * Returns what the optimized `plan` reads from each of its streaming
relations of Data Source
+ * V2: what the planner takes from the projections and the filters over each
occurrence of the
+ * relation, or all its columns and no filter where nothing is projected or
filtered over it. A
+ * relation that the optimized plan does not have, or any relation if the
optimizer fails, has no
+ * entry, so that its scan reads all its columns and gets no filter.
+ *
+ * @param optimizer the optimizer of the micro-batches of the query
+ * @param conf the configuration that the micro-batches of the query run with
+ * @param plan the plan of the query, with the current time of a micro-batch
as the micro-batches
+ * have it, which the optimizer does not fold, so that no filter
on it is pushed down
+ */
+ def reads(
+ optimizer: RuleExecutor[LogicalPlan],
+ conf: SQLConf,
+ plan: LogicalPlan): Map[StreamingRelationV2, Reads] = {
+ SQLConf.withExistingConf(conf) {
+ try {
+ val occurrences = mutable.HashMap.empty[StreamingRelationV2,
mutable.Buffer[Reads]]
+ // Visits the nodes as the planner plans them: the projections and the
filters over a
+ // relation are planned with its scan, and the other nodes on their
own.
+ def visit(node: LogicalPlan): Unit = node match {
+ case PhysicalOperation(project, filters, relation:
StreamingRelationV2) =>
+ occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) +=
Reads(
+ AttributeSet((project ++ filters).flatMap(_.references)),
+ filters.filter(f => f.deterministic &&
!SubqueryExpression.hasSubquery(f)))
+ case _ =>
+ node.children.foreach(visit)
+ }
+ visit(optimizer.execute(plan))
Review Comment:
This runs the full `SparkOptimizer` once more on the whole query at every
query start and restart. It includes `V2ScanRelationPushDown` and
`PruneFileSourcePartitions` for the static sides, which can build batch scans,
spawn Python workers, and list partitions only to collect the projections and
the filters over the streaming relations.
For example, a stream-static join with a large partitioned Hive or DSv2
table now lists the partitions or builds a batch scan, in addition to a full
extra optimizer pass, on the query thread before the first batch. This cost
didn't exist before, and each micro-batch repeats it anyway. Could we use a
narrower pass (e.g., only the operator optimization rules, or excluding the
early scan pushdown rules) to avoid these side effects?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,39 +124,59 @@ class ColumnarScanBuilder(
override def pruneColumns(requiredSchema: StructType): Unit = {
Review Comment:
Since `ColumnarScanBuilder` still implements
`SupportsPushDownRequiredColumns`, every streaming query over a native columnar
source triggers the extra startup optimization (`prunesColumns` is `true` in
`buildScan`). However, the streaming read schema of `ColumnarScan` is always
the full schema, so `buildScan` returns `None` for `prunedOutput` and the
result is thrown away. It is wasted work until the stream readers support the
pruning.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala:
##########
@@ -0,0 +1,241 @@
+/*
+ * 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.streaming.runtime
+
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.spark.internal.{Logging, LogKeys}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, ExpressionSet, SubqueryExpression}
+import org.apache.spark.sql.catalyst.planning.PhysicalOperation
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownRequiredColumns,
SupportsPushDownV2Filters}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.v2.PushDownUtils
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Pushes the top-level columns and the filters that the optimized plan of a
streaming query reads
+ * from each Data Source V2 streaming source down to the scan of the source, as
+ * `V2ScanRelationPushDown` does for the scans of batch queries.
+ *
+ * The stream of a source is created from its scan, and has to exist before
any micro-batch is
+ * planned, so the scan is built when the query starts. The query is optimized
once at that point,
+ * by the optimizer of its micro-batches and with their configuration, and the
scan of each source
+ * gets what the planner takes from the optimized plan for the relation of the
source (see
+ * `DataSourceV2Strategy`): the projections and the filters over each
occurrence of the relation.
+ *
+ * The plan of the query does not change, so each micro-batch is optimized as
without these
+ * pushdowns: the optimizer keeps deciding which columns every operator gets,
and in which order,
+ * including the stateful ones, and Spark still evaluates all the filters over
the relation. Only
+ * the scan node of the physical plan outputs the columns that the scan reads,
and the planner
+ * checks that the micro-batch reads no other column of the relation.
+ *
+ * The other operations that the scans of batch queries can apply, such as
aggregates and limits,
+ * apply to the data of all the micro-batches of a streaming query, which the
scan of a micro-batch
+ * does not read, so they are not pushed down.
+ */
+private[sql] object StreamingV2ScanPushDown extends Logging {
+
+ /**
+ * What the optimized plan of a streaming query reads from a streaming
relation.
+ *
+ * @param requiredColumns the columns of the relation that the plan reads
+ * @param filters the deterministic filters without subqueries that the plan
applies over every
+ * occurrence of the relation, which its scan may apply too
+ */
+ case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression])
+
+ /**
+ * Returns what the optimized `plan` reads from each of its streaming
relations of Data Source
+ * V2: what the planner takes from the projections and the filters over each
occurrence of the
+ * relation, or all its columns and no filter where nothing is projected or
filtered over it. A
+ * relation that the optimized plan does not have, or any relation if the
optimizer fails, has no
+ * entry, so that its scan reads all its columns and gets no filter.
+ *
+ * @param optimizer the optimizer of the micro-batches of the query
+ * @param conf the configuration that the micro-batches of the query run with
+ * @param plan the plan of the query, with the current time of a micro-batch
as the micro-batches
+ * have it, which the optimizer does not fold, so that no filter
on it is pushed down
+ */
+ def reads(
+ optimizer: RuleExecutor[LogicalPlan],
+ conf: SQLConf,
+ plan: LogicalPlan): Map[StreamingRelationV2, Reads] = {
+ SQLConf.withExistingConf(conf) {
+ try {
+ val occurrences = mutable.HashMap.empty[StreamingRelationV2,
mutable.Buffer[Reads]]
+ // Visits the nodes as the planner plans them: the projections and the
filters over a
+ // relation are planned with its scan, and the other nodes on their
own.
+ def visit(node: LogicalPlan): Unit = node match {
+ case PhysicalOperation(project, filters, relation:
StreamingRelationV2) =>
+ occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) +=
Reads(
+ AttributeSet((project ++ filters).flatMap(_.references)),
+ filters.filter(f => f.deterministic &&
!SubqueryExpression.hasSubquery(f)))
+ case _ =>
+ node.children.foreach(visit)
+ }
+ visit(optimizer.execute(plan))
+ occurrences.map { case (relation, reads) =>
+ // The scan of a relation reads the rows of all its occurrences, so
it can only apply the
+ // filters that all of them apply.
+ val filters = reads.map(r => ExpressionSet(r.filters)).reduce {
(left, right) =>
+ ExpressionSet(left.filter(right.contains))
+ }
+ relation -> Reads(
+ AttributeSet.fromAttributeSets(reads.map(_.requiredColumns)),
filters.toSeq)
+ }.toMap
+ } catch {
+ case NonFatal(e) =>
+ logWarning(log"Could not optimize the streaming query to push down
what it reads to " +
+ log"the scans of its sources, which read all their columns with no
filter.", e)
+ Map.empty
+ }
+ }
+ }
+
+ /**
+ * Builds the scan of a streaming source, and creates its stream with
`createStream`. If its
+ * scan builder supports it, and `conf` enables it, the scan only reads the
columns of `output`
+ * that the query reads, and gets the filters that the query applies to all
the rows that it reads
+ * from the source (see `reads`). `reads` is only evaluated for such a scan
builder. Returns the
+ * scan, its stream, and the columns of `output` that the scan reads, in the
order of its read
+ * schema, unless these are `output` itself.
+ *
+ * A pruned scan must read the required columns and only columns of
`output`, each once, as
+ * `SupportsPushDownRequiredColumns` requires. A scan that does not is
replaced by a scan of all
+ * the columns, as without pruning, and its stream is stopped. Spark asks a
scan for its read
+ * schema once the stream is created, since a scan may only know then what
its stream reads.
+ *
+ * @param output the output of the relation of the source, with the metadata
columns that the
+ * query reads, which the analyzer added to it, so that
connectors implementing
+ * `SupportsMetadataColumns` include them in `readSchema()`
+ * @param conf the configuration of the query
+ * @param reads what the query reads from the relation, or None to read all
its columns
+ */
+ def buildScan[S <: SparkDataStream](
+ table: SupportsRead,
+ options: CaseInsensitiveStringMap,
+ output: Seq[AttributeReference],
+ conf: SQLConf,
+ reads: => Option[Reads])(
+ createStream: Scan => S): (Scan, S, Option[Seq[AttributeReference]]) = {
+ val builder = table.newScanBuilder(options)
+ val prunesColumns = conf.streamingV2ColumnPruningEnabled &&
+ builder.isInstanceOf[SupportsPushDownRequiredColumns]
+ val pushesFilters = conf.streamingV2FilterPushdownEnabled &&
supportsFilters(builder)
+ val known = if (prunesColumns || pushesFilters) reads else None
+ val requiredOutput = known.filter(_ => prunesColumns)
+ .map(r => output.filter(r.requiredColumns.contains))
+ .getOrElse(output)
+ val filters = known.filter(_ => pushesFilters)
+ .map(r => DataSourceStrategy.normalizeExprs(r.filters, output))
+ .getOrElse(Nil)
+
+ def build(builder: ScanBuilder, columns: Seq[AttributeReference]): Scan = {
+ if (filters.nonEmpty) {
+ // The plan of the query keeps the filters, which Spark evaluates
after the scan, so the
+ // scan may apply any of them.
+ PushDownUtils.pushFilters(builder, filters, partitionFields = None)
+ }
+ builder match {
+ case pruning: SupportsPushDownRequiredColumns =>
+ pruning.pruneColumns(DataTypeUtils.fromAttributes(columns))
+ case _ =>
+ }
+ builder.build()
+ }
+ def buildFullScan(builder: ScanBuilder): (Scan, S,
Option[Seq[AttributeReference]]) = {
+ val scan = build(builder, output)
+ (scan, createStream(scan), None)
+ }
+
+ if (requiredOutput.length == output.length) {
+ buildFullScan(builder)
+ } else {
+ val scan = build(builder, requiredOutput)
+ val stream = createStream(scan)
+ val (readSchema, scanOutput) = try {
+ val readSchema = scan.readSchema()
+ (readSchema, prunedScanOutput(readSchema, output, requiredOutput))
+ } catch {
+ case e: Throwable =>
Review Comment:
The fallback looks incomplete.
- If `readSchema()` of the pruned scan throws, the stream is stopped and the
query fails instead of falling back to a full scan, while an invalid read
schema falls back below. A connector whose `readSchema()` throws for a schema
which it cannot honor now fails at query start, although it worked with all the
columns before.
- The fallback creates a second stream with the same `metadataPath` after
the first stream may have initialized it. A source which persists state on
creation (e.g., the initial offsets or a creation timestamp under
`metadataPath`) runs its initialization twice for one source id.
- `case e: Throwable` catches the fatal errors too.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala:
##########
@@ -478,14 +478,21 @@ case class StreamingDataSourceV2Relation(
/**
* A specialization of [[DataSourceV2ScanRelation]] with the streaming bit set
to true, as well
* as start and end offsets for Microbatch processing.
+ *
+ * @param output the columns of the source that the query references, which
the plan of each
+ * micro-batch is optimized with
+ * @param prunedOutput the columns of `output` that the scan reads, in the
order of its read schema,
+ * if it was built to read fewer of them. The scan node of
the physical plan
+ * outputs these, so the plan must not read the other
columns of `output`.
*/
case class StreamingDataSourceV2ScanRelation(
relation: StreamingDataSourceV2Relation,
scan: Scan,
output: Seq[AttributeReference],
stream: SparkDataStream,
startOffset: Option[Offset] = None,
- endOffset: Option[Offset] = None)
+ endOffset: Option[Offset] = None,
+ prunedOutput: Option[Seq[AttributeReference]] = None)
Review Comment:
`stringArgs` (`stringArgsVal`) and `simpleString` don't include this new
`prunedOutput`, so the streaming logical plans and the `explain` output still
show all the columns even when the scan reads only some of them. When we debug
the new `INTERNAL_ERROR` (`does not read the columns ...`) or a pruning issue,
only the physical `MicroBatchScanExec` shows which columns the scan actually
reads.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,39 +124,59 @@ class ColumnarScanBuilder(
override def pruneColumns(requiredSchema: StructType): Unit = {
// Only top-level columns are pruned by the data source. Spark prunes the
nested fields of the
- // columns it reads.
+ // columns it reads. A batch scan asks its planner to prune them when it
is planned, while a
+ // streaming scan, which Spark also prunes, has no batch planner and reads
all the columns.
val required =
StructType(requiredSchema.fieldNames.flatMap(name => schema.find(_.name
== name)))
- if (required.length < schema.length &&
batchPlanner.pruneColumns(required)) {
- readSchema = required
+ if (required.length < schema.length) {
+ this.requiredSchema = Some(required)
}
}
// The scan takes over the planner, so it is built once.
private lazy val scan = new ColumnarScan(
- shortName, schema, readSchema, dataSource, Option(planner), pushedLimit)
+ shortName, schema, requiredSchema, dataSource, Option(planner),
pushedLimit)
override def build(): Scan = scan
}
/**
+ * The read schema of a scan of a micro-batch stream is only known once the
stream is created, as
+ * `StreamingV2ScanPushDown` asks for it: it is the full schema, while the
read schema that the
+ * scan reports before is the one of a batch, whose planner it creates to
prune the columns.
+ *
+ * @param requiredSchema the top-level columns that the query needs, if fewer
than those of the
+ * full schema. A batch scan asks its planner to prune
the other columns,
+ * while a streaming scan reads all the columns of the
full schema, which is
+ * the schema that a stream reader gets.
* @param pushedDownPlanner the planner that the operations were pushed down
to, if any. Otherwise,
- * the scan creates one when its batch is planned.
+ * the scan creates one when its batch is planned or
pruned.
*/
class ColumnarScan(
shortName: String,
fullSchema: StructType,
- prunedSchema: StructType,
+ requiredSchema: Option[StructType],
dataSource: ColumnarDataSource,
pushedDownPlanner: Option[ColumnarBatchPlanner],
pushedLimit: Option[Int])
extends Scan with SupportsMetadata {
+ private lazy val planner =
pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))
+
+ // The columns that the batch reads: the required ones if the planner prunes
the others.
+ private lazy val batchSchema =
requiredSchema.filter(planner.pruneColumns).getOrElse(fullSchema)
+
// Spark may call `toBatch` more than once for the same scan; the partitions
are planned once.
- private lazy val batch = new ColumnarBatchScan(
- pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema)))
+ private lazy val batch = {
+ // The planner prunes the columns before it plans the partitions.
+ val _ = batchSchema
+ new ColumnarBatchScan(planner)
+ }
+
+ // Whether the scan is the one of a micro-batch stream, whose read schema is
the full schema.
+ @volatile private var streaming = false
- override def readSchema(): StructType = prunedSchema
+ override def readSchema(): StructType = if (streaming) fullSchema else
batchSchema
Review Comment:
`readSchema()` depends on the call order via the mutable `streaming` flag
which `toMicroBatchStream` sets. If anything (explain, logging, or a future
change of `buildScan`) calls `readSchema()` or `getMetaData()` on a streaming
scan before its stream is created, `batchSchema` is evaluated and calls
`dataSource.batchPlanner(fullSchema)`. For a streaming-only
`ColumnarDataSource`, it can throw or allocate planner resources which are
never released, which the builder comment says the data source may not support.
It also reports a pruned batch schema which the stream reader won't produce.
--
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]
