VaibhaveS opened a new issue, #12187:
URL: https://github.com/apache/arrow-rs/issues/12187
### Is your feature request related to a problem or challenge?
We write Parquet from Arrow `RecordBatch`es that come with a per-row
selection mask. The batches come from a change stream. Within a snapshot, a
later change can replace or delete a row that an earlier batch holds, so only
the rows whose bit is set may be written. The batches are immutable and can be
large.
Today the public API needs a filtered copy first:
```rust
let selected = arrow_select::filter::filter_record_batch(&batch, &mask)?;
if selected.num_rows() > 0 {
writer.write(&selected)?;
}
```
`filter_record_batch` copies every selected value into a new batch. The
writer then reads that copy to encode it. With large batches, each writer also
holds one extra batch in memory while it writes.
Writing one zero-copy `RecordBatch::slice` per run of consecutive selected
rows avoids the copy. But every `write` call has a fixed cost per column, so it
gets slow when unselected rows are scattered.
The writer already has the pieces. `LevelInfoBuilder::write` builds levels
for any row range, the leaf levels keep `non_null_indices`, and the column
writers gather values by index. A selection could be applied while building
levels, straight from the original arrays.
### Describe the solution you'd like
A way to tell `ArrowWriter` (and `AsyncArrowWriter`) to write only the
selected rows of a batch, without materializing a filtered `RecordBatch`. For
example:
```rust
impl<W: Write + Send> ArrowWriter<W> {
/// Writes only the rows of `batch` that `selection` selects.
pub fn write_with_selection(
&mut self,
batch: &RecordBatch,
selection: &RowSelection,
) -> Result<()>;
}
```
The name and the selection type are up to you. A `BooleanArray` or
`BooleanBuffer` mask, the reader's `RowSelection`
(`RowSelection::from_boolean_buffer` keeps a bitmap without copying it), or an
iterator of row ranges would all work for us. What matters to us:
- The file holds the same rows, statistics, page indexes and row counts as
writing `filter_record_batch(&batch, &mask)`.
- No filtered copy of the batch is built.
- Selected rows count toward the row group limits, so one call can span row
groups, as `write` does today.
A selection-aware `compute_leaves` would also help callers that encode
columns in parallel with `ArrowColumnWriter`.
### Describe alternatives you've considered
- `filter_record_batch`, then `write`. This works today, at the cost of the
copy and its memory.
- One `RecordBatch::slice` per run of selected rows. There is no copy, but
in our tests it used 2.3 to 6.8 times the CPU of filter + write when 10% to 99%
of rows were unselected at random.
- The low-level `SerializedFileWriter` column writers. They lose the Arrow
to Parquet type mapping and the Arrow buffer fast paths that `ArrowWriter` has.
### Additional context
We prototyped this on parquet 59.3.2 to see if it was worth asking for. Four
writer threads encoded 41.9 million rows of a 9-column table (Int64, Int32,
Decimal128(12, 2), Float64, Utf8 and Timestamp columns), with data page v2 and
no compression. Rows were unselected at random. Compared with
`filter_record_batch` + `write`, writing the selection directly used:
| Unselected rows | 8,192-row batches | 1,048,576-row batches |
| --- | --- | --- |
| 0% | the same CPU | the same CPU |
| 10% | 6% less CPU | 14% less CPU, 0.43 GB less peak RSS |
| 50% | 7% less CPU | 13% less CPU, 0.36 GB less peak RSS |
| 99% | 21% more CPU | 34% more CPU |
The prototype lost at 99% for two reasons that an implementation can avoid:
- `WritePrimitive::map` converts every value of the array before the column
writer gathers the selected ones, for example `Decimal128` to `i64`. With a
selection, only the selected values need converting. Storing our decimal column
as `Int64` instead removed about two thirds of the gap.
- The prototype scanned the mask several times per column to build levels.
Finding the selected rows once per batch and sharing them across columns avoids
that.
Related: #9957 asks for index-based writes at the column-writer level. This
request is for the `ArrowWriter` API.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
