HyukjinKwon commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4461833598
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/Columnar.scala:
##########
@@ -589,6 +590,9 @@ case class ApplyColumnarRulesAndInsertTransitions(
case write: DataWritingCommandExec
if write.cmd.isInstanceOf[V1WriteCommand] &&
conf.plannedWriteEnabled =>
write.child.supportsColumnar
+ // A V2 write whose data writers take columnar batches writes those of
a query that only
+ // supports columnar output.
+ case write: V2TableWriteExec if write.writesColumnar => true
Review Comment:
Good point, fixed in 732fe67. `ApplyColumnarRulesAndInsertTransitions` now
only lets a V2 write take the batches of its query when the session has no
columnar rules, and Spark has none of its own. With the columnar rules of an
extension, the query gets a `ColumnarToRowExec` as before, which the extension
may replace with its own, and the write takes rows. The doc of `writesColumnar`
and the guide say so. The new test "write rows when an extension has columnar
rules" applies the transitions to the physical plan of a copy without columnar
rules and with the rule of an extension, and checks that the write takes the
batches of the scan in the first case, and rows through a transition in the
other.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,39 +131,60 @@ 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. The required columns are only recorded here: a batch
scan asks its planner
+ // to prune the others, and a streaming scan, which has no planner, passes
the required
+ // columns to its stream reader.
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 only these columns: they
are the schema that
+ * its 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))
+
+ // 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)
- override def readSchema(): StructType = prunedSchema
+ // The columns that the stream reads.
+ private def streamSchema: StructType = requiredSchema.getOrElse(fullSchema)
Review Comment:
Right: until SPARK-60734 (#59943) lands, the micro-batch and continuous
executions pass all the columns of the relation to `pruneColumns`, so
`create_stream_reader` gets all of them. The header allows that, since its
schema is "the columns to read ... or only the top-level columns of it that the
query needs".
#59943 makes both executions pass the columns that a streaming query reads,
taken from the plan of the query optimized when it starts. Its version of
`ColumnarScan` still gives the stream reader the full schema, since master's
stream reader does not take the columns to read yet. Whichever of #59931 and
#59943 lands second connects the two: `ColumnarScan` then passes the pruned
columns to `create_stream_reader`, with a streaming test of a native data
source that reads some of the columns. #59943 already has such a test for the
columnar adapter, "streaming read of some of the columns" in
`ColumnarDataSourceSuite`.
--
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]
