Avoid temporary heap allocation when copying off-heap UTF8String into Arrow

[sunchao]

What is the problem the feature request solves?

When a Spark row contains an off-heap UTF8String, Comet's ordinary Spark-row-to-Arrow string writer copies the string into a temporary JVM byte array, then copies that array into the Arrow value buffer. The first allocation and copy can be avoided while preserving Arrow's ownership of the output.

This is relevant to row-to-columnar boundaries carrying large strings, such as text or JSON stored in a STRING column. The benefit depends on the actual input representation: the same temporary payload allocation does not occur when getByteBuffer() can wrap an existing on-heap byte array.

Current execution path

At Comet main 451c9996:

1. RowArrowReader.loadNextBatch creates an Arrow writer and writes each InternalRow into the current batch.
2. StringWriter.setValue obtains UTF8String, calls getByteBuffer(), and passes the buffer to VarCharVector.setSafe().
3. In Spark 4.1.1 UTF8String.getByteBuffer, an existing byte-array backing can be wrapped directly. An off-heap backing takes the getBytes() branch, which allocates byte[numBytes] and copies the payload.
4. Arrow Java 18.3.0 BaseVariableWidthVector.setSafe ensures capacity, updates validity/offsets, and copies the bytes into its own value buffer. This is the Arrow Java version pinned by the inspected Comet source.

For an off-heap string of length L:

Current:
  Spark-owned off-heap bytes
      -> temporary JVM byte[L]
      -> Arrow-owned value buffer

Desired:
  Spark-owned off-heap bytes
      -> Arrow-owned value buffer

For example, converting 32 strings of 1 MiB each currently allocates approximately 32 MiB of temporary byte-array payloads and copies 64 MiB of payload across these two copy operations, excluding array headers, wrapper objects, and Arrow buffer growth. Removing the intermediate payload would leave the necessary 32 MiB copy into Arrow. These figures are arithmetic from the code path, not measured peak heap, RSS, elapsed time, or a claimed production speedup. The arrays are created per value; their aggregate allocation must not be confused with simultaneous live memory.

Describe the potential solution

Add a narrowly scoped path in the ordinary StringWriter that copies off-heap UTF-8 bytes directly into Arrow-owned storage, without allocating a payload-sized JVM byte array.

Possible implementation approaches

A. Reserve the destination range and copy into its address. Arrow's public setValueLengthSafe can grow the buffers, fill missing offsets, and establish the value's byte range. Spark's UTF8String.writeToMemory can then copy the source bytes into a sufficiently sized destination address. Complete the validity/value bookkeeping through Arrow's APIs.

This avoids the temporary payload and can also avoid a per-value buffer wrapper. It requires careful sequencing: reserve capacity first, retrieve the destination address after any reallocation, check the destination range and integer arithmetic, copy before advancing the source iterator, and make the value visible only with consistent offsets and validity. A failed write must not publish a partially constructed batch.

B. Wrap the source address in a non-owning direct buffer and keep setSafe(). Arrow provides MemoryUtil.directBuffer(address, capacity), and its byte-buffer setter has a direct-memory copy path. This keeps more of Arrow's existing setter bookkeeping, at the cost of a small wrapper allocation. Validate wrapper ownership and supported-JDK behavior: this API uses reflective direct-buffer construction and can be unavailable. A fast path must preserve the current supported configurations and a safe fallback where necessary.

These are implementation options, not a prescribed patch. Prefer the smallest maintainable approach that demonstrably removes the payload-sized allocation and preserves the existing vector contract. Check the required APIs across Comet's supported Spark/JDK profiles before choosing one.

Ownership and scope

• One copy into Arrow remains intentional. Spark's row backing may be reused or released after conversion; the completed Arrow batch must own its bytes independently. This issue does not propose retaining a pointer into a Spark row or transferring ownership of Spark-managed memory.
• Keep null handling, byte content, row order, and existing UTF-8 ingress behavior unchanged. This is not a string decoding, normalization, or validation change.
• Start with ordinary StringWriter / VarCharVector. LargeStringWriter has the same staging pattern, but support for its 64-bit offsets should be a follow-up unless the same helper can cover it with explicit tests and no extra design work.
• Keep native Parquet decoding, columnar-to-row conversion, binary columns, Python transport, shuffle formats, and batch-size policy outside this issue. The change applies to whichever existing callers use this writer; it does not add new eligible input operators or data types.

Acceptance criteria

Correctness and lifecycle

• Compare output bytes with the existing writer for off-heap strings, including nonzero source offsets, empty strings, nulls, embedded zero bytes, ASCII, and multibyte UTF-8.
• Cover mixed null/non-null values and values large enough to grow the data buffer; independently force validity/offset capacity growth. Verify offsets, validity, value count, and previously written values after growth.
• Reuse or overwrite the source row backing immediately after each write and verify the resulting Arrow values remain unchanged. Include a multi-batch reader case where an earlier exported batch remains readable while a later batch is produced.
• Retain on-heap whole-array and sliced-array controls; preserve their current semantics and performance. Exercise any fallback path explicitly.
• Verify failed conversion/allocator exhaustion does not emit an incomplete batch and normal reader/vector cleanup releases owned buffers. Run the relevant supported Spark/JDK build and conversion checks.

Performance evidence

• Add a focused allocation benchmark using explicitly off-heap-backed UTF8String inputs; do not infer source placement from a SQL query alone. Prepare input outside the measured region and consume/validate the output.
• Cover short values and larger values such as 1 KiB, 64 KiB, and 1 MiB; include a varied-size batch, on-heap controls, and both preallocated and forced-growth destinations.
• Report allocated JVM bytes per row or per batch, allocation profiles, and conversion throughput against the same baseline. Demonstrate that the off-heap path no longer allocates a temporary byte[] proportional to the string length, and report any tradeoff for small strings or heap-backed input.
• Separate the direct-copy saving from unavoidable Arrow growth copies and output allocations. An end-to-end query benchmark is useful additional evidence, but the initial success criterion is removal of this specific allocation with correct output and no material regression in the controls.

Additional context

The inspected StringWriter already has a TODO asking how to pass off-heap UTF8String data to Arrow without the extra copy. The scope here is to resolve that staging allocation with explicit ownership and benchmark evidence.

RowArrowReader documents why fresh Arrow buffers are needed across exported batches. That lifetime guarantee must remain intact.

Related but distinct work includes #3518: conversion from Spark ColumnarBatch, which concerns a columnar input boundary, and #5310: invalid UTF-8 handling at JVM/native ingress, which concerns string semantics. This proposal is limited to eliminating an intermediate allocation in the existing row-to-Arrow string writer.

The copy/allocation mechanism and available APIs above were checked in source. No performance improvement has yet been measured for a proposed implementation.