fix: serialize large-offset string and binary vectors with 32-bit offsets

[andygrove]

Which issue does this PR close?

Closes #5488.

Rationale for this change

Utils.getFieldVector accepts a fixed list of Arrow vector classes that leaves out LargeVarCharVector and LargeVarBinaryVector, so Utils.serializeBatches throws Unsupported Arrow Vector for serialize on a batch that holds either. Comet reads these vectors on purpose elsewhere. A PyArrow UDF can return large_string or large_binary (pandas 3 backs string columns with them), CometPlainVector reads 64-bit offsets, and Utils.toArrowType maps LargeUtf8 to StringType.

The gap is latent on main, and I have not reproduced it with a query. serializeBatches has two callers. getByteArrayRdd in operators.scala has no callers. CometBroadcastExchangeExec only broadcasts a plan whose children are all CometNativeExec. CometMapInBatchExec, the operator that keeps the Python worker's vectors, is not one of those. EliminateRedundantTransitions also inserts it after CometExecRule has run, so its output reaches the rest of the plan through a row transition. This change makes the serialization boundary hold once such output can reach it.

The issue suggested two fixes: accept the large vectors in getFieldVector, or narrow their offsets to 32 bits before serializing. This PR narrows them, because of what happens to the bytes after serializeBatches:

• On the build side of a broadcast join, CometBatchRDD decodes the bytes with Utils.decodeBatches, and CometArrowStream.wrapColumnarBatchRDD feeds the batches to the native plan. reconcileStreamSchema advertises the first batch's actual Arrow types for the whole stream, and ColumnarBatchArrowReader loads every later batch's buffers under that schema. If each batch kept its own offset width, a broadcast that holds both widths would load one width's offset buffers under the other's type.
• Utils.coalesceBroadcastBatches appends every batch into a root built from the first batch's schema. When a later batch doesn't match, it leaves the broadcast uncoalesced, so mixed widths would also lose coalescing.
• CometNativeColumnarToRowExec hands decoded broadcast batches to native code through NativeUtil.exportBatch, which checks each vector with the same getFieldVector. Accepting large vectors there would widen what every JVM-to-native export accepts, not just serialization.
• The column's Spark type is StringType or BinaryType, which Comet carries as Utf8 or Binary everywhere else, including what the cache serializer writes when it converts a batch. Writing that form gives every reader the type it plans for. Without narrowing, native ScanExec would cast LargeUtf8 back to the planned Utf8 in every task that reads the broadcast. Narrowing copies the column once, when it is serialized.

What changes are included in this PR?

• Utils.getBatchFieldVectorsWithProviders builds the vectors serializeBatches writes. It now replaces a LargeVarCharVector or LargeVarBinaryVector with a VarCharVector or VarBinaryVector copy, made by a new private Utils.narrowOffsets.
  • The copy moves the validity and data buffers in bulk, rewrites the offsets, and keeps the field's name, nullability and metadata.
  • A column whose data does not fit 32-bit offsets (more than 2 GiB in one batch) fails with a SparkException that names the column. The size is checked before anything is allocated.
  • Like a materialized ConstantColumnVector in the same method, the copy is released when serializeBatches clears its root. The original vector stays with its owner.
• getFieldVector and isSupportedFieldVector are unchanged and still agree. The FFI export in NativeUtil.exportBatch still rejects large-offset vectors. isArrowBacked still answers false for them, so the cache serializer keeps converting such a batch. The comments explaining that answer are updated, since writing such a batch no longer fails.
• Nested large-offset children, such as a struct of large_string, already pass getFieldVector and are written as they are. This PR leaves them alone.

How are these changes tested?

New tests in UtilsSuite:

• serializeBatches writes large-offset string and binary vectors with 32-bit offsets: a batch with a LargeVarCharVector and a LargeVarBinaryVector column, whose values include a null, an empty value and a multi-byte character, is serialized and then decoded with Utils.decodeBatches, as the broadcast consumers do. The decoded columns are VarCharVector and VarBinaryVector and hold the original values.
• coalesceBroadcastBatches appends a large-offset batch to a 32-bit one: a large-offset batch followed by a VarCharVector batch coalesces into a single batch, without falling back, and keeps all four values in order.
• serializeBatches refuses a large-offset column too big for 32-bit offsets: offsets that claim more than Int.MaxValue bytes fail with the new message, before anything is copied.

