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

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions