HyukjinKwon commented on code in PR #59943:
URL: https://github.com/apache/spark/pull/59943#discussion_r4461623514
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala:
##########
@@ -261,6 +260,13 @@ class MicroBatchExecution(
}
import
org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._
+ // The columns of the Data Source V2 sources that the query needs, which
their scans read.
+ val requiredAttributes =
+ if (sparkSession.sessionState.conf.streamingV2ColumnPruningEnabled) {
Review Comment:
Done in 29864de: the micro-batch execution reads the confs from
`streamConf`, and the continuous one from the session of the stream.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala:
##########
@@ -93,19 +99,13 @@ class ContinuousExecution(
log"from DataSourceV2 named '${MDC(STREAMING_DATA_SOURCE_NAME,
sourceName)}' " +
log"${MDC(STREAMING_DATA_SOURCE_DESCRIPTION, dsStr)}")
// TODO: operator pushdown.
- // Passes the full output schema (not a pruned subset) so that
connectors
- // implementing SupportsMetadataColumns can include metadata columns
in readSchema().
- val scanBuilder = table.newScanBuilder(options)
- scanBuilder match {
- case r: SupportsPushDownRequiredColumns =>
- r.pruneColumns(output.toStructType)
- case _ =>
- }
- val scan = scanBuilder.build()
- val stream = scan.toContinuousStream(metadataPath)
+ // The output includes the metadata columns that the query reads, so
that connectors
+ // implementing SupportsMetadataColumns include them in readSchema().
+ val (scan, stream, scanOutput) = StreamingV2ColumnPruning.buildScan(
Review Comment:
Done in 29864de: "read only the columns that a continuous query needs" runs
a continuous query over a test table of a `ContinuousMemoryStream`, whose scans
only read the columns they are asked for. It checks the read schema and the
rows. The fallback is the same `buildScan` for both modes, which the
micro-batch tests cover.
##########
sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingV2ColumnPruningSuite.scala:
##########
@@ -0,0 +1,237 @@
+/*
+ * 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 org.scalatest.BeforeAndAfter
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.catalog.{Identifier, InMemoryBaseTable,
InMemoryTableCatalog}
+import
org.apache.spark.sql.execution.datasources.v2.StreamingDataSourceV2ScanRelation
+import org.apache.spark.sql.internal.SQLConf
+
+class StreamingV2ColumnPruningSuite extends StreamTest with BeforeAndAfter {
+
+ 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): StreamAction = Execute { _ =>
+ sql(s"INSERT INTO $tbl VALUES $values")
+ }
+
+ 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
+ }
+
+ 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"),
+ CheckAnswer(Row("b")),
+ StopStream,
+ insert("(3, 'c', 30), (4, 'd', 5)"),
+ StartStream(),
+ ProcessAllAvailable(),
+ assertReadSchema("data", "extra"),
+ CheckAnswer(Row("b"), Row("c")))
+ }
+ }
+
+ 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 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 all the columns of the children of a union") {
+ withSourceTable() {
+ val source = spark.readStream.table(tbl)
+ // A union maps the columns of its children by position, and its output
is the one of its
+ // first child, so the query does not reference the columns of its other
children.
+ val df = source.union(source).select("data")
+ testStream(df)(
+ ProcessAllAvailable(),
+ assertReadSchema("id", "data", "extra"),
+ 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 df = spark.readStream.table(tbl).select("data")
+ testStream(df)(
+ 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")))
+ }
+ }
+
+ 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")
+ val readSchemas = AssertOnQuery { query =>
+ val schemas = query.logicalPlan.collect {
+ case relation: StreamingDataSourceV2ScanRelation =>
+ relation.scan.readSchema().fieldNames.toSeq
+ }
+ assert(schemas.sortBy(_.mkString(",")) === Seq(Seq("id", "data"),
Seq("id", "extra")))
+ true
+ }
+ testStream(df)(
+ ProcessAllAvailable(),
+ readSchemas,
+ CheckAnswer(Row("a", 10), Row("b", 20)),
+ StopStream,
+ insert("(3, 'c', 30)"),
Review Comment:
Done in 29864de: "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"
and "keep the state of a streaming aggregation ..." restart from pruning
disabled to enabled and back, with new rows whose keys match the state, over a
scan that returns the pruned columns in another order. Since the plan keeps the
unpruned relation and only the scan node is pruned, the stateful operators get
the same input in both modes, so the conf does not need to be pinned in the
offset log.
##########
docs/streaming/ss-migration-guide.md:
##########
@@ -38,6 +38,11 @@ Please refer [Migration Guide: SQL, Datasets and
DataFrame](../sql-migration-gui
To recover, write to a new table location with a new checkpoint location, or
rewrite the
committed files with the partition columns of the table.
+- Since Spark 4.4, the scan of a Data Source V2 streaming source whose scan
builder implements
+ `SupportsPushDownRequiredColumns` only reads the top-level columns that the
streaming query
+ needs, as the scans of batch queries do. Before, it read all the columns of
the source. To
Review Comment:
Done in 29864de. Since the columns now come from the optimized plan, the
parity holds for the top-level columns: `withColumn("x",
lit(1)).select("data")`, chained selects, aliases and `SELECT *` read only
`data`, as in the new "read only the columns that the optimized query reads"
tests. The guide and the conf doc now say that Spark takes the columns from the
plan of the query optimized when it starts, as it does for batch scans.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ColumnPruning.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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 org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, NamedExpression}
+import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, Distinct,
LogicalPlan, SetOperation, UnionBase}
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan,
SupportsPushDownRequiredColumns}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Prunes the top-level columns that the scans of Data Source V2 streaming
sources read to those
+ * that a streaming query needs, as `V2ScanRelationPushDown` does for batch
scans. The scan of a
+ * streaming source is built once, when the query starts, from its analyzed
plan, so the required
+ * columns are taken from that plan: the optimizer prunes the columns of each
micro-batch further,
+ * so stateful operators keep getting the same columns as without this pruning.
+ */
+private[sql] object StreamingV2ColumnPruning {
+
+ /**
+ * Returns the attributes of its sources that a query needs: those in its
output, those that its
+ * operators reference, and every column of the children of the operators
that need all of them,
+ * such as a distinct, or that map the columns of a child to their output by
position rather
+ * than by attribute, such as a union of any kind or the definition of a
common table expression.
+ */
+ def requiredAttributes(plan: LogicalPlan): AttributeSet = {
+ var required = plan.outputSet
+ plan.foreach { node =>
+ required ++= node.references
+ val needsAllColumns = node match {
+ case _: UnionBase | _: SetOperation | _: Distinct | _: CTERelationDef
=> true
+ case _ => mapsColumnsByPosition(node)
+ }
+ if (needsAllColumns) {
+ node.children.foreach(child => required ++= child.outputSet)
+ }
+ }
+ required
+ }
+
+ /**
+ * Returns whether a node outputs attributes that are neither of its
children, nor produced or
+ * defined by its expressions, which it then maps from the columns of its
children by position.
+ */
+ private def mapsColumnsByPosition(node: LogicalPlan): Boolean =
node.children.nonEmpty && {
+ val known = AttributeSet(node.children.flatMap(_.output)) ++
node.producedAttributes ++
+ AttributeSet(node.expressions.collect { case named: NamedExpression =>
named.toAttribute })
+ node.output.exists(attribute => !known.contains(attribute))
+ }
+
+ /**
+ * Builds the scan of a streaming source that reads the columns of `output`
in `required`, when
+ * its scan builder supports column pruning, and creates its stream with
`createStream`.
+ * Returns the scan, its stream and its output: the columns of `output` that
it reads, in the
+ * order of its read schema when the scan prunes them, and otherwise
`output` as it is. A scan
+ * that does not read the required columns, as named in `output`, is
replaced by one of all the
+ * columns, as without pruning.
+ *
+ * The read schema of the scan is asked once its stream is created, since a
scan may only know
+ * then what it reads for a stream.
+ *
+ * @param output the output of the source, with the metadata columns that
the query reads, which
+ * the analyzer added to it
+ * @param required the attributes that the query needs, or None to read all
the columns
+ */
+ def buildScan[S <: SparkDataStream](
+ table: SupportsRead,
+ options: CaseInsensitiveStringMap,
+ output: Seq[AttributeReference],
+ required: Option[AttributeSet])(
+ createStream: Scan => S): (Scan, S, Seq[AttributeReference]) = {
+ def build(columns: Seq[AttributeReference]): (Scan, Boolean) = {
+ val builder = table.newScanBuilder(options)
+ val prunes = builder match {
+ case pruning: SupportsPushDownRequiredColumns =>
+ pruning.pruneColumns(DataTypeUtils.fromAttributes(columns))
+ true
+ case _ => false
+ }
+ (builder.build(), prunes)
+ }
+ val requiredOutput = required.map(r =>
output.filter(r.contains)).getOrElse(output)
+ val (scan, prunes) = build(requiredOutput)
+ val stream = createStream(scan)
+ if (!prunes || requiredOutput.length == output.length) {
+ (scan, stream, output)
+ } else {
+ val byName = output.map(attribute => attribute.name -> attribute).toMap
+ val readSchema = scan.readSchema()
+ val scanOutput = readSchema.fieldNames.toSeq.flatMap(byName.get)
+ if (scanOutput.length == readSchema.length &&
requiredOutput.forall(scanOutput.contains)) {
+ (scan, stream, scanOutput)
+ } else {
+ stream.stop()
Review Comment:
Done in 29864de: the warning names the table, the required columns, the read
schema and the reason, which also covers a field that does not map to the
output and a field that appears twice. The doc of `buildScan` lists the cases.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ColumnPruning.scala:
##########
@@ -0,0 +1,117 @@
+/*
+ * 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 org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, NamedExpression}
+import org.apache.spark.sql.catalyst.plans.logical.{CTERelationDef, Distinct,
LogicalPlan, SetOperation, UnionBase}
+import org.apache.spark.sql.catalyst.types.DataTypeUtils
+import org.apache.spark.sql.connector.catalog.SupportsRead
+import org.apache.spark.sql.connector.read.{Scan,
SupportsPushDownRequiredColumns}
+import org.apache.spark.sql.connector.read.streaming.SparkDataStream
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Prunes the top-level columns that the scans of Data Source V2 streaming
sources read to those
+ * that a streaming query needs, as `V2ScanRelationPushDown` does for batch
scans. The scan of a
+ * streaming source is built once, when the query starts, from its analyzed
plan, so the required
+ * columns are taken from that plan: the optimizer prunes the columns of each
micro-batch further,
+ * so stateful operators keep getting the same columns as without this pruning.
+ */
+private[sql] object StreamingV2ColumnPruning {
+
+ /**
+ * Returns the attributes of its sources that a query needs: those in its
output, those that its
+ * operators reference, and every column of the children of the operators
that need all of them,
+ * such as a distinct, or that map the columns of a child to their output by
position rather
+ * than by attribute, such as a union of any kind or the definition of a
common table expression.
+ */
+ def requiredAttributes(plan: LogicalPlan): AttributeSet = {
+ var required = plan.outputSet
+ plan.foreach { node =>
+ required ++= node.references
Review Comment:
The pass is gone in 29864de. The columns that each occurrence of a relation
reads are collected once and combined with `AttributeSet.fromAttributeSets`.
--
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]