Without the change to Utils.scala, all three fail with SparkException: Unsupported Arrow Vector for serialize: class org.apache.arrow.vector.LargeVarCharVector.

With the change, on the default profile (Spark 4.1, Scala 2.13), I ran UtilsSuite, NativeUtilSuite, CometInMemoryCacheSuite and the CometJoinSuite tests whose names contain Broadcast. All 94 tests pass.

[sunchao]

Summary

• Prior state and problem: serializeBatches rejected top-level large-offset string and binary vectors that Comet could otherwise read.
• Design approach: Normalize these columns to 32-bit offsets before serialization, giving broadcast consumers consistent schemas.
• Correctness / compatibility analysis: The conversion preserves values, nulls and field metadata. Spark sources for supported versions support the resulting logical types. One introduced P2 cleanup leak remains when a later column fails conversion.
• Key design decisions: The conversion is localized to serialization, preserving the existing FFI acceptance boundary.
• Implementation sketch: narrowOffsets checks the payload size, copies validity and data buffers, and rewrites rebased offsets.
• Behavioral changes worth calling out: Affected columns incur one copy and temporary off-heap allocation. Payloads exceeding Int.MaxValue are rejected before allocation for that column.
• Suggested improvements: Track temporary vectors across batch assembly and release them on failure. Add a two-column cleanup regression test.

Reviewed the entire diff from 8805e3f066440752ac98d5b48e197b43e4fb9ebd to 30efe2b5cbbf1ad45297f6ee35e84fb1f96f90de. Routed skills: review-comet-pr, review-comet-ffi-pr, and review-comet-memory-pr. The PR is non-draft and has no existing discussion or review concerns.

Exact-head CI: 24 checks passed and 14 were skipped. Required checks passed. The Linux Spark 4.1 execution job passed 1,116 tests, including all three new tests. Spark SQL, PyArrow UDF, Iceberg and macOS jobs were skipped.

Local validation used verbatim extracted helpers with Arrow 18.3.0, Scala 2.13.17, JDK 17 and a thin Spark vector adapter. Seven IPC round trips passed, covering empty/all-null values, binary data, metadata and nonzero starting offsets. The failure probe retained 72 bytes at head versus zero at base. A full local JVM/native build and an end-to-end SQL reproduction were not run.

[sunchao]

[P2] Could we release previously created copies when conversion of a later column fails? With two large-offset columns, the first can successfully allocate its narrowed copy before the second triggers the new oversized-payload exception. The enclosing map then abandons the first copy before serializeBatches creates its root, so root.clear() cannot release it. Closing the input batch only releases the original vectors. Each such failure therefore leaks the earlier columns' copied off-heap buffers. Track newly materialized vectors and close them when batch assembly fails, while preserving caller-owned vectors.

Evidence: A bounded probe used the exact helper bodies with Arrow 18.3.0: column 0 contained small, and column 1 used the PR's existing Int.MaxValue + 1L offset fixture without allocating a large payload. Each failed call retained 24 temporary bytes. After three calls and closing every input vector, the allocator still reported 72 bytes. Replacing batch dispatch with the base revision's implementation retained zero bytes. Reproduction and logs: /tmp/comet6371-current-review/run_probe.py, OffsetReviewProbe.log, and BaseOffsetReviewProbe.log.

[comphead]

Thanks for the thorough write-up. The bulk copy, the rebasing by the first offset and the size check before any allocation read correctly to me, and the coalesceBroadcastBatches test nicely shows why narrowing is needed. I haven't run this locally. Two small test suggestions inline.

[comphead]

Do the new tests reach the non-zero start path? fill always writes from index 0, so start is 0 in all three, and dropping the - start below or the start source index in the data copy would still pass. Would a case like hello, null, wörld with offsets [2, 7, 7, 13] over xxhellowörld be worth adding, for example via loadFieldBuffers? A zero-row vector also skips this code.

