diff --git a/parquet/benches/arrow_writer.rs b/parquet/benches/arrow_writer.rs index 55bd8e980217..54fdd27397e9 100644 --- a/parquet/benches/arrow_writer.rs +++ b/parquet/benches/arrow_writer.rs @@ -31,8 +31,8 @@ use std::sync::Arc; use arrow::datatypes::*; use arrow::util::bench_util::{create_f16_array, create_f32_array, create_f64_array}; use arrow::{record_batch::RecordBatch, util::data_gen::*}; -use arrow_array::builder::{ListBuilder, StringBuilder}; -use arrow_array::{RecordBatchOptions, StringArray}; +use arrow_array::builder::{FixedSizeBinaryBuilder, ListBuilder, StringBuilder}; +use arrow_array::{Int32Array, LargeBinaryArray, RecordBatchOptions, StringArray, StringViewArray}; use parquet::errors::Result; use parquet::file::properties::{CdcOptions, WriterProperties, WriterVersion}; @@ -1005,10 +1005,224 @@ fn bench_delta_byte_array_writers(c: &mut Criterion) { } } +const NDV_BATCH_SIZE: usize = 65_536; +const NDV_NUM_BATCHES: usize = 10; +const NDV_LOW_CARDINALITY: usize = 100; + +// Generates NDV_NUM_BATCHES random batches for high-cardinality cases. +fn ndv_random_batches(schema: Arc) -> Vec { + (0..NDV_NUM_BATCHES) + .map(|_| create_random_batch(schema.clone(), NDV_BATCH_SIZE, 0.0, 0.5).unwrap()) + .collect() +} + +// Low-cardinality helpers: cycle through NDV_LOW_CARDINALITY distinct values. +fn make_int32_batches_low() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + (0..NDV_NUM_BATCHES) + .map(|batch_idx| { + let values: Vec = (0..NDV_BATCH_SIZE) + .map(|row| ((batch_idx * NDV_BATCH_SIZE + row) % NDV_LOW_CARDINALITY) as i32) + .collect(); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(values)) as _], + ) + .unwrap() + }) + .collect() +} + +fn make_utf8_batches_low() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "label", + DataType::Utf8, + false, + )])); + (0..NDV_NUM_BATCHES) + .map(|batch_idx| { + let values: Vec = (0..NDV_BATCH_SIZE) + .map(|row| { + format!( + "label_{}", + (batch_idx * NDV_BATCH_SIZE + row) % NDV_LOW_CARDINALITY + ) + }) + .collect(); + let array = StringArray::from(values.iter().map(|s| s.as_str()).collect::>()); + RecordBatch::try_new(schema.clone(), vec![Arc::new(array) as _]).unwrap() + }) + .collect() +} + +fn make_large_binary_batches_low() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::LargeBinary, + false, + )])); + (0..NDV_NUM_BATCHES) + .map(|batch_idx| { + let payloads: Vec> = (0..NDV_BATCH_SIZE) + .map(|row| { + format!( + "payload_{}", + (batch_idx * NDV_BATCH_SIZE + row) % NDV_LOW_CARDINALITY + ) + .into_bytes() + }) + .collect(); + let refs: Vec<&[u8]> = payloads.iter().map(|p| p.as_slice()).collect(); + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(LargeBinaryArray::from(refs)) as _], + ) + .unwrap() + }) + .collect() +} + +// Utf8View is not yet tracked by NDV (skipped in update_distinct_values_seen), so these +// benchmarks measure base write cost only — useful as a no-overhead baseline. +fn make_utf8_view_batches_low() -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "label", + DataType::Utf8View, + false, + )])); + (0..NDV_NUM_BATCHES) + .map(|batch_idx| { + let values: Vec = (0..NDV_BATCH_SIZE) + .map(|row| { + format!( + "label_{}", + (batch_idx * NDV_BATCH_SIZE + row) % NDV_LOW_CARDINALITY + ) + }) + .collect(); + let array = StringViewArray::from_iter_values(values.iter().map(|s| s.as_str())); + RecordBatch::try_new(schema.clone(), vec![Arc::new(array) as _]).unwrap() + }) + .collect() +} + +fn make_fsb_batches_low(byte_width: i32) -> Vec { + let schema = Arc::new(Schema::new(vec![Field::new( + "id", + DataType::FixedSizeBinary(byte_width), + false, + )])); + let width = byte_width as usize; + (0..NDV_NUM_BATCHES) + .map(|batch_idx| { + let mut builder = FixedSizeBinaryBuilder::new(byte_width); + for row in 0..NDV_BATCH_SIZE { + let distinct_id = (batch_idx * NDV_BATCH_SIZE + row) % NDV_LOW_CARDINALITY; + let mut buf = vec![0u8; width]; + let encoded = distinct_id.to_le_bytes(); + buf[..encoded.len().min(width)] + .copy_from_slice(&encoded[..encoded.len().min(width)]); + builder.append_value(&buf).unwrap(); + } + RecordBatch::try_new(schema.clone(), vec![Arc::new(builder.finish()) as _]).unwrap() + }) + .collect() +} + +fn write_ndv_batches(b: &mut criterion::Bencher, batches: &[RecordBatch]) { + let props = WriterProperties::builder() + .set_write_row_group_number_distinct_values(true) + .build(); + let schema = batches[0].schema(); + b.iter(|| { + let mut file = Empty::default(); + let mut writer = + ArrowWriter::try_new(&mut file, schema.clone(), Some(props.clone())).unwrap(); + for batch in batches { + writer.write(black_box(batch)).unwrap(); + } + black_box(writer.close()).unwrap(); + }); +} + +fn bench_ndv(c: &mut Criterion) { + let int32_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let bool_schema = Arc::new(Schema::new(vec![Field::new( + "flag", + DataType::Boolean, + false, + )])); + let utf8_schema = Arc::new(Schema::new(vec![Field::new( + "label", + DataType::Utf8, + false, + )])); + let large_binary_schema = Arc::new(Schema::new(vec![Field::new( + "payload", + DataType::LargeBinary, + false, + )])); + let fsb16_schema = Arc::new(Schema::new(vec![Field::new( + "id", + DataType::FixedSizeBinary(16), + false, + )])); + let utf8_view_schema = Arc::new(Schema::new(vec![Field::new( + "label", + DataType::Utf8View, + false, + )])); + + let cases: Vec<(&str, Vec)> = vec![ + ("int32_low_cardinality", make_int32_batches_low()), + ("int32_high_cardinality", ndv_random_batches(int32_schema)), + ("bool", ndv_random_batches(bool_schema)), + ("utf8_low_cardinality", make_utf8_batches_low()), + ("utf8_high_cardinality", ndv_random_batches(utf8_schema)), + ( + "large_binary_low_cardinality", + make_large_binary_batches_low(), + ), + ( + "large_binary_high_cardinality", + ndv_random_batches(large_binary_schema), + ), + ("fsb16_low_cardinality", make_fsb_batches_low(16)), + ("fsb16_high_cardinality", ndv_random_batches(fsb16_schema)), + ("utf8_view_low_cardinality", make_utf8_view_batches_low()), + ( + "utf8_view_high_cardinality", + ndv_random_batches(utf8_view_schema), + ), + ]; + + for (name, batches) in cases { + let total_bytes: u64 = batches + .iter() + .flat_map(|b| b.columns()) + .map(|c| c.get_array_memory_size() as u64) + .sum(); + + let mut group = c.benchmark_group(format!("ndv/{name}")); + group.throughput(Throughput::Bytes(total_bytes)); + group.bench_function("ndv_on", |b| write_ndv_batches(b, &batches)); + group.finish(); + } +} + criterion_group!( benches, bench_all_writers, bench_small_delta_byte_array_writers, - bench_delta_byte_array_writers + bench_delta_byte_array_writers, + bench_ndv, ); criterion_main!(benches);