sunchao commented on code in PR #5557:
URL: https://github.com/apache/datafusion-comet/pull/5557#discussion_r3890174485
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -329,6 +348,55 @@ private[python] trait CometArrowPythonRunnerBase
private[python] object CometArrowPythonRunnerBase {
+ /**
+ * Supply logical Arrow vectors to the serializer for the duration of the
body.
+ *
+ * Plain Comet vectors already expose their logical values and remain
borrowed.
+ * Dictionary-backed shuffle columns expose only their integer indices
through getValueVector,
+ * so materialize those columns first. The temporary decoded vectors own
their buffers and are
+ * closed after the synchronous write, including schema and serialization
failures.
+ */
+ private[python] def withMaterializedInputVectors[T](
+ columns: Seq[CometDecodedVector],
+ allocator: BufferAllocator)(body: Seq[FieldVector] => T): T = {
+ val materialized = new ArrayList[FieldVector]()
+ try {
+ val vectors = columns.map {
+ case dictionaryVector: CometDictionaryVector =>
Review Comment:
I split the dictionary-input fix into #5560. I kept that fix scoped to the
top-level `CometDictionaryVector` shape that the JVM shuffle produces today:
nested list, struct, and map builders explicitly disable dictionary encoding,
and the container import paths do not surface a dictionary-backed child to this
runner.
Given that there is no production producer for the hand-built nested shape,
adding recursive materialization here would introduce ownership logic with no
end-to-end path to test. #5560 covers the reachable JVM-shuffle shape with
real-worker tests and allocator/refcount checks. If nested dictionaries become
producible, I agree they should be handled recursively with a provider-aware
test at that boundary rather than allowed to reach Arrow's null-provider
failure. The dictionary code is now outside this feature PR.
##########
spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala:
##########
@@ -379,17 +449,27 @@ private[python] object CometArrowPythonRunnerBase {
buffers.add(structValidity)
buffers.addAll(sourceBatch.getBuffers)
- val wrappedBatch = new ArrowRecordBatch(
- numRows,
- nodes,
- buffers,
- sourceBatch.getBodyCompression,
- sourceBatch.getVariadicBufferCounts,
- true)
+ val widenedOffsets = if (useLargeVarTypes) new ArrayList[ArrowBuf]()
else null
Review Comment:
Done on the narrowed feature branch in `1876303b9`. `widenedOffsets` is now
always an `ArrayList`, cleanup no longer uses a null sentinel, and
`serializeBatch` requires the Boolean at every call site. The final narrow head
passed 21/21 focused JVM tests and all 128 real-worker cases.
--
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]
