HyukjinKwon commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4461625749
########## 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: Done in 29864de: "read no column of the source when the query needs none, also after a restart" runs `groupBy().count()`, with a restart, and `select(lit(1))`. It checks the empty read schema, the output of the scan node and the rows. ########## 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) Review Comment: Fixed in 29864de. The relation now keeps `output` in the plan, so `ColumnPruning` puts the projection in the original order over it, as before this PR. Only the scan node of the physical plan outputs the columns in the order of the read schema, under that projection, so the join stores the same layout with and without pruning. "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" is your scenario: a source of `(k, v, unused)` returns `(v, k)` once pruned. It stores the unmatched left row `(1, 100)` with pruning disabled, restarts with it enabled, adds `rk = 1`, and expects 100, and then does the same the other way around. With the relation output in the order of the read schema, as in the previous version, it gets 1. ########## 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: both confs are read from `streamConf`, the clone taken when the query is created. ########## 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: "streaming read of some of the columns" in `ColumnarDataSourceSuite` streams `select("name")` of a columnar source. It checks the rows, and that the scan node outputs both columns, as the stream reader reads. If the `streaming` flag regressed, the scan node would output only `name` over batches with both columns. -- 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]
