HyukjinKwon commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4461621560
########## 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: Done in 29864de. The relation now keeps `output` in the plan of the query, so the optimizer projects its columns in their order, as before this PR. Only the scan node of the physical plan outputs the columns in the order of the read schema (`StreamingDataSourceV2ScanRelation.prunedOutput`), and the projection over it restores the order. So a stateful operator gets the same input with and without pruning. The new table property `simulate-reordered-column-pruning` of the in-memory table returns the pruned columns in the reverse order. "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" stores unmatched rows with pruning disabled, restarts with it enabled, and back, and checks the joined rows. With the relation output in the order of the read schema, as in the previous version, the join returns 1 instead of 100. "keep the state of a streaming aggregation ..." does the same for an aggregation. ########## 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 fallback logs a warning with the table, the required columns, the read schema and why it does not follow `SupportsPushDownRequiredColumns`: it has a column that the relation does not output, has a column more than once, or misses a required column. The check now also rejects a column that appears twice. "read all the columns when a scan does not read the required ones" checks the warning. ########## 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 Review Comment: Done in 29864de: `PushDownUtils.toOutputAttrs` also takes the attributes now, and the streaming scan maps its read schema with it, once the names are checked. ########## 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)) { Review Comment: Done in 29864de: the check uses an `AttributeSet`. ########## 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 Review Comment: Done in 29864de, as suggested: `val _ = batchSchema`, with a comment that the planner prunes the columns before it plans the partitions. -- 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]
