peterxcli commented on code in PR #5051:
URL: https://github.com/apache/datafusion-comet/pull/5051#discussion_r3677859989
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala:
##########
@@ -74,4 +74,45 @@ object CometArrowConverters extends Logging {
}
}
}
+
+ /**
+ * Copy a Spark `ColumnarBatch` whose columns are not Arrow-backed (e.g.
+ * `On/OffHeapColumnVector` from Spark's vectorized Parquet reader, or a
third-party connector's
+ * vectors) into a freshly allocated Arrow `ColumnarBatch` of `CometVector`s.
+ *
+ * The input batch is not consumed or closed; the caller owns the returned
batch and must close
+ * it. Values are copied element-wise, since Spark's `ColumnVector`
implementations do not
+ * expose Arrow buffers.
+ */
+ def columnarBatchToArrowBatch(
Review Comment:
A `ColumnarBatch` is a container:
```scala
class ColumnarBatch {
ColumnVector[] columns;
}
```
Spark 4.1.2’s relevant `ColumnVector` implementations are:
Producer | Vectors inside ColumnarBatch
-- | --
Vectorized Parquet | OnHeapColumnVector, OffHeapColumnVector,
ConstantColumnVector
Vectorized ORC | OrcAtomicColumnVector, OrcArrayColumnVector,
OrcMapColumnVector, OrcStructColumnVector; sometimes heap/off-heap or constant
vectors
Spark Arrow/Python paths | ArrowColumnVector
Comet | CometVector
External columnar sources | Any implementation extending Spark’s ColumnVector
[SparkColumnarArrowReader](https://github.com/apache/datafusion-comet/blob/63c5a916b2e5626b040dc1f097384ab1a443bf91/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala#L27-L101)
does not inspect those concrete classes. It wraps each column in
`ColumnarArray`, and `ArrowWriter` reads through Spark’s standard accessors
such as `getLong`, `getUTF8String`, `getArray`, and `getStruct`. Therefore, it
supports any implementation that correctly follows the [spark's
ColumnVector](https://github.com/apache/spark/blob/v4.1.2/sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnVector.java)
`ColumnVector` contract for the declared data type.
The cache writer could reuse it approximately like this:
```scala
override def convertColumnarBatchToCachedBatch(
input: RDD[ColumnarBatch],
schema: Seq[Attribute],
storageLevel: StorageLevel,
conf: SQLConf): RDD[CachedBatch] = {
val sparkSchema = toStructType(schema)
val batchSize = conf.columnBatchSize
val timeZoneId = conf.sessionLocalTimeZone
input.mapPartitions { batches =>
// Arrow Schema is not serializable, so construct it inside the task.
val arrowSchema = Utils.toArrowSchema(sparkSchema, timeZoneId)
val arrowBatches = CometArrowStream.readerBatchIter(
"CometCacheWrite",
allocator =>
new SparkColumnarArrowReader(
allocator,
arrowSchema,
batches,
batchSize))
encodeBatches(arrowBatches, schema)
}
}
```
That gives this path:
```
Spark ColumnarBatch
→ SparkColumnarArrowReader
→ Arrow-backed ColumnarBatch containing CometVector
→ existing encodeBatches/getBatchFieldVectors
→ CometCachedBatch</code>
```
_Assist with LLM_
--
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]