[comphead]

Nit: the isArrowBacked rejects large-offset Arrow vectors test above builds the same LargeVarCharVector and LargeVarBinaryVector by hand. Would it make sense to reuse fill and cometVector there, so the suite has one way to build them?

[parthchandra]

Notes on the large-offset narrowing (confirmed it errors cleanly past 2GB rather than silently truncating):

• [+1 to sunchao's P2] spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala:600 — when the dataLength > Int.MaxValue throw fires on, say, the third column of a batch, the narrowed copies already built for the first two are never placed in a VectorSchemaRoot, so serializeBatches never runs root.clear() on them and they leak from CometArrowAllocator (same for the materialized ConstantColumnVector copies above). This PR widens a pre-existing gap by adding a new throwing conversion. Could getBatchFieldVectorsWithProviders close whatever it already built if a later column fails?

• Utils.scala:575 (the narrowOffsets scaladoc) — a struct or list holding a large_string child still serializes with the child's 64-bit offsets untouched, i.e. the same mixed-width schema drift this PR fixes for top-level columns, one level down. Reasonable to scope this PR to the top level, but please file a tracking issue for the nested case and reference it here; nothing tests or errors on it today.

• Utils.scala:617 (the validity copy) — the data offsets are rebased by start, but validity is copied from bit 0. For an Arrow Java vector that's fine (logical index 0 always sits at validity bit 0 even when offset[0] is nonzero). Just confirming no path hands narrowOffsets a vector whose validity is bit-offset relative to its first value, which would silently misalign nulls. If it can't happen, a one-line note would help.

• Utils.scala:623 (the data copy) — this is a full copy of up to nearly 2 GiB into CometArrowAllocator (which no memory budget sees), briefly holding both the original large column and the narrowed copy. It matches the existing ConstantColumnVector path, so no change requested - just confirming the transient peak is understood and acceptable.

[sunchao]

Summary

• Prior state and problem: serializeBatches rejected top-level large-offset string and binary vectors that Comet could otherwise read.
• Design approach: Normalize these columns to 32-bit offsets before serialization, keeping broadcast schemas consistent.
• Correctness / compatibility analysis: The conversion preserves values, nulls and field metadata. Spark sources for all five supported profiles confirm the resulting logical-type mappings. The existing P2 cleanup leak remains unresolved and was reproduced again.
• Key design decisions: One private helper localizes conversion while preserving existing FFI and cache acceptance rules. Nested large-offset children remain unchanged.
• Implementation sketch: narrowOffsets checks payload size, copies validity and data buffers, rewrites rebased offsets, and sets the value count.
• Behavioral changes worth calling out: Each affected column incurs a buffer copy and temporary off-heap allocation. Payloads exceeding Int.MaxValue are rejected before allocating that column's copy.
• Suggested improvements: Address the existing blocker by releasing previously materialized vectors when batch assembly fails, while preserving caller-owned vectors. Cover this with a cleanup regression test.

Reviewed the full diff from 8805e3f066440752ac98d5b48e197b43e4fb9ebd to 30efe2b5cbbf1ad45297f6ee35e84fb1f96f90de. The PR is non-draft. Routed skills: review-comet-pr, review-comet-ffi-pr, and review-comet-memory-pr.

Exact-head CI: 24 checks succeeded and 14 were skipped. Required Checks passed. The Linux Spark 4.1 execution job passed 1,116 tests, including all three new tests. Spark SQL, PyArrow UDF, Iceberg and macOS jobs were skipped.

Local validation: verified helper extraction against head and base, then compiled bounded probes using Arrow 18.3.0, Scala 2.13.17, JDK 17 and a thin Spark vector adapter. Eleven IPC round trips passed, including empty vectors, binary data, metadata, nonzero offsets and sliced validity. Three failed batch assemblies retained 72 temporary bytes at head versus zero with base dispatch. No full local JVM/native build, end-to-end PyArrow query or Spark SQL suite was run.

No additional introduced P1/P2 issues found within this review. The existing blocker is not duplicated as a new finding.

[andygrove]

@Visorgood you worked on JVM columnar shuffle in #6110. If you have time, would you be able to review this one? Thanks.

[sunchao]

Summary

• Prior state and problem: serializeBatches rejected top-level LargeVarCharVector and LargeVarBinaryVector columns that Comet could otherwise read. The reported gap is at the serialization boundary. An end-to-end query triggering it has not been demonstrated.
• Design approach: Normalize these columns to VarCharVector and VarBinaryVector before serialization, keeping offset widths consistent across broadcast batches.
• Correctness: Buffer copying, offset rebasing and setLastSet preserve values and nulls in the exercised cases. Arrow 18.3.0 source confirms that Java slicing rebases validity bits appropriately. The existing P2 cleanup leak remains unresolved: a later column's conversion failure abandons copies already created for earlier columns.
• Compatibility analysis: Checked ArrowUtils sources for Spark 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0. All support the resulting Utf8/Binary mappings. Spark 3.5 and later also map the large variants to the same logical types. No additional P1/P2 compatibility issue was identified.
• Key design decisions: Narrowing stays within batch serialization. getFieldVector, FFI export acceptance and the cache's direct-write eligibility remain unchanged. Nested large-offset children retain their previous behavior.
• Implementation sketch: narrowOffsets checks the rebased payload length before allocation, preserves field name, nullability and metadata, copies validity and data buffers, writes 32-bit offsets, and sets the value count. Three added tests exercise round trips, mixed-width coalescing and oversized-payload rejection.
• Performance: Each affected column adds one payload copy and an offset rewrite, temporarily retaining both original and narrowed buffers. These allocations use the unbounded CometArrowAllocator. Ordinary vectors retain their existing path. Normalization enables mixed-width coalescing, but no throughput improvement was measured.
• Design: Normalizing at serialization gives downstream readers a consistent schema without widening the general FFI contract. This is a focused approach. Batch-level cleanup must still account for partially completed conversion.
• Abstraction & complexity: One private helper handles both variable-width types without adding a framework or configuration. Its local exception cleanup handles its own copy, but cannot release copies created by earlier iterations of the enclosing map.
• Behavioral changes worth calling out: Previously rejected top-level large-offset columns now serialize using 32-bit offsets. Payloads exceeding Int.MaxValue bytes fail explicitly. Source vectors remain caller-owned. The change does not establish a new Python-to-broadcast execution route.
• Suggested improvements: Address the already-reported cleanup blocker by tracking newly materialized vectors and releasing them if batch assembly fails, while preserving borrowed vectors. A two-column regression test should assert that allocator usage returns to its starting value after failure and input cleanup.

Reviewed the entire two-file, one-commit diff from 8805e3f066440752ac98d5b48e197b43e4fb9ebd to 30efe2b5cbbf1ad45297f6ee35e84fb1f96f90de. The PR is non-draft. Routed skills: review-comet-pr, review-comet-ffi-pr and review-comet-memory-pr. Existing reviews, conversation comments, inline comments and threads were read.

Exact-head CI: 24 checks succeeded and 14 were skipped, with no failed checks. Required Checks passed. The Linux Spark 4.1 execution log confirms testing the requested head merged into the requested base, with 1,116 tests passing, including all three added tests. Spark SQL, PyArrow UDF, Iceberg and macOS jobs were skipped.

Local validation: Recompiled probes after verifying extracted helper bodies against the requested head and base dispatch. With Arrow 18.3.0, Scala 2.13.17 and JDK 17, eleven IPC round trips passed, covering empty/all-null vectors, binary data, metadata, nonzero offsets and sliced validity. Three bounded failed batch assemblies retained 72 temporary bytes at head versus zero with base dispatch, confirming the existing P2 blocker. The probes use a thin Spark vector adapter. No full local JVM/native build, end-to-end PyArrow query, Spark SQL suite or performance benchmark was run.

No additional introduced P1/P2 issues found within this review. The existing blocker is not duplicated as a new finding.