dongjoon-hyun commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4420266595
########## 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: **1. [Medium] The union test passes without the `UnionBase` case** `source.union(source)` is a self-union, so `DeduplicateRelations` puts a `Project` of aliases of all the columns over its second child. The references of that `Project` already require `id`, `data` and `extra`. So this test still reads all the columns when `_: UnionBase` is removed from `requiredAttributes` (`StreamingV2ColumnPruning.scala` L48). For the same reason, the comment above (L126-127) does not hold for a self-union. The `UnionBase` case is meant for a union of two reads that have their own attributes, where nothing references the columns of the second read. Could the test union two reads of the table instead? ```scala val df = spark.readStream.table(tbl).union(spark.readStream.table(tbl)).select("data") ``` Without the case, the scan of the second read would then be pruned to no columns. `Distinct`, `SetOperation`, `CTERelationDef` and `mapsColumnsByPosition` have no test either. ########## 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: **3. [Low] An operator missing from this list makes the scan drop a column that it needs** This lists the operators that need all the columns of their children, and treats every other operator as safe. So an operator that the list misses fails in the unsafe direction. The scan drops a column that the plan of a micro-batch needs later, and the query fails at planning or reads the wrong column, instead of just reading too much. Two kinds of operators are at risk: - An operator that keeps the attributes of its children but maps them by position, as `Union` does. - An operator that an optimizer rule rewrites with new references to the columns of its child, as `ReplaceDistinctWithAggregate` does for `Distinct` (which the list covers). Analyzed plans can also hold positions. An example is the `unrequiredChildIndex` of the `Generate` that `ResolveChangelogTable` builds, which is safe today only because its child is an `Aggregate`. Could this be inverted? Only the operators known to map their columns by attribute (`Project`, `Filter`, `Aggregate`, `Join`, `EventTimeWatermark`, `Deduplicate`, ...) would let a column be pruned, and every other operator would keep all the columns of its children. A missed operator would then only cost some pruning. ########## 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: **5. [Low] No test reads some of the columns of a columnar stream** `ColumnarScanBuilder.pruneColumns` now gets the pruned columns of streaming queries too. A stream then returns the right columns only because `readSchema()` returns `fullSchema` once `streaming` is set. The streaming read test of `ColumnarDataSourceSuite` reads all the columns, so it still passes if this regresses. Could the suite have a streaming counterpart of its pruning test (L76-79 of `ColumnarDataSourceSuite.scala`)? For example: ```scala val query = spark.readStream.format(format).option("end", "5").load().select("name") .writeStream.format("memory").queryName("pruned_stream").start() try query.processAllAvailable() finally query.stop() checkAnswer(spark.table("pruned_stream"), (0L until 5L).map(id => Row(s"name$id"))) ``` ########## 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: **4. [Low] The first stream is not stopped when the check or the second scan fails** The stream is created before the read schema is checked. If `scan.readSchema()` (L106), `build(output)` or `createStream(fullScan)` (L112-113) throws, nothing stops `stream`. It is not in `sources` yet, so `StreamExecution` does not stop it either when the query fails, and whatever the stream holds stays open. Could the check and the rebuild run in a `try` that stops `stream` on a failure? ########## 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: **7. [Low] The fallback to all the columns is silent** When the read schema does not have the required columns, the scan is replaced without a log. That means the connector does not follow `SupportsPushDownRequiredColumns`, and it costs a second scan and a second stream at every start. Could this log a warning with the source and the missing columns? The check also compares only the names. A read schema that has a column twice passes it (`scanOutput.length == readSchema.length`), and the relation then outputs the same attribute twice. Could it also require distinct names? ########## 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: **6. [Low] The output follows the order of the read schema, not of `output`** Before this PR, `ColumnPruning` put `Project(output.filter(required))` over the relation, so an operator right above it got the columns in the order of `output`. Now the relation itself is pruned and takes the order of `readSchema()`, and no `Project` is added when the read schema has exactly the required columns. Some connectors return the pruned columns in another order than they were asked for, e.g. with their key columns first. Over such a scan, a stateful operator (for example a side of a stream-stream join) gets its input in another layout. A query restarted from an existing checkpoint then fails the state schema check. If the reordered columns have the same types, the check passes instead, because `StateSchemaCompatibilityChecker` ignores the names, and the query reads the stored rows with the fields swapped. `SupportsPushDownRequiredColumns` does not say that the order is kept. Would it be safer to keep the order of `requiredOutput`, for example with a `Project` in that order over the relation, or by falling back when the orders differ? ########## 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: **9. [Nit] `forall(scanOutput.contains)` is quadratic in the number of columns** `Seq.contains` compares the whole `AttributeReference` (name, exprId, data type, metadata and qualifier) with each element. So a source with thousands of columns makes millions of comparisons when the query starts. An `AttributeSet` does the same check in linear time: ```scala if (scanOutput.length == readSchema.length && AttributeSet(requiredOutput).subsetOf(AttributeSet(scanOutput))) { ``` ########## 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: **8. [Nit] This remapping repeats `PushDownUtils.toOutputAttrs`** Batch scans map the read schema back to the attributes of the relation by name in `PushDownUtils.toOutputAttrs`. That function also takes the data types and the CHAR/VARCHAR metadata of the read schema, while this code keeps the attributes of `output`. So the same scan can get different attributes in batch and streaming queries. `toOutputAttrs` only uses `relation.output`, so it could take the attributes as a parameter and be shared, once the names are checked as they are here. ########## 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: **10. [Nit] `batchSchema` is evaluated only for its side effect** The bare `batchSchema` makes the planner prune the columns before the batch is planned, but it reads like a leftover. If it were removed, the batch would read all the columns while `readSchema()` reports the pruned ones. Could the dependency be explicit? For example: ```scala private lazy val batch = { // The planner prunes the columns before it plans the partitions. val _ = batchSchema new ColumnarBatchScan(planner) } ``` ########## 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: **2. [Low] `ColumnarScan.readSchema()` depends on whether `toMicroBatchStream` was called** `StreamingV2ColumnPruning.buildScan` asks for the read schema only after it creates the stream, so this works today. But the result of `readSchema()` now depends on the order of the calls, and the `Scan` API does not define that order. If anything asks a streaming `ColumnarScan` for its read schema before `toMicroBatchStream`, three things go wrong: - It evaluates `batchSchema`, which creates a batch planner for a data source that may not support batch scans (L91-92). - It asks that planner to prune the columns. - It reports the pruned columns, while the stream reader still reads `fullSchema`. If the flag stays, could the doc of `ColumnarScan` say that the read schema of a streaming scan is only valid once its stream is created? Then a later caller does not rely on the other order. -- 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]
