gengliangwang commented on code in PR #59933:
URL: https://github.com/apache/spark/pull/59933#discussion_r4426218792
##########
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:
**[P2] Stream column pruning never happens in a real streaming query**
MicroBatchExecution (and ContinuousExecution) always call
`pruneColumns(output.toStructType)` with the full relation output, so
`streamSchema` is always the full schema in production. The stream reader is
still created with every column, although the header, the docs and
`ColumnarDataSource.streamReader` promise that only the needed columns are
passed. The new `streaming` flag and the requiredSchema plumbing for streams
are only exercised by unit tests that call pruneColumns by hand.
**Failure scenario:** readStream.format(<native>).load().select("mod") calls
create_stream_reader with every column of the source rather than only `mod`.
That contradicts the documented contract, and the unit tests pass anyway
because they call pruneColumns themselves.
##########
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:
**[P2] Columnar V2 write skips the ColumnarToRowExec that ColumnarRule
extensions replace**
When `writesColumnar` is true, `ensureOutputsRowBased` no longer puts a
ColumnarToRowExec above a query that only outputs columnar batches.
ColumnarRule extensions (Gluten/Velox, Comet) usually swap that transition for
their own in postColumnarTransitions. With no transition in the plan, their
batches go straight to `ArrowBatchingDataWriter.writeBatch`, which only copies
`ArrowColumnVector` columns and otherwise falls back to `batch.getRow(i)`.
**Failure scenario:** With Gluten or Comet installed, run INSERT INTO
native_ds SELECT * FROM t, where t is read by the extension's columnar scan.
The child is columnar-only, so no C2R is inserted and the extension has none to
replace. Its opaque batches (e.g. IndicatorVector) reach writeBatch, and the
row-by-row fallback through their ColumnVector getters either throws or is very
slow. Before this change, the same query was written through the extension's
C2R.
--
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]
