HyukjinKwon commented on code in PR #59932:
URL: https://github.com/apache/spark/pull/59932#discussion_r4469603470
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -351,17 +366,20 @@ class NativeDataSourceStreamReader(
}
override def partitions(start: String, end: String): Array[InputPartition] =
{
- val partitions = call(name, "plan a scan of") {
- library.streamPartitions(handle.get, start, end)
+ if (plannedRange != (start, end)) {
Review Comment:
Fixed in #59931 (9681e3e), merged here in 9f94cbf. The calls of
`NativeDataSourceStreamReader` on the driver (`initial_offset`,
`latest_offset`, `stream_partitions`, `commit_offset`, and the close in `stop`)
are now `synchronized` on the reader, and the range that it planned last keeps
its failure as well as its partitions, so that all the scans of a micro-batch
get the result of a single `stream_partitions` call, and a failed call is not
repeated on the handle. "plan each micro-batch once, also when a query reads a
source twice" now plans 20 ranges from two threads at the same time, and checks
that each range has one call and that both threads get the same partitions
(without the lock, it fails with the partitions of two calls), and that
planning a failed range again rethrows the same error without calling the
library.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -123,53 +134,80 @@ 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 (!pruned) {
+ // The columns are pruned at most once, since the scan asks the reader
once, as the header
+ // promises, though Spark prunes the columns of a scan builder once
anyway.
+ pruned = true
+ if (required.length < schema.length) {
+ this.requiredSchema = Some(required)
+ }
+ } else if (!required.fieldNames.toSet.subsetOf(
+ this.requiredSchema.getOrElse(schema).fieldNames.toSet)) {
+ throw SparkException.internalError(
+ "The columns of a scan were pruned again, with columns that it does
not read.")
}
}
// 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
}
/**
+ * A scan of a [[ColumnarDataSource]]. It reads the top-level columns that the
query needs, the
+ * columns of its read schema: a batch scan asks its planner to prune the
other columns when it
+ * plans its partitions, and drops them from the batches of the reader if the
planner does not,
+ * while the stream reader of a streaming scan reads only these columns.
+ *
+ * @param requiredSchema the top-level columns that the query needs, if fewer
than those of the
+ * full schema
* @param pushedDownPlanner the planner that the operations were pushed down
to, if any. Otherwise,
* the scan creates one when its batch is planned.
*/
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)))
+ // Spark may call `toBatch` more than once for the same scan; the partitions
are planned once,
+ // after the planner prunes the columns, if it does. A failure is kept too,
so that the planner
+ // is asked at most once, as the header promises.
+ private lazy val batch: Try[ColumnarBatchScan] = Try {
+ val planner =
pushedDownPlanner.getOrElse(dataSource.batchPlanner(fullSchema))
+ // The ordinals in the full schema of the columns that the scan reads, if
the planner does
+ // not prune the others.
+ val projection = requiredSchema.filterNot(planner.pruneColumns)
Review Comment:
Added in #59931 (9681e3e), merged here. "push down a limit" also runs
`select("id").limit(3)` with the default library: it checks the rows, the
pushed limit, and that the scan reads with the reader factory of the library
rather than with one that drops columns, so the planner that got the limit
pruned the columns too. Here, the order-violation checks of the test library
see push_limit and then prune_columns on the same reader. "push down a limit
and the required columns" in `ColumnarDataSourceSuite` combines a projection
with a limit, for a planner that prunes and one that does not, and checks which
reader factory the batch uses.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -802,6 +855,7 @@ class NativePartitionReader(
throw error(name, "read", s"The column $column has an offset of
${array.offset}, but " +
"Spark only reads arrays with an offset of 0. Copy sliced arrays
before exporting them.")
}
+ checkBuffers(array, s"The column $column")
Review Comment:
Fixed in #59931 (9681e3e), merged here: `checkBuffers` and `checkValidity`
take their owner by name, so the message is only built when a check fails.
--
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]
