Repository navigation
Broadcast serialization keeps 64-bit offsets for large_string and large_binary children of nested columns #6933
Description
Describe the bug
Utils.serializeBatches writes a struct, list or map column as it stands, so a large_string or large_binary child keeps its 64-bit offsets. #6371 narrows top-level LargeVarCharVector and LargeVarBinaryVector columns to 32-bit offsets before they are written. getFieldVector accepts a StructVector, ListVector or MapVector without looking at its children, though, so a nested child is written with the offsets it arrived with.
That matters when the batches of one broadcast disagree on the width. It is the case #6371 fixes for top-level columns, one level down:
coalesceBroadcastBatchescannot append astruct<c: large_string>batch to astruct<c: string>one, in either order, so the broadcast is left uncoalesced.CometArrowStream.reconcileStreamSchematakes the input stream's schema from the first batch, children included, andColumnarBatchArrowReaderloads every later batch's buffers under it. A later batch of the other width then has its offsets read at the wrong width. With the 64-bit batch first, reading the second batch's child from the reader's root fails withIndexOutOfBoundsException: index: 17179869184. With the 32-bit batch first, the second batch silently reads["", "cccc"]where it holds["cccc", "dddd"]. That root is what gets exported to the native plan.
When every batch agrees, the stream still carries Struct<c: LargeUtf8> where the plan expects Struct<c: Utf8>, so native ScanExec would cast the column in every task that reads the broadcast.
Like #5488, this is latent on main. CometBroadcastExchangeExec only broadcasts a plan whose children are all CometNativeExec, and CometMapInBatchExec, the operator that keeps a Python worker's vectors, is not one of them. The gap is at the serialization boundary, ready for when such output can reach it.
Steps to reproduce
Verified on main at 474b5ec (Spark 4.1) with a throwaway suite:
- Build two
StructVectorcolumns with one child each: aLargeVarCharVectorholdingaaaaandbbbb, and aVarCharVectorholdingccccanddddd. Wrap each in aCometVectorand aColumnarBatch. - Call
Utils.serializeBatcheson the two batches. Decoding the first withUtils.decodeBatchesgives back aLargeVarCharVectorchild. - Call
Utils.coalesceBroadcastBatcheson the two buffers. It returns a batch count of 0, the uncoalesced fallback. - Decode both buffers and pass the batches through
CometArrowStream.reconcileStreamSchema, with the schemaUtils.toArrowSchemagives forst struct<c: string>, then through aColumnarBatchArrowReaderover the reconciled schema. Read the struct child of each batch the reader loads. The second batch fails or reads wrong values as above, depending on which width comes first.
Expected behavior
A nested large_string or large_binary child is written with 32-bit offsets, as #6371 does for top-level columns, so every batch of a broadcast has the same schema and native readers get the Utf8 or Binary types they plan for. A child holding more data than 32-bit offsets can address can still be refused.
Additional context
Raised by @parthchandra in review of #6371 (#6371 (comment)), which scopes its fix to top-level columns.
The cache serializer is not affected. CachedBatchIpc.matchesReaderLayout compares nested types too, and converts a batch whose layout does not match.
Two ways to close this:
- Narrow nested children in
getBatchFieldVectorsWithProviders, rebuilding a struct, list or map vector around narrowed children and sharing every other buffer. - Convert a batch whose nested types differ from the Spark schema's Arrow types, as the cache serializer does.
serializeBatchesis not given the schema today, so it would need passing in.
Activity
fix: release the vectors a failed batch assembly made, and test offse…