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

parquet/column/writer/
encoder.rs
