HyukjinKwon commented on code in PR #59934:
URL: https://github.com/apache/spark/pull/59934#discussion_r4472101021
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnarTable.scala:
##########
@@ -222,9 +338,52 @@ abstract class ColumnarOnlyPartitionReaderFactory extends
PartitionReaderFactory
override def supportColumnarReads(partition: InputPartition): Boolean = true
}
-class ColumnarStreamPartitionReaderFactory(reader: ColumnarStreamReader)
+object ColumnarProjection {
+ /**
+ * Returns `factory`, or a factory whose readers only return the columns at
`projection` of the
+ * batches of the readers of `factory`, if any.
+ */
+ def apply(
+ factory: PartitionReaderFactory,
+ projection: Option[Array[Int]]): PartitionReaderFactory = projection
match {
+ case Some(ordinals) => new ColumnarProjectionReaderFactory(factory,
ordinals)
+ case None => factory
+ }
+}
+
+/**
+ * Reads only the columns at `ordinals` of the batches of the readers of
`factory`, for a scan
+ * whose partitions return more columns than the scan reads.
+ */
+class ColumnarProjectionReaderFactory(factory: PartitionReaderFactory,
ordinals: Array[Int])
extends ColumnarOnlyPartitionReaderFactory {
override def createColumnarReader(partition: InputPartition):
PartitionReader[ColumnarBatch] = {
- reader.read(partition)
+ new ColumnarProjectionReader(factory.createColumnarReader(partition),
ordinals)
}
}
+
+private class ColumnarProjectionReader(
+ reader: PartitionReader[ColumnarBatch],
+ ordinals: Array[Int])
+ extends PartitionReader[ColumnarBatch] {
+
+ override def next(): Boolean = reader.next()
+
+ override def get(): ColumnarBatch = new ProjectedColumnarBatch(reader.get(),
ordinals)
Review Comment:
Fixed in #59931 (6ab9036), which is merged here in 8a28f1d:
`ColumnarProjectionReader.next()` builds the projected batch once, and `get()`
returns it until the next call of `next()`. The test of the projection checks
that `reader.get() eq reader.get()`.
--
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]
