diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index b6896459fddc..1339a0bfa5f4 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -143,6 +143,11 @@ name = "external_metadata" required-features = ["arrow", "async"] path = "./examples/external_metadata.rs" +[[example]] +name = "chunk_probe_writer" +required-features = ["arrow"] +path = "./examples/chunk_probe_writer.rs" + [[example]] name = "read_parquet" required-features = ["arrow"] diff --git a/parquet/examples/chunk_probe_writer.rs b/parquet/examples/chunk_probe_writer.rs new file mode 100644 index 000000000000..f5c0b39ce2a1 --- /dev/null +++ b/parquet/examples/chunk_probe_writer.rs @@ -0,0 +1,728 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! An adaptive Parquet writer that measures encodings instead of guessing them. +//! +//! The writer's default encoding choices are made ahead of time from the schema +//! and a handful of size limits. They cannot know whether a particular column's +//! actual values compress better as a dictionary, as deltas, or as plain +//! values. This example shows how a caller can find out by measurement, using +//! [`ArrowRowGroupWriterFactory::create_column_writer`], which builds a writer +//! for one leaf column of one row group at properties of the caller's choosing. +//! +//! For each row group, and for each column that is due to decide, the writer +//! encodes a probe of that column's leading rows once per candidate set of +//! writer properties, through throwaway single-column writers. Closing +//! one yields a `ColumnCloseResult` whose metadata carries the compressed size +//! that candidate actually achieved. The smallest wins, the probe chunks are +//! discarded, and the row group is then written for real through one ordinary +//! column writer per column at that column's current choice. +//! +//! The cost model is deliberately simple. A column that is deciding encodes its +//! probe K + 1 times: K throwaway passes plus the real one. The probe is one +//! data page worth of rows rather than a whole row group, so the extra work is +//! bounded by the page size and not by the data. A column that has settled costs +//! nothing extra at all: no probe writer is built for it, so no page store is +//! allocated for it either. That is what addressing one column at a time buys +//! over building every column writer at once. +//! +//! Growing the probe instead, a page at a time until the leader is decisive, is +//! a natural next step and works. It is left out here because it did not pay: +//! measured over TPC-H and ClickBench it bought under a point of compression for +//! 10 to 30 percent more write CPU, since a chunk's compressed size can only be +//! read once the chunk is closed, so each extra page re-encodes the ones before +//! it. A caller whose data rewards a longer look can grow the probe the same +//! way; this example keeps one page. +//! +//! # Deciding, and how often to decide again +//! +//! A race that comes back close decides nothing. A leader that has not beaten +//! the runner-up by [`DECISIVE_MARGIN`] has won by an amount as likely to be an +//! accident of this row group as a property of the column, so it is used for +//! this row group but held provisionally, and the column races again on the very +//! next one. A column whose leader is decisive doubles the interval before its +//! next race, up to [`MAX_RERACE_INTERVAL`] row groups, so a column with an +//! obvious answer stops paying to be asked it, while a column whose data keeps +//! drifting is asked often. +//! +//! Closer still, inside [`NEAR_TIE_BAND`], the sizes are not telling the +//! candidates apart at all. There the tie is broken toward the page a reader +//! decodes most cheaply rather than toward a nominally smaller chunk, and it is +//! broken deterministically, never by a coin flip. +//! +//! # Dictionary candidates +//! +//! A dictionary is not a special case here, because a probe measures it rather +//! than estimating it. Closing a probe writer produces a complete column chunk, +//! and the `compressed_size` on its `ColumnCloseResult` is the size of the whole +//! chunk: the dictionary page, then every data page that indexes into it. A +//! dictionary candidate is therefore charged the full cost of the dictionary it +//! built, and cheap looking `RLE_DICTIONARY` data pages cannot flatter it. So +//! "is a dictionary worthwhile for this column" needs no ratio or cardinality +//! heuristic; it is decided the same way as every other candidate, by the +//! smallest measured chunk winning. +//! +//! What a probe cannot see is the rest of the row group. A prefix's distinct +//! value count is only a lower bound on the row group's, so a prefix can flatter +//! a dictionary that the full data would not sustain. Two things keep that from +//! being a problem. The first is the writer's own fallback: a column chunk may +//! carry at most one dictionary page, so when the dictionary being built passes +//! `dictionary_page_size_limit`, the writer emits that page for the values +//! indexed so far, flushes the data pages that reference it, and encodes the +//! remainder of the chunk with the fallback encoding. The chunk stays valid and +//! the choice degrades instead of failing. The second is the re-race above, +//! which is what moves a column off a dictionary in later row groups once the +//! data has drifted away from what the probe suggested. +//! +//! ```text +//! cargo run --example chunk_probe_writer --features arrow +//! ``` + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::fs::File; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Instant; + +use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; + +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; +use parquet::arrow::arrow_writer::{ + ArrowColumnWriter, ArrowLeafColumn, ArrowWriter, compute_leaves, +}; +use parquet::basic::{Compression, Encoding, Type as PhysicalType}; +use parquet::errors::Result; +use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder, WriterPropertiesPtr}; +use parquet::schema::types::{ColumnDescPtr, ColumnPath}; + +/// Total rows in the generated dataset. +const ROWS: usize = 400_000; +/// Rows per row group. Eight row groups, so re-racing has something to do. +const ROW_GROUP_ROWS: usize = 50_000; +/// Rows per record batch handed to the writers. +const BATCH_ROWS: usize = 10_000; +/// Rows per data page, shared by both writers so their output is comparable. +/// This is also the length of a probe, so what a probe measures is one whole +/// data page: enough for a dictionary candidate to have built a real dictionary +/// and for the compressor to have something to chew on. +const DATA_PAGE_ROWS: usize = 10_000; + +/// How much smaller than the runner-up the leading candidate's chunk must be +/// for the race to count as decided, and so for the column to be allowed to +/// skip row groups before racing again. +const DECISIVE_MARGIN: f64 = 0.10; + +/// How close two candidates must be before the difference between them stops +/// counting as a size difference at all and the tie is broken on decoding cost. +/// +/// Deliberately much narrower than [`DECISIVE_MARGIN`], and measured rather than +/// guessed. A band as wide as the decisive margin gives away real wins: on +/// TPC-H, delta beats plain or dictionary by 5 to 10 percent on many columns, so +/// a 10 percent band hands those columns to the cheaper-to-decode candidate and +/// costs several percent of the whole file. At 2 percent the tie-break is free +/// on that data, and still catches the case it is for, where the sizes are close +/// enough that picking the smaller one is reading noise. +const NEAR_TIE_BAND: f64 = 0.02; + +/// The most row groups a column may go without racing again. +/// +/// A column that keeps winning decisively doubles its way up to this cadence. +/// Racing again is also what corrects a dictionary that a short probe made look +/// better than the whole column can sustain, so the cadence is capped rather +/// than allowed to grow without limit. +const MAX_RERACE_INTERVAL: usize = 8; + +// --------------------------------------------------------------------------- +// Measurement: what each arm of the comparison costs +// --------------------------------------------------------------------------- + +/// A pass-through allocator that counts live bytes and keeps a high water mark, +/// so each arm below can report the heap it actually needed. +/// +/// Counting in the allocator, rather than reading RSS, keeps the number about +/// what this program asked for instead of what the OS chose to reclaim. +struct CountingAllocator; + +/// Bytes handed out and not yet freed. +static LIVE: AtomicUsize = AtomicUsize::new(0); +/// The largest `LIVE` has reached since the last [`reset_peak`]. +static PEAK: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + let ptr = unsafe { System.alloc(layout) }; + if !ptr.is_null() { + let live = LIVE.fetch_add(layout.size(), Ordering::Relaxed) + layout.size(); + PEAK.fetch_max(live, Ordering::Relaxed); + } + ptr + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + LIVE.fetch_sub(layout.size(), Ordering::Relaxed); + unsafe { System.dealloc(ptr, layout) }; + } +} + +#[global_allocator] +static ALLOCATOR: CountingAllocator = CountingAllocator; + +/// Drops the high water mark to the bytes live right now, and returns that +/// baseline. Without this an arm's peak would be swamped by the input batches, +/// which are much larger than anything the writers themselves allocate. +fn reset_peak() -> usize { + let live = LIVE.load(Ordering::Relaxed); + PEAK.store(live, Ordering::Relaxed); + live +} + +/// How far above `baseline` the live bytes rose since [`reset_peak`]. +fn peak_over(baseline: usize) -> usize { + PEAK.load(Ordering::Relaxed).saturating_sub(baseline) +} + +/// Bytes as mebibytes, for the report. +fn mib(bytes: usize) -> f64 { + bytes as f64 / (1024.0 * 1024.0) +} + +// --------------------------------------------------------------------------- +// Candidates: the property sets a column races against each other +// --------------------------------------------------------------------------- + +/// One candidate way of encoding a column. +/// +/// Each variant is only a recipe for a few per-column settings on a +/// [`WriterPropertiesBuilder`]. Nothing here reaches inside the writer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Candidate { + /// Dictionary encoding, which the writer would pick by default. + Dictionary, + /// No dictionary, values written as they are. + Plain, + /// No dictionary, delta encoded. Only for physical types that have a delta + /// encoding. + Delta, +} + +impl Candidate { + fn name(self) -> &'static str { + match self { + Candidate::Dictionary => "dictionary", + Candidate::Plain => "plain", + Candidate::Delta => "delta", + } + } + + /// What a reader pays to decode this candidate's data pages, cheapest + /// first. `PLAIN` is a memcpy, and `RLE_DICTIONARY` is a bit-unpacked index + /// followed by a lookup in a dictionary page the reader decodes once per + /// chunk. The `DELTA_*` encodings are dearer: every value comes out of a + /// bit-packed miniblock through a running sum, and `DELTA_BYTE_ARRAY` also + /// reassembles each value from a prefix of the one before it. + fn decode_rank(self) -> u8 { + match self { + Candidate::Dictionary | Candidate::Plain => 0, + Candidate::Delta => 1, + } + } + + /// The delta encoding for a physical type, if it has one. DOUBLE does not, + /// which is why a float column races only two candidates. + fn delta_encoding(physical: PhysicalType) -> Option { + match physical { + PhysicalType::INT32 | PhysicalType::INT64 => Some(Encoding::DELTA_BINARY_PACKED), + PhysicalType::BYTE_ARRAY => Some(Encoding::DELTA_BYTE_ARRAY), + _ => None, + } + } + + /// The candidates worth racing for one column. + fn for_column(descr: &ColumnDescPtr) -> Vec { + let mut out = vec![Candidate::Dictionary, Candidate::Plain]; + if Self::delta_encoding(descr.physical_type()).is_some() { + out.push(Candidate::Delta); + } + out + } + + /// Applies this candidate's settings for `column` to `builder`. + fn apply( + self, + builder: WriterPropertiesBuilder, + column: ColumnPath, + physical: PhysicalType, + ) -> WriterPropertiesBuilder { + match self { + Candidate::Dictionary => builder.set_column_dictionary_enabled(column, true), + Candidate::Plain => builder + .set_column_dictionary_enabled(column.clone(), false) + .set_column_encoding(column, Encoding::PLAIN), + Candidate::Delta => { + let encoding = Self::delta_encoding(physical) + .expect("a Delta candidate is only built for a type that has one"); + builder + .set_column_dictionary_enabled(column.clone(), false) + .set_column_encoding(column, encoding) + } + } + } +} + +/// What the writer knows about one column as it works through the row groups. +struct ColumnState { + path: ColumnPath, + physical: PhysicalType, + candidates: Vec, + /// The candidate in force: the last race's winner, or the writer's own + /// default until a race has run. + current: Candidate, + /// The next row group this column races on. + next_race: usize, + /// Row groups to wait before racing again. Doubles after every decisive + /// race, up to [`MAX_RERACE_INTERVAL`], and drops back to one row group as + /// soon as a race comes back short of [`DECISIVE_MARGIN`]. + race_interval: usize, +} + +impl ColumnState { + fn new(descr: &ColumnDescPtr) -> Self { + let candidates = Candidate::for_column(descr); + Self { + path: descr.path().clone(), + physical: descr.physical_type(), + current: candidates[0], + candidates, + next_race: 0, + race_interval: 1, + } + } + + /// Whether this column races on `row_group`. Every column races on the + /// first one; after that the cadence below decides. + fn is_racing(&self, row_group: usize) -> bool { + row_group >= self.next_race + } + + /// Records the outcome of one race, and scales the cadence to it. + /// + /// A decisive winner is worth asking about less often, so the interval + /// doubles: a column whose answer is obvious pays for a probe on a + /// vanishing fraction of its row groups. A race the leader did not win + /// decisively is not an answer to keep, so the interval drops back to the + /// minimum and the column races on the next row group. + fn record(&mut self, row_group: usize, winner: Candidate, margin: f64) { + self.current = winner; + self.race_interval = if margin >= DECISIVE_MARGIN { + (self.race_interval * 2).min(MAX_RERACE_INTERVAL) + } else { + 1 + }; + self.next_race = row_group + self.race_interval; + } +} + +// --------------------------------------------------------------------------- +// The adaptive writer +// --------------------------------------------------------------------------- + +/// Writes `groups` to `path`, choosing each column's encoding by measurement. +/// +/// Returns the candidate each column ended on, in column order. +fn write_with_probes( + schema: &SchemaRef, + groups: &[Vec], + props: &WriterProperties, + path: &Path, +) -> Result> { + let arrow_writer = + ArrowWriter::try_new(File::create(path)?, schema.clone(), Some(props.clone()))?; + // Taking the file writer and the row group writer factory apart from the + // ArrowWriter is what gives access to individual column writers. + let (mut file_writer, factory) = arrow_writer.into_serialized_writer()?; + + let mut columns: Vec = file_writer + .schema_descr() + .columns() + .iter() + .map(ColumnState::new) + .collect(); + + for (row_group, group) in groups.iter().enumerate() { + // Step 1: cut the probe, the leading data page of this row group, and + // flatten it once into leaf columns. Every probe writer below is fed + // from it. A row group in which nothing is racing cuts no probe at all. + let probe: Vec> = if columns.iter().any(|c| c.is_racing(row_group)) { + probe_page(group) + .iter() + .map(|batch| columns_of(schema, batch)) + .collect::>()? + } else { + Vec::new() + }; + + // Step 2: race the candidates, one column at a time. A column that is + // not due is skipped entirely, so nothing at all is built for it. + for (idx, column) in columns.iter_mut().enumerate() { + if !column.is_racing(row_group) { + continue; + } + + // One throwaway writer per candidate, all of them for this one leaf + // column, each at that candidate's properties. + let mut probes: Vec = column + .candidates + .iter() + .map(|candidate| { + let props = candidate_props(props, column, *candidate); + factory.create_column_writer(row_group, idx, &props) + }) + .collect::>()?; + + for leaves in &probe { + for writer in &mut probes { + writer.write(&leaves[idx])?; + } + } + + // Closing a probe yields a `ColumnCloseResult` carrying the + // compressed size that candidate actually achieved, counting every + // page of the chunk: for a dictionary candidate that is the + // dictionary page as well as the data pages indexing into it, so + // each candidate is charged its full cost. The column chunks + // themselves are thrown away; only the sizes matter. + let mut sizes: Vec = Vec::with_capacity(probes.len()); + for writer in probes { + let chunk = writer.close()?; + sizes.push(chunk.close().metadata.compressed_size() as u64); + } + + let (winner, margin) = decide(&column.candidates, &sizes); + column.record(row_group, winner, margin); + } + + // Step 3: write the whole row group for real, through one ordinary + // column writer per column at that column's current candidate. The + // probe rows are encoded a second time here, which is the price of the + // measurement. + let mut writers: Vec = Vec::with_capacity(columns.len()); + for (idx, column) in columns.iter().enumerate() { + let props = candidate_props(props, column, column.current); + writers.push(factory.create_column_writer(row_group, idx, &props)?); + } + + for batch in group { + for (idx, column) in columns_of(schema, batch)?.into_iter().enumerate() { + writers[idx].write(&column)?; + } + } + + // Step 4: hand the finished column chunks to the file writer. They are + // ordinary chunks: the same page store and, with the encryption feature + // on, the same encryptor as the default write path would have used. + let mut rg = file_writer.next_row_group()?; + for writer in writers { + writer.close()?.append_to_row_group(&mut rg)?; + } + rg.close()?; + } + + file_writer.close()?; + Ok(columns.iter().map(|c| c.current).collect()) +} + +/// How much smaller `best` is than `other`, as a fraction of `other`. +fn margin_over(best: u64, other: u64) -> f64 { + if other == 0 { + 0.0 + } else { + (other - best) as f64 / other as f64 + } +} + +/// The winning candidate and the margin the leader won by, from one round of +/// measured sizes. +fn decide(candidates: &[Candidate], sizes: &[u64]) -> (Candidate, f64) { + let mut order: Vec = (0..sizes.len()).collect(); + order.sort_by_key(|&k| (sizes[k], k)); + let best = sizes[order[0]]; + // Every column races at least a dictionary and a plain candidate, so there + // is always a runner-up for the leader to be measured against. + let margin = margin_over(best, sizes[order[1]]); + + // Candidates within NEAR_TIE_BAND of the leader are a near-tie: the sizes + // are not telling them apart, and picking the nominally smallest would be + // reading noise. The tie goes to the cheapest page to decode instead, which + // spends it where it is nearly free, on bytes the reader barely pays for. + // + // Deliberately not a coin flip. This example promises that the same input + // produces the same file, and a random tie-break would make a column's + // encoding depend on the run rather than on the data, so two writes of one + // dataset could differ and neither would be reproducible. + let winner = order + .iter() + .copied() + .take_while(|&k| margin_over(best, sizes[k]) < NEAR_TIE_BAND) + .min_by_key(|&k| (candidates[k].decode_rank(), sizes[k], k)) + .expect("the leader is always within the margin of itself"); + (candidates[winner], margin) +} + +/// The properties for one column encoded with `candidate`: the file's own +/// properties, with that one column's dictionary and encoding settings +/// overridden. Every other column's settings are irrelevant here, because a +/// writer built for a single column only ever consults its own. +fn candidate_props( + base: &WriterProperties, + column: &ColumnState, + candidate: Candidate, +) -> WriterPropertiesPtr { + let builder = candidate.apply( + base.clone().into_builder(), + column.path.clone(), + column.physical, + ); + Arc::new(builder.build()) +} + +/// The columns of one record batch, flattened into the order the column +/// writers expect. A nested field would expand into several entries here, one +/// per column of the Parquet schema. +fn columns_of(schema: &SchemaRef, batch: &RecordBatch) -> Result> { + let mut out = Vec::new(); + for (field, array) in schema.fields().iter().zip(batch.columns()) { + out.extend(compute_leaves(field.as_ref(), array)?); + } + Ok(out) +} + +/// The leading [`DATA_PAGE_ROWS`] rows of a row group, cut at a record batch +/// boundary with a slice of the batch that straddles the cut. Slicing an Arrow +/// array copies nothing, so cutting the probe costs no data movement, and a row +/// group shorter than a page simply yields all of it. +fn probe_page(group: &[RecordBatch]) -> Vec { + let mut out = Vec::new(); + let mut taken = 0usize; + for batch in group { + if taken >= DATA_PAGE_ROWS { + break; + } + let take = (DATA_PAGE_ROWS - taken).min(batch.num_rows()); + out.push(batch.slice(0, take)); + taken += take; + } + out +} + +// --------------------------------------------------------------------------- +// The dataset, and the stock writer to compare against +// --------------------------------------------------------------------------- + +/// splitmix64, so every run generates byte-identical data. +fn mix(state: &mut u64) -> u64 { + *state = state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut z = *state; + z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ (z >> 31) +} + +/// Three columns whose right answers differ: a low-cardinality string column +/// that wants a dictionary, a monotonic timestamp column that wants deltas, and +/// a column of unique floats that wants neither. +fn dataset() -> (SchemaRef, Vec) { + let schema = Arc::new(Schema::new(vec![ + Field::new("category", DataType::Utf8, false), + Field::new( + "event_time", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new("measurement", DataType::Float64, false), + ])); + + let mut seed = 0x5eedu64; + let mut batches = Vec::new(); + let mut row = 0usize; + while row < ROWS { + let n = BATCH_ROWS.min(ROWS - row); + let category: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..n).map(|i| format!("region-{:02}", (row + i) % 24)), + )); + // Strictly increasing microseconds, which is what delta encoding is + // built for and what a dictionary is useless against. + let event_time: ArrayRef = Arc::new( + Int64Array::from_iter_values( + (0..n).map(|i| 1_700_000_000_000_000i64 + (row + i) as i64 * 1_000), + ) + .reinterpret_cast::(), + ); + let measurement: ArrayRef = Arc::new(Float64Array::from_iter_values( + (0..n).map(|_| mix(&mut seed) as f64 / 1e6), + )); + batches.push( + RecordBatch::try_new(schema.clone(), vec![category, event_time, measurement]).unwrap(), + ); + row += n; + } + (schema, batches) +} + +/// Groups the batches into fixed-size row groups. +fn row_groups(batches: &[RecordBatch]) -> Vec> { + batches + .chunks(ROW_GROUP_ROWS / BATCH_ROWS) + .map(|c| c.to_vec()) + .collect() +} + +/// A stock [`ArrowWriter`] at the same properties, cutting row groups in the +/// same places, so the two files differ only in encoding choices. +fn write_stock( + schema: &SchemaRef, + groups: &[Vec], + props: &WriterProperties, + path: &Path, +) -> Result<()> { + let mut writer = + ArrowWriter::try_new(File::create(path)?, schema.clone(), Some(props.clone()))?; + for group in groups { + for batch in group { + writer.write(batch)?; + } + writer.flush()?; + } + writer.close()?; + Ok(()) +} + +/// Reads `path` back and checks it row for row against the source batches. +fn verify(path: &Path, source: &[RecordBatch], schema: &SchemaRef) -> Result<()> { + let reader = ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)? + .with_batch_size(BATCH_ROWS) + .build()?; + let read: Vec = reader.collect::>()?; + let read = arrow_select::concat::concat_batches(schema, &read).unwrap(); + let want = arrow_select::concat::concat_batches(schema, source).unwrap(); + assert_eq!(read, want, "{} did not read back exactly", path.display()); + Ok(()) +} + +/// The encodings each column chunk of the first row group actually used. +fn encodings(path: &Path) -> Result> { + let reader = ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)?; + let rg = reader.metadata().row_group(0); + Ok((0..rg.num_columns()) + .map(|i| { + let c = rg.column(i); + let mut names: Vec = c.encodings().map(|e| format!("{e}")).collect(); + names.sort(); + names.join("+") + }) + .collect()) +} + +fn main() -> Result<()> { + let (schema, batches) = dataset(); + let groups = row_groups(&batches); + + // Identical properties for both writers. Only the per-column encoding + // settings the probe writer adds on top may differ. + let props = WriterProperties::builder() + .set_compression(Compression::UNCOMPRESSED) + .set_data_page_row_count_limit(DATA_PAGE_ROWS) + .set_max_row_group_row_count(Some(ROW_GROUP_ROWS)) + .build(); + + let dir = std::env::temp_dir(); + let probed_path = dir.join("chunk_probe_writer_probed.parquet"); + let stock_path = dir.join("chunk_probe_writer_stock.parquet"); + + // Both arms are measured from here, after the input batches have been + // materialized, so what is reported is the cost of writing and not of + // generating the data. Each arm runs once. The byte totals are + // deterministic; the time and memory figures are indicative only, and the + // arm that runs second can look slightly cheaper for reusing heap that the + // first one freed. Wall clock is a stand-in for CPU time, which the standard + // library cannot read; both arms are single threaded, so on an unloaded + // machine the two track each other closely. + // + // The probe arm can come out ahead on time as well as bytes, which is not a + // contradiction: a probe pass costs one page worth of rows per candidate, + // while a column the stock writer starts encoding as a dictionary and then + // falls back on has cost a dictionary build over the whole row group. The + // encodings printed below show where that happened. + let baseline = reset_peak(); + let started = Instant::now(); + write_stock(&schema, &groups, &props, &stock_path)?; + let stock_wall = started.elapsed(); + let stock_peak = peak_over(baseline); + + let baseline = reset_peak(); + let started = Instant::now(); + let chosen = write_with_probes(&schema, &groups, &props, &probed_path)?; + let probed_wall = started.elapsed(); + let probed_peak = peak_over(baseline); + + verify(&probed_path, &batches, &schema)?; + verify(&stock_path, &batches, &schema)?; + + let probed_bytes = std::fs::metadata(&probed_path)?.len(); + let stock_bytes = std::fs::metadata(&stock_path)?.len(); + let delta = 100.0 * (probed_bytes as f64 - stock_bytes as f64) / stock_bytes as f64; + + let probed_encodings = encodings(&probed_path)?; + let stock_encodings = encodings(&stock_path)?; + println!( + "{ROWS} rows in {} row groups, both read back exactly\n", + groups.len() + ); + println!( + " {:<18} {:>10} {:>6} {:>10}", + "arm", "bytes", "wall s", "peak alloc" + ); + println!( + " {:<18} {stock_bytes:>10} {:>6.2} {:>7.1} MB", + "stock ArrowWriter", + stock_wall.as_secs_f64(), + mib(stock_peak) + ); + println!( + " {:<18} {probed_bytes:>10} {:>6.2} {:>7.1} MB ({delta:+.1}% bytes)\n", + "probe writer", + probed_wall.as_secs_f64(), + mib(probed_peak) + ); + println!( + " {:<14} {:<12} {:<24} encodings (stock)", + "column", "chosen", "encodings (probed)" + ); + for (i, field) in schema.fields().iter().enumerate() { + let (probed, stock) = (&probed_encodings[i], &stock_encodings[i]); + println!( + " {:<14} {:<12} {probed:<24} {stock}", + field.name(), + chosen[i].name() + ); + } + + std::fs::remove_file(probed_path)?; + std::fs::remove_file(stock_path)?; + Ok(()) +} diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index ac38a11f1017..af7e14ccc9a1 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1359,17 +1359,144 @@ impl ArrowRowGroupWriterFactory { Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema)) } + /// The [`PageStoreFactory`] this factory allocates column chunk buffers + /// from, as set by [`Self::with_page_store_factory`]. + /// + /// A caller that builds some of a row group's column chunks by another + /// route can pass this to that route as well, so that every chunk in the + /// row group is buffered by the same page store, whether that spills to a + /// temporary file, to object storage or to the heap. + pub fn page_store_factory(&self) -> &Arc { + &self.page_store_factory + } + /// Create column writers for a new row group, with the given row group index pub fn create_column_writers(&self, row_group_index: usize) -> Result> { - let mut writers = Vec::with_capacity(self.arrow_schema.fields.len()); + let props = Arc::clone(&self.props); + let writers = self.create_selected_column_writers(row_group_index, &props, &|_| true)?; + Ok(writers + .into_iter() + .map(|writer| writer.expect("every leaf column was selected")) + .collect()) + } + + /// Create a writer for one leaf column of the given row group, using + /// `props` in place of the file writer's properties. + /// + /// `column_index` indexes the parquet *leaf* columns of the schema, in + /// schema order: the order of [`SchemaDescriptor::columns`], which is also + /// the order in which [`compute_leaves`] yields [`ArrowLeafColumn`]s and + /// the order of the writers returned by [`Self::create_column_writers`]. + /// + /// It is not an index into the Arrow schema's top level fields. A struct, + /// list or map field expands to one leaf column per primitive it contains, + /// so the Arrow fields `[i, struct, t]` have the four leaf columns + /// `[i, s.a, s.b, t]` and the field `t` is column 3, not column 2. An index + /// at or past the number of leaf columns is an error. + /// + /// The writer allocates its page store from the same [`PageStoreFactory`] + /// (and, with the `encryption` feature, is wired to the same file + /// encryptor) as [`Self::create_column_writers`]. Only the encoding level + /// properties differ, so the [`ArrowColumnChunk`] it produces is still + /// appendable to a row group of the same file writer. This makes it + /// possible to encode the same rows several times with different + /// properties and keep only the chunk that came out smallest, or to build + /// the rest of the row group's chunks by some other route. + /// + /// `props` must describe the same schema as the file writer: properties + /// that change the physical layout of the file, such as the writer version + /// or the schema itself, will produce chunks that the file writer cannot + /// accept. Per-column encoding, compression, dictionary and statistics + /// settings are the intended use. + /// + /// ``` + /// # use std::sync::Arc; + /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch}; + /// # use parquet::arrow::ArrowSchemaConverter; + /// # use parquet::arrow::arrow_writer::{ArrowRowGroupWriterFactory, compute_leaves}; + /// # use parquet::basic::Encoding; + /// # use parquet::file::properties::WriterProperties; + /// # use parquet::file::writer::SerializedFileWriter; + /// # fn main() -> parquet::errors::Result<()> { + /// let values: ArrayRef = Arc::new(Int32Array::from_iter_values(0..1024)); + /// let batch = RecordBatch::try_from_iter([("id", values)])?; + /// let schema = batch.schema(); + /// + /// let props = Arc::new(WriterProperties::builder().build()); + /// let parquet_schema = ArrowSchemaConverter::new().convert(&schema)?; + /// let mut buffer = Vec::new(); + /// let mut file_writer = + /// SerializedFileWriter::new(&mut buffer, parquet_schema.root_schema_ptr(), props.clone())?; + /// let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&schema)); + /// + /// // Encoding properties for this column chunk only. + /// let column_props = Arc::new( + /// WriterProperties::builder() + /// .set_dictionary_enabled(false) + /// .set_encoding(Encoding::DELTA_BINARY_PACKED) + /// .build(), + /// ); + /// // The single leaf column of this schema, written at those properties. + /// let mut writer = factory.create_column_writer(0, 0, &column_props)?; + /// + /// for (field, column) in schema.fields().iter().zip(batch.columns()) { + /// for leaf in compute_leaves(field, column)? { + /// writer.write(&leaf)?; + /// } + /// } + /// + /// let mut row_group = file_writer.next_row_group()?; + /// writer.close()?.append_to_row_group(&mut row_group)?; + /// row_group.close()?; + /// file_writer.close()?; + /// # Ok(()) + /// # } + /// ``` + pub fn create_column_writer( + &self, + row_group_index: usize, + column_index: usize, + props: &WriterPropertiesPtr, + ) -> Result { + let num_columns = self.schema.num_columns(); + if column_index >= num_columns { + return Err(general_err!( + "column_index {column_index} is out of range: the schema has {num_columns} leaf columns, so the valid indices are 0..{num_columns}" + )); + } + let mut writers = self + .create_selected_column_writers(row_group_index, props, &|leaf| leaf == column_index)?; + Ok(writers[column_index] + .take() + .expect("the requested leaf column was selected")) + } + + /// Create a column writer for each leaf column `select` accepts, with one + /// entry per leaf column of the schema in leaf order, holding `None` + /// wherever `select` returned `false` so that the entries stay aligned with + /// that order. Nothing is allocated for an unselected column: in particular + /// no page store is created for it, which matters when the page store + /// factory allocates something more expensive than a heap buffer. + /// + /// This is the only place column writers are constructed. + /// [`Self::create_column_writers`] and [`Self::create_column_writer`] are + /// its two degenerate cases, selecting every column and exactly one. + fn create_selected_column_writers( + &self, + row_group_index: usize, + props: &WriterPropertiesPtr, + select: &dyn Fn(usize) -> bool, + ) -> Result>> { + let mut writers = Vec::with_capacity(self.schema.num_columns()); let mut leaves = self.schema.columns().iter(); let column_factory = self.column_writer_factory(row_group_index); for field in &self.arrow_schema.fields { column_factory.get_arrow_column_writer( field.data_type(), - &self.props, + props, &mut leaves, &mut writers, + select, )?; } Ok(writers) @@ -1461,38 +1588,49 @@ impl ArrowColumnWriterFactory { } /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the - /// output ColumnDesc to `leaves` and the column writers to `out` + /// output ColumnDesc to `leaves` and the column writers to `out`. + /// + /// One entry is appended to `out` per leaf column, holding `None` for a + /// leaf `select` rejects so that the entries stay aligned with the leaf + /// order either way. fn get_arrow_column_writer( &self, data_type: &ArrowDataType, props: &WriterPropertiesPtr, leaves: &mut Iter<'_, ColumnDescPtr>, - out: &mut Vec, + out: &mut Vec>, + select: &dyn Fn(usize) -> bool, ) -> Result<()> { let write_distinct_values = props.write_row_group_number_distinct_values(); // Instantiate writers for normal columns - let col = |desc: &ColumnDescPtr| -> Result { + let col = |desc: &ColumnDescPtr| -> Result> { + if !select(out.len()) { + return Ok(None); + } let page_writer = self.create_page_writer(desc, out.len())?; let chunk = page_writer.buffer.clone(); let writer = get_column_writer(desc.clone(), props.clone(), page_writer); - Ok(ArrowColumnWriter { + Ok(Some(ArrowColumnWriter { chunk, writer: ArrowColumnWriterImpl::Column(writer), distinct_values_seen: write_distinct_values.then(HashSet::new), - }) + })) }; // Instantiate writers for byte arrays (e.g. Utf8, Binary, etc) - let bytes = |desc: &ColumnDescPtr| -> Result { + let bytes = |desc: &ColumnDescPtr| -> Result> { + if !select(out.len()) { + return Ok(None); + } let page_writer = self.create_page_writer(desc, out.len())?; let chunk = page_writer.buffer.clone(); let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer); - Ok(ArrowColumnWriter { + Ok(Some(ArrowColumnWriter { chunk, writer: ArrowColumnWriterImpl::ByteArray(writer), distinct_values_seen: write_distinct_values.then(HashSet::new), - }) + })) }; match data_type { @@ -1511,17 +1649,17 @@ impl ArrowColumnWriterFactory { | ArrowDataType::FixedSizeList(f, _) | ArrowDataType::ListView(f) | ArrowDataType::LargeListView(f) => { - self.get_arrow_column_writer(f.data_type(), props, leaves, out)? + self.get_arrow_column_writer(f.data_type(), props, leaves, out, select)? } ArrowDataType::Struct(fields) => { for field in fields { - self.get_arrow_column_writer(field.data_type(), props, leaves, out)? + self.get_arrow_column_writer(field.data_type(), props, leaves, out, select)? } } ArrowDataType::Map(f, _) => match f.data_type() { ArrowDataType::Struct(f) => { - self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?; - self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)? + self.get_arrow_column_writer(f[0].data_type(), props, leaves, out, select)?; + self.get_arrow_column_writer(f[1].data_type(), props, leaves, out, select)? } _ => unreachable!("invalid map type"), }, @@ -1537,7 +1675,7 @@ impl ArrowColumnWriterFactory { _ => out.push(col(leaves.next().unwrap())?), }, ArrowDataType::RunEndEncoded(_, value_field) => { - self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)? + self.get_arrow_column_writer(value_field.data_type(), props, leaves, out, select)? } _ => { return Err(ParquetError::NYI(format!( @@ -2322,6 +2460,439 @@ mod tests { ); } + /// Writes `batch` as a single row group and returns the finished file. + /// + /// With `column_props` set, every leaf column's writer is built one at a + /// time through [`ArrowRowGroupWriterFactory::create_column_writer`] at + /// those properties. With `None`, the row group is built in one call to + /// [`ArrowRowGroupWriterFactory::create_column_writers`] at the file + /// writer's own properties, which is the ordinary write path. + fn write_via_row_group_factory( + batch: &RecordBatch, + file_props: WriterProperties, + column_props: Option, + ) -> Bytes { + let schema = batch.schema(); + let file_props = Arc::new(file_props); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let mut writer = SerializedFileWriter::new( + &mut buf, + parquet_schema.root_schema_ptr(), + file_props.clone(), + ) + .unwrap(); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + let mut col_writers: Vec = match column_props { + Some(props) => { + let props = Arc::new(props); + (0..parquet_schema.num_columns()) + .map(|column| factory.create_column_writer(0, column, &props).unwrap()) + .collect() + } + None => factory.create_column_writers(0).unwrap(), + }; + + let mut writers = col_writers.iter_mut(); + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + writers.next().unwrap().write(&leaf).unwrap(); + } + } + + let mut rg = writer.next_row_group().unwrap(); + for chunk in col_writers { + chunk.close().unwrap().append_to_row_group(&mut rg).unwrap(); + } + rg.close().unwrap(); + writer.close().unwrap(); + Bytes::from(buf) + } + + /// Building every column one at a time at the file writer's own properties + /// must reproduce `create_column_writers`, and different properties must + /// really reach the encoders while still producing an appendable chunk. + #[test] + fn create_column_writer_honours_per_column_properties() { + let array: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..4096).map(|i| format!("value-{:04}", i % 97)), + )); + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + + let props = || { + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_dictionary_page_size_limit(4096) + .build() + }; + + // Handing the primitive the file writer's own properties must + // reproduce `create_column_writers` byte for byte. + let batched = write_via_row_group_factory(&batch, props(), None); + let per_column = write_via_row_group_factory(&batch, props(), Some(props())); + assert_eq!(batched, per_column); + + // Per-column properties really do change the encoding, and the + // resulting chunk is still appendable to the same file writer. + let pinned = write_via_row_group_factory( + &batch, + props(), + Some( + WriterProperties::builder() + .set_statistics_enabled(EnabledStatistics::Page) + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .build(), + ), + ); + assert_ne!(batched, pinned); + + let reader = SerializedFileReader::new(pinned.clone()).unwrap(); + let column = reader.metadata().row_group(0).column(0); + assert!(column.dictionary_page_offset().is_none()); + let encodings: Vec<_> = column.encodings().collect(); + assert!( + encodings.contains(&Encoding::DELTA_BYTE_ARRAY), + "expected the requested encoding, got {encodings:?}" + ); + + // Every variant reads back to the original rows. + for data in [batched, per_column, pinned] { + let read = ParquetRecordBatchReader::try_new(data, 1024) + .unwrap() + .collect::>>() + .unwrap(); + let read = arrow_select::concat::concat_batches(&batch.schema(), &read).unwrap(); + assert_eq!(read, batch); + } + } + + /// `column_index` must address the schema's leaf columns, including across + /// a nested field that expands into several leaves, and the chunks the + /// primitive produces must be indistinguishable from the ones the + /// all-columns call produces. + #[test] + fn create_column_writer_addresses_leaf_columns() { + let ints: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); + let inner_a: ArrayRef = Arc::new(Int32Array::from_iter_values((0..64).map(|i| i * 2))); + let inner_b: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..64).map(|i| format!("s{i}")), + )); + let structs: ArrayRef = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("a", ArrowDataType::Int32, false)), + inner_a, + ), + ( + Arc::new(Field::new("b", ArrowDataType::Utf8, false)), + inner_b, + ), + ])); + let tail: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..64).map(|i| format!("t{}", i % 4)), + )); + let batch = RecordBatch::try_from_iter([("i", ints), ("s", structs), ("t", tail)]).unwrap(); + + let schema = batch.schema(); + let props = Arc::new(WriterProperties::builder().build()); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let writer = + SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) + .unwrap(); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + // Three Arrow fields, but four leaf columns: the struct field `s` + // expands into two of them, so the last field is column 3, not 2. + assert_eq!(schema.fields().len(), 3); + assert_eq!(parquet_schema.num_columns(), 4); + let paths: Vec = parquet_schema + .columns() + .iter() + .map(|c| c.path().string()) + .collect(); + assert_eq!(paths, ["i", "s.a", "s.b", "t"]); + + // One writer per leaf column, built one at a time, against one writer + // per leaf column built in a single call. + let mut one_at_a_time: Vec = (0..parquet_schema.num_columns()) + .map(|column| factory.create_column_writer(0, column, &props).unwrap()) + .collect(); + let mut all = factory.create_column_writers(0).unwrap(); + assert_eq!(all.len(), 4); + + let mut leaf_idx = 0usize; + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + one_at_a_time[leaf_idx].write(&leaf).unwrap(); + all[leaf_idx].write(&leaf).unwrap(); + leaf_idx += 1; + } + } + assert_eq!(leaf_idx, 4); + + // A writer built for the wrong leaf would have the wrong descriptor and + // so encode to a different size: `s.b` and `t` are both strings, and + // only the right pairing matches. + for (idx, (ours, theirs)) in one_at_a_time.into_iter().zip(all).enumerate() { + let ours = ours.close().unwrap(); + let theirs = theirs.close().unwrap(); + assert_eq!( + ours.close().metadata.compressed_size(), + theirs.close().metadata.compressed_size(), + "leaf column {idx} encoded differently from the all-columns writer" + ); + assert_eq!( + ours.close().metadata.num_values(), + theirs.close().metadata.num_values() + ); + assert_eq!( + ours.close().metadata.column_path(), + theirs.close().metadata.column_path(), + "leaf column {idx} was built for the wrong column" + ); + } + } + + /// A writer built by the primitive must allocate its page store from the + /// factory's own [`PageStoreFactory`]. + #[test] + fn create_column_writer_uses_the_configured_page_store() { + let a: ArrayRef = Arc::new(Int32Array::from_iter_values(0..128)); + let b: ArrayRef = Arc::new(Int32Array::from_iter_values((0..128).map(|i| i * 3))); + let batch = RecordBatch::try_from_iter([("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let props = Arc::new(WriterProperties::builder().build()); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let writer = + SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) + .unwrap(); + + let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let store_factory: Arc = + Arc::new(RecordingPageStoreFactory { puts: puts.clone() }); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)) + .with_page_store_factory(Arc::clone(&store_factory)); + + // The getter hands back exactly the factory that was configured, so a + // caller producing chunks by another route can share it. + assert!(Arc::ptr_eq(factory.page_store_factory(), &store_factory)); + + // Build a writer for the second leaf column only, and feed it only its + // own leaf. + let mut column_writer = factory.create_column_writer(0, 1, &props).unwrap(); + let mut leaf_idx = 0usize; + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + if leaf_idx == 1 { + column_writer.write(&leaf).unwrap(); + } + leaf_idx += 1; + } + } + column_writer.close().unwrap(); + + assert!( + puts.load(std::sync::atomic::Ordering::Relaxed) > 0, + "the column's pages did not go through the configured page store" + ); + } + + /// An index at or past the number of leaf columns is a clear error naming + /// the valid range, not a panic. + #[test] + fn create_column_writer_rejects_an_out_of_range_column() { + let a: ArrayRef = Arc::new(Int32Array::from_iter_values(0..8)); + let b: ArrayRef = Arc::new(Int32Array::from_iter_values(8..16)); + let batch = RecordBatch::try_from_iter([("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let props = Arc::new(WriterProperties::builder().build()); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let writer = + SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) + .unwrap(); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + // The last valid index works. + factory.create_column_writer(0, 1, &props).unwrap(); + + let err = factory.create_column_writer(0, 2, &props).unwrap_err(); + let message = err.to_string(); + assert!( + message.contains("column_index 2 is out of range"), + "unexpected error: {message}" + ); + assert!( + message.contains("0..2"), + "the error must name the valid range: {message}" + ); + } + + /// Several writers for the *same* leaf column, under different properties, + /// can be alive at once and each closes independently. This is the probe + /// pattern: encode the same rows every way and keep the smallest chunk. + #[test] + fn create_column_writer_probes_one_column_under_several_properties() { + let array: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..2048).map(|i| format!("value-{:04}", i % 53)), + )); + let batch = RecordBatch::try_from_iter([("col", array)]).unwrap(); + let schema = batch.schema(); + + let file_props = Arc::new(WriterProperties::builder().build()); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let mut writer = SerializedFileWriter::new( + &mut buf, + parquet_schema.root_schema_ptr(), + file_props.clone(), + ) + .unwrap(); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + let candidates = [ + Arc::new(WriterProperties::builder().build()), + Arc::new( + WriterProperties::builder() + .set_dictionary_enabled(false) + .set_encoding(Encoding::DELTA_BYTE_ARRAY) + .build(), + ), + Arc::new( + WriterProperties::builder() + .set_dictionary_enabled(false) + .set_encoding(Encoding::PLAIN) + .build(), + ), + ]; + + // All of the probe writers for column 0 exist at the same time. + let mut probes: Vec = candidates + .iter() + .map(|props| factory.create_column_writer(0, 0, props).unwrap()) + .collect(); + + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + for probe in &mut probes { + probe.write(&leaf).unwrap(); + } + } + } + + // Each one closes on its own, and the candidates really did encode + // differently. + let sizes: Vec = probes + .into_iter() + .map(|probe| probe.close().unwrap().close().metadata.compressed_size()) + .collect(); + assert_eq!(sizes.len(), candidates.len()); + assert!( + sizes.iter().collect::>().len() == sizes.len(), + "the candidate properties all encoded to the same size: {sizes:?}" + ); + + // The probe chunks are thrown away; the winner is written again for + // real and the file still reads back. + let winner = sizes + .iter() + .enumerate() + .min_by_key(|(_, size)| **size) + .map(|(idx, _)| idx) + .unwrap(); + let mut real = factory + .create_column_writer(0, 0, &candidates[winner]) + .unwrap(); + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + real.write(&leaf).unwrap(); + } + } + let mut rg = writer.next_row_group().unwrap(); + real.close().unwrap().append_to_row_group(&mut rg).unwrap(); + rg.close().unwrap(); + writer.close().unwrap(); + + let read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024) + .unwrap() + .collect::>>() + .unwrap(); + let read = arrow_select::concat::concat_batches(&schema, &read).unwrap(); + assert_eq!(read, batch); + } + + /// With the `encryption` feature on, the writers the primitive builds must + /// be wired to the file writer's encryptor: a file written entirely through + /// it must be readable only with the matching decryption properties. + #[cfg(feature = "encryption")] + #[test] + fn create_column_writer_uses_the_file_encryptor() { + use crate::arrow::arrow_reader::ArrowReaderOptions; + use crate::encryption::decrypt::FileDecryptionProperties; + use crate::encryption::encrypt::FileEncryptionProperties; + + let key: Vec = b"0123456789012345".to_vec(); + let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..256)); + let batch = RecordBatch::try_from_iter([("v", array)]).unwrap(); + let schema = batch.schema(); + + let props = Arc::new( + WriterProperties::builder() + .with_file_encryption_properties( + FileEncryptionProperties::builder(key.clone()) + .build() + .unwrap(), + ) + .build(), + ); + let parquet_schema = ArrowSchemaConverter::new().convert(&schema).unwrap(); + let mut buf = Vec::with_capacity(1024); + let mut writer = + SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) + .unwrap(); + let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + + let mut writers: Vec = (0..parquet_schema.num_columns()) + .map(|column| factory.create_column_writer(0, column, &props).unwrap()) + .collect(); + let mut leaf_idx = 0usize; + for (field, column) in schema.fields().iter().zip(batch.columns()) { + for leaf in compute_leaves(field.as_ref(), column).unwrap() { + writers[leaf_idx].write(&leaf).unwrap(); + leaf_idx += 1; + } + } + let mut rg = writer.next_row_group().unwrap(); + for chunk in writers { + chunk.close().unwrap().append_to_row_group(&mut rg).unwrap(); + } + rg.close().unwrap(); + writer.close().unwrap(); + let data = Bytes::from(buf); + + // Without the key the pages cannot be read back. + assert!( + ParquetRecordBatchReader::try_new(data.clone(), 1024).is_err(), + "the chunks were written unencrypted" + ); + + let options = ArrowReaderOptions::new().with_file_decryption_properties( + FileDecryptionProperties::builder(key).build().unwrap(), + ); + let read = ParquetRecordBatchReaderBuilder::try_new_with_options(data, options) + .unwrap() + .build() + .unwrap() + .collect::>>() + .unwrap(); + let read = arrow_select::concat::concat_batches(&schema, &read).unwrap(); + assert_eq!(read, batch); + } #[test] fn arrow_writer() { // define schema