Skip to content

refactor: centralize Spark-to-Arrow batch encoding in ArrowWriter #5317

Description

@peterxcli

What problem does this solve?

Comet currently has several Spark-to-Arrow conversion paths:

  • RowArrowReader
  • SparkColumnarArrowReader
  • CometArrowConverters
  • ArrowWriter / ArrowFieldWriter

These paths have different ownership and lifecycle requirements, so the readers and converters should remain separate. However, they share the same lower-level responsibility: writing Spark rows or ColumnVector slices into Arrow vectors.

PR #5051 introduced CometArrowConverters.writeColumns to share the column-copy loop between SparkColumnarArrowReader and the cache conversion path. This removes duplication, but it also places low-level Arrow encoding inside a higher-level converter that is otherwise responsible for allocating independently owned batches.

The current structure also converts each source column into a ColumnarArray before dispatching to ArrowWriter. This hides the original ColumnVector, slice offset, and length from the field writer, making it difficult to implement specialized bulk-copy paths such as #5299.

Row count is also currently derived indirectly from column writes. This previously caused zero-column batches with numRows > 0 to be emitted with zero rows, as fixed in #4795.

Proposed change

Centralize Spark-to-Arrow value encoding in ArrowWriter and ArrowFieldWriter, while keeping batching, ownership, and transport logic in the existing readers and converters.

A possible API shape is:

class ArrowWriter {
  def write(row: InternalRow): Unit

  def writeColumns(
      batch: ColumnarBatch,
      startRow: Int,
      numRows: Int): Unit

  def finish(rowCount: Int): Unit
}

abstract class ArrowFieldWriter {
  def writeColumnSlice(
      source: ColumnVector,
      startRow: Int,
      numRows: Int): Unit
}

writeColumnSlice should:

  1. Retain access to the original ColumnVector, startRow, and numRows.
  2. Dispatch to an optimized implementation when available.
  3. Fall back to the existing ColumnarArray plus writeCol / writeColNoNull behavior.
  4. Keep capacity checks and non-resizing writes coupled inside the writer implementation.

finish(rowCount) should set the logical row count explicitly instead of deriving it from the last column written. This naturally handles zero-column batches.

After this change:

  • SparkColumnarArrowReader remains responsible for selecting input batches and slices.
  • RowArrowReader remains responsible for consuming and batching InternalRows.
  • CometArrowConverters remains responsible for allocating fresh roots, transferring ownership, and cleaning up on failure.
  • ColumnarBatchArrowReader remains the separate Arrow-backed retain/transfer path.
  • ArrowWriter becomes the single implementation point for Spark-value-to-Arrow-buffer encoding.

Why is this useful?

This should:

  • remove the remaining duplicated column-writing behavior;
  • make the distinction between encoding and ownership clearer;
  • centralize row-count and capacity invariants;
  • make the conversion code easier to review and maintain;
  • provide the correct extension point for perf: explore bulk copies for no-null fixed-width Arrow writes #5299, where eligible no-null fixed-width Spark vectors may be copied into Arrow buffers in bulk.

Non-goals

This issue should not:

  • merge the existing ArrowReader implementations;
  • change Arrow C Stream or FFI ownership behavior;
  • change which component owns or closes input/output batches;
  • unconditionally convert Arrow-backed CometVector batches;
  • implement the bulk-copy optimization from perf: explore bulk copies for no-null fixed-width Arrow writes #5299;
  • introduce reflection or access Spark private fields.

Acceptance criteria

  • SparkColumnarArrowReader and fresh-batch conversion use one shared writer API.
  • The low-level column-copy loop no longer lives in CometArrowConverters.
  • Logical row count is passed explicitly when finishing a batch.
  • Zero-column batches preserve their nonzero row count.
  • The writer receives the original ColumnVector, slice offset, and row count before falling back to ColumnarArray.
  • Stable-reader and independently owned batch lifecycles remain unchanged.
  • Arrow-backed batches continue to use their existing non-element-wise path.
  • Variable-width, nullable, nested, dictionary, and third-party vectors continue to use safe fallback behavior.
  • Tests cover nonzero slice offsets, split batches, zero-column batches, nullable/no-null columns, nested fallback, and cleanup when conversion throws.
  • Existing Spark-to-Arrow benchmarks show no material regression.

Related work

Metadata

Metadata

Assignees

Labels

area:ffiArrow FFI / JNI boundaryenhancementNew feature or request

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions