HyukjinKwon commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4461619803
########## 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") Review Comment: Done in 29864de. The pass over the analyzed plan and its `UnionBase` case are gone: the columns now come from the plan optimized when the query starts, where `ColumnPruning` prunes the children of a union by position. "read only the columns of each child of a union" now unions two reads of the table, whose scans each read only `data`, and two reads of different columns, which read `id` and `extra`. It also checks that a self-union has one scan, which reads the columns of both children. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala: ########## @@ -123,39 +124,56 @@ 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 } /** + * @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 { - // 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 planner = pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema)) - override def readSchema(): StructType = prunedSchema + // 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, + // after the columns are pruned. + private lazy val batch = { + batchSchema + new ColumnarBatchScan(planner) + } + + // Whether the scan is the one of a micro-batch stream, which Spark asks for its read schema + // after it creates the stream. + @volatile private var streaming = false + + override def readSchema(): StructType = if (streaming) fullSchema else batchSchema Review Comment: Done in 29864de: the doc of `ColumnarScan` now says that the read schema of the scan of a micro-batch stream is only known once its stream is created, which is when `StreamingV2ScanPushDown.buildScan` asks for it, and `buildScan` documents that order. ########## 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 Review Comment: The list is gone in 29864de. The required columns now come from the plan of the query optimized once when it starts, with the optimizer and the configuration of its micro-batches, and the scan reads what the planner takes from it for the relation, with `PhysicalOperation` as `DataSourceV2Strategy` does. So `ColumnPruning` and the other rules decide what each operator needs, as for batch scans, including rewrites like `ReplaceDistinctWithAggregate`. The plan of each micro-batch still has the unpruned relation, and is optimized as before. If a micro-batch ever read a column that the scan does not read, the planner fails with an internal error that names `spark.sql.streaming.v2ColumnPruning.enabled`, rather than reading another column ("fail to plan a micro-batch that reads a column that its scan does not read"). ########## 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) Review Comment: Done in 29864de: if asking for the read schema or checking it throws, the stream is stopped before the error is rethrown, and a failure to stop it is added as suppressed. When the check fails, the stream is stopped before the scan of all the columns is built, so a failure there leaves nothing open. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala: ########## @@ -165,11 +183,12 @@ class ColumnarScan( override def toBatch: Batch = batch override def toMicroBatchStream(checkpointLocation: String): MicroBatchStream = { + streaming = true Review Comment: Done in 29864de: "streaming read of some of the columns" in `ColumnarDataSourceSuite` streams `select("name")` of the columnar source. It checks the rows, and that the scan node reads both columns, as the stream reader does. -- 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]
