HyukjinKwon commented on code in PR #59943: URL: https://github.com/apache/spark/pull/59943#discussion_r4465617292
########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala: ########## @@ -0,0 +1,241 @@ +/* + * 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 scala.collection.mutable +import scala.util.control.NonFatal + +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExpressionSet, SubqueryExpression} +import org.apache.spark.sql.catalyst.planning.PhysicalOperation +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.SupportsRead +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} +import org.apache.spark.sql.connector.read.streaming.SparkDataStream +import org.apache.spark.sql.execution.datasources.DataSourceStrategy +import org.apache.spark.sql.execution.datasources.v2.PushDownUtils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Pushes the top-level columns and the filters that the optimized plan of a streaming query reads + * from each Data Source V2 streaming source down to the scan of the source, as + * `V2ScanRelationPushDown` does for the scans of batch queries. + * + * The stream of a source is created from its scan, and has to exist before any micro-batch is + * planned, so the scan is built when the query starts. The query is optimized once at that point, + * by the optimizer of its micro-batches and with their configuration, and the scan of each source + * gets what the planner takes from the optimized plan for the relation of the source (see + * `DataSourceV2Strategy`): the projections and the filters over each occurrence of the relation. + * + * The plan of the query does not change, so each micro-batch is optimized as without these + * pushdowns: the optimizer keeps deciding which columns every operator gets, and in which order, + * including the stateful ones, and Spark still evaluates all the filters over the relation. Only + * the scan node of the physical plan outputs the columns that the scan reads, and the planner + * checks that the micro-batch reads no other column of the relation. + * + * The other operations that the scans of batch queries can apply, such as aggregates and limits, + * apply to the data of all the micro-batches of a streaming query, which the scan of a micro-batch + * does not read, so they are not pushed down. + */ +private[sql] object StreamingV2ScanPushDown extends Logging { + + /** + * What the optimized plan of a streaming query reads from a streaming relation. + * + * @param requiredColumns the columns of the relation that the plan reads + * @param filters the deterministic filters without subqueries that the plan applies over every + * occurrence of the relation, which its scan may apply too + */ + case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression]) + + /** + * Returns what the optimized `plan` reads from each of its streaming relations of Data Source + * V2: what the planner takes from the projections and the filters over each occurrence of the + * relation, or all its columns and no filter where nothing is projected or filtered over it. A + * relation that the optimized plan does not have, or any relation if the optimizer fails, has no + * entry, so that its scan reads all its columns and gets no filter. + * + * @param optimizer the optimizer of the micro-batches of the query + * @param conf the configuration that the micro-batches of the query run with + * @param plan the plan of the query, with the current time of a micro-batch as the micro-batches + * have it, which the optimizer does not fold, so that no filter on it is pushed down + */ + def reads( + optimizer: RuleExecutor[LogicalPlan], + conf: SQLConf, + plan: LogicalPlan): Map[StreamingRelationV2, Reads] = { + SQLConf.withExistingConf(conf) { + try { + val occurrences = mutable.HashMap.empty[StreamingRelationV2, mutable.Buffer[Reads]] + // Visits the nodes as the planner plans them: the projections and the filters over a + // relation are planned with its scan, and the other nodes on their own. + def visit(node: LogicalPlan): Unit = node match { + case PhysicalOperation(project, filters, relation: StreamingRelationV2) => + occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) += Reads( + AttributeSet((project ++ filters).flatMap(_.references)), + filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f))) + case _ => + node.children.foreach(visit) + } + visit(optimizer.execute(plan)) + occurrences.map { case (relation, reads) => + // The scan of a relation reads the rows of all its occurrences, so it can only apply the + // filters that all of them apply. + val filters = reads.map(r => ExpressionSet(r.filters)).reduce { (left, right) => + ExpressionSet(left.filter(right.contains)) + } + relation -> Reads( + AttributeSet.fromAttributeSets(reads.map(_.requiredColumns)), filters.toSeq) + }.toMap + } catch { + case NonFatal(e) => + logWarning(log"Could not optimize the streaming query to push down what it reads to " + + log"the scans of its sources, which read all their columns with no filter.", e) + Map.empty + } + } + } + + /** + * Builds the scan of a streaming source, and creates its stream with `createStream`. If its + * scan builder supports it, and `conf` enables it, the scan only reads the columns of `output` + * that the query reads, and gets the filters that the query applies to all the rows that it reads + * from the source (see `reads`). `reads` is only evaluated for such a scan builder. Returns the + * scan, its stream, and the columns of `output` that the scan reads, in the order of its read + * schema, unless these are `output` itself. + * + * A pruned scan must read the required columns and only columns of `output`, each once, as + * `SupportsPushDownRequiredColumns` requires. A scan that does not is replaced by a scan of all + * the columns, as without pruning, and its stream is stopped. Spark asks a scan for its read + * schema once the stream is created, since a scan may only know then what its stream reads. + * + * @param output the output of the relation of the source, with the metadata columns that the + * query reads, which the analyzer added to it, so that connectors implementing + * `SupportsMetadataColumns` include them in `readSchema()` + * @param conf the configuration of the query + * @param reads what the query reads from the relation, or None to read all its columns + */ + def buildScan[S <: SparkDataStream]( + table: SupportsRead, + options: CaseInsensitiveStringMap, + output: Seq[AttributeReference], + conf: SQLConf, + reads: => Option[Reads])( + createStream: Scan => S): (Scan, S, Option[Seq[AttributeReference]]) = { + val builder = table.newScanBuilder(options) + val prunesColumns = conf.streamingV2ColumnPruningEnabled && + builder.isInstanceOf[SupportsPushDownRequiredColumns] + val pushesFilters = conf.streamingV2FilterPushdownEnabled && supportsFilters(builder) + val known = if (prunesColumns || pushesFilters) reads else None + val requiredOutput = known.filter(_ => prunesColumns) + .map(r => output.filter(r.requiredColumns.contains)) + .getOrElse(output) + val filters = known.filter(_ => pushesFilters) + .map(r => DataSourceStrategy.normalizeExprs(r.filters, output)) + .getOrElse(Nil) + + def build(builder: ScanBuilder, columns: Seq[AttributeReference]): Scan = { + if (filters.nonEmpty) { + // The plan of the query keeps the filters, which Spark evaluates after the scan, so the + // scan may apply any of them. + PushDownUtils.pushFilters(builder, filters, partitionFields = None) Review Comment: Added to the migration guide entry in 9ef62c1: a source that applies the filters reads fewer rows, so `numInputRows`, `inputRowsPerSecond` and `processedRowsPerSecond` only count the rows that pass the filters, and `spark.sql.streaming.v2FilterPushdown.enabled` restores the previous behavior. The user-facing section of the PR description says it too. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/StreamingV2ScanPushDown.scala: ########## @@ -0,0 +1,241 @@ +/* + * 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 scala.collection.mutable +import scala.util.control.NonFatal + +import org.apache.spark.internal.{Logging, LogKeys} +import org.apache.spark.sql.catalyst.expressions.{AttributeReference, AttributeSet, Expression, ExpressionSet, SubqueryExpression} +import org.apache.spark.sql.catalyst.planning.PhysicalOperation +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.catalyst.rules.RuleExecutor +import org.apache.spark.sql.catalyst.streaming.StreamingRelationV2 +import org.apache.spark.sql.catalyst.types.DataTypeUtils +import org.apache.spark.sql.connector.catalog.SupportsRead +import org.apache.spark.sql.connector.read.{Scan, ScanBuilder, SupportsPushDownFilters, SupportsPushDownRequiredColumns, SupportsPushDownV2Filters} +import org.apache.spark.sql.connector.read.streaming.SparkDataStream +import org.apache.spark.sql.execution.datasources.DataSourceStrategy +import org.apache.spark.sql.execution.datasources.v2.PushDownUtils +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters +import org.apache.spark.sql.types.StructType +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +/** + * Pushes the top-level columns and the filters that the optimized plan of a streaming query reads + * from each Data Source V2 streaming source down to the scan of the source, as + * `V2ScanRelationPushDown` does for the scans of batch queries. + * + * The stream of a source is created from its scan, and has to exist before any micro-batch is + * planned, so the scan is built when the query starts. The query is optimized once at that point, + * by the optimizer of its micro-batches and with their configuration, and the scan of each source + * gets what the planner takes from the optimized plan for the relation of the source (see + * `DataSourceV2Strategy`): the projections and the filters over each occurrence of the relation. + * + * The plan of the query does not change, so each micro-batch is optimized as without these + * pushdowns: the optimizer keeps deciding which columns every operator gets, and in which order, + * including the stateful ones, and Spark still evaluates all the filters over the relation. Only + * the scan node of the physical plan outputs the columns that the scan reads, and the planner + * checks that the micro-batch reads no other column of the relation. + * + * The other operations that the scans of batch queries can apply, such as aggregates and limits, + * apply to the data of all the micro-batches of a streaming query, which the scan of a micro-batch + * does not read, so they are not pushed down. + */ +private[sql] object StreamingV2ScanPushDown extends Logging { + + /** + * What the optimized plan of a streaming query reads from a streaming relation. + * + * @param requiredColumns the columns of the relation that the plan reads + * @param filters the deterministic filters without subqueries that the plan applies over every + * occurrence of the relation, which its scan may apply too + */ + case class Reads(requiredColumns: AttributeSet, filters: Seq[Expression]) + + /** + * Returns what the optimized `plan` reads from each of its streaming relations of Data Source + * V2: what the planner takes from the projections and the filters over each occurrence of the + * relation, or all its columns and no filter where nothing is projected or filtered over it. A + * relation that the optimized plan does not have, or any relation if the optimizer fails, has no + * entry, so that its scan reads all its columns and gets no filter. + * + * @param optimizer the optimizer of the micro-batches of the query + * @param conf the configuration that the micro-batches of the query run with + * @param plan the plan of the query, with the current time of a micro-batch as the micro-batches + * have it, which the optimizer does not fold, so that no filter on it is pushed down + */ + def reads( + optimizer: RuleExecutor[LogicalPlan], + conf: SQLConf, + plan: LogicalPlan): Map[StreamingRelationV2, Reads] = { + SQLConf.withExistingConf(conf) { + try { + val occurrences = mutable.HashMap.empty[StreamingRelationV2, mutable.Buffer[Reads]] + // Visits the nodes as the planner plans them: the projections and the filters over a + // relation are planned with its scan, and the other nodes on their own. + def visit(node: LogicalPlan): Unit = node match { + case PhysicalOperation(project, filters, relation: StreamingRelationV2) => + occurrences.getOrElseUpdate(relation, mutable.Buffer.empty) += Reads( + AttributeSet((project ++ filters).flatMap(_.references)), + filters.filter(f => f.deterministic && !SubqueryExpression.hasSubquery(f))) + case _ => + node.children.foreach(visit) + } + visit(optimizer.execute(plan)) + occurrences.map { case (relation, reads) => + // The scan of a relation reads the rows of all its occurrences, so it can only apply the + // filters that all of them apply. + val filters = reads.map(r => ExpressionSet(r.filters)).reduce { (left, right) => + ExpressionSet(left.filter(right.contains)) + } + relation -> Reads( + AttributeSet.fromAttributeSets(reads.map(_.requiredColumns)), filters.toSeq) + }.toMap + } catch { + case NonFatal(e) => + logWarning(log"Could not optimize the streaming query to push down what it reads to " + + log"the scans of its sources, which read all their columns with no filter.", e) + Map.empty + } + } + } + + /** + * Builds the scan of a streaming source, and creates its stream with `createStream`. If its + * scan builder supports it, and `conf` enables it, the scan only reads the columns of `output` + * that the query reads, and gets the filters that the query applies to all the rows that it reads + * from the source (see `reads`). `reads` is only evaluated for such a scan builder. Returns the + * scan, its stream, and the columns of `output` that the scan reads, in the order of its read + * schema, unless these are `output` itself. + * + * A pruned scan must read the required columns and only columns of `output`, each once, as + * `SupportsPushDownRequiredColumns` requires. A scan that does not is replaced by a scan of all + * the columns, as without pruning, and its stream is stopped. Spark asks a scan for its read + * schema once the stream is created, since a scan may only know then what its stream reads. + * + * @param output the output of the relation of the source, with the metadata columns that the + * query reads, which the analyzer added to it, so that connectors implementing + * `SupportsMetadataColumns` include them in `readSchema()` + * @param conf the configuration of the query + * @param reads what the query reads from the relation, or None to read all its columns + */ + def buildScan[S <: SparkDataStream]( + table: SupportsRead, + options: CaseInsensitiveStringMap, + output: Seq[AttributeReference], + conf: SQLConf, + reads: => Option[Reads])( + createStream: Scan => S): (Scan, S, Option[Seq[AttributeReference]]) = { + val builder = table.newScanBuilder(options) + val prunesColumns = conf.streamingV2ColumnPruningEnabled && + builder.isInstanceOf[SupportsPushDownRequiredColumns] + val pushesFilters = conf.streamingV2FilterPushdownEnabled && supportsFilters(builder) + val known = if (prunesColumns || pushesFilters) reads else None + val requiredOutput = known.filter(_ => prunesColumns) + .map(r => output.filter(r.requiredColumns.contains)) + .getOrElse(output) + val filters = known.filter(_ => pushesFilters) + .map(r => DataSourceStrategy.normalizeExprs(r.filters, output)) + .getOrElse(Nil) + + def build(builder: ScanBuilder, columns: Seq[AttributeReference]): Scan = { + if (filters.nonEmpty) { + // The plan of the query keeps the filters, which Spark evaluates after the scan, so the + // scan may apply any of them. + PushDownUtils.pushFilters(builder, filters, partitionFields = None) + } + builder match { + case pruning: SupportsPushDownRequiredColumns => + pruning.pruneColumns(DataTypeUtils.fromAttributes(columns)) + case _ => + } + builder.build() + } + def buildFullScan(builder: ScanBuilder): (Scan, S, Option[Seq[AttributeReference]]) = { + val scan = build(builder, output) + (scan, createStream(scan), None) + } + + if (requiredOutput.length == output.length) { + buildFullScan(builder) + } else { + val scan = build(builder, requiredOutput) + val stream = createStream(scan) + val (readSchema, scanOutput) = try { + val readSchema = scan.readSchema() + (readSchema, prunedScanOutput(readSchema, output, requiredOutput)) + } catch { + case e: Throwable => Review Comment: Fixed in 9ef62c1: - `buildScan` asks the pruned scan for its read schema before it creates the stream, and falls back to a scan of all the columns when `readSchema()` throws, as when it is invalid. New test: "read all the columns of a source whose pruned scan has no read schema". - Since the decision comes first, the source only creates the stream of the scan that it keeps: no stream to stop, and no second stream for the same `metadataPath`. That test counts the streams that the scans of its table create, and "read all the columns when a scan does not read the required ones" now checks that the in-memory table creates one stream. - The `Throwable` case is gone; the fallback only catches `NonFatal`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala: ########## @@ -123,39 +124,59 @@ 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 } /** + * The read schema of a scan of a micro-batch stream is only known once the stream is created, as + * `StreamingV2ScanPushDown` asks for it: it is the full schema, while the read schema that the + * scan reports before is the one of a batch, whose planner it creates to prune the columns. + * + * @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 { + private lazy val planner = pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema)) + + // 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. - private lazy val batch = new ColumnarBatchScan( - pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))) + private lazy val batch = { + // The planner prunes the columns before it plans the partitions. + val _ = batchSchema + new ColumnarBatchScan(planner) + } + + // Whether the scan is the one of a micro-batch stream, whose read schema is the full schema. + @volatile private var streaming = false - override def readSchema(): StructType = prunedSchema + override def readSchema(): StructType = if (streaming) fullSchema else batchSchema Review Comment: Fixed in 9ef62c1: the `streaming` flag is gone. `ColumnarScan.readSchema()` is the required columns whatever Spark calls first, and neither it nor `getMetaData()` creates a planner. Only `toBatch` does, and asks it to prune the columns; if it does not, the batch drops the other columns from the batches of the reader (and releases them with the batch). A stream drops them from the batches of its stream reader. So a streaming-only `ColumnarDataSource` never gets `batchPlanner` for a stream. `ColumnarDataSourceSuite` checks a batch read whose planner does not prune, and `NativeDataSourceSuite` checks the read-only library, which does not prune. #59931 has the same flag, so I will make the same change there. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousExecution.scala: ########## @@ -78,6 +77,13 @@ class ContinuousExecution( val v2ToRelationMap = MutableMap[StreamingRelationV2, StreamingDataSourceV2ScanRelation]() var nextSourceId = 0 import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Implicits._ + // The columns and the filters that the query reads from each Data Source V2 source, which the + // scan of the source gets if it supports them. They come from the plan of the query optimized + // once, as it is optimized when it runs. The plan is only optimized for a scan builder that + // supports a pushdown. + val streamConf = sparkSessionForStream.sessionState.conf + lazy val sourceReads = StreamingV2ScanPushDown.reads( Review Comment: Done in 9ef62c1: `StreamingV2ScanPushDown.reads(session, plan, offsetLogMetadata)` now clones the configuration of the stream, applies the values that the offset log pins, disables CBO and replaces the values of the time, so both executions call it in one line. `ContinuousExecution` passes no offset-log metadata, since it does not pin configurations from it. `reads` also runs with the session of the stream active, as the micro-batches do; the other overrides of `runStream` (AQE, the cluster keys of the distribution, the stateful shuffle partitions, the real-time mode) are about physical planning and the state, not the optimized plan. `MicroBatchExecution` keeps the batch timestamp substitution of master as it is. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Relation.scala: ########## @@ -478,14 +478,21 @@ case class StreamingDataSourceV2Relation( /** * A specialization of [[DataSourceV2ScanRelation]] with the streaming bit set to true, as well * as start and end offsets for Microbatch processing. + * + * @param output the columns of the source that the query references, which the plan of each + * micro-batch is optimized with + * @param prunedOutput the columns of `output` that the scan reads, in the order of its read schema, + * if it was built to read fewer of them. The scan node of the physical plan + * outputs these, so the plan must not read the other columns of `output`. */ case class StreamingDataSourceV2ScanRelation( relation: StreamingDataSourceV2Relation, scan: Scan, output: Seq[AttributeReference], stream: SparkDataStream, startOffset: Option[Offset] = None, - endOffset: Option[Offset] = None) + endOffset: Option[Offset] = None, + prunedOutput: Option[Seq[AttributeReference]] = None) Review Comment: Done in 9ef62c1: `simpleString` and `stringArgs` of `StreamingDataSourceV2ScanRelation` show the pruned output and the pushed filters, e.g. `StreamingDataSourceV2ScanRelation[id#1L, data#2, extra#3] testcat.t, PrunedOutput: [data#2, extra#3], PushedFilters: [(extra#3 > 10)]`. The tests of the column and filter checks assert them. -- 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]
