apache / apache/arrow-rs

Expose next_column_with_factory and add write_batch_with_indices for streaming column writes

Open
#9,957 0 comments 5 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
3.6k
Forks
1.3k
Avg merge
2d 18h
Merged PRs (30d)
169

Description

**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**

I want to write **a large record batch into a parquet file consisting of a single (arbitrarily large) row group**. A typical use case is log data — a single parquet file can reach several GB to tens of GB.

The ideal memory profile is:

- columns processed strictly sequentially, only one leaf column being encoded at any time
- for each column, pages are flushed to the file sink as soon as they are full
- peak memory is on the order of one encoder's in-progress page, **independent of total row count and column count**

This is the natural shape of the problem: a single row group means no cross-row-group buffering; sequential columns means no cross-column buffering; page-level streaming means no cross-page buffering. For files in the several-GB to tens-of-GB range, the absence of this streaming capability means writer peak memory scales with the total row group size instead of converging on a single page.

Neither of the existing public paths supports this:

**`ArrowWriter`** — `ArrowWriter::write(batch)` pushes data row-wise into every column's encoder. `ArrowColumnWriter` internally uses `ArrowPageWriter`, which buffers every compressed page into a `SharedColumnChunk` until `close()` is called, then appends the whole column chunk to the file sink via `ArrowColumnChunk::append_to_row_group`. Peak memory = Σ(compressed bytes of every leaf column's chunk) + encoder state for every column. For a 10 GB target file with a wide schema, even with 5× compression the combined compressed column chunks can reach the 2 GB range. There is no knob to make `ArrowPageWriter` flush pages as they are produced.

**`SerializedFileWriter` + `next_column()` + `SerializedColumnWriter::typed::()`** — this is the natural low-level path. `SerializedPageWriter` (the default page writer under `next_column()`) already flushes each compressed page directly to the file sink, which is exactly the behavior I want. But two problems block this path:

***Problem 1: `SerializedColumnWriter` hardcodes the generic encoder.***

```rust
pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl>;

pub enum ColumnWriter<'a> {
ByteArrayColumnWriter(ColumnWriterImpl<'a, ByteArrayType>), ...
}
```

Going through `next_column()` locks in `ColumnValueEncoderImpl`. To use `ByteArrayEncoder` (the specialized zero-copy byte array encoder that `ArrowColumnWriter` uses internally), I need to construct `GenericColumnWriter` myself via a factory — but `SerializedRowGroupWriter::next_column_with_factory` is `pub(crate)`:

```rust
pub(crate) fn next_column_with_factory<'b, F, C>(
&'b mut self,
factory: F,
) -> Result>
where
F: FnOnce(
ColumnDescPtr,
WriterPropertiesPtr,
Box,
OnCloseColumnChunk<'b>,
) -> Result,
```

The signature is already fully generic in the return type `C`; the only barrier is `pub(crate)`.

***Problem 2: `write_batch` requires a dense value array, discarding null-handling information I already have.***

Arrow arrays naturally produce `(dense values, non-null indices)` pairs when computing definition/repetition levels (see `ArrayLevels::non_null_indices()`). The public API is:

```rust
pub fn write_batch(
&mut self,
values: &E::Values,
def_levels: Option<&[i16]>,
rep_levels: Option<&[i16]>,
) -> Result
```

There is no way for external callers to say "take values at these indices". The internal `write_batch_internal` has exactly this entry point:

```rust
pub(crate) fn write_batch_internal(
&mut self,
values: &E::Values,
value_indices: Option<&[usize]>, // <-- this
def_levels: LevelDataRef<'_>,
rep_levels: LevelDataRef<'_>,
min: Option<&E::T>,
max: Option<&E::T>,
distinct_count: Option,
) -> Result
```

This is exactly the path `ArrowColumnWriter` uses internally via `write_primitive` to avoid materializing a gathered copy. External callers can only gather into a scratch `Vec` and call `write_batch`, or go through `ArrowColumnWriter` and accept the column-chunk buffering.

**Describe the solution you'd like**

Two visibility changes + one thin wrapper method. No implementation changes, no behavior changes.

***1. `SerializedRowGroupWriter::next_column_with_factory` → `pub`***

```diff
impl<'a, W: Write> SerializedRowGroupWriter<'a, W> {
- pub(crate) fn next_column_with_factory<'b, F, C>(&'b mut self, factory: F) -> Result>
+ pub fn next_column_with_factory<'b, F, C>(&'b mut self, factory: F) -> Result>
}
```

This lets callers plug in their own `GenericColumnWriter` (including `GenericColumnWriter`) with the file's `SerializedPageWriter` already correctly wired up by arrow-rs internally.

***2. New `pub fn write_batch_with_indices` on `GenericColumnWriter`***

```rust
impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
pub fn write_batch_with_indices(
&mut self,
values: &E::Values,
indices: &[usize],
def_levels: Option<&[i16]>,
rep_levels: Option<&[i16]>,
) -> Result {
self.write_batch_internal(
values, Some(indices),
LevelDataRef::from(def_levels),
LevelDataRef::from(rep_levels),
None, None, None,
)
}

pub fn write_batch_with_indices_and_statistics(
&mut self,
values: &E::Values,
indices: &[usize],
def_levels: Option<&[i16]>,
rep_levels: Option<&[i16]>,
min: Option<&E::T>,
max: Option<&E::T>,
distinct_count: Option,
) -> Result { ... }
}
```

Exposes the existing `value_indices` code path. `write_batch` already delegates to `write_batch_internal` with `value_indices: None`; the new API just exposes the `Some(indices)` variant as a public wrapper.

**Describe alternatives you've considered**

*Use `ArrowColumnWriter` as-is.* This is the right default for most callers, but not when the file is designed to contain a single large row group (as in log workloads reaching several GB to tens of GB per file), columns are processed strictly sequentially, and peak memory is dominated by per-column-chunk buffering rather than encoding cost. In this shape, the natural memory ceiling is one encoder's in-progress page, not one column chunk.

*Pre-gather into a dense `Vec` before `write_batch`.* Works for fixed-width types, but wastes `K × sizeof(T)` of memory that the `value_indices` path would save. For `ByteArrayType` it wastes `Vec` (~32 bytes per row).

*Reimplement parquet encoding downstream.* The Arrow → Parquet dispatch (`write_leaf`), level computation, and byte array encoders are non-trivial and evolve with the format spec. Not a viable long-term maintenance path.

**Additional context**

The `SerializedFileWriter` + `SerializedPageWriter` infrastructure already supports page-streamed column writes end-to-end — what's missing is just the plumbing to connect it with Arrow input and the existing specialized encoders. Happy to submit a PR if this direction is acceptable.

Contributor guide

Open the contributing guide

Research direction

Start by reading SerializedRowGroupWriter::next_column_with_factory and GenericColumnWriter::write_batch_internal, including the existing write_batch wrapper and value_indices path. Trace how SerializedPageWriter is supplied and how ArrowColumnWriter uses write_primitive. Done means the factory entry point and indexed batch-writing API are publicly usable without changing encoding behavior or requiring gathered values.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data-engineering
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
58/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.