From d7700cfadab52a77ab0fb9b4bf3924f8ab44f2bd Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 10 Apr 2026 01:37:05 +0300 Subject: [PATCH 01/26] feat(parquet): dictionary fallback heuristics In parquet, add `dictionary_fallback` option to `ColumnProperties`. Its value type is defined as the `DictionaryFallback` enum with the varians `OnDataPageSize` (the previous behavior, the default) and `OnUnfavorableCompression`. The latter replicates the behavior of parquet-java, which falls back to non-dictionary encoding when the estimated compressed size of the data page with dictionary encoding is not smaller than the estimated size of the data page with plain encoding. --- parquet/src/arrow/arrow_writer/byte_array.rs | 34 +++++++++- parquet/src/arrow/arrow_writer/mod.rs | 44 ++++++++++++- parquet/src/column/writer/encoder.rs | 35 +++++++++- parquet/src/column/writer/mod.rs | 33 +++++++--- parquet/src/encodings/encoding/mod.rs | 5 +- .../src/encodings/encoding/plain_counter.rs | 64 +++++++++++++++++++ parquet/src/file/properties.rs | 64 +++++++++++++++++++ 7 files changed, 262 insertions(+), 17 deletions(-) create mode 100644 parquet/src/encodings/encoding/plain_counter.rs diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index f56f9570adfb..c0bfc9c1d519 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -21,10 +21,12 @@ use crate::column::writer::encoder::{ ColumnValueEncoder, DataPageValues, DictionaryPage, create_bloom_filter, }; use crate::data_type::{AsBytes, ByteArray, Int32Type}; -use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder}; +use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder, PlainDataSizeCounter}; use crate::encodings::rle::RleEncoder; use crate::errors::{ParquetError, Result}; -use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion}; +use crate::file::properties::{ + DictionaryFallback, EnabledStatistics, WriterProperties, WriterVersion, +}; use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::ColumnDescPtr; @@ -421,6 +423,7 @@ impl DictEncoder { pub struct ByteArrayEncoder { fallback: FallbackEncoder, dict_encoder: Option, + plain_data_size_counter: Option, statistics_enabled: EnabledStatistics, min_value: Option, max_value: Option, @@ -446,6 +449,17 @@ impl ColumnValueEncoder for ByteArrayEncoder { .dictionary_enabled(descr.path()) .then(DictEncoder::default); + let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { + DictionaryFallback::OnPageSizeLimit => None, + DictionaryFallback::OnUnfavorableCompression => { + if dictionary.is_some() { + Some(PlainDataSizeCounter::new(descr)) + } else { + None + } + } + }; + let fallback = FallbackEncoder::new(descr, props)?; let (bloom_filter, bloom_filter_target_fpp) = create_bloom_filter(props, descr)?; @@ -460,6 +474,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { bloom_filter, bloom_filter_target_fpp, dict_encoder: dictionary, + plain_data_size_counter, min_value: None, max_value: None, geo_stats_accumulator, @@ -510,6 +525,11 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(self.dict_encoder.as_ref()?.estimated_dict_page_size()) } + fn uncompressed_data_size(&self) -> Option { + let counter = self.plain_data_size_counter.as_ref()?; + Some(counter.uncompressed_data_size()) + } + /// Returns an estimate of the data page size in bytes /// /// This includes: @@ -582,7 +602,15 @@ where } match &mut encoder.dict_encoder { - Some(dict_encoder) => dict_encoder.encode(values, indices), + Some(dict_encoder) => { + dict_encoder.encode(values, indices); + if let Some(counter) = encoder.plain_data_size_counter.as_mut() { + for idx in indices { + let value = values.value(*idx); + counter.update_byte_array(value.as_ref()); + } + } + } None => encoder.fallback.encode(values, indices), } } diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 8422263b1f63..d11f086557be 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1691,7 +1691,7 @@ mod tests { use crate::data_type::AsBytes; use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader}; use crate::file::properties::{ - BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion, + BloomFilterPosition, DictionaryFallback, EnabledStatistics, ReaderProperties, WriterVersion, }; use crate::file::serialized_reader::ReadOptionsBuilder; use crate::file::{ @@ -4827,6 +4827,48 @@ mod tests { assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); } + #[test] + fn test_dict_page_size_decided_by_compression_fallback() { + let array = Arc::new(Int64Array::from_iter(0..1024 * 1024)); + let schema = Arc::new(Schema::new(vec![ + Field::new("col0", arrow_schema::DataType::Int64, false), + Field::new("col1", arrow_schema::DataType::Int64, false), + ])); + let batch = + arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_page_size_limit(1024 * 1024) + .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4) + .set_column_dictionary_fallback( + ColumnPath::from("col1"), + DictionaryFallback::OnUnfavorableCompression, + ) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let col0_meta = metadata.row_group(0).column(0); + let col1_meta = metadata.row_group(0).column(1); + + let get_dict_page_size = move |meta: &ColumnChunkMetaData| { + let mut reader = + SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap(); + let page = reader.get_next_page().unwrap().unwrap(); + match page { + Page::DictionaryPage { buf, .. } => buf.len(), + _ => panic!("expected DictionaryPage"), + } + }; + + assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024); + assert_eq!(get_dict_page_size(col1_meta), 8192); + } + struct WriteBatchesShape { num_batches: usize, rows_per_batch: usize, diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index ec1afca58335..a808cb6286e5 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -25,9 +25,9 @@ use crate::column::writer::{ }; use crate::data_type::DataType; use crate::data_type::private::ParquetValueType; -use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder}; +use crate::encodings::encoding::{DictEncoder, Encoder, PlainDataSizeCounter, get_encoder}; use crate::errors::{ParquetError, Result}; -use crate::file::properties::{EnabledStatistics, WriterProperties}; +use crate::file::properties::{DictionaryFallback, EnabledStatistics, WriterProperties}; use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; use crate::geospatial::statistics::GeospatialStatistics; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; @@ -109,6 +109,12 @@ pub trait ColumnValueEncoder { /// + fn estimated_data_page_size(&self) -> usize; + /// Returns the estimated size of plainly encoded data, in bytes, + /// that would be written without a dictionary. + /// If there is no dictionary, or the data size statistic is not available, + /// returns `None`. + fn uncompressed_data_size(&self) -> Option; + /// Flush the dictionary page for this column chunk if any. Any subsequent calls to /// [`Self::write`] will not be dictionary encoded /// @@ -132,6 +138,7 @@ pub trait ColumnValueEncoder { pub struct ColumnValueEncoderImpl { encoder: Box>, dict_encoder: Option>, + plain_data_size_counter: Option, descr: ColumnDescPtr, num_values: usize, statistics_enabled: EnabledStatistics, @@ -176,7 +183,13 @@ impl ColumnValueEncoderImpl { } match &mut self.dict_encoder { - Some(encoder) => encoder.put(slice), + Some(encoder) => { + encoder.put(slice)?; + if let Some(counter) = self.plain_data_size_counter.as_mut() { + counter.update(slice); + } + Ok(()) + } _ => self.encoder.put(slice), } } @@ -197,6 +210,16 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let dict_supported = props.dictionary_enabled(descr.path()) && has_dictionary_support(T::get_physical_type(), props); let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); + let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { + DictionaryFallback::OnPageSizeLimit => None, + DictionaryFallback::OnUnfavorableCompression => { + if dict_encoder.is_some() { + Some(PlainDataSizeCounter::new(descr)) + } else { + None + } + } + }; // Set either main encoder or fallback encoder. let encoder = get_encoder( @@ -215,6 +238,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Ok(Self { encoder, dict_encoder, + plain_data_size_counter, descr: descr.clone(), num_values: 0, statistics_enabled, @@ -277,6 +301,11 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Some(self.dict_encoder.as_ref()?.dict_encoded_size()) } + fn uncompressed_data_size(&self) -> Option { + let counter = self.plain_data_size_counter.as_ref()?; + Some(counter.uncompressed_data_size()) + } + fn estimated_data_page_size(&self) -> usize { match &self.dict_encoder { Some(encoder) => encoder.estimated_data_encoded_size(), diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 46f90d3f7762..944ce8df17bc 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -746,18 +746,33 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { /// Returns true if we need to fall back to non-dictionary encoding. /// - /// We can only fall back if dictionary encoder is set and we have exceeded dictionary - /// size. - #[inline] + /// The behavior is governed by the `dictionary_fallback` column property. fn should_dict_fallback(&self) -> bool { - match self.encoder.estimated_dict_page_size() { - Some(size) => { - size >= self - .props - .column_dictionary_page_size_limit(self.descr.path()) + let dict_size = match self.encoder.estimated_dict_page_size() { + Some(size) => size, + None => return false, + }; + + // First check: dictionary size exceeds limit + if dict_size + >= self + .props + .column_dictionary_page_size_limit(self.descr.path()) + { + return true; + } + + // Second check, if enabled: the compression heuristic. + // For similar logic in parquet-java, + // see DictionaryValuesWriter.isCompressionSatisfying + if let Some(raw_size) = self.encoder.uncompressed_data_size() { + let encoded_size = self.encoder.estimated_data_page_size(); + if encoded_size + dict_size >= raw_size { + return true; } - None => false, } + + false } /// Returns true if there is enough data for a data page, false otherwise. diff --git a/parquet/src/encodings/encoding/mod.rs b/parquet/src/encodings/encoding/mod.rs index eeabcf4ba5ce..21b2caed6212 100644 --- a/parquet/src/encodings/encoding/mod.rs +++ b/parquet/src/encodings/encoding/mod.rs @@ -29,10 +29,13 @@ use crate::util::bit_util::{BitWriter, num_required_bits}; use byte_stream_split_encoder::{ByteStreamSplitEncoder, VariableWidthByteStreamSplitEncoder}; use bytes::Bytes; -pub use dict_encoder::DictEncoder; mod byte_stream_split_encoder; mod dict_encoder; +mod plain_counter; + +pub use dict_encoder::DictEncoder; +pub use plain_counter::PlainDataSizeCounter; // ---------------------------------------------------------------------- // Encoders diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs new file mode 100644 index 000000000000..b9498e8ec874 --- /dev/null +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -0,0 +1,64 @@ +use crate::basic::Type; +use crate::data_type::private::ParquetValueType; +use crate::schema::types::ColumnDescriptor; + +/// A helper to estimate the size of plain encoding of the values +/// that were written to the dictionary encoder. +/// +/// This is used to enhance the dictionary fallback heuristic with the logic +/// that the writer should fall back to the plain encoding when at a certain point, +/// e.g. after encoding the first batch, the total size of unencoded data +/// is calculated as smaller than `(encodedSize + dictionarySize)`. +pub struct PlainDataSizeCounter { + raw_data_byte_size: usize, + // Cached type length to improve performance for fixed-length types. + type_length: usize, +} + +impl PlainDataSizeCounter { + pub fn new(desc: &ColumnDescriptor) -> Self { + Self { + raw_data_byte_size: 0, + type_length: desc.type_length() as usize, + } + } + + /// Updates the counter with the given slice. + pub fn update(&mut self, values: &[T]) { + let raw_size = match T::PHYSICAL_TYPE { + Type::BOOLEAN => 1 * values.len(), + Type::INT32 | Type::FLOAT => 4 * values.len(), + Type::INT64 | Type::DOUBLE => 8 * values.len(), + Type::INT96 => 12 * values.len(), + Type::BYTE_ARRAY => { + // For variable-length types, the length prefix and the actual data are are encoded. + values + .iter() + .map(|value| { + let (base_size, num_elements) = value.dict_encoding_size(); + base_size + num_elements + }) + .sum() + } + Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), + }; + self.raw_data_byte_size = self.raw_data_byte_size.saturating_add(raw_size); + } + + /// Like `update`, but specialized for byte array data exposed by Arrow + /// array accessors. + #[cfg(feature = "arrow")] + #[inline] + pub fn update_byte_array(&mut self, value: &[u8]) { + // For variable-length types, the length prefix and the actual data are are encoded. + let raw_size = std::mem::size_of::() + value.len(); + self.raw_data_byte_size = self.raw_data_byte_size.saturating_add(raw_size); + } + + /// Returns the total size in bytes of values passed to `update` as if they were written + /// in plain encoding, without a dictionary. + #[inline] + pub fn uncompressed_data_size(&self) -> usize { + self.raw_data_byte_size + } +} diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index 65630cfed218..333d25fad3aa 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -35,6 +35,8 @@ pub const DEFAULT_WRITER_VERSION: WriterVersion = WriterVersion::PARQUET_1_0; pub const DEFAULT_COMPRESSION: Compression = Compression::UNCOMPRESSED; /// Default value for [`WriterProperties::dictionary_enabled`] pub const DEFAULT_DICTIONARY_ENABLED: bool = true; +/// Default value for [`WriterProperties::dictionary_fallback`] +pub const DEFAULT_DICTIONARY_FALLBACK: DictionaryFallback = DictionaryFallback::OnPageSizeLimit; /// Default value for [`WriterProperties::dictionary_page_size_limit`] pub const DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT: usize = DEFAULT_PAGE_SIZE; /// Default value for [`WriterProperties::data_page_row_count_limit`] @@ -491,6 +493,17 @@ impl WriterProperties { .unwrap_or(DEFAULT_DICTIONARY_ENABLED) } + /// Returns the dictionary fallback behavior for a column. + /// + /// For more details see [`WriterPropertiesBuilder::set_dictionary_fallback`] + pub fn dictionary_fallback(&self, col: &ColumnPath) -> DictionaryFallback { + self.column_properties + .get(col) + .and_then(|c| c.dictionary_fallback()) + .or_else(|| self.default_column_properties.dictionary_fallback()) + .unwrap_or(DEFAULT_DICTIONARY_FALLBACK) + } + /// Returns which statistics are written for a column. /// /// For more details see [`WriterPropertiesBuilder::set_statistics_enabled`] @@ -915,6 +928,16 @@ impl WriterPropertiesBuilder { self } + /// Sets default behavior to control dictionary fallback for all columns + /// (defaults to [`OnPageSizeLimit`] via [`DEFAULT_DICTIONARY_FALLBACK`]). + /// + /// [`OnPageSizeLimit`]: DictionaryFallback::OnPageSizeLimit + pub fn set_dictionary_fallback(mut self, behavior: DictionaryFallback) -> Self { + self.default_column_properties + .set_dictionary_fallback(behavior); + self + } + /// Sets best effort maximum dictionary page size, in bytes (defaults to `1024 * 1024` /// via [`DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT`]). /// @@ -1073,6 +1096,18 @@ impl WriterPropertiesBuilder { self } + /// Sets dictionary fallback behavior for a specific column. + /// + /// Takes precedence over [`Self::set_dictionary_fallback`]. + pub fn set_column_dictionary_fallback( + mut self, + col: ColumnPath, + behavior: DictionaryFallback, + ) -> Self { + self.get_mut_props(col).set_dictionary_fallback(behavior); + self + } + /// Sets dictionary page size limit for a specific column. /// /// Takes precedence over [`Self::set_dictionary_page_size_limit`]. @@ -1163,6 +1198,24 @@ impl From for WriterPropertiesBuilder { } } +/// Controls the behavior of the writer in deciding whether to fall back to non-dictionary encoding. +#[derive(Debug, Clone, Copy, Eq, PartialEq)] +#[non_exhaustive] +pub enum DictionaryFallback { + /// Fall back to non-dictionary encoding only if the dictionary page size limit is exceeded. + OnPageSizeLimit, + /// Fall back to non-dictionary encoding if the dictionary page size limit is exceeded + /// or if the dictionary encoding upon encoding a write batch is larger than the plain + /// encoding for the same data. + OnUnfavorableCompression, +} + +impl Default for DictionaryFallback { + fn default() -> Self { + DEFAULT_DICTIONARY_FALLBACK + } +} + /// Controls the level of statistics to be computed by the writer and stored in /// the parquet file. /// @@ -1278,6 +1331,7 @@ struct ColumnProperties { data_page_size_limit: Option, dictionary_page_size_limit: Option, dictionary_enabled: Option, + dictionary_fallback: Option, statistics_enabled: Option, write_page_header_statistics: Option, /// bloom filter related properties @@ -1323,6 +1377,11 @@ impl ColumnProperties { self.dictionary_page_size_limit = Some(value); } + /// Sets the dictionary fallback behavior for this column. + fn set_dictionary_fallback(&mut self, value: DictionaryFallback) { + self.dictionary_fallback = Some(value); + } + /// Sets the statistics level for this column. fn set_statistics_enabled(&mut self, enabled: EnabledStatistics) { self.statistics_enabled = Some(enabled); @@ -1392,6 +1451,11 @@ impl ColumnProperties { self.dictionary_page_size_limit } + /// Returns optional dictionary fallback behavior for this column. + fn dictionary_fallback(&self) -> Option { + self.dictionary_fallback + } + /// Returns optional data page size limit for this column. fn data_page_size_limit(&self) -> Option { self.data_page_size_limit From 3aacbe5cc61de6261804fcfc2598af4e23776373 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 13 Apr 2026 12:50:49 +0300 Subject: [PATCH 02/26] test: fallback with OnUnfavorableCompression --- parquet/src/arrow/arrow_writer/mod.rs | 68 +++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index d11f086557be..d8e24aad944e 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -2572,6 +2572,74 @@ mod tests { ); } + #[test] + fn arrow_writer_dictionary_fallback_on_unfavorable_compression() { + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let mut builder = StringBuilder::with_capacity(100, 329 * 10_000); + + // Generate an array of 10 unique 10 character strings. + // This results in a dictionary encoding larger than the plain encoded data, + // which should trigger a fallback to PLAIN encoding. + for i in 0..10 { + let value = i + .to_string() + .repeat(10) + .chars() + .take(10) + .collect::(); + + builder.append_value(value); + } + + let array = Arc::new(builder.finish()); + + let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + + let file = tempfile::tempfile().unwrap(); + + // Set dictionary fallback to trigger fallback to PLAIN encoding on unfavorable compression + let props = WriterProperties::builder() + .set_dictionary_fallback(DictionaryFallback::OnUnfavorableCompression) + .set_data_page_size_limit(1) + .set_write_batch_size(1) + .build(); + + let mut writer = + ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props)) + .expect("Unable to write file"); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let options = ReadOptionsBuilder::new().with_page_index().build(); + let reader = + SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(); + + let column = reader.metadata().row_group(0).columns(); + + assert_eq!(column.len(), 1); + + // We should write one row before falling back to PLAIN encoding so there should still be a + // dictionary page. + assert!( + column[0].dictionary_page_offset().is_some(), + "Expected a dictionary page" + ); + + assert!(reader.metadata().offset_index().is_some()); + let offset_indexes = &reader.metadata().offset_index().unwrap()[0]; + + let page_locations = offset_indexes[0].page_locations.clone(); + + // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes + // so we expect one dictionary encoded page and then a page per row thereafter. + assert_eq!( + page_locations.len(), + 10, + "Expected 10 pages but got {page_locations:#?}" + ); + } + #[test] fn arrow_writer_float_nans() { let f16_field = Field::new("a", DataType::Float16, false); From 83ded362c90bbd70019b8366c856f17485f7407f Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 13 Apr 2026 16:37:27 +0300 Subject: [PATCH 03/26] refactor(parquet): simplify dict_encoding_size The method as it is used now does not need to return the tuple of base size and number of elements, instead just compute the encoded size as appropriate for the type. --- parquet/src/data_type.rs | 31 +++++++++++++------ .../src/encodings/encoding/dict_encoder.rs | 5 +-- .../src/encodings/encoding/plain_counter.rs | 8 +---- 3 files changed, 23 insertions(+), 21 deletions(-) diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index d8c7b9201389..a331412428f7 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -721,10 +721,8 @@ pub(crate) mod private { fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result; - /// Return the encoded size for a type - fn dict_encoding_size(&self) -> (usize, usize) { - (std::mem::size_of::(), 1) - } + /// Return the size in bytes for the value encoded in the dictionary. + fn dict_encoding_size(&self) -> usize; /// Return the number of variable length bytes in a given slice of data /// @@ -803,6 +801,10 @@ pub(crate) mod private { Ok(values_read) } + fn dict_encoding_size(&self) -> usize { + 1 + } + #[inline] fn as_i64(&self) -> Result { Ok(*self as i64) @@ -887,6 +889,10 @@ pub(crate) mod private { Ok(num_values) } + fn dict_encoding_size(&self) -> usize { + std::mem::size_of::() + } + #[inline] fn as_i64(&$self) -> Result { $as_i64 @@ -984,6 +990,10 @@ pub(crate) mod private { Ok(num_values) } + fn dict_encoding_size(&self) -> usize { + 12 + } + #[inline] fn as_any(&self) -> &dyn std::any::Any { self @@ -1071,9 +1081,8 @@ pub(crate) mod private { Ok(num_values) } - #[inline] - fn dict_encoding_size(&self) -> (usize, usize) { - (std::mem::size_of::(), self.len()) + fn dict_encoding_size(&self) -> usize { + 4 + self.len() } #[inline] @@ -1171,9 +1180,11 @@ pub(crate) mod private { Ok(num_values) } - #[inline] - fn dict_encoding_size(&self) -> (usize, usize) { - (std::mem::size_of::(), self.len()) + fn dict_encoding_size(&self) -> usize { + // The encoding of fixed-length byte arrays only encodes the bytes + // without a length prefix. In practice, the length of fixed-length + // column values is taken from the column metadata and this call is avoided. + self.len() } #[inline] diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 79a1f247670c..0cd8d1f0a885 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -49,12 +49,9 @@ impl Storage for KeyStorage { } fn push(&mut self, value: &Self::Value) -> Self::Key { - let (base_size, num_elements) = value.dict_encoding_size(); - let unique_size = match T::get_physical_type() { - Type::BYTE_ARRAY => base_size + num_elements, Type::FIXED_LEN_BYTE_ARRAY => self.type_length, - _ => base_size, + _ => value.dict_encoding_size(), }; self.size_in_bytes += unique_size; diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index b9498e8ec874..f7a80f357e1f 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -32,13 +32,7 @@ impl PlainDataSizeCounter { Type::INT96 => 12 * values.len(), Type::BYTE_ARRAY => { // For variable-length types, the length prefix and the actual data are are encoded. - values - .iter() - .map(|value| { - let (base_size, num_elements) = value.dict_encoding_size(); - base_size + num_elements - }) - .sum() + values.iter().map(|value| value.dict_encoding_size()).sum() } Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), }; From 674a1b05d48fc05f120f60be100ee62c2fb4b520 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 13 Apr 2026 20:25:08 +0300 Subject: [PATCH 04/26] fix: unset plain data counter when flushing dict It's not really necessary because the counter is only used while dictionary encoding is enabled, but it's good to safeguard against future refactoring. --- parquet/src/arrow/arrow_writer/byte_array.rs | 2 ++ parquet/src/column/writer/encoder.rs | 2 ++ 2 files changed, 4 insertions(+) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index c0bfc9c1d519..7cad49568332 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -550,6 +550,8 @@ impl ColumnValueEncoder for ByteArrayEncoder { )); } + self.plain_data_size_counter = None; + Ok(Some(encoder.flush_dict_page())) } _ => Ok(None), diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index a808cb6286e5..00855193e622 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -322,6 +322,8 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { )); } + self.plain_data_size_counter = None; + let buf = encoder.write_dict()?; Ok(Some(DictionaryPage { From 13e042ea2a8d58fd1f122e091c85f1b84c372079 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 13 Apr 2026 20:53:11 +0300 Subject: [PATCH 05/26] chore: fix clippy --- parquet/src/encodings/encoding/plain_counter.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index f7a80f357e1f..49545085a24e 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -26,7 +26,7 @@ impl PlainDataSizeCounter { /// Updates the counter with the given slice. pub fn update(&mut self, values: &[T]) { let raw_size = match T::PHYSICAL_TYPE { - Type::BOOLEAN => 1 * values.len(), + Type::BOOLEAN => values.len(), Type::INT32 | Type::FLOAT => 4 * values.len(), Type::INT64 | Type::DOUBLE => 8 * values.len(), Type::INT96 => 12 * values.len(), From becdcef9ae8f5bfc45cf8845e8172ba1d2e360e9 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 13 Apr 2026 20:58:24 +0300 Subject: [PATCH 06/26] chore: license on plain_counter --- parquet/src/encodings/encoding/plain_counter.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index 49545085a24e..327bbe5cacfc 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -1,3 +1,20 @@ +// 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. + use crate::basic::Type; use crate::data_type::private::ParquetValueType; use crate::schema::types::ColumnDescriptor; From 1de7b012e763afc065cf275223a60cb3d91ec16e Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 14 Apr 2026 19:19:01 +0300 Subject: [PATCH 07/26] refactor(parquet): dict_encoded_size cleanup In dict_encoded_size, use size_of and symbolic constants instead of hardcoded values. For the boolean type, backstop with a panic since any returned value would be bogus for this method's contract (as plain encoding for BOOLEAN is bit-packed). The dictionary encoding is never used for boolean values. --- parquet/src/data_type.rs | 9 ++++++--- parquet/src/encodings/encoding/mod.rs | 4 +--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index a331412428f7..82468bb0a7c1 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -57,6 +57,9 @@ const MICROSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MICROSECONDS; const NANOSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * NANOSECONDS; impl Int96 { + /// Size of an INT96 value in bytes. + const SIZE_IN_BYTES: usize = 12; + /// Creates new INT96 type struct with no data set. pub fn new() -> Self { Self { value: [0; 3] } @@ -802,7 +805,7 @@ pub(crate) mod private { } fn dict_encoding_size(&self) -> usize { - 1 + panic!("Dictionary encoding should not be used for BOOLEAN type") } #[inline] @@ -991,7 +994,7 @@ pub(crate) mod private { } fn dict_encoding_size(&self) -> usize { - 12 + Self::SIZE_IN_BYTES } #[inline] @@ -1082,7 +1085,7 @@ pub(crate) mod private { } fn dict_encoding_size(&self) -> usize { - 4 + self.len() + std::mem::size_of::() + self.len() } #[inline] diff --git a/parquet/src/encodings/encoding/mod.rs b/parquet/src/encodings/encoding/mod.rs index 21b2caed6212..848c88b92289 100644 --- a/parquet/src/encodings/encoding/mod.rs +++ b/parquet/src/encodings/encoding/mod.rs @@ -813,7 +813,6 @@ mod tests { #[test] fn test_bool() { BoolType::test(Encoding::PLAIN, TEST_SET_SIZE, -1); - BoolType::test(Encoding::PLAIN_DICTIONARY, TEST_SET_SIZE, -1); BoolType::test(Encoding::RLE, TEST_SET_SIZE, -1); } @@ -881,8 +880,7 @@ mod tests { assert_eq!(encoder.dict_encoded_size(), expected_size); } - // Only 2 variations of values 1 byte each - run_test::(-1, &[true, false, true, false, true], 2); + // Dictionary encoding is not supported for BoolType, so we don't test it here. run_test::(-1, &[1i32, 2i32, 3i32, 4i32, 5i32], 20); run_test::(-1, &[1i64, 2i64, 3i64, 4i64, 5i64], 40); run_test::(-1, &[1f32, 2f32, 3f32, 4f32, 5f32], 20); From 701ff2bffc19394479c1699079f79ee1abc1d378 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 14 Apr 2026 22:31:16 +0300 Subject: [PATCH 08/26] chore(parquet): rename uncompressed_data_size To plain_encoded_data_size as suggested in review. --- parquet/src/arrow/arrow_writer/byte_array.rs | 4 ++-- parquet/src/column/writer/encoder.rs | 6 +++--- parquet/src/column/writer/mod.rs | 2 +- parquet/src/encodings/encoding/plain_counter.rs | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 7cad49568332..75b5f154cec8 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -525,9 +525,9 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(self.dict_encoder.as_ref()?.estimated_dict_page_size()) } - fn uncompressed_data_size(&self) -> Option { + fn plain_encoded_data_size(&self) -> Option { let counter = self.plain_data_size_counter.as_ref()?; - Some(counter.uncompressed_data_size()) + Some(counter.plain_encoded_data_size()) } /// Returns an estimate of the data page size in bytes diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 00855193e622..ff32d8f4bc68 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -113,7 +113,7 @@ pub trait ColumnValueEncoder { /// that would be written without a dictionary. /// If there is no dictionary, or the data size statistic is not available, /// returns `None`. - fn uncompressed_data_size(&self) -> Option; + fn plain_encoded_data_size(&self) -> Option; /// Flush the dictionary page for this column chunk if any. Any subsequent calls to /// [`Self::write`] will not be dictionary encoded @@ -301,9 +301,9 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Some(self.dict_encoder.as_ref()?.dict_encoded_size()) } - fn uncompressed_data_size(&self) -> Option { + fn plain_encoded_data_size(&self) -> Option { let counter = self.plain_data_size_counter.as_ref()?; - Some(counter.uncompressed_data_size()) + Some(counter.plain_encoded_data_size()) } fn estimated_data_page_size(&self) -> usize { diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 944ce8df17bc..a57b305f3461 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -765,7 +765,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { // Second check, if enabled: the compression heuristic. // For similar logic in parquet-java, // see DictionaryValuesWriter.isCompressionSatisfying - if let Some(raw_size) = self.encoder.uncompressed_data_size() { + if let Some(raw_size) = self.encoder.plain_encoded_data_size() { let encoded_size = self.encoder.estimated_data_page_size(); if encoded_size + dict_size >= raw_size { return true; diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index 327bbe5cacfc..dadbcf5c61ed 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -69,7 +69,7 @@ impl PlainDataSizeCounter { /// Returns the total size in bytes of values passed to `update` as if they were written /// in plain encoding, without a dictionary. #[inline] - pub fn uncompressed_data_size(&self) -> usize { + pub fn plain_encoded_data_size(&self) -> usize { self.raw_data_byte_size } } From 1b6dd3756c447111baf78af969875b7790909070 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Wed, 15 Apr 2026 13:16:24 +0300 Subject: [PATCH 09/26] test: compression fallback wins Modify the compression fallback test to illustrate the benefit in an admittedly differently contrived case. The heuristic borrowed from parquet-java is still not ideal for all cases, so we'll need more configurability. --- parquet/src/arrow/arrow_writer/mod.rs | 79 ++++++++++++++------------- 1 file changed, 42 insertions(+), 37 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index d8e24aad944e..35b3e04dff73 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4857,6 +4857,15 @@ mod tests { assert_eq!(chunk_page_stats, file_page_stats); } + fn get_dict_page_size(meta: &ColumnChunkMetaData, data: Bytes) -> usize { + let mut reader = SerializedPageReader::new(Arc::new(data), meta, 0, None).unwrap(); + let page = reader.get_next_page().unwrap().unwrap(); + match page { + Page::DictionaryPage { buf, .. } => buf.len(), + _ => panic!("expected DictionaryPage"), + } + } + #[test] fn test_different_dict_page_size_limit() { let array = Arc::new(Int64Array::from_iter(0..1024 * 1024)); @@ -4881,60 +4890,56 @@ mod tests { let col0_meta = metadata.row_group(0).column(0); let col1_meta = metadata.row_group(0).column(1); - let get_dict_page_size = move |meta: &ColumnChunkMetaData| { - let mut reader = - SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap(); - let page = reader.get_next_page().unwrap().unwrap(); - match page { - Page::DictionaryPage { buf, .. } => buf.len(), - _ => panic!("expected DictionaryPage"), - } - }; - - assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024); - assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4); + assert_eq!(get_dict_page_size(col0_meta, data.clone()), 1024 * 1024); + assert_eq!(get_dict_page_size(col1_meta, data.clone()), 1024 * 1024 * 4); } #[test] fn test_dict_page_size_decided_by_compression_fallback() { - let array = Arc::new(Int64Array::from_iter(0..1024 * 1024)); - let schema = Arc::new(Schema::new(vec![ - Field::new("col0", arrow_schema::DataType::Int64, false), - Field::new("col1", arrow_schema::DataType::Int64, false), - ])); - let batch = - arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap(); + // Generate values that are well dispersed across a range approximating (0..256 * 1024) + let array = Arc::new(Int32Array::from_iter( + (0i32..1024 * 1024).map(|x| x.wrapping_mul(163019) % 262139), + )); + let schema = Arc::new(Schema::new(vec![Field::new( + "col0", + arrow_schema::DataType::Int32, + false, + )])); + let batch = arrow_array::RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + + let props = WriterProperties::builder() + .set_dictionary_page_size_limit(1024 * 1024) + .build(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap(); + writer.write(&batch).unwrap(); + let data = Bytes::from(writer.into_inner().unwrap()); + + // println!("file length, dictionary: {}", data.len()); + + let mut metadata = ParquetMetaDataReader::new(); + metadata.try_parse(&data).unwrap(); + let metadata = metadata.finish().unwrap(); + let full_dict_meta = metadata.row_group(0).column(0); + assert_eq!(get_dict_page_size(full_dict_meta, data.clone()), 1_048_576); let props = WriterProperties::builder() .set_dictionary_page_size_limit(1024 * 1024) - .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4) .set_column_dictionary_fallback( - ColumnPath::from("col1"), + ColumnPath::from("col0"), DictionaryFallback::OnUnfavorableCompression, ) .build(); - let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap(); writer.write(&batch).unwrap(); let data = Bytes::from(writer.into_inner().unwrap()); + // println!("file length, fallback: {}", data.len()); + let mut metadata = ParquetMetaDataReader::new(); metadata.try_parse(&data).unwrap(); let metadata = metadata.finish().unwrap(); - let col0_meta = metadata.row_group(0).column(0); - let col1_meta = metadata.row_group(0).column(1); - - let get_dict_page_size = move |meta: &ColumnChunkMetaData| { - let mut reader = - SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap(); - let page = reader.get_next_page().unwrap().unwrap(); - match page { - Page::DictionaryPage { buf, .. } => buf.len(), - _ => panic!("expected DictionaryPage"), - } - }; - - assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024); - assert_eq!(get_dict_page_size(col1_meta), 8192); + let fallback_meta = metadata.row_group(0).column(0); + assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 4096); } struct WriteBatchesShape { From da73778eebc80299aaced8f123840bc173785b56 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 07:10:36 +0300 Subject: [PATCH 10/26] feat: DictionaryFallback::OnUnfavorableAfter Change OnUnfavorableCompression to the variant carrying a value specicifying the minimal sample length, to give the user more control over when to fall back to the plain encoding. --- parquet/src/arrow/arrow_writer/byte_array.rs | 2 +- parquet/src/arrow/arrow_writer/mod.rs | 6 +++--- parquet/src/column/writer/encoder.rs | 2 +- parquet/src/column/writer/mod.rs | 21 +++++++++++++++----- parquet/src/file/properties.rs | 8 ++++---- 5 files changed, 25 insertions(+), 14 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 75b5f154cec8..864aebafade1 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -451,7 +451,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { DictionaryFallback::OnPageSizeLimit => None, - DictionaryFallback::OnUnfavorableCompression => { + DictionaryFallback::OnUnfavorableAfter(_) => { if dictionary.is_some() { Some(PlainDataSizeCounter::new(descr)) } else { diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 35b3e04dff73..6d14a7115e31 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -2600,7 +2600,7 @@ mod tests { // Set dictionary fallback to trigger fallback to PLAIN encoding on unfavorable compression let props = WriterProperties::builder() - .set_dictionary_fallback(DictionaryFallback::OnUnfavorableCompression) + .set_dictionary_fallback(DictionaryFallback::OnUnfavorableAfter(1)) .set_data_page_size_limit(1) .set_write_batch_size(1) .build(); @@ -4926,7 +4926,7 @@ mod tests { .set_dictionary_page_size_limit(1024 * 1024) .set_column_dictionary_fallback( ColumnPath::from("col0"), - DictionaryFallback::OnUnfavorableCompression, + DictionaryFallback::OnUnfavorableAfter(8192), ) .build(); let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap(); @@ -4939,7 +4939,7 @@ mod tests { metadata.try_parse(&data).unwrap(); let metadata = metadata.finish().unwrap(); let fallback_meta = metadata.row_group(0).column(0); - assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 4096); + assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 8192); } struct WriteBatchesShape { diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index ff32d8f4bc68..5979441b12f8 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -212,7 +212,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { DictionaryFallback::OnPageSizeLimit => None, - DictionaryFallback::OnUnfavorableCompression => { + DictionaryFallback::OnUnfavorableAfter(_) => { if dict_encoder.is_some() { Some(PlainDataSizeCounter::new(descr)) } else { diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index a57b305f3461..0c884b4c7e8a 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -43,7 +43,7 @@ use crate::file::metadata::{ OffsetIndexBuilder, PageEncodingStats, }; use crate::file::properties::{ - EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion, + DictionaryFallback, EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion, }; use crate::file::statistics::{Statistics, ValueStatistics}; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; @@ -337,6 +337,7 @@ pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> { descr: ColumnDescPtr, props: WriterPropertiesPtr, statistics_enabled: EnabledStatistics, + dict_fallback_sample_len: usize, page_writer: Box, codec: Compression, @@ -379,6 +380,13 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { let compressor = create_codec(codec, &codec_options).unwrap(); let encoder = E::try_new(&descr, props.as_ref()).unwrap(); + let dict_fallback_sample_len = match props.dictionary_fallback(descr.path()) { + DictionaryFallback::OnUnfavorableAfter(sample_len) if encoder.has_dictionary() => { + sample_len + } + _ => 0, + }; + let statistics_enabled = props.statistics_enabled(descr.path()); let mut encodings = BTreeSet::new(); @@ -416,6 +424,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { descr, props, statistics_enabled, + dict_fallback_sample_len, page_writer, codec, compressor, @@ -765,10 +774,12 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { // Second check, if enabled: the compression heuristic. // For similar logic in parquet-java, // see DictionaryValuesWriter.isCompressionSatisfying - if let Some(raw_size) = self.encoder.plain_encoded_data_size() { - let encoded_size = self.encoder.estimated_data_page_size(); - if encoded_size + dict_size >= raw_size { - return true; + if self.encoder.num_values() >= self.dict_fallback_sample_len { + if let Some(raw_size) = self.encoder.plain_encoded_data_size() { + let encoded_size = self.encoder.estimated_data_page_size(); + if encoded_size + dict_size >= raw_size { + return true; + } } } diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index 333d25fad3aa..2873b1cfa583 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -1204,10 +1204,10 @@ impl From for WriterPropertiesBuilder { pub enum DictionaryFallback { /// Fall back to non-dictionary encoding only if the dictionary page size limit is exceeded. OnPageSizeLimit, - /// Fall back to non-dictionary encoding if the dictionary page size limit is exceeded - /// or if the dictionary encoding upon encoding a write batch is larger than the plain - /// encoding for the same data. - OnUnfavorableCompression, + /// Fall back to non-dictionary encoding if the dictionary page size limit is exceeded, + /// or if the dictionary encoding upon encoding at least the given number of values + /// is larger than the plain encoding for the same data. + OnUnfavorableAfter(usize), } impl Default for DictionaryFallback { From b3927380dfd2220e20c1c06b0d0d76b29c5143f4 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 14:16:56 +0300 Subject: [PATCH 11/26] refactor: more compact plain data counter init --- parquet/src/arrow/arrow_writer/byte_array.rs | 21 ++++++++++---------- parquet/src/column/writer/encoder.rs | 17 ++++++++-------- 2 files changed, 18 insertions(+), 20 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 864aebafade1..0d46c9abf83e 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -445,19 +445,18 @@ impl ColumnValueEncoder for ByteArrayEncoder { where Self: Sized, { - let dictionary = props + let dict_encoder = props .dictionary_enabled(descr.path()) .then(DictEncoder::default); - let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { - DictionaryFallback::OnPageSizeLimit => None, - DictionaryFallback::OnUnfavorableAfter(_) => { - if dictionary.is_some() { - Some(PlainDataSizeCounter::new(descr)) - } else { - None - } - } + let plain_data_size_counter = if dict_encoder.is_some() + && matches!( + props.dictionary_fallback(descr.path()), + DictionaryFallback::OnUnfavorableAfter(_) + ) { + Some(PlainDataSizeCounter::new(descr)) + } else { + None }; let fallback = FallbackEncoder::new(descr, props)?; @@ -473,7 +472,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { statistics_enabled, bloom_filter, bloom_filter_target_fpp, - dict_encoder: dictionary, + dict_encoder, plain_data_size_counter, min_value: None, max_value: None, diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 5979441b12f8..9dc9128a0d38 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -210,15 +210,14 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let dict_supported = props.dictionary_enabled(descr.path()) && has_dictionary_support(T::get_physical_type(), props); let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); - let plain_data_size_counter = match props.dictionary_fallback(descr.path()) { - DictionaryFallback::OnPageSizeLimit => None, - DictionaryFallback::OnUnfavorableAfter(_) => { - if dict_encoder.is_some() { - Some(PlainDataSizeCounter::new(descr)) - } else { - None - } - } + let plain_data_size_counter = if dict_encoder.is_some() + && matches!( + props.dictionary_fallback(descr.path()), + DictionaryFallback::OnUnfavorableAfter(_) + ) { + Some(PlainDataSizeCounter::new(descr)) + } else { + None }; // Set either main encoder or fallback encoder. From bd914bb556e135bc0d8771e8036e31c7dd1413c6 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 17:51:54 +0300 Subject: [PATCH 12/26] test(parquet): fix up compression fallback --- parquet/src/arrow/arrow_writer/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 6d14a7115e31..64162cad420d 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4914,7 +4914,7 @@ mod tests { writer.write(&batch).unwrap(); let data = Bytes::from(writer.into_inner().unwrap()); - // println!("file length, dictionary: {}", data.len()); + println!("file length, dictionary: {}", data.len()); let mut metadata = ParquetMetaDataReader::new(); metadata.try_parse(&data).unwrap(); @@ -4933,13 +4933,13 @@ mod tests { writer.write(&batch).unwrap(); let data = Bytes::from(writer.into_inner().unwrap()); - // println!("file length, fallback: {}", data.len()); + println!("file length, fallback: {}", data.len()); let mut metadata = ParquetMetaDataReader::new(); metadata.try_parse(&data).unwrap(); let metadata = metadata.finish().unwrap(); let fallback_meta = metadata.row_group(0).column(0); - assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 8192); + assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 32_768); } struct WriteBatchesShape { From 898e7e582701ef0643809a3ff9a159bc61fa4615 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Thu, 16 Apr 2026 21:25:44 +0300 Subject: [PATCH 13/26] test(parquet): test dictionary_fallback property --- parquet/src/file/properties.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index 2873b1cfa583..e6345356eafe 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -1649,6 +1649,10 @@ mod tests { props.dictionary_enabled(&ColumnPath::from("col")), DEFAULT_DICTIONARY_ENABLED ); + assert_eq!( + props.dictionary_fallback(&ColumnPath::from("col")), + DEFAULT_DICTIONARY_FALLBACK + ); assert_eq!( props.statistics_enabled(&ColumnPath::from("col")), DEFAULT_STATISTICS_ENABLED @@ -1730,11 +1734,16 @@ mod tests { .set_encoding(Encoding::DELTA_BINARY_PACKED) .set_compression(Compression::GZIP(Default::default())) .set_dictionary_enabled(false) + .set_dictionary_fallback(DictionaryFallback::OnUnfavorableAfter(1024)) .set_statistics_enabled(EnabledStatistics::None) // specific column settings .set_column_encoding(ColumnPath::from("col"), Encoding::RLE) .set_column_compression(ColumnPath::from("col"), Compression::SNAPPY) .set_column_dictionary_enabled(ColumnPath::from("col"), true) + .set_column_dictionary_fallback( + ColumnPath::from("col"), + DictionaryFallback::OnUnfavorableAfter(2048), + ) .set_column_statistics_enabled(ColumnPath::from("col"), EnabledStatistics::Chunk) .set_column_bloom_filter_enabled(ColumnPath::from("col"), true) .set_column_bloom_filter_ndv(ColumnPath::from("col"), 100_u64) @@ -1764,6 +1773,10 @@ mod tests { Compression::GZIP(Default::default()) ); assert!(!props.dictionary_enabled(&ColumnPath::from("a"))); + assert_eq!( + props.dictionary_fallback(&ColumnPath::from("a")), + DictionaryFallback::OnUnfavorableAfter(1024) + ); assert_eq!( props.statistics_enabled(&ColumnPath::from("a")), EnabledStatistics::None @@ -1778,6 +1791,10 @@ mod tests { Compression::SNAPPY ); assert!(props.dictionary_enabled(&ColumnPath::from("col"))); + assert_eq!( + props.dictionary_fallback(&ColumnPath::from("col")), + DictionaryFallback::OnUnfavorableAfter(2048) + ); assert_eq!( props.statistics_enabled(&ColumnPath::from("col")), EnabledStatistics::Chunk From 1e8117343c45c2ebc28120606ac5cf6735b1ee00 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 17 Apr 2026 00:24:37 +0300 Subject: [PATCH 14/26] chore: suggestions from code review Co-authored-by: Ed Seidl --- parquet/src/arrow/arrow_writer/byte_array.rs | 3 +-- parquet/src/column/writer/encoder.rs | 3 +-- parquet/src/data_type.rs | 2 +- parquet/src/encodings/encoding/plain_counter.rs | 2 +- 4 files changed, 4 insertions(+), 6 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 0d46c9abf83e..5792fe4d4a69 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -525,8 +525,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { } fn plain_encoded_data_size(&self) -> Option { - let counter = self.plain_data_size_counter.as_ref()?; - Some(counter.plain_encoded_data_size()) + Some(self.plain_data_size_counter.as_ref()?.plain_encoded_data_size()) } /// Returns an estimate of the data page size in bytes diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 9dc9128a0d38..e95d6911a9d3 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -301,8 +301,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { } fn plain_encoded_data_size(&self) -> Option { - let counter = self.plain_data_size_counter.as_ref()?; - Some(counter.plain_encoded_data_size()) + Some(self.plain_data_size_counter.as_ref()?.plain_encoded_data_size()) } fn estimated_data_page_size(&self) -> usize { diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index 82468bb0a7c1..57839e98c386 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -58,7 +58,7 @@ const NANOSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * NANOSECONDS; impl Int96 { /// Size of an INT96 value in bytes. - const SIZE_IN_BYTES: usize = 12; + pub const SIZE_IN_BYTES: usize = std::mem::size_of::<[u32; 3]>(); /// Creates new INT96 type struct with no data set. pub fn new() -> Self { diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index dadbcf5c61ed..ce36e8fb07df 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -46,7 +46,7 @@ impl PlainDataSizeCounter { Type::BOOLEAN => values.len(), Type::INT32 | Type::FLOAT => 4 * values.len(), Type::INT64 | Type::DOUBLE => 8 * values.len(), - Type::INT96 => 12 * values.len(), + Type::INT96 => Int96::SIZE_IN_BYTES * values.len(), Type::BYTE_ARRAY => { // For variable-length types, the length prefix and the actual data are are encoded. values.iter().map(|value| value.dict_encoding_size()).sum() From 34b819d457fa3b5c2d219605d8a6c489d1b133f2 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 17 Apr 2026 14:37:26 +0300 Subject: [PATCH 15/26] chore: add missing import The web-edited suggestions did not include the import. --- parquet/src/encodings/encoding/plain_counter.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index ce36e8fb07df..681bb006627c 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -16,6 +16,7 @@ // under the License. use crate::basic::Type; +use crate::data_type::Int96; use crate::data_type::private::ParquetValueType; use crate::schema::types::ColumnDescriptor; From 692018eaf06cf01c632f6d20cc225ac1ba235ec9 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 17 Apr 2026 14:32:41 +0300 Subject: [PATCH 16/26] chore: rename dict_encoding_size Rename the ParquetValueType::dict_encoding_size trait method to plain_encoded_size, to better reflect its usage since it's also used to calculate the comparative size of the plain encoding for the dictionary fallback heuristic. --- parquet/src/data_type.rs | 14 +++++++------- parquet/src/encodings/encoding/dict_encoder.rs | 2 +- parquet/src/encodings/encoding/plain_counter.rs | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index 57839e98c386..bf218e91bd2f 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -724,8 +724,8 @@ pub(crate) mod private { fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result; - /// Return the size in bytes for the value encoded in the dictionary. - fn dict_encoding_size(&self) -> usize; + /// Return the size in bytes for the value encoded in the plain encoding. + fn plain_encoded_size(&self) -> usize; /// Return the number of variable length bytes in a given slice of data /// @@ -804,7 +804,7 @@ pub(crate) mod private { Ok(values_read) } - fn dict_encoding_size(&self) -> usize { + fn plain_encoded_size(&self) -> usize { panic!("Dictionary encoding should not be used for BOOLEAN type") } @@ -892,7 +892,7 @@ pub(crate) mod private { Ok(num_values) } - fn dict_encoding_size(&self) -> usize { + fn plain_encoded_size(&self) -> usize { std::mem::size_of::() } @@ -993,7 +993,7 @@ pub(crate) mod private { Ok(num_values) } - fn dict_encoding_size(&self) -> usize { + fn plain_encoded_size(&self) -> usize { Self::SIZE_IN_BYTES } @@ -1084,7 +1084,7 @@ pub(crate) mod private { Ok(num_values) } - fn dict_encoding_size(&self) -> usize { + fn plain_encoded_size(&self) -> usize { std::mem::size_of::() + self.len() } @@ -1183,7 +1183,7 @@ pub(crate) mod private { Ok(num_values) } - fn dict_encoding_size(&self) -> usize { + fn plain_encoded_size(&self) -> usize { // The encoding of fixed-length byte arrays only encodes the bytes // without a length prefix. In practice, the length of fixed-length // column values is taken from the column metadata and this call is avoided. diff --git a/parquet/src/encodings/encoding/dict_encoder.rs b/parquet/src/encodings/encoding/dict_encoder.rs index 0cd8d1f0a885..28abcc2e989b 100644 --- a/parquet/src/encodings/encoding/dict_encoder.rs +++ b/parquet/src/encodings/encoding/dict_encoder.rs @@ -51,7 +51,7 @@ impl Storage for KeyStorage { fn push(&mut self, value: &Self::Value) -> Self::Key { let unique_size = match T::get_physical_type() { Type::FIXED_LEN_BYTE_ARRAY => self.type_length, - _ => value.dict_encoding_size(), + _ => value.plain_encoded_size(), }; self.size_in_bytes += unique_size; diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs index 681bb006627c..34665a0fd730 100644 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ b/parquet/src/encodings/encoding/plain_counter.rs @@ -50,7 +50,7 @@ impl PlainDataSizeCounter { Type::INT96 => Int96::SIZE_IN_BYTES * values.len(), Type::BYTE_ARRAY => { // For variable-length types, the length prefix and the actual data are are encoded. - values.iter().map(|value| value.dict_encoding_size()).sum() + values.iter().map(|value| value.plain_encoded_size()).sum() } Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), }; From 4511917287bd55db451bab045fc3abfc642cad18 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 17 Apr 2026 19:20:05 +0300 Subject: [PATCH 17/26] fix: rework plain_counter to DictFallbackCounter Fix the counting logic that mistakenly relied on values that are reset with every flushed page. Move the fallback decision logic to the DictFallbackCounter implementation which supersedes PlainDataSizeCounter. --- parquet/src/arrow/arrow_writer/byte_array.rs | 41 ++++--- parquet/src/column/writer/encoder.rs | 53 +++++---- parquet/src/column/writer/mod.rs | 22 +--- .../src/encodings/encoding/dict_fallback.rs | 107 ++++++++++++++++++ parquet/src/encodings/encoding/mod.rs | 4 +- .../src/encodings/encoding/plain_counter.rs | 76 ------------- 6 files changed, 169 insertions(+), 134 deletions(-) create mode 100644 parquet/src/encodings/encoding/dict_fallback.rs delete mode 100644 parquet/src/encodings/encoding/plain_counter.rs diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 5792fe4d4a69..9aadbfc8a9d0 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -21,7 +21,7 @@ use crate::column::writer::encoder::{ ColumnValueEncoder, DataPageValues, DictionaryPage, create_bloom_filter, }; use crate::data_type::{AsBytes, ByteArray, Int32Type}; -use crate::encodings::encoding::{DeltaBitPackEncoder, Encoder, PlainDataSizeCounter}; +use crate::encodings::encoding::{DeltaBitPackEncoder, DictFallbackCounter, Encoder}; use crate::encodings::rle::RleEncoder; use crate::errors::{ParquetError, Result}; use crate::file::properties::{ @@ -423,7 +423,7 @@ impl DictEncoder { pub struct ByteArrayEncoder { fallback: FallbackEncoder, dict_encoder: Option, - plain_data_size_counter: Option, + dict_fallback_counter: Option, statistics_enabled: EnabledStatistics, min_value: Option, max_value: Option, @@ -449,14 +449,11 @@ impl ColumnValueEncoder for ByteArrayEncoder { .dictionary_enabled(descr.path()) .then(DictEncoder::default); - let plain_data_size_counter = if dict_encoder.is_some() - && matches!( - props.dictionary_fallback(descr.path()), - DictionaryFallback::OnUnfavorableAfter(_) - ) { - Some(PlainDataSizeCounter::new(descr)) - } else { - None + let dict_fallback_counter = match props.dictionary_fallback(descr.path()) { + DictionaryFallback::OnUnfavorableAfter(min_sample_len) if dict_encoder.is_some() => { + Some(DictFallbackCounter::new(descr, min_sample_len)) + } + _ => None, }; let fallback = FallbackEncoder::new(descr, props)?; @@ -473,7 +470,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { bloom_filter, bloom_filter_target_fpp, dict_encoder, - plain_data_size_counter, + dict_fallback_counter, min_value: None, max_value: None, geo_stats_accumulator, @@ -524,8 +521,14 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(self.dict_encoder.as_ref()?.estimated_dict_page_size()) } - fn plain_encoded_data_size(&self) -> Option { - Some(self.plain_data_size_counter.as_ref()?.plain_encoded_data_size()) + fn is_dict_encoding_unfavorable(&self) -> bool { + match (&self.dict_encoder, &self.dict_fallback_counter) { + (Some(encoder), Some(counter)) => { + let dict_size = encoder.estimated_dict_page_size(); + counter.is_dict_encoding_unfavorable(dict_size) + } + _ => false, + } } /// Returns an estimate of the data page size in bytes @@ -548,7 +551,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { )); } - self.plain_data_size_counter = None; + self.dict_fallback_counter = None; Ok(Some(encoder.flush_dict_page())) } @@ -561,7 +564,13 @@ impl ColumnValueEncoder for ByteArrayEncoder { let max_value = self.max_value.take(); match &mut self.dict_encoder { - Some(encoder) => Ok(encoder.flush_data_page(min_value, max_value)), + Some(encoder) => { + let data_page = encoder.flush_data_page(min_value, max_value); + if let Some(counter) = self.dict_fallback_counter.as_mut() { + counter.count_dict_encoded_data(data_page.buf.len()); + } + Ok(data_page) + } _ => self.fallback.flush_data_page(min_value, max_value), } } @@ -604,7 +613,7 @@ where match &mut encoder.dict_encoder { Some(dict_encoder) => { dict_encoder.encode(values, indices); - if let Some(counter) = encoder.plain_data_size_counter.as_mut() { + if let Some(counter) = encoder.dict_fallback_counter.as_mut() { for idx in indices { let value = values.value(*idx); counter.update_byte_array(value.as_ref()); diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index e95d6911a9d3..0cf6d0a1fbc3 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -25,7 +25,7 @@ use crate::column::writer::{ }; use crate::data_type::DataType; use crate::data_type::private::ParquetValueType; -use crate::encodings::encoding::{DictEncoder, Encoder, PlainDataSizeCounter, get_encoder}; +use crate::encodings::encoding::{DictEncoder, DictFallbackCounter, Encoder, get_encoder}; use crate::errors::{ParquetError, Result}; use crate::file::properties::{DictionaryFallback, EnabledStatistics, WriterProperties}; use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator}; @@ -109,11 +109,13 @@ pub trait ColumnValueEncoder { /// + fn estimated_data_page_size(&self) -> usize; - /// Returns the estimated size of plainly encoded data, in bytes, - /// that would be written without a dictionary. + /// Returns true if the estimated size of plainly encoded data, in bytes, + /// would be smaller than the size of data encoded with a dictionary. /// If there is no dictionary, or the data size statistic is not available, - /// returns `None`. - fn plain_encoded_data_size(&self) -> Option; + /// returns false. + /// For similar logic in parquet-java, + /// see `DictionaryValuesWriter.isCompressionSatisfying`. + fn is_dict_encoding_unfavorable(&self) -> bool; /// Flush the dictionary page for this column chunk if any. Any subsequent calls to /// [`Self::write`] will not be dictionary encoded @@ -138,7 +140,7 @@ pub trait ColumnValueEncoder { pub struct ColumnValueEncoderImpl { encoder: Box>, dict_encoder: Option>, - plain_data_size_counter: Option, + dict_fallback_counter: Option, descr: ColumnDescPtr, num_values: usize, statistics_enabled: EnabledStatistics, @@ -185,8 +187,8 @@ impl ColumnValueEncoderImpl { match &mut self.dict_encoder { Some(encoder) => { encoder.put(slice)?; - if let Some(counter) = self.plain_data_size_counter.as_mut() { - counter.update(slice); + if let Some(counter) = self.dict_fallback_counter.as_mut() { + counter.update_values(slice); } Ok(()) } @@ -210,14 +212,11 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { let dict_supported = props.dictionary_enabled(descr.path()) && has_dictionary_support(T::get_physical_type(), props); let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone())); - let plain_data_size_counter = if dict_encoder.is_some() - && matches!( - props.dictionary_fallback(descr.path()), - DictionaryFallback::OnUnfavorableAfter(_) - ) { - Some(PlainDataSizeCounter::new(descr)) - } else { - None + let dict_fallback_counter = match props.dictionary_fallback(descr.path()) { + DictionaryFallback::OnUnfavorableAfter(min_sample_len) if dict_encoder.is_some() => { + Some(DictFallbackCounter::new(descr, min_sample_len)) + } + _ => None, }; // Set either main encoder or fallback encoder. @@ -237,7 +236,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Ok(Self { encoder, dict_encoder, - plain_data_size_counter, + dict_fallback_counter, descr: descr.clone(), num_values: 0, statistics_enabled, @@ -300,8 +299,14 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Some(self.dict_encoder.as_ref()?.dict_encoded_size()) } - fn plain_encoded_data_size(&self) -> Option { - Some(self.plain_data_size_counter.as_ref()?.plain_encoded_data_size()) + fn is_dict_encoding_unfavorable(&self) -> bool { + match (&self.dict_encoder, &self.dict_fallback_counter) { + (Some(encoder), Some(counter)) => { + let dict_size = encoder.dict_encoded_size(); + counter.is_dict_encoding_unfavorable(dict_size) + } + _ => false, + } } fn estimated_data_page_size(&self) -> usize { @@ -320,7 +325,7 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { )); } - self.plain_data_size_counter = None; + self.dict_fallback_counter = None; let buf = encoder.write_dict()?; @@ -336,7 +341,13 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { fn flush_data_page(&mut self) -> Result> { let (buf, encoding) = match &mut self.dict_encoder { - Some(encoder) => (encoder.write_indices()?, Encoding::RLE_DICTIONARY), + Some(encoder) => { + let buf = encoder.write_indices()?; + if let Some(counter) = self.dict_fallback_counter.as_mut() { + counter.count_dict_encoded_data(buf.len()); + } + (buf, Encoding::RLE_DICTIONARY) + } _ => (self.encoder.flush_buffer()?, self.encoder.encoding()), }; diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index 0c884b4c7e8a..e9cb61db4fa5 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -43,7 +43,7 @@ use crate::file::metadata::{ OffsetIndexBuilder, PageEncodingStats, }; use crate::file::properties::{ - DictionaryFallback, EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion, + EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion, }; use crate::file::statistics::{Statistics, ValueStatistics}; use crate::schema::types::{ColumnDescPtr, ColumnDescriptor}; @@ -337,7 +337,6 @@ pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> { descr: ColumnDescPtr, props: WriterPropertiesPtr, statistics_enabled: EnabledStatistics, - dict_fallback_sample_len: usize, page_writer: Box, codec: Compression, @@ -380,13 +379,6 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { let compressor = create_codec(codec, &codec_options).unwrap(); let encoder = E::try_new(&descr, props.as_ref()).unwrap(); - let dict_fallback_sample_len = match props.dictionary_fallback(descr.path()) { - DictionaryFallback::OnUnfavorableAfter(sample_len) if encoder.has_dictionary() => { - sample_len - } - _ => 0, - }; - let statistics_enabled = props.statistics_enabled(descr.path()); let mut encodings = BTreeSet::new(); @@ -424,7 +416,6 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { descr, props, statistics_enabled, - dict_fallback_sample_len, page_writer, codec, compressor, @@ -772,15 +763,8 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } // Second check, if enabled: the compression heuristic. - // For similar logic in parquet-java, - // see DictionaryValuesWriter.isCompressionSatisfying - if self.encoder.num_values() >= self.dict_fallback_sample_len { - if let Some(raw_size) = self.encoder.plain_encoded_data_size() { - let encoded_size = self.encoder.estimated_data_page_size(); - if encoded_size + dict_size >= raw_size { - return true; - } - } + if self.encoder.is_dict_encoding_unfavorable() { + return true; } false diff --git a/parquet/src/encodings/encoding/dict_fallback.rs b/parquet/src/encodings/encoding/dict_fallback.rs new file mode 100644 index 000000000000..6248156c57f5 --- /dev/null +++ b/parquet/src/encodings/encoding/dict_fallback.rs @@ -0,0 +1,107 @@ +// 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. + +use crate::basic::Type; +use crate::data_type::Int96; +use crate::data_type::private::ParquetValueType; +use crate::schema::types::ColumnDescriptor; + +// TODO: it's possible to tighten the worst-case estimate for fallback encodings +// other than PLAIN, e.g. by estimating the bit widths of delta encoded values, +// possibly making use of the min/max statistics if they are available. +// This would require more complex logic in the counter, and it's not clear if +// the improvement in the heuristic would be worth it. +// The size of the plain encoding should be a reasonable bottom estimate +// for any fallback encoding. + +/// A helper to estimate the favorability of the dictionary encoding +/// compared to a pessimistic estimate of the size of data encoded without +/// the dictionary. +/// +/// This is used to enhance the dictionary fallback heuristic with the logic +/// that the writer should fall back to a non-dictionary encoding when, +/// after encoding a prescribed minimum number of values, the worst case on +/// the size of data encoded without the dictionary is calculated as smaller +/// than `(encodedSize + dictionarySize)`. +pub struct DictFallbackCounter { + // Estimated size of the data encoded without the dictionary, in bytes. + raw_data_size: usize, + // Size of the data encoded with the dictionary, in bytes. + encoded_data_size: usize, + // Number of values passed to the counter. + num_values: usize, + // Minimum number of values to sample before + // the counter can return a favorable estimate for fallback. + min_sample_len: usize, + // Cached type length to improve performance for fixed-length types. + type_length: usize, +} + +impl DictFallbackCounter { + pub fn new(desc: &ColumnDescriptor, min_sample_len: usize) -> Self { + Self { + raw_data_size: 0, + encoded_data_size: 0, + num_values: 0, + min_sample_len, + type_length: desc.type_length() as usize, + } + } + + /// Updates the counter with the given slice of values. + pub fn update_values(&mut self, values: &[T]) { + let raw_size = match T::PHYSICAL_TYPE { + Type::BOOLEAN => values.len(), + Type::INT32 | Type::FLOAT => 4 * values.len(), + Type::INT64 | Type::DOUBLE => 8 * values.len(), + Type::INT96 => Int96::SIZE_IN_BYTES * values.len(), + Type::BYTE_ARRAY => { + // For variable-length types, the length prefix and the actual data are are encoded. + values.iter().map(|value| value.plain_encoded_size()).sum() + } + Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), + }; + self.raw_data_size = self.raw_data_size.saturating_add(raw_size); + self.num_values += values.len(); + } + + /// Like `update_values`, but specialized for byte array data exposed by Arrow + /// array accessors. Updates the counter with the single given byte array value. + #[cfg(feature = "arrow")] + #[inline] + pub fn update_byte_array(&mut self, value: &[u8]) { + let raw_size = std::mem::size_of::() + value.len(); + self.raw_data_size = self.raw_data_size.saturating_add(raw_size); + self.num_values += 1; + } + + /// Increments the counter of the size of dictionary encoded data + /// by the given amount in bytes. + #[inline] + pub fn count_dict_encoded_data(&mut self, encoded_len: usize) { + self.encoded_data_size = self.encoded_data_size.saturating_add(encoded_len); + } + + /// Returns true if the estimated size of plainly encoded data, in bytes, + /// would not exceed the size of data encoded with a dictionary, + /// as counted by the `count_dict_encoded_data` calls made on this counter. + #[inline] + pub fn is_dict_encoding_unfavorable(&self, dict_encoded_size: usize) -> bool { + self.num_values >= self.min_sample_len + && self.raw_data_size <= dict_encoded_size.saturating_add(self.encoded_data_size) + } +} diff --git a/parquet/src/encodings/encoding/mod.rs b/parquet/src/encodings/encoding/mod.rs index 848c88b92289..6d5f5fade88a 100644 --- a/parquet/src/encodings/encoding/mod.rs +++ b/parquet/src/encodings/encoding/mod.rs @@ -32,10 +32,10 @@ use bytes::Bytes; mod byte_stream_split_encoder; mod dict_encoder; -mod plain_counter; +mod dict_fallback; pub use dict_encoder::DictEncoder; -pub use plain_counter::PlainDataSizeCounter; +pub use dict_fallback::DictFallbackCounter; // ---------------------------------------------------------------------- // Encoders diff --git a/parquet/src/encodings/encoding/plain_counter.rs b/parquet/src/encodings/encoding/plain_counter.rs deleted file mode 100644 index 34665a0fd730..000000000000 --- a/parquet/src/encodings/encoding/plain_counter.rs +++ /dev/null @@ -1,76 +0,0 @@ -// 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. - -use crate::basic::Type; -use crate::data_type::Int96; -use crate::data_type::private::ParquetValueType; -use crate::schema::types::ColumnDescriptor; - -/// A helper to estimate the size of plain encoding of the values -/// that were written to the dictionary encoder. -/// -/// This is used to enhance the dictionary fallback heuristic with the logic -/// that the writer should fall back to the plain encoding when at a certain point, -/// e.g. after encoding the first batch, the total size of unencoded data -/// is calculated as smaller than `(encodedSize + dictionarySize)`. -pub struct PlainDataSizeCounter { - raw_data_byte_size: usize, - // Cached type length to improve performance for fixed-length types. - type_length: usize, -} - -impl PlainDataSizeCounter { - pub fn new(desc: &ColumnDescriptor) -> Self { - Self { - raw_data_byte_size: 0, - type_length: desc.type_length() as usize, - } - } - - /// Updates the counter with the given slice. - pub fn update(&mut self, values: &[T]) { - let raw_size = match T::PHYSICAL_TYPE { - Type::BOOLEAN => values.len(), - Type::INT32 | Type::FLOAT => 4 * values.len(), - Type::INT64 | Type::DOUBLE => 8 * values.len(), - Type::INT96 => Int96::SIZE_IN_BYTES * values.len(), - Type::BYTE_ARRAY => { - // For variable-length types, the length prefix and the actual data are are encoded. - values.iter().map(|value| value.plain_encoded_size()).sum() - } - Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), - }; - self.raw_data_byte_size = self.raw_data_byte_size.saturating_add(raw_size); - } - - /// Like `update`, but specialized for byte array data exposed by Arrow - /// array accessors. - #[cfg(feature = "arrow")] - #[inline] - pub fn update_byte_array(&mut self, value: &[u8]) { - // For variable-length types, the length prefix and the actual data are are encoded. - let raw_size = std::mem::size_of::() + value.len(); - self.raw_data_byte_size = self.raw_data_byte_size.saturating_add(raw_size); - } - - /// Returns the total size in bytes of values passed to `update` as if they were written - /// in plain encoding, without a dictionary. - #[inline] - pub fn plain_encoded_data_size(&self) -> usize { - self.raw_data_byte_size - } -} From 289e4e1fa2b1074563f475c172e912cdf1d0cfb8 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Fri, 17 Apr 2026 20:01:11 +0300 Subject: [PATCH 18/26] test: adjust dict fallback tests as per review --- parquet/src/arrow/arrow_writer/mod.rs | 52 ++++++++++++++++----------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 64162cad420d..78c9fb0f16e0 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1669,6 +1669,7 @@ mod tests { use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder}; use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY}; + use crate::basic::PageType; use crate::column::page::{Page, PageReader}; use crate::file::metadata::thrift::PageHeader; use crate::file::page_index::column_index::ColumnIndexMetaData; @@ -2601,7 +2602,7 @@ mod tests { // Set dictionary fallback to trigger fallback to PLAIN encoding on unfavorable compression let props = WriterProperties::builder() .set_dictionary_fallback(DictionaryFallback::OnUnfavorableAfter(1)) - .set_data_page_size_limit(1) + .set_data_page_row_count_limit(2) .set_write_batch_size(1) .build(); @@ -2611,7 +2612,9 @@ mod tests { writer.write(&batch).unwrap(); writer.close().unwrap(); - let options = ReadOptionsBuilder::new().with_page_index().build(); + let options = ReadOptionsBuilder::new() + .with_encoding_stats_as_mask(false) + .build(); let reader = SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap(); @@ -2619,25 +2622,29 @@ mod tests { assert_eq!(column.len(), 1); - // We should write one row before falling back to PLAIN encoding so there should still be a - // dictionary page. + // check page encoding stats, should be one dict page, one dict encoded page, and 5 + // plain encoded pages + let stats = column[0].page_encoding_stats().unwrap(); + println!("pes: {stats:?}"); assert!( - column[0].dictionary_page_offset().is_some(), - "Expected a dictionary page" - ); - - assert!(reader.metadata().offset_index().is_some()); - let offset_indexes = &reader.metadata().offset_index().unwrap()[0]; - - let page_locations = offset_indexes[0].page_locations.clone(); - - // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes - // so we expect one dictionary encoded page and then a page per row thereafter. - assert_eq!( - page_locations.len(), - 10, - "Expected 10 pages but got {page_locations:#?}" + stats + .iter() + .any(|s| s.page_type == PageType::DICTIONARY_PAGE) ); + let num_dict_encoded: i32 = stats + .iter() + .filter(|s| { + s.page_type == PageType::DATA_PAGE && s.encoding == Encoding::RLE_DICTIONARY + }) + .map(|s| s.count) + .sum(); + assert_eq!(num_dict_encoded, 1); + let num_plain_encoded: i32 = stats + .iter() + .filter(|s| s.page_type == PageType::DATA_PAGE && s.encoding == Encoding::PLAIN) + .map(|s| s.count) + .sum(); + assert_eq!(num_plain_encoded, 5); } #[test] @@ -4926,7 +4933,7 @@ mod tests { .set_dictionary_page_size_limit(1024 * 1024) .set_column_dictionary_fallback( ColumnPath::from("col0"), - DictionaryFallback::OnUnfavorableAfter(8192), + DictionaryFallback::OnUnfavorableAfter(32_768), ) .build(); let mut writer = ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap(); @@ -4939,7 +4946,10 @@ mod tests { metadata.try_parse(&data).unwrap(); let metadata = metadata.finish().unwrap(); let fallback_meta = metadata.row_group(0).column(0); - assert_eq!(get_dict_page_size(fallback_meta, data.clone()), 32_768); + assert_eq!( + get_dict_page_size(fallback_meta, data.clone()), + 32_768 * std::mem::size_of::() + ); } struct WriteBatchesShape { From 3ff12c8a1903b4b9d711d0343a37518ccd5f6ea4 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Sat, 18 Apr 2026 01:53:58 +0300 Subject: [PATCH 19/26] feat: make dict fallback decision one-off Once the sample size reaches the configured minimum, shut down the fallback counter even if the dictionary encoding is still favorable, to avoid the accounting overhead. --- parquet/src/arrow/arrow_writer/byte_array.rs | 4 +++- parquet/src/column/writer/encoder.rs | 4 +++- parquet/src/encodings/encoding/dict_fallback.rs | 16 ++++++++++++---- parquet/src/file/properties.rs | 4 +++- 4 files changed, 21 insertions(+), 7 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 9aadbfc8a9d0..2b32132dce88 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -567,7 +567,9 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(encoder) => { let data_page = encoder.flush_data_page(min_value, max_value); if let Some(counter) = self.dict_fallback_counter.as_mut() { - counter.count_dict_encoded_data(data_page.buf.len()); + if !counter.continue_with_dict_encoded_page(data_page.buf.len()) { + self.dict_fallback_counter = None; + } } Ok(data_page) } diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 0cf6d0a1fbc3..2d563a1b41d2 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -344,7 +344,9 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Some(encoder) => { let buf = encoder.write_indices()?; if let Some(counter) = self.dict_fallback_counter.as_mut() { - counter.count_dict_encoded_data(buf.len()); + if !counter.continue_with_dict_encoded_page(buf.len()) { + self.dict_fallback_counter = None; + } } (buf, Encoding::RLE_DICTIONARY) } diff --git a/parquet/src/encodings/encoding/dict_fallback.rs b/parquet/src/encodings/encoding/dict_fallback.rs index 6248156c57f5..2a4954f4555f 100644 --- a/parquet/src/encodings/encoding/dict_fallback.rs +++ b/parquet/src/encodings/encoding/dict_fallback.rs @@ -89,16 +89,24 @@ impl DictFallbackCounter { self.num_values += 1; } - /// Increments the counter of the size of dictionary encoded data - /// by the given amount in bytes. + /// If the number of values accounted so far reaches or exceeds the minimum + /// sample length, returns false to indicate that the counting + /// should be stopped. + /// Otherwise, increments the total counted size of dictionary encoded data + /// by the given amount in bytes and returns true. #[inline] - pub fn count_dict_encoded_data(&mut self, encoded_len: usize) { + pub fn continue_with_dict_encoded_page(&mut self, encoded_len: usize) -> bool { + if self.num_values >= self.min_sample_len { + // The sample size has been reached, no need to proceed with counting. + return false; + } self.encoded_data_size = self.encoded_data_size.saturating_add(encoded_len); + true } /// Returns true if the estimated size of plainly encoded data, in bytes, /// would not exceed the size of data encoded with a dictionary, - /// as counted by the `count_dict_encoded_data` calls made on this counter. + /// as counted by the `continue_with_dict_encoded_page` calls made on this counter. #[inline] pub fn is_dict_encoding_unfavorable(&self, dict_encoded_size: usize) -> bool { self.num_values >= self.min_sample_len diff --git a/parquet/src/file/properties.rs b/parquet/src/file/properties.rs index e6345356eafe..fa320c62ce0d 100644 --- a/parquet/src/file/properties.rs +++ b/parquet/src/file/properties.rs @@ -1206,7 +1206,9 @@ pub enum DictionaryFallback { OnPageSizeLimit, /// Fall back to non-dictionary encoding if the dictionary page size limit is exceeded, /// or if the dictionary encoding upon encoding at least the given number of values - /// is larger than the plain encoding for the same data. + /// is larger than the plain encoding for the same data. The latter check is performed once + /// per column chunk, so the encoding efficiency may still degrade with subsequent pages in + /// the same column chunk. OnUnfavorableAfter(usize), } From 1fcea02d7e54048d8e86c108d96cffb015531cc7 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 20 Apr 2026 18:06:22 +0300 Subject: [PATCH 20/26] refactor: revert counter disabling, rename --- parquet/src/arrow/arrow_writer/byte_array.rs | 4 +- parquet/src/column/writer/encoder.rs | 20 ++++----- .../src/encodings/encoding/dict_fallback.rs | 41 ++++++++++++------- 3 files changed, 36 insertions(+), 29 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 2b32132dce88..66a2942e9599 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -567,9 +567,7 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(encoder) => { let data_page = encoder.flush_data_page(min_value, max_value); if let Some(counter) = self.dict_fallback_counter.as_mut() { - if !counter.continue_with_dict_encoded_page(data_page.buf.len()) { - self.dict_fallback_counter = None; - } + counter.commit_page(&data_page); } Ok(data_page) } diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 2d563a1b41d2..47e7f50b493f 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -341,26 +341,24 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { fn flush_data_page(&mut self) -> Result> { let (buf, encoding) = match &mut self.dict_encoder { - Some(encoder) => { - let buf = encoder.write_indices()?; - if let Some(counter) = self.dict_fallback_counter.as_mut() { - if !counter.continue_with_dict_encoded_page(buf.len()) { - self.dict_fallback_counter = None; - } - } - (buf, Encoding::RLE_DICTIONARY) - } + Some(encoder) => (encoder.write_indices()?, Encoding::RLE_DICTIONARY), _ => (self.encoder.flush_buffer()?, self.encoder.encoding()), }; - Ok(DataPageValues { + let page = DataPageValues { buf, encoding, num_values: std::mem::take(&mut self.num_values), min_value: self.min_value.take(), max_value: self.max_value.take(), variable_length_bytes: self.variable_length_bytes.take(), - }) + }; + + if let Some(counter) = self.dict_fallback_counter.as_mut() { + counter.commit_page(&page); + } + + Ok(page) } fn flush_geospatial_statistics(&mut self) -> Option> { diff --git a/parquet/src/encodings/encoding/dict_fallback.rs b/parquet/src/encodings/encoding/dict_fallback.rs index 2a4954f4555f..b6c49fb78c11 100644 --- a/parquet/src/encodings/encoding/dict_fallback.rs +++ b/parquet/src/encodings/encoding/dict_fallback.rs @@ -15,7 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::basic::Type; +use crate::basic::{Encoding, Type}; +use crate::column::writer::encoder::DataPageValues; use crate::data_type::Int96; use crate::data_type::private::ParquetValueType; use crate::schema::types::ColumnDescriptor; @@ -89,27 +90,37 @@ impl DictFallbackCounter { self.num_values += 1; } - /// If the number of values accounted so far reaches or exceeds the minimum - /// sample length, returns false to indicate that the counting - /// should be stopped. - /// Otherwise, increments the total counted size of dictionary encoded data - /// by the given amount in bytes and returns true. + /// Increments the total counted size of dictionary encoded data + /// for a page encoded with the dictionary encoding. + pub fn commit_page(&mut self, page: &DataPageValues) + where + T: ParquetValueType, + { + assert_eq!( + page.encoding, + Encoding::RLE_DICTIONARY, + "should only be used with the dictionary encoder" + ); + self.encoded_data_size = self.encoded_data_size.saturating_add(page.buf.len()); + } + + /// If the number of dictionary encoded values accounted so far + /// reaches or exceeds the configured minimum, returns true to indicate + /// that the counting should be stopped, otherwise returns false. #[inline] - pub fn continue_with_dict_encoded_page(&mut self, encoded_len: usize) -> bool { - if self.num_values >= self.min_sample_len { - // The sample size has been reached, no need to proceed with counting. - return false; - } - self.encoded_data_size = self.encoded_data_size.saturating_add(encoded_len); - true + pub fn min_sample_len_reached(&self) -> bool { + self.num_values >= self.min_sample_len } /// Returns true if the estimated size of plainly encoded data, in bytes, /// would not exceed the size of data encoded with a dictionary, - /// as counted by the `continue_with_dict_encoded_page` calls made on this counter. + /// as counted by the `commit_page` calls made on this counter. + /// + /// This method does not return true until the minimum sample length has been reached, + /// as indicated by the `min_sample_len_reached` method. #[inline] pub fn is_dict_encoding_unfavorable(&self, dict_encoded_size: usize) -> bool { - self.num_values >= self.min_sample_len + self.min_sample_len_reached() && self.raw_data_size <= dict_encoded_size.saturating_add(self.encoded_data_size) } } From 389d038d9fa24264f4ba83438e15508822f8d1cc Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 20 Apr 2026 18:07:02 +0300 Subject: [PATCH 21/26] test: verify encoded data after dict fallback --- parquet/src/arrow/arrow_writer/mod.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 78c9fb0f16e0..275fd7ef7f25 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -1699,6 +1699,7 @@ mod tests { reader::{FileReader, SerializedFileReader}, statistics::Statistics, }; + use crate::record::RowAccessor; #[test] fn arrow_writer() { @@ -2595,7 +2596,7 @@ mod tests { let array = Arc::new(builder.finish()); - let batch = RecordBatch::try_new(schema, vec![array]).unwrap(); + let batch = RecordBatch::try_new(schema, vec![array.clone()]).unwrap(); let file = tempfile::tempfile().unwrap(); @@ -2645,6 +2646,17 @@ mod tests { .map(|s| s.count) .sum(); assert_eq!(num_plain_encoded, 5); + + // Read back the values and confirm they match the original array. + let rows: Vec<_> = reader + .get_row_iter(None) + .unwrap() + .map(|r| r.unwrap()) + .collect(); + assert_eq!(rows.len(), array.len()); + for (i, row) in rows.iter().enumerate() { + assert_eq!(row.get_string(0).unwrap(), array.value(i)); + } } #[test] From 8873bc1d0f796935815407788bd4849346ec759e Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Mon, 20 Apr 2026 19:55:14 +0300 Subject: [PATCH 22/26] perf(parquet): eager cutoff for fallback counter As soon as the minimum number of values to sample has been reached for to reach a decision on dictionary fallback under the `OnUnfavorableAfter` policy, the fallback counter can be disabled to avoid the overhead of further accounting. In previous code, this did not happen until the next page flush, but the write batch size can be much smaller and this is the granularity at which should_dict_fallback is checked. --- parquet/src/arrow/arrow_writer/byte_array.rs | 24 +++++++++------- parquet/src/column/writer/encoder.rs | 28 ++++++++++++------- parquet/src/column/writer/mod.rs | 10 +++++-- .../src/encodings/encoding/dict_fallback.rs | 23 ++++++++------- 4 files changed, 52 insertions(+), 33 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/byte_array.rs b/parquet/src/arrow/arrow_writer/byte_array.rs index 66a2942e9599..36b8945fdf5d 100644 --- a/parquet/src/arrow/arrow_writer/byte_array.rs +++ b/parquet/src/arrow/arrow_writer/byte_array.rs @@ -521,16 +521,6 @@ impl ColumnValueEncoder for ByteArrayEncoder { Some(self.dict_encoder.as_ref()?.estimated_dict_page_size()) } - fn is_dict_encoding_unfavorable(&self) -> bool { - match (&self.dict_encoder, &self.dict_fallback_counter) { - (Some(encoder), Some(counter)) => { - let dict_size = encoder.estimated_dict_page_size(); - counter.is_dict_encoding_unfavorable(dict_size) - } - _ => false, - } - } - /// Returns an estimate of the data page size in bytes /// /// This includes: @@ -542,6 +532,20 @@ impl ColumnValueEncoder for ByteArrayEncoder { } } + fn is_dict_encoding_unfavorable(&self) -> Option { + match (&self.dict_encoder, &self.dict_fallback_counter) { + (Some(encoder), Some(counter)) => { + let dict_size = encoder.estimated_dict_page_size(); + counter.is_dict_encoding_unfavorable(dict_size) + } + _ => None, + } + } + + fn disable_dict_fallback_accounting(&mut self) { + self.dict_fallback_counter = None; + } + fn flush_dict_page(&mut self) -> Result> { match self.dict_encoder.take() { Some(encoder) => { diff --git a/parquet/src/column/writer/encoder.rs b/parquet/src/column/writer/encoder.rs index 47e7f50b493f..d1ce1a024a78 100644 --- a/parquet/src/column/writer/encoder.rs +++ b/parquet/src/column/writer/encoder.rs @@ -109,13 +109,17 @@ pub trait ColumnValueEncoder { /// + fn estimated_data_page_size(&self) -> usize; - /// Returns true if the estimated size of plainly encoded data, in bytes, + /// Returns `Some(true)` if the estimated size of plainly encoded data, in bytes, /// would be smaller than the size of data encoded with a dictionary. + /// If the estimate does not show a clear benefit, returns `Some(false)`. /// If there is no dictionary, or the data size statistic is not available, - /// returns false. + /// or it is not yet possible to make an estimate due to insufficient + /// collected data, returns `None`. /// For similar logic in parquet-java, /// see `DictionaryValuesWriter.isCompressionSatisfying`. - fn is_dict_encoding_unfavorable(&self) -> bool; + fn is_dict_encoding_unfavorable(&self) -> Option; + + fn disable_dict_fallback_accounting(&mut self); /// Flush the dictionary page for this column chunk if any. Any subsequent calls to /// [`Self::write`] will not be dictionary encoded @@ -299,21 +303,25 @@ impl ColumnValueEncoder for ColumnValueEncoderImpl { Some(self.dict_encoder.as_ref()?.dict_encoded_size()) } - fn is_dict_encoding_unfavorable(&self) -> bool { + fn estimated_data_page_size(&self) -> usize { + match &self.dict_encoder { + Some(encoder) => encoder.estimated_data_encoded_size(), + _ => self.encoder.estimated_data_encoded_size(), + } + } + + fn is_dict_encoding_unfavorable(&self) -> Option { match (&self.dict_encoder, &self.dict_fallback_counter) { (Some(encoder), Some(counter)) => { let dict_size = encoder.dict_encoded_size(); counter.is_dict_encoding_unfavorable(dict_size) } - _ => false, + _ => None, } } - fn estimated_data_page_size(&self) -> usize { - match &self.dict_encoder { - Some(encoder) => encoder.estimated_data_encoded_size(), - _ => self.encoder.estimated_data_encoded_size(), - } + fn disable_dict_fallback_accounting(&mut self) { + self.dict_fallback_counter = None; } fn flush_dict_page(&mut self) -> Result> { diff --git a/parquet/src/column/writer/mod.rs b/parquet/src/column/writer/mod.rs index e9cb61db4fa5..bfb94acb6fc6 100644 --- a/parquet/src/column/writer/mod.rs +++ b/parquet/src/column/writer/mod.rs @@ -747,7 +747,7 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { /// Returns true if we need to fall back to non-dictionary encoding. /// /// The behavior is governed by the `dictionary_fallback` column property. - fn should_dict_fallback(&self) -> bool { + fn should_dict_fallback(&mut self) -> bool { let dict_size = match self.encoder.estimated_dict_page_size() { Some(size) => size, None => return false, @@ -763,8 +763,12 @@ impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> { } // Second check, if enabled: the compression heuristic. - if self.encoder.is_dict_encoding_unfavorable() { - return true; + if let Some(should_fallback) = self.encoder.is_dict_encoding_unfavorable() { + // The decision on the efficiency is made after processing + // the requisite number of values, so disable further accounting + // to avoid the overhead. + self.encoder.disable_dict_fallback_accounting(); + return should_fallback; } false diff --git a/parquet/src/encodings/encoding/dict_fallback.rs b/parquet/src/encodings/encoding/dict_fallback.rs index b6c49fb78c11..5fe16146bc90 100644 --- a/parquet/src/encodings/encoding/dict_fallback.rs +++ b/parquet/src/encodings/encoding/dict_fallback.rs @@ -107,20 +107,23 @@ impl DictFallbackCounter { /// If the number of dictionary encoded values accounted so far /// reaches or exceeds the configured minimum, returns true to indicate /// that the counting should be stopped, otherwise returns false. - #[inline] - pub fn min_sample_len_reached(&self) -> bool { + fn min_sample_len_reached(&self) -> bool { self.num_values >= self.min_sample_len } - /// Returns true if the estimated size of plainly encoded data, in bytes, + /// Returns `Some(true)` if the estimated size of plainly encoded data, in bytes, /// would not exceed the size of data encoded with a dictionary, - /// as counted by the `commit_page` calls made on this counter. - /// - /// This method does not return true until the minimum sample length has been reached, - /// as indicated by the `min_sample_len_reached` method. + /// as counted by the `commit_page` calls made on this counter and the provided size of + /// the encoded dictionary page. + /// This method returns `None` until the minimum number of values given in + /// `DictFallbackCounter::new` has been processed. The third alternative, + /// `Some(false)`, indicates that the sample size is sufficient, but the dictionary encoding + /// is not-unfavorable, that is, the collected metrics show no clear advantage in falling + /// back to plain (or other, presumably more efficient) encoding. #[inline] - pub fn is_dict_encoding_unfavorable(&self, dict_encoded_size: usize) -> bool { - self.min_sample_len_reached() - && self.raw_data_size <= dict_encoded_size.saturating_add(self.encoded_data_size) + pub fn is_dict_encoding_unfavorable(&self, dict_encoded_size: usize) -> Option { + self.min_sample_len_reached().then_some( + self.raw_data_size <= dict_encoded_size.saturating_add(self.encoded_data_size), + ) } } From 8ba290b9c3b93c30f66344ddf124b7231ba5d641 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Apr 2026 16:45:11 +0300 Subject: [PATCH 23/26] test: assert efficiency of dictionary fallback --- parquet/src/arrow/arrow_writer/mod.rs | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 275fd7ef7f25..2913a5faa7c6 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -4933,7 +4933,7 @@ mod tests { writer.write(&batch).unwrap(); let data = Bytes::from(writer.into_inner().unwrap()); - println!("file length, dictionary: {}", data.len()); + let file_length_dict = data.len(); let mut metadata = ParquetMetaDataReader::new(); metadata.try_parse(&data).unwrap(); @@ -4952,7 +4952,7 @@ mod tests { writer.write(&batch).unwrap(); let data = Bytes::from(writer.into_inner().unwrap()); - println!("file length, fallback: {}", data.len()); + let file_length_fallback = data.len(); let mut metadata = ParquetMetaDataReader::new(); metadata.try_parse(&data).unwrap(); @@ -4962,6 +4962,14 @@ mod tests { get_dict_page_size(fallback_meta, data.clone()), 32_768 * std::mem::size_of::() ); + + let compression_ratio = file_length_fallback as f64 / file_length_dict as f64; + assert!( + compression_ratio < 0.9, + "File encoded with dictionary fallback encoding does not result in sufficient compression, + got {file_length_fallback} vs {file_length_dict} ({:.2}%)", + compression_ratio * 100.0 + ); } struct WriteBatchesShape { From 1827d048da775e906083b2ad26e0b7d221baef7f Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Apr 2026 19:42:14 +0300 Subject: [PATCH 24/26] test: remove page stats printout Moved the printout to an assertion failure message below. Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> --- parquet/src/arrow/arrow_writer/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 2913a5faa7c6..70cf6dc978ed 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -2626,11 +2626,11 @@ mod tests { // check page encoding stats, should be one dict page, one dict encoded page, and 5 // plain encoded pages let stats = column[0].page_encoding_stats().unwrap(); - println!("pes: {stats:?}"); assert!( stats .iter() - .any(|s| s.page_type == PageType::DICTIONARY_PAGE) + .any(|s| s.page_type == PageType::DICTIONARY_PAGE), + "stats are {stats:?}" ); let num_dict_encoded: i32 = stats .iter() From 87a19c4e7d709f5bb660de02b10201affdda2f2d Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Apr 2026 23:13:20 +0300 Subject: [PATCH 25/26] docs: clarify purpose of plain_encoded_size Explain why it's OK to panic in the implementation for bool. --- parquet/src/data_type.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index bf218e91bd2f..34c3ce00e699 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -725,6 +725,10 @@ pub(crate) mod private { fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result; /// Return the size in bytes for the value encoded in the plain encoding. + /// + /// This method is only used with the dictionary encoding. Since the writer + /// does not use the dictionary encoding for BOOLEAN type, this method's + /// implementation for bool will panic if called. fn plain_encoded_size(&self) -> usize; /// Return the number of variable length bytes in a given slice of data From 14ab09e794ff4e7489a47bcc19bd19974b3e6686 Mon Sep 17 00:00:00 2001 From: Mikhail Zabaluev Date: Tue, 21 Apr 2026 23:17:45 +0300 Subject: [PATCH 26/26] refactor: consistently panic on bool Panic also in DictFallbackCounter::update_values when the type is bool. --- parquet/src/data_type.rs | 2 +- parquet/src/encodings/encoding/dict_fallback.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/parquet/src/data_type.rs b/parquet/src/data_type.rs index 34c3ce00e699..2649bf29d060 100644 --- a/parquet/src/data_type.rs +++ b/parquet/src/data_type.rs @@ -809,7 +809,7 @@ pub(crate) mod private { } fn plain_encoded_size(&self) -> usize { - panic!("Dictionary encoding should not be used for BOOLEAN type") + panic!("dictionary encoding should not be used for BOOLEAN type") } #[inline] diff --git a/parquet/src/encodings/encoding/dict_fallback.rs b/parquet/src/encodings/encoding/dict_fallback.rs index 5fe16146bc90..9d9f9c438a21 100644 --- a/parquet/src/encodings/encoding/dict_fallback.rs +++ b/parquet/src/encodings/encoding/dict_fallback.rs @@ -66,7 +66,6 @@ impl DictFallbackCounter { /// Updates the counter with the given slice of values. pub fn update_values(&mut self, values: &[T]) { let raw_size = match T::PHYSICAL_TYPE { - Type::BOOLEAN => values.len(), Type::INT32 | Type::FLOAT => 4 * values.len(), Type::INT64 | Type::DOUBLE => 8 * values.len(), Type::INT96 => Int96::SIZE_IN_BYTES * values.len(), @@ -75,6 +74,7 @@ impl DictFallbackCounter { values.iter().map(|value| value.plain_encoded_size()).sum() } Type::FIXED_LEN_BYTE_ARRAY => self.type_length * values.len(), + Type::BOOLEAN => panic!("dictionary encoding should not be used for BOOLEAN type"), }; self.raw_data_size = self.raw_data_size.saturating_add(raw_size); self.num_values += values.len();