xudong963 opened a new issue, #12027:
URL: https://github.com/apache/arrow-rs/issues/12027
**Describe the bug**
Since #10364 (parquet 59.0.0) `ArrowWriter` converts Decimal32/64/128/256,
Float16, FixedSizeBinary and Interval columns to FIXED_LEN_BYTE_ARRAY by
filling one contiguous `Vec<u8>` per batch and column, turning it into `Bytes`,
and handing out `split_to` slices as the `FixedLenByteArray` values
(`chunk_array_slice` / `chunk_contiguous_vec` in
`parquet/src/arrow/arrow_writer/mod.rs`).
The generic column writer tracks min/max by cloning the value (`update_min`
/ `update_max` / `update_stat` in `parquet/src/column/writer/mod.rs`, `*cur =
val.clone()`). For `ByteArray` / `FixedLenByteArray` a clone is a refcount bump
on the shared `Bytes`, so the running min/max are views into the batch arena.
That is fine while the column chunk is open: at most two arenas per column are
pinned at any time.
The problem is what happens when the column chunk is closed.
`build_column_metadata` puts `min_column_value.clone()` /
`max_column_value.clone()` into the `ColumnChunkMetaData` statistics and runs
them through `truncate_statistics`. For BYTE_ARRAY and for truncatable FLBA
logical types this rebuilds the values via `truncate_min_value` /
`truncate_max_value`, which return `data.to_vec()` even when nothing is
truncated, so those statistics end up owning their bytes. Decimal and Float16
are exempt from truncation (`can_truncate_value()` is false for them, from
#5075 / #5154), so they take the `stats => stats` arm and are returned
unchanged: the chunk statistics still reference the arena.
Those `ColumnChunkMetaData` live in `SerializedFileWriter::row_groups` until
the file is closed, and `write_metadata` moves them into the `ParquetMetaData`
returned from `close()`. So every closed row group of a Decimal/Float16 column
permanently pins one or two batch-sized arenas (min and max can come from
different batches), and a 16-byte statistic keeps `batch_rows * byte_width`
bytes alive: 8192 x 16 B = 128 KiB per statistic with a typical batch size.
`ParquetMetaData::memory_size()` reports the 16 bytes, not the 128 KiB.
Anything that retains `ParquetMetaData` for the duration of a write session
turns this into unbounded growth. DataFusion's `ParquetSink` keeps the metadata
of every written file until the sink finishes. In our case (two `Decimal128(38,
8)` columns, ~190 row groups per file, 8192-row batches) that is ~95 MiB of
pinned arenas per 0.85 GiB output file, about 1 GiB of RSS per 10 GB written,
and a 64 GiB container is OOM-killed after ~500 GiB of output. Heap profiles
attribute the growth to `chunk_array_slice` under
`get_decimal_array_slice::<Decimal128Type>`. parquet 58.x is not affected
because each value was its own allocation, so a cloned statistic pinned only
its 16 bytes.
**To Reproduce**
parquet 60.1.0, counting global allocator. 64 batches of 8192 ascending
`Decimal128(38, 8)` values are written through `ArrowWriter` into a `Vec<u8>`;
`close()` consumes the writer and the sink, only the returned `ParquetMetaData`
is kept.
```text
64 batches x 8192 rows; Decimal128 arena per batch = 131072 B
Decimal128(38,8), stats=Chunk row_group_size= 1024 row_groups= 512
metadata.memory_size()= 340936 B live heap after close= 8725375 B
live heap after drop(metadata)= 63 B
Decimal128(38,8), stats=None row_group_size= 1024 row_groups= 512
metadata.memory_size()= 324552 B live heap after close= 324391 B
Int64, stats=Chunk row_group_size= 1024 row_groups= 512
metadata.memory_size()= 324552 B live heap after close= 324391 B
Decimal128(38,8), stats=Chunk row_group_size=32768 row_groups= 16
metadata.memory_size()= 11976 B live heap after close= 4206375 B
live heap after drop(metadata)= -25 B
Decimal128(38,8), stats=None row_group_size=32768 row_groups= 16
metadata.memory_size()= 11464 B live heap after close= 11303 B
Int64, stats=Chunk row_group_size=32768 row_groups= 16
metadata.memory_size()= 11464 B live heap after close= 11303 B
```
With 1024-row row groups all 64 arenas (8 MiB) survive `close()`. With
32768-row row groups (4 batches each) every one of the 16 row groups pins the
arena of its first batch (min) and its last batch (max): 16 x 2 x 128 KiB = 4
MiB. Dropping the metadata releases everything. The same column with statistics
disabled, or an Int64 column, retains only the metadata itself.
<details>
<summary>reproducer</summary>
`Cargo.toml`
```toml
[package]
name = "flba-stats-repro"
version = "0.1.0"
edition = "2021"
[dependencies]
arrow-array = "60.1.0"
arrow-schema = "60.1.0"
parquet = "60.1.0"
```
`src/main.rs`
```rust
use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicIsize, Ordering::Relaxed};
use std::sync::Arc;
use arrow_array::{ArrayRef, Decimal128Array, Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use parquet::arrow::ArrowWriter;
use parquet::file::properties::{EnabledStatistics, WriterProperties};
struct Counting;
static LIVE: AtomicIsize = AtomicIsize::new(0);
unsafe impl GlobalAlloc for Counting {
unsafe fn alloc(&self, l: Layout) -> *mut u8 {
LIVE.fetch_add(l.size() as isize, Relaxed);
unsafe { System.alloc(l) }
}
unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
LIVE.fetch_sub(l.size() as isize, Relaxed);
unsafe { System.dealloc(p, l) }
}
}
#[global_allocator]
static GLOBAL: Counting = Counting;
const ROWS_PER_BATCH: usize = 8192;
const BATCHES: usize = 64;
fn run(name: &str, data_type: DataType, stats: EnabledStatistics,
row_group_size: usize) {
let schema = Arc::new(Schema::new(vec![Field::new("v",
data_type.clone(), false)]));
let props = WriterProperties::builder()
.set_max_row_group_row_count(Some(row_group_size))
.set_statistics_enabled(stats)
.build();
let baseline = LIVE.load(Relaxed);
let mut writer = ArrowWriter::try_new(Vec::<u8>::new(), schema.clone(),
Some(props)).unwrap();
for b in 0..BATCHES {
let values = (0..ROWS_PER_BATCH).map(|i| (b * ROWS_PER_BATCH + i) as
i64);
let array: ArrayRef = match &data_type {
DataType::Decimal128(p, s) => Arc::new(
Decimal128Array::from_iter_values(values.map(|v| v as i128))
.with_precision_and_scale(*p, *s)
.unwrap(),
),
DataType::Int64 =>
Arc::new(Int64Array::from_iter_values(values)),
_ => unreachable!(),
};
let batch = RecordBatch::try_new(schema.clone(),
vec![array]).unwrap();
writer.write(&batch).unwrap();
}
// close() consumes the writer and drops the Vec<u8> sink; only
`metadata` survives.
let metadata = writer.close().unwrap();
let live = LIVE.load(Relaxed) - baseline;
println!(
"{name:<32} row_group_size={row_group_size:>5} row_groups={:>4}
metadata.memory_size()={:>8} B live heap after close={:>9} B",
metadata.num_row_groups(),
metadata.memory_size(),
live
);
drop(metadata);
println!("live heap after drop(metadata)={} B", LIVE.load(Relaxed) -
baseline);
}
fn main() {
println!("{BATCHES} batches x {ROWS_PER_BATCH} rows; Decimal128 arena
per batch = {} B\n", ROWS_PER_BATCH * 16);
for rg in [1024, 32768] {
run("Decimal128(38,8), stats=Chunk", DataType::Decimal128(38, 8),
EnabledStatistics::Chunk, rg);
run("Decimal128(38,8), stats=None", DataType::Decimal128(38, 8),
EnabledStatistics::None, rg);
run("Int64, stats=Chunk", DataType::Int64, EnabledStatistics::Chunk,
rg);
}
}
```
</details>
**Expected behavior**
Statistics stored in `ColumnChunkMetaData` / `ParquetMetaData` own their
bytes and do not alias the writer's input buffers, so the memory retained by
the returned metadata matches `memory_size()` and does not depend on the batch
size the file was written with.
**Additional context**
- Affected: parquet 59.0.0 through 60.1.0 and current main
(`chunk_contiguous_vec` in `arrow_writer/mod.rs`; `stats => stats` in
`truncate_statistics` and `min_column_value.clone()` in `build_column_metadata`
in `column/writer/mod.rs`). Only FLBA columns with a Decimal or Float16 logical
type are affected through `ArrowWriter`: BYTE_ARRAY columns go through the
Arrow byte-array encoder, whose `compute_min_max` copies with `to_vec()`, and
the other FLBA logical types are copied by the truncation path. The low-level
`ColumnWriter` API has the same behaviour for any caller that passes
`FixedLenByteArray` values sliced from a shared `Bytes`.
- Possible fix: in `truncate_statistics`, rebuild Decimal/Float16 FLBA
min/max with `to_vec()` without truncating (two small allocations per column
chunk). Alternatives are copying at `update_stat` time, or not sharing the
arena in `chunk_contiguous_vec`, which would give up the gains from #10364.
--
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]
