spark-test commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4425284122
########## 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: **Nit (P3):** This restart inserts only the new key `3`, so no row restored from the join state reaches the output, and every restart in the suite keeps pruning on in both runs. Since `spark.sql.streaming.v2ColumnPruning.enabled` is read at each start and is not pinned in the offset log, every existing checkpoint over a pruning-capable source moves to pruned scans on its first restart after the upgrade, and turning the conf off moves it back. Restarting a stream-stream join and a streaming aggregation across the two values, with new rows whose keys match stored state, would verify that the input layout of the stateful operators is the same in both modes. **Verification:** - **Compatibility:** A stream-stream join and a streaming aggregation restarted with streaming column pruning enabled from a checkpoint written with it disabled, and the reverse, produce the same results as a run without the toggle when new rows match keys held in state. ########## 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: **Nit (P3):** The conf is read from the caller's live session, but `logicalPlan` is initialized lazily on the query thread after `start()` has returned: `runStream` calls `startLatch.countDown()` before it forces `logicalPlan`, and the offset log is listed before this line. If a user turns the conf off around `start()` to work around a connector and restores it afterwards, the query can still build pruned scans, and the reverse can happen too. The other per-query reads in this initializer go through `streamConf`, the clone taken when the query is created (line 242); reading `streamConf.streamingV2ColumnPruningEnabled` here would bind the value at creation. Continuous queries read the conf during construction, so only micro-batch is affected. **Verification:** - **Behavior:** A micro-batch query started with streaming column pruning disabled keeps reading all source columns even when the session conf is set back to true right after start() returns, and a query started with pruning enabled keeps pruning when the conf is then disabled. ########## 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: **Nit (P3):** `AttributeSet.++` builds a new set from `baseSet ++ other.baseSet`, so each `required ++= ...` here and at line 52 copies everything accumulated so far, which makes the pass O(nodes x attributes) at every query start. This only matters for wide plans with many sources, such as a union of many wide streams, and the fix is local: collect the output set, the references of each node and the child output sets of the all-columns nodes, then combine them once with `AttributeSet.fromAttributeSets`, as `QueryPlan.inputSet` does. The membership and order stay the same. **Verification:** - **Behavior:** The required attributes, and therefore the read schemas of every existing pruning scenario, are unchanged. ########## 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: **Nit (P3):** This path is on by default for continuous queries, but nothing tests it: `InMemoryBaseTable` declares only `MICRO_BATCH_READ`, the continuous-capable built-in sources do not prune, and `StreamingV2ColumnPruningSuite` runs only micro-batch queries. A mistake that only continuous mode hits, such as binding the scan relation to `output` instead of `scanOutput`, taking the required columns from another plan, or breaking the fallback and its stream stop, would leave every test green. A small test source that supports `CONTINUOUS_READ` and column pruning could cover a pruned projection and the fallback. **Verification:** - **Behavior:** A continuous query over a source that supports continuous reads and column pruning reads only the columns it needs and returns the right rows. - **Regression:** In continuous mode, a scan whose read schema lacks a required column is replaced by a scan of all the columns, and the replaced stream is stopped. ########## 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: **Nit (P3):** "as the scans of batch queries do" does not hold for common query shapes. The required columns come from the references of every node of the analyzed plan, in which the analyzer has not collapsed Projects, so `spark.readStream.table(t).withColumn("x", lit(1)).select("data")` reads every column of `t`, while `V2ScanRelationPushDown` prunes the batch scan of the same query to `data`. Chained selects, `SELECT *` subqueries and the time and session window rewrites behave the same way, and the conf doc in `SQLConf` makes the same claim ("as for batch scans"). Could both texts say that the scan reads the top-level columns that the analyzed query references? **Verification:** - **Inspection:** The migration-guide entry and the conf documentation describe the pruned columns as those that the analyzed streaming query references and no longer state that streaming scans prune as batch scans do. ########## 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: **Nit (P3):** The scan is replaced without any log, both when the read schema lacks a required column and when it has a field that does not map to `output`, such as a name in another case or an extra internal column: `byName` is an exact-case map, so `scanOutput.length == readSchema.length` then fails. Such a connector loses pruning and costs an extra builder, scan and stream at every start, and nothing tells its owner. A warning that names the source, the missing required columns and the unmapped read-schema fields would make this diagnosable; a warning that lists only the missing columns would show nothing for an extra field. The doc of this method (lines 72-74) also mentions only the first trigger. **Verification:** - **Behavior:** Starting a query over a scan whose read schema lacks a required column, or has a field that does not map to the output, logs one warning that names the source and those columns, and the query still reads all the columns. ########## 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) Review Comment: **Nit (P3):** When the query needs no source column, as `spark.readStream.table(t).groupBy().count()` or `select(lit(1))` do, `requiredOutput` is empty: `pruneColumns` gets an empty struct, the `forall` at line 108 accepts an empty read schema, and both relations get no output. The engine never built an empty streaming scan before, and every test in `StreamingV2ColumnPruningSuite` needs at least one source column. A regression on this path, in `buildScan`, the per-batch `LocalRelation` substitution or a reader, would go unnoticed. A global count and a constant projection, including a restart, would cover it. **Verification:** - **Behavior:** groupBy().count() and a constant projection over a pruning-capable streaming source read no source column and return the right counts and rows across micro-batches and after a restart. -- 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]
