parquet/column/writer/
encoder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use bytes::Bytes;
19
20use crate::basic::{ConvertedType, Encoding, LogicalType, Type};
21use crate::bloom_filter::Sbbf;
22use crate::column::writer::{
23    compare_greater, fallback_encoding, has_dictionary_support, is_nan, update_max, update_min,
24};
25use crate::data_type::DataType;
26use crate::data_type::private::ParquetValueType;
27use crate::encodings::encoding::{DictEncoder, Encoder, get_encoder_with_properties};
28use crate::errors::{ParquetError, Result};
29use crate::file::properties::{EnabledStatistics, ResolvedColumnProperties, WriterProperties};
30use crate::geospatial::accumulator::{GeoStatsAccumulator, try_new_geo_stats_accumulator};
31use crate::geospatial::statistics::GeospatialStatistics;
32use crate::schema::types::{BasicTypeInfo, ColumnDescPtr};
33
34/// A collection of [`ParquetValueType`] encoded by a [`ColumnValueEncoder`]
35pub trait ColumnValues {
36    /// The number of values in this collection
37    fn len(&self) -> usize;
38}
39
40#[cfg(feature = "arrow")]
41impl ColumnValues for dyn arrow_array::Array {
42    fn len(&self) -> usize {
43        arrow_array::Array::len(self)
44    }
45}
46
47impl<T: ParquetValueType> ColumnValues for [T] {
48    fn len(&self) -> usize {
49        self.len()
50    }
51}
52
53/// The encoded data for a dictionary page
54pub struct DictionaryPage {
55    pub buf: Bytes,
56    pub num_values: usize,
57    pub is_sorted: bool,
58}
59
60/// The encoded values for a data page, with optional statistics
61pub struct DataPageValues<T> {
62    pub buf: Bytes,
63    pub num_values: usize,
64    pub encoding: Encoding,
65    pub min_value: Option<T>,
66    pub max_value: Option<T>,
67    pub nan_count: Option<u64>,
68    pub variable_length_bytes: Option<i64>,
69}
70
71/// A generic encoder of [`ColumnValues`] to data and dictionary pages used by
72/// [`super::GenericColumnWriter`]
73pub trait ColumnValueEncoder {
74    /// The underlying value type of [`Self::Values`]
75    ///
76    /// Note: this avoids needing to fully qualify `<Self::Values as ColumnValues>::T`
77    type T: ParquetValueType;
78
79    /// The values encoded by this encoder
80    type Values: ColumnValues + ?Sized;
81
82    /// Create a new [`ColumnValueEncoder`]
83    ///
84    /// `column_props` holds the settings in `props` that apply specifically to
85    /// `descr`, already resolved by the caller.
86    #[expect(
87        private_interfaces,
88        reason = "this trait is not nameable outside the crate"
89    )]
90    fn try_new(
91        descr: &ColumnDescPtr,
92        props: &WriterProperties,
93        column_props: &ResolvedColumnProperties,
94    ) -> Result<Self>
95    where
96        Self: Sized;
97
98    /// Write the corresponding values to this [`ColumnValueEncoder`]
99    fn write(&mut self, values: &Self::Values, offset: usize, len: usize) -> Result<()>;
100
101    /// Write the values at the indexes in `indices` to this [`ColumnValueEncoder`]
102    fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()>;
103
104    /// Returns the largest `k` such that the first `k` values in
105    /// `values[offset..offset + len]` encode to at most `byte_budget`
106    /// bytes — i.e. how many values fit in a single page byte budget.
107    ///
108    /// Returns `len` if every value fits. Returns at least 1 if a single
109    /// value alone exceeds the budget, matching parquet's "at least one
110    /// value per data page" rule.
111    ///
112    /// `None` means "no cheap estimate available"; the caller stays on
113    /// the batched fast path and lets the post-write
114    /// `should_add_data_page` check handle bounding.
115    ///
116    /// Implementations should short-circuit aggressively: the typical
117    /// case is "everything fits, return `len`", and the next-most-common
118    /// case is "one wide value, return 1." The variable-width walk only
119    /// needs to be precise when the chunk is genuinely near the budget.
120    fn count_values_within_byte_budget(
121        _values: &Self::Values,
122        _offset: usize,
123        _len: usize,
124        _byte_budget: usize,
125    ) -> Option<usize> {
126        None
127    }
128
129    /// As [`Self::count_values_within_byte_budget`] but using gather
130    /// `indices` rather than a contiguous range. Returns the number of
131    /// `indices` that fit, not the maximum index value.
132    fn count_values_within_byte_budget_gather(
133        _values: &Self::Values,
134        _indices: &[usize],
135        _byte_budget: usize,
136    ) -> Option<usize> {
137        None
138    }
139
140    /// Returns the number of buffered values
141    fn num_values(&self) -> usize;
142
143    /// Returns true if this encoder has a dictionary page
144    fn has_dictionary(&self) -> bool;
145
146    /// Returns true if the encoder compresses each value against the value
147    /// immediately before it, within the current page.
148    ///
149    /// For such encodings a page boundary is not free: flushing discards the
150    /// previous value, so the first value of the next page is stored in full.
151    /// [`GenericColumnWriter::should_add_data_page`] uses this to exempt a
152    /// page's mandatory first value from the data page byte limit.
153    ///
154    /// Per encoding:
155    /// * `DELTA_BYTE_ARRAY`: true. Each value is stored as the length of the
156    ///   prefix it shares with its predecessor plus the remaining suffix.
157    /// * Everything else: false, the default. `PLAIN` and
158    ///   `DELTA_LENGTH_BYTE_ARRAY` store a value at the same cost wherever it
159    ///   lands, and a dictionary outlives the pages that index into it, so no
160    ///   page boundary makes a value more expensive.
161    ///
162    /// [`GenericColumnWriter::should_add_data_page`]: crate::column::writer::GenericColumnWriter::should_add_data_page
163    fn compresses_against_previous_value(&self) -> bool {
164        false
165    }
166
167    /// Returns the estimated total memory usage of the encoder
168    ///
169    fn estimated_memory_size(&self) -> usize;
170
171    /// Returns an estimate of the encoded size of dictionary page size in bytes, or `None` if no dictionary
172    fn estimated_dict_page_size(&self) -> Option<usize>;
173
174    /// Returns an estimate of the encoded data page size in bytes
175    ///
176    /// This should include:
177    /// <already_written_encoded_byte_size> + <estimated_encoded_size_of_unflushed_bytes>
178    fn estimated_data_page_size(&self) -> usize;
179
180    /// Flush the dictionary page for this column chunk if any. Any subsequent calls to
181    /// [`Self::write`] will not be dictionary encoded
182    ///
183    /// Note: [`Self::flush_data_page`] must be called first, as this will error if there
184    /// are any pending page values
185    fn flush_dict_page(&mut self) -> Result<Option<DictionaryPage>>;
186
187    /// Flush the next data page for this column chunk
188    fn flush_data_page(&mut self) -> Result<DataPageValues<Self::T>>;
189
190    /// Flushes bloom filter if enabled and returns it, otherwise returns `None`. Subsequent writes
191    /// will *not* be tracked by the bloom filter as it is empty since. This should be called once
192    /// near the end of encoding.
193    fn flush_bloom_filter(&mut self) -> Option<Sbbf>;
194
195    /// Computes [`GeospatialStatistics`], if any, and resets internal state such that any internal
196    /// accumulator is prepared to accumulate statistics for the next column chunk.
197    fn flush_geospatial_statistics(&mut self) -> Option<Box<GeospatialStatistics>>;
198}
199
200pub struct ColumnValueEncoderImpl<T: DataType> {
201    encoder: Box<dyn Encoder<T>>,
202    dict_encoder: Option<DictEncoder<T>>,
203    descr: ColumnDescPtr,
204    num_values: usize,
205    statistics_enabled: EnabledStatistics,
206    min_value: Option<T::T>,
207    max_value: Option<T::T>,
208    nan_count: Option<u64>,
209    bloom_filter: Option<Sbbf>,
210    bloom_filter_target_fpp: f64,
211    variable_length_bytes: Option<i64>,
212    geo_stats_accumulator: Option<Box<dyn GeoStatsAccumulator>>,
213}
214
215impl<T: DataType> ColumnValueEncoderImpl<T> {
216    fn is_floating_point_column(&self) -> bool {
217        matches!(self.descr.physical_type(), Type::FLOAT | Type::DOUBLE)
218            || self.descr.logical_type_ref() == Some(&LogicalType::Float16)
219    }
220
221    fn write_slice(&mut self, slice: &[T::T]) -> Result<()> {
222        if self.statistics_enabled != EnabledStatistics::None
223            // INTERVAL, Geometry, and Geography have undefined sort order, so don't write min/max stats for them
224            && self.descr.converted_type() != ConvertedType::INTERVAL
225        {
226            if let Some(accumulator) = self.geo_stats_accumulator.as_deref_mut() {
227                update_geo_stats_accumulator(accumulator, slice.iter());
228            } else if let Some((min, max, nan_count)) =
229                get_min_max(self.descr.get_basic_info(), slice.iter())
230            {
231                update_min(&self.descr, &min, &mut self.min_value);
232                update_max(&self.descr, &max, &mut self.max_value);
233                if self.is_floating_point_column() {
234                    *self.nan_count.get_or_insert(0) += nan_count;
235                }
236            }
237
238            if let Some(var_bytes) = T::T::variable_length_bytes(slice) {
239                *self.variable_length_bytes.get_or_insert(0) += var_bytes;
240            }
241        }
242
243        // While a dictionary is in use the filter is populated from its distinct values
244        // in `flush_dict_page`, so each value is hashed once rather than once per row.
245        match &mut self.dict_encoder {
246            Some(encoder) => encoder.put(slice),
247            _ => {
248                if let Some(bloom_filter) = &mut self.bloom_filter {
249                    for value in slice {
250                        bloom_filter.insert(value);
251                    }
252                }
253                self.encoder.put(slice)
254            }
255        }
256    }
257}
258
259impl<T: DataType> ColumnValueEncoder for ColumnValueEncoderImpl<T> {
260    type T = T::T;
261
262    type Values = [T::T];
263
264    fn flush_bloom_filter(&mut self) -> Option<Sbbf> {
265        let mut sbbf = self.bloom_filter.take()?;
266        sbbf.fold_to_target_fpp(self.bloom_filter_target_fpp);
267        Some(sbbf)
268    }
269
270    #[expect(
271        private_interfaces,
272        reason = "this trait is not nameable outside the crate"
273    )]
274    fn try_new(
275        descr: &ColumnDescPtr,
276        props: &WriterProperties,
277        column_props: &ResolvedColumnProperties,
278    ) -> Result<Self> {
279        let dict_supported =
280            column_props.dictionary_enabled && has_dictionary_support(T::get_physical_type());
281        let dict_encoder = dict_supported.then(|| DictEncoder::new(descr.clone()));
282
283        // Set either main encoder or fallback encoder.
284        let encoder = get_encoder_with_properties(
285            column_props
286                .encoding
287                .unwrap_or_else(|| fallback_encoding(T::get_physical_type(), props)),
288            descr,
289            column_props,
290        )?;
291
292        let statistics_enabled = column_props.statistics_enabled;
293
294        let (bloom_filter, bloom_filter_target_fpp) = create_bloom_filter(column_props)?;
295
296        let geo_stats_accumulator = try_new_geo_stats_accumulator(descr);
297
298        Ok(Self {
299            encoder,
300            dict_encoder,
301            descr: descr.clone(),
302            num_values: 0,
303            statistics_enabled,
304            bloom_filter,
305            bloom_filter_target_fpp,
306            min_value: None,
307            max_value: None,
308            nan_count: None,
309            variable_length_bytes: None,
310            geo_stats_accumulator,
311        })
312    }
313
314    fn write(&mut self, values: &[T::T], offset: usize, len: usize) -> Result<()> {
315        self.num_values += len;
316
317        let slice = values.get(offset..offset + len).ok_or_else(|| {
318            general_err!(
319                "Expected to write {} values, but have only {}",
320                len,
321                values.len() - offset
322            )
323        })?;
324
325        self.write_slice(slice)
326    }
327
328    fn write_gather(&mut self, values: &Self::Values, indices: &[usize]) -> Result<()> {
329        self.num_values += indices.len();
330        let slice: Vec<_> = indices.iter().map(|idx| values[*idx].clone()).collect();
331        self.write_slice(&slice)
332    }
333
334    fn count_values_within_byte_budget(
335        values: &[T::T],
336        offset: usize,
337        len: usize,
338        byte_budget: usize,
339    ) -> Option<usize> {
340        // Clamp so that a caller-supplied `len` that overruns the input
341        // (e.g. a level/value mismatch the encoder will reject later)
342        // returns an estimate instead of panicking here.
343        let end = (offset + len).min(values.len());
344        let start = offset.min(end);
345        count_within_budget::<T>(
346            end - start,
347            byte_budget,
348            values[start..end].iter().map(Some),
349        )
350    }
351
352    fn count_values_within_byte_budget_gather(
353        values: &[T::T],
354        indices: &[usize],
355        byte_budget: usize,
356    ) -> Option<usize> {
357        // `values.get` yields `None` for an out-of-range index (defensive
358        // against a level/value mismatch the encoder rejects later); such a
359        // position is counted but contributes no bytes.
360        count_within_budget::<T>(
361            indices.len(),
362            byte_budget,
363            indices.iter().map(|&i| values.get(i)),
364        )
365    }
366
367    fn num_values(&self) -> usize {
368        self.num_values
369    }
370
371    fn has_dictionary(&self) -> bool {
372        self.dict_encoder.is_some()
373    }
374
375    fn compresses_against_previous_value(&self) -> bool {
376        // While dictionary encoding is active `self.encoder` is unused: the
377        // data page holds RLE indices, which carry no cross-value state.
378        self.dict_encoder.is_none() && self.encoder.encoding() == Encoding::DELTA_BYTE_ARRAY
379    }
380
381    fn estimated_memory_size(&self) -> usize {
382        let encoder_size = self.encoder.estimated_memory_size();
383
384        let dict_encoder_size = self
385            .dict_encoder
386            .as_ref()
387            .map(|encoder| encoder.estimated_memory_size())
388            .unwrap_or_default();
389
390        let bloom_filter_size = self
391            .bloom_filter
392            .as_ref()
393            .map(|bf| bf.estimated_memory_size())
394            .unwrap_or_default();
395
396        encoder_size + dict_encoder_size + bloom_filter_size
397    }
398
399    fn estimated_dict_page_size(&self) -> Option<usize> {
400        Some(self.dict_encoder.as_ref()?.dict_encoded_size())
401    }
402
403    fn estimated_data_page_size(&self) -> usize {
404        match &self.dict_encoder {
405            Some(encoder) => encoder.estimated_data_encoded_size(),
406            _ => self.encoder.estimated_data_encoded_size(),
407        }
408    }
409
410    fn flush_dict_page(&mut self) -> Result<Option<DictionaryPage>> {
411        match self.dict_encoder.take() {
412            Some(encoder) => {
413                if self.num_values != 0 {
414                    return Err(general_err!(
415                        "Must flush data pages before flushing dictionary"
416                    ));
417                }
418
419                if let Some(bloom_filter) = &mut self.bloom_filter {
420                    for value in encoder.uniques() {
421                        bloom_filter.insert(value);
422                    }
423                }
424
425                let buf = encoder.write_dict()?;
426
427                Ok(Some(DictionaryPage {
428                    buf,
429                    num_values: encoder.num_entries(),
430                    is_sorted: encoder.is_sorted(),
431                }))
432            }
433            _ => Ok(None),
434        }
435    }
436
437    fn flush_data_page(&mut self) -> Result<DataPageValues<T::T>> {
438        let (buf, encoding) = match &mut self.dict_encoder {
439            Some(encoder) => (encoder.write_indices()?, Encoding::RLE_DICTIONARY),
440            _ => (self.encoder.flush_buffer()?, self.encoder.encoding()),
441        };
442
443        Ok(DataPageValues {
444            buf,
445            encoding,
446            num_values: std::mem::take(&mut self.num_values),
447            min_value: self.min_value.take(),
448            max_value: self.max_value.take(),
449            nan_count: self.nan_count.take(),
450            variable_length_bytes: self.variable_length_bytes.take(),
451        })
452    }
453
454    fn flush_geospatial_statistics(&mut self) -> Option<Box<GeospatialStatistics>> {
455        self.geo_stats_accumulator.as_mut().map(|a| a.finish())?
456    }
457}
458
459// Get min and max values for all values in `iter`.
460//
461// For floating point we need to compare NaN values until we encounter a non-NaN
462// value which then becomes the new min/max. After this, only non-NaN values are
463// evaluated. If all values are NaN, then the min/max NaNs as determined by
464// IEEE 754 total order are returned.
465fn get_min_max<'a, T, I>(basic_type_info: &BasicTypeInfo, mut iter: I) -> Option<(T, T, u64)>
466where
467    T: ParquetValueType + 'a,
468    I: Iterator<Item = &'a T>,
469{
470    let first = iter.next()?;
471    let mut min_max_nan = is_nan(basic_type_info, first);
472    let mut nan_count = min_max_nan as u64;
473
474    let mut min = first;
475    let mut max = first;
476    for val in iter {
477        match (min_max_nan, is_nan(basic_type_info, val)) {
478            // skip NaNs if we've encounter non-NaN
479            (false, true) => {
480                nan_count += 1;
481            }
482            // if min/max are NaN, check for non-NaN and reset
483            (true, false) => {
484                min = val;
485                max = val;
486                min_max_nan = false;
487            }
488            // both are NaN or non-NaN, so do the comparison
489            (_, val_is_nan) => {
490                nan_count += val_is_nan as u64;
491                // we've already initialized min and max, so a single value can't be both
492                // extremes
493                if compare_greater(basic_type_info, min, val) {
494                    min = val;
495                } else if compare_greater(basic_type_info, val, max) {
496                    max = val;
497                }
498            }
499        }
500    }
501
502    Some((min.clone(), max.clone(), nan_count))
503}
504
505/// Creates a bloom filter sized for the column's configured NDV, returning the filter
506/// and the target FPP for folding.
507pub(crate) fn create_bloom_filter(
508    column_props: &ResolvedColumnProperties,
509) -> Result<(Option<Sbbf>, f64)> {
510    match column_props.bloom_filter_properties.as_ref() {
511        Some(bf_props) => Ok((
512            Some(Sbbf::new_with_ndv_fpp(bf_props.ndv(), bf_props.fpp())?),
513            bf_props.fpp(),
514        )),
515        None => Ok((None, 0.0)),
516    }
517}
518
519fn update_geo_stats_accumulator<'a, T, I>(bounder: &mut dyn GeoStatsAccumulator, iter: I)
520where
521    T: ParquetValueType + 'a,
522    I: Iterator<Item = &'a T>,
523{
524    if bounder.is_valid() {
525        for val in iter {
526            bounder.update_wkb(val.as_bytes());
527        }
528    }
529}
530
531/// Plain-encoded byte cost of a single value of type `T::T`.
532///
533/// Derived from [`ParquetValueType::dict_encoding_size`] (which returns
534/// `(per-value overhead, value-bytes)`) so we don't add a parallel
535/// per-value-size hook to the trait. Mirrors the dispatch in
536/// `KeyStorage::push` (`encodings/encoding/dict_encoder.rs`).
537///
538/// Placed at the end of the module deliberately. Inserting it above the
539/// `ColumnValueEncoder` trait shifts the trait and `ColumnValueEncoderImpl`
540/// within the compiled module enough to perturb downstream code placement,
541/// which measurably regresses unrelated arrow-writer string benchmarks
542/// (~5-9% on `string` / `string_and_binary_view`). Defining it last keeps
543/// the hot encoder code at the offsets it has on `main`.
544#[inline]
545fn plain_encoded_byte_size<T: DataType>(value: &T::T) -> usize {
546    let (overhead, bytes) = value.dict_encoding_size();
547    match <T::T as ParquetValueType>::PHYSICAL_TYPE {
548        // Plain BYTE_ARRAY = 4-byte length prefix + payload.
549        Type::BYTE_ARRAY => overhead + bytes,
550        // Plain FLBA = raw bytes only; `dict_encoding_size`'s length prefix
551        // is irrelevant here, so the encoder passes `type_length` directly.
552        Type::FIXED_LEN_BYTE_ARRAY => bytes,
553        // Numeric/bool are short-circuited by the caller via
554        // `mem::size_of`, so this is unreachable in practice; fall back to
555        // `overhead` defensively.
556        _ => overhead,
557    }
558}
559
560/// How many leading values fit in `byte_budget` bytes, shared by the two
561/// `ColumnValueEncoder::count_values_within_byte_budget*` methods (one walks a
562/// contiguous slice, the other gathers by index).
563///
564/// `n` is the answer when everything fits; `vals` yields each candidate value,
565/// or `None` for a position that should still be counted but contributes no
566/// bytes (an out-of-range gather index). The boundary value that crosses the
567/// budget is included in the count so the caller's page-flush check trips on
568/// this mini-batch rather than leaving a sliver for the next page; this also
569/// catches a lone outlier wherever it lands among small values.
570///
571/// Defined at the end of the module alongside `plain_encoded_byte_size` for
572/// the same reason — see that function's note on code placement and the
573/// `string` / `string_and_binary_view` benchmarks.
574#[inline]
575fn count_within_budget<'a, T: DataType>(
576    n: usize,
577    byte_budget: usize,
578    vals: impl Iterator<Item = Option<&'a T::T>>,
579) -> Option<usize>
580where
581    T::T: 'a,
582{
583    // Fixed-size physical types have a constant per-value byte cost, so the
584    // answer is one division — no walk needed.
585    let phys = <T::T as ParquetValueType>::PHYSICAL_TYPE;
586    if phys != Type::BYTE_ARRAY && phys != Type::FIXED_LEN_BYTE_ARRAY {
587        let per = std::mem::size_of::<T::T>().max(1);
588        return Some((byte_budget / per).max(1).min(n));
589    }
590    // Variable-width: accumulate, exit at the first value past the budget.
591    let mut cum: usize = 0;
592    for (i, v) in vals.enumerate() {
593        if let Some(v) = v {
594            cum = cum.saturating_add(plain_encoded_byte_size::<T>(v));
595        }
596        if cum > byte_budget {
597            return Some(i + 1);
598        }
599    }
600    Some(n)
601}
parquet/column/writer/encoder.rs

parquet/column/writer/
encoder.rs
