From e4800749755351897fc4fed613fe4fc224e8c0a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 30 Aug 2026 17:06:47 +0000 Subject: [PATCH 1/2] feat(parquet): create a subset of column writers with per-row-group properties `ArrowRowGroupWriterFactory` can only build column writers for every leaf column of a row group, at the properties the file writer was created with. This adds two accessors: * `create_column_writer(row_group_index, column_index, props)` builds the writer for one leaf column with the given `WriterProperties`. `column_index` indexes the parquet leaf columns in schema order, the order of `SchemaDescriptor::columns()` and of the `ArrowLeafColumn`s that `compute_leaves` yields, not the Arrow schema's top level fields; an index past the last leaf column is an error naming the valid range. The writer shares the factory's page store factory and, with the `encryption` feature, its file encryptor, so the chunk it produces is appendable to the same file writer. * `page_store_factory()` returns the configured `PageStoreFactory`, so a caller that builds some of a row group's column chunks by another route can buffer them through the same page store. Both this accessor and `create_column_writers` go through one private construction path, of which they are the two degenerate cases: one column, or every column at the file writer's own properties. The default write path is unchanged. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01MwnUAoPDMQaYaQPVP2iUcz --- parquet/src/arrow/arrow_writer/mod.rs | 601 +++++++++++++++++++++++++- 1 file changed, 586 insertions(+), 15 deletions(-) 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 From b99c55ff7087a47aceccfc940f397d4f6f828ef9 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 30 Aug 2026 17:06:47 +0000 Subject: [PATCH 2/2] feat(parquet): add the chunk_probe_writer example An adaptive writer built on `create_column_writer`: for each row group, each column that is due to race encodes a probe of its leading rows through several throwaway single-column writers, one per candidate property set (dictionary, delta where the physical type has one, and plain). The compressed sizes on the returned `ColumnCloseResult`s pick a winner, the probe chunks are dropped, and the row group is then written through one ordinary column writer per column at its current choice. Two small mechanisms make the measurement pay for itself only where it is worth it. `DECISIVE_MARGIN` is the fraction by which the leading chunk must beat the runner-up for the race to count as decided. Closer than `NEAR_TIE_BAND`, the sizes are not telling the candidates apart at all, and the tie is broken toward the page a reader decodes most cheaply, dictionary and plain ahead of the delta encodings, rather than toward a nominally smaller chunk. The band is a separate and much narrower constant than `DECISIVE_MARGIN` for a measured reason: at 10 percent it gives away real wins, since delta beats plain or dictionary by 5 to 10 percent on many TPC-H columns, and at 2 percent the tie-break is free on that data. The tie-break is deterministic and not random, because the example promises that the same input produces the same file. The interval before a column races again scales with the margin it won by, doubling after each decisive race up to `MAX_RERACE_INTERVAL` row groups and dropping back to one row group after a race the leader did not win decisively. A settled column then costs nothing at all: no probe writer is built for it, so no page store is allocated for it either. The probe itself stays one data page long. Growing it a page at a time until the leader is decisive was tried and 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. The module documentation records that result and leaves adaptive probe length to callers whose data rewards it. The example generates a deterministic dataset whose best encoding differs per column, writes it both with the probe writer and with a stock `ArrowWriter` at identical properties and row group boundaries, verifies that both read back exactly, and prints the total bytes, elapsed time and peak allocation of each arm alongside the encodings each column ended on. The two measured numbers are there because the technique trades work for size, so the report should show both sides of that trade. Neither needs a dependency: elapsed time is `std::time::Instant`, standing in for CPU time, which the standard library cannot read, and peak allocation comes from a short counting `#[global_allocator]` in the example itself, sampled as each arm's increment over the bytes live when it started so that the input batches do not drown the signal. Both arms are measured after the input is materialized and each runs once, so the byte totals are deterministic while the time and memory figures are indicative. The module documentation also spells out how a dictionary candidate is handled, since that is the case where a measured comparison differs most from a guess: a probe's compressed size covers the dictionary page as well as the data pages indexing into it, so a dictionary is charged its full cost and wins only by being smallest; and a dictionary that a short probe flattered is corrected by the writer's own single-dictionary-page fallback and by the next re-race. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01MwnUAoPDMQaYaQPVP2iUcz --- parquet/Cargo.toml | 5 + parquet/examples/chunk_probe_writer.rs | 728 +++++++++++++++++++++++++ 2 files changed, 733 insertions(+) create mode 100644 parquet/examples/chunk_probe_writer.rs 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(()) +}