Skip to main content

parquet/column/writer/
mod.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
18//! Contains column writer API.
19
20use bytes::Bytes;
21use half::f16;
22
23use crate::bloom_filter::Sbbf;
24use crate::file::page_index::column_index::ColumnIndexMetaData;
25use crate::file::page_index::offset_index::OffsetIndexMetaData;
26use std::cmp::Ordering;
27use std::collections::{BTreeSet, VecDeque};
28use std::str;
29
30use crate::basic::{
31    BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, LogicalType, PageType,
32    SortOrder, Type,
33};
34use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter};
35use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues};
36use crate::compression::{Codec, CodecOptionsBuilder, create_codec};
37use crate::data_type::private::ParquetValueType;
38use crate::data_type::*;
39use crate::encodings::levels::LevelEncoder;
40#[cfg(feature = "encryption")]
41use crate::encryption::encrypt::get_column_crypto_metadata;
42use crate::errors::{ParquetError, Result};
43use crate::file::metadata::{
44    ColumnChunkMetaData, ColumnChunkMetaDataBuilder, ColumnIndexBuilder, LevelHistogram,
45    OffsetIndexBuilder, PageEncodingStats,
46};
47use crate::file::properties::{
48    EnabledStatistics, ResolvedColumnProperties, WriterProperties, WriterPropertiesPtr,
49    WriterVersion,
50};
51use crate::file::statistics::{Statistics, ValueStatistics};
52use crate::schema::types::{BasicTypeInfo, ColumnDescPtr, ColumnDescriptor};
53
54mod byte_budget_chunker;
55pub(crate) mod encoder;
56pub(crate) mod min_max;
57
58use byte_budget_chunker::{ByteBudgetChunker, SubBatchStrategy};
59
60macro_rules! downcast_writer {
61    ($e:expr, $i:ident, $b:expr) => {
62        match $e {
63            Self::BoolColumnWriter($i) => $b,
64            Self::Int32ColumnWriter($i) => $b,
65            Self::Int64ColumnWriter($i) => $b,
66            Self::Int96ColumnWriter($i) => $b,
67            Self::FloatColumnWriter($i) => $b,
68            Self::DoubleColumnWriter($i) => $b,
69            Self::ByteArrayColumnWriter($i) => $b,
70            Self::FixedLenByteArrayColumnWriter($i) => $b,
71        }
72    };
73}
74
75/// Column writer for a Parquet type.
76///
77/// See [`get_column_writer`] to create instances of this type
78pub enum ColumnWriter<'a> {
79    /// Column writer for boolean type
80    BoolColumnWriter(ColumnWriterImpl<'a, BoolType>),
81    /// Column writer for int32 type
82    Int32ColumnWriter(ColumnWriterImpl<'a, Int32Type>),
83    /// Column writer for int64 type
84    Int64ColumnWriter(ColumnWriterImpl<'a, Int64Type>),
85    /// Column writer for int96 (timestamp) type
86    Int96ColumnWriter(ColumnWriterImpl<'a, Int96Type>),
87    /// Column writer for float type
88    FloatColumnWriter(ColumnWriterImpl<'a, FloatType>),
89    /// Column writer for double type
90    DoubleColumnWriter(ColumnWriterImpl<'a, DoubleType>),
91    /// Column writer for byte array type
92    ByteArrayColumnWriter(ColumnWriterImpl<'a, ByteArrayType>),
93    /// Column writer for fixed length byte array type
94    FixedLenByteArrayColumnWriter(ColumnWriterImpl<'a, FixedLenByteArrayType>),
95}
96
97impl ColumnWriter<'_> {
98    /// Returns the estimated total memory usage
99    #[cfg(feature = "arrow")]
100    pub(crate) fn memory_size(&self) -> usize {
101        downcast_writer!(self, typed, typed.memory_size())
102    }
103
104    /// Returns the estimated total encoded bytes for this column writer
105    #[cfg(feature = "arrow")]
106    pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
107        downcast_writer!(self, typed, typed.get_estimated_total_bytes())
108    }
109
110    /// Finalize the currently buffered values as a data page.
111    ///
112    /// This is used by content-defined chunking to force a page boundary at
113    /// content-determined positions.
114    #[cfg(feature = "arrow")]
115    pub(crate) fn add_data_page(&mut self) -> Result<()> {
116        downcast_writer!(self, typed, typed.add_data_page())
117    }
118
119    /// Sets a pre-computed distinct count on this column writer.
120    ///
121    /// See [`GenericColumnWriter::set_distinct_count_override`] for details.
122    #[cfg(feature = "arrow")]
123    pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
124        downcast_writer!(self, typed, typed.set_distinct_count_override(count))
125    }
126
127    /// Close this [`ColumnWriter`], returning the metadata for the column chunk.
128    pub fn close(self) -> Result<ColumnCloseResult> {
129        downcast_writer!(self, typed, typed.close())
130    }
131}
132
133/// Create a specific column writer corresponding to column descriptor `descr`.
134pub fn get_column_writer<'a>(
135    descr: ColumnDescPtr,
136    props: WriterPropertiesPtr,
137    page_writer: Box<dyn PageWriter + 'a>,
138) -> ColumnWriter<'a> {
139    match descr.physical_type() {
140        Type::BOOLEAN => {
141            ColumnWriter::BoolColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
142        }
143        Type::INT32 => {
144            ColumnWriter::Int32ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
145        }
146        Type::INT64 => {
147            ColumnWriter::Int64ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
148        }
149        Type::INT96 => {
150            ColumnWriter::Int96ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
151        }
152        Type::FLOAT => {
153            ColumnWriter::FloatColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
154        }
155        Type::DOUBLE => {
156            ColumnWriter::DoubleColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
157        }
158        Type::BYTE_ARRAY => {
159            ColumnWriter::ByteArrayColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
160        }
161        Type::FIXED_LEN_BYTE_ARRAY => ColumnWriter::FixedLenByteArrayColumnWriter(
162            ColumnWriterImpl::new(descr, props, page_writer),
163        ),
164    }
165}
166
167/// Gets a typed column writer for the specific type `T`, by "up-casting" `col_writer` of
168/// non-generic type to a generic column writer type `ColumnWriterImpl`.
169///
170/// # Panics
171///
172/// Panics if actual enum value for `col_writer` does not match the type `T`.
173pub fn get_typed_column_writer<T: DataType>(col_writer: ColumnWriter) -> ColumnWriterImpl<T> {
174    T::get_column_writer(col_writer).unwrap_or_else(|| {
175        panic!(
176            "Failed to convert column writer into a typed column writer for `{}` type",
177            T::get_physical_type()
178        )
179    })
180}
181
182/// Similar to `get_typed_column_writer` but returns a reference.
183pub fn get_typed_column_writer_ref<'a, 'b: 'a, T: DataType>(
184    col_writer: &'b ColumnWriter<'a>,
185) -> &'b ColumnWriterImpl<'a, T> {
186    T::get_column_writer_ref(col_writer).unwrap_or_else(|| {
187        panic!(
188            "Failed to convert column writer into a typed column writer for `{}` type",
189            T::get_physical_type()
190        )
191    })
192}
193
194/// Similar to `get_typed_column_writer` but returns a reference.
195pub fn get_typed_column_writer_mut<'a, 'b: 'a, T: DataType>(
196    col_writer: &'a mut ColumnWriter<'b>,
197) -> &'a mut ColumnWriterImpl<'b, T> {
198    T::get_column_writer_mut(col_writer).unwrap_or_else(|| {
199        panic!(
200            "Failed to convert column writer into a typed column writer for `{}` type",
201            T::get_physical_type()
202        )
203    })
204}
205
206/// Metadata for a column chunk of a Parquet file.
207///
208/// Note this structure is returned by [`ColumnWriter::close`].
209#[derive(Debug, Clone)]
210pub struct ColumnCloseResult {
211    /// The total number of bytes written
212    pub bytes_written: u64,
213    /// The total number of rows written
214    pub rows_written: u64,
215    /// Metadata for this column chunk
216    pub metadata: ColumnChunkMetaData,
217    /// Optional bloom filter for this column
218    pub bloom_filter: Option<Sbbf>,
219    /// Optional column index, for filtering
220    pub column_index: Option<ColumnIndexMetaData>,
221    /// Optional offset index, identifying page locations
222    pub offset_index: Option<OffsetIndexMetaData>,
223}
224
225impl ColumnCloseResult {
226    /// Rewrite the page offsets for a dictionary-first on-disk layout.
227    ///
228    /// A writer that buffers the whole column chunk and splices it later (the
229    /// Arrow path) may accept the data pages *before* the dictionary page so the
230    /// data pages can stream straight through, then emit the dictionary page
231    /// first at splice. The offsets recorded during encoding therefore assume a
232    /// data-pages-first layout; call this with the serialized length of the
233    /// dictionary page to move it to offset 0 and shift every data page after
234    /// it. A `dictionary_len` of 0 (no dictionary page) leaves the result
235    /// unchanged.
236    pub fn update_dictionary_location(mut self, dictionary_len: usize) -> Result<Self> {
237        if dictionary_len > 0 {
238            self.metadata = self
239                .metadata
240                .into_builder()
241                .set_dictionary_page_offset(Some(0))
242                .set_data_page_offset(dictionary_len as i64)
243                .build()?;
244            if let Some(offset_index) = self.offset_index.as_mut() {
245                let mut offset = dictionary_len as i64;
246                for location in &mut offset_index.page_locations {
247                    location.offset = offset;
248                    offset += location.compressed_page_size as i64;
249                }
250            }
251        }
252        Ok(self)
253    }
254}
255
256// Metrics per page
257#[derive(Default)]
258struct PageMetrics {
259    num_buffered_values: u32,
260    num_buffered_rows: u32,
261    /// Encoded bytes that the data page byte limit does not apply to,
262    /// because they belong to the page's mandatory first value and cannot be
263    /// moved elsewhere. Zero unless that value alone exceeded the limit
264    /// *and* the encoding compresses against the preceding value; see
265    /// [`ColumnValueEncoder::compresses_against_previous_value`].
266    page_size_exemption: usize,
267    num_page_nulls: u64,
268    num_page_nans: Option<u64>,
269    repetition_level_histogram: Option<LevelHistogram>,
270    definition_level_histogram: Option<LevelHistogram>,
271}
272
273impl PageMetrics {
274    fn new() -> Self {
275        Default::default()
276    }
277
278    /// Initialize the repetition level histogram
279    fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
280        self.repetition_level_histogram = LevelHistogram::try_new(max_level);
281        self
282    }
283
284    /// Initialize the definition level histogram
285    fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
286        self.definition_level_histogram = LevelHistogram::try_new(max_level);
287        self
288    }
289
290    /// Resets the state of this `PageMetrics` to the initial state.
291    /// If histograms have been initialized their contents will be reset to zero.
292    fn new_page(&mut self) {
293        self.num_buffered_values = 0;
294        self.num_buffered_rows = 0;
295        self.page_size_exemption = 0;
296        self.num_page_nulls = 0;
297        self.num_page_nans = None;
298        self.repetition_level_histogram
299            .as_mut()
300            .map(LevelHistogram::reset);
301        self.definition_level_histogram
302            .as_mut()
303            .map(LevelHistogram::reset);
304    }
305}
306
307// Metrics per column writer
308#[derive(Default)]
309struct ColumnMetrics<T: Default> {
310    total_bytes_written: u64,
311    total_rows_written: u64,
312    total_uncompressed_size: u64,
313    total_compressed_size: u64,
314    total_num_values: u64,
315    dictionary_page_offset: Option<u64>,
316    data_page_offset: Option<u64>,
317    min_column_value: Option<T>,
318    max_column_value: Option<T>,
319    num_column_nulls: u64,
320    num_column_nans: Option<u64>,
321    column_distinct_count: Option<u64>,
322    variable_length_bytes: Option<i64>,
323    repetition_level_histogram: Option<LevelHistogram>,
324    definition_level_histogram: Option<LevelHistogram>,
325}
326
327impl<T: Default> ColumnMetrics<T> {
328    fn new() -> Self {
329        Default::default()
330    }
331
332    /// Initialize the repetition level histogram
333    fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
334        self.repetition_level_histogram = LevelHistogram::try_new(max_level);
335        self
336    }
337
338    /// Initialize the definition level histogram
339    fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
340        self.definition_level_histogram = LevelHistogram::try_new(max_level);
341        self
342    }
343
344    /// Sum `page_histogram` into `chunk_histogram`
345    fn update_histogram(
346        chunk_histogram: &mut Option<LevelHistogram>,
347        page_histogram: Option<&LevelHistogram>,
348    ) {
349        if let (Some(page_hist), Some(chunk_hist)) = (page_histogram, chunk_histogram) {
350            chunk_hist.add(page_hist);
351        }
352    }
353
354    /// Sum the provided PageMetrics histograms into the chunk histograms. Does nothing if
355    /// page histograms are not initialized.
356    fn update_from_page_metrics(&mut self, page_metrics: &PageMetrics) {
357        ColumnMetrics::<T>::update_histogram(
358            &mut self.definition_level_histogram,
359            page_metrics.definition_level_histogram.as_ref(),
360        );
361        ColumnMetrics::<T>::update_histogram(
362            &mut self.repetition_level_histogram,
363            page_metrics.repetition_level_histogram.as_ref(),
364        );
365    }
366
367    /// Sum the provided page variable_length_bytes into the chunk variable_length_bytes
368    fn update_variable_length_bytes(&mut self, variable_length_bytes: Option<i64>) {
369        if let Some(var_bytes) = variable_length_bytes {
370            *self.variable_length_bytes.get_or_insert(0) += var_bytes;
371        }
372    }
373}
374
375/// Borrowed view of level data, analogous to `&str` for `LevelData`'s `String`.
376///
377/// `LevelDataRef` can be constructed from `LevelData` and directly from an existing
378/// `&[i16]` without allocating.
379///
380/// The variants are different physical representations of the same logical
381/// sequence of levels.
382#[derive(Debug, Clone, Copy)]
383pub(crate) enum LevelDataRef<'a> {
384    Absent,
385    Materialized(&'a [i16]),
386    Uniform { value: i16, count: usize },
387}
388
389impl<'a> From<&'a [i16]> for LevelDataRef<'a> {
390    fn from(levels: &'a [i16]) -> Self {
391        Self::Materialized(levels)
392    }
393}
394
395impl<'a> From<Option<&'a [i16]>> for LevelDataRef<'a> {
396    fn from(levels: Option<&'a [i16]>) -> Self {
397        levels.map_or(Self::Absent, Self::from)
398    }
399}
400
401impl LevelDataRef<'_> {
402    pub(crate) fn len(self) -> usize {
403        match self {
404            Self::Absent => 0,
405            Self::Materialized(values) => values.len(),
406            Self::Uniform { count, .. } => count,
407        }
408    }
409
410    pub(crate) fn first(self) -> Option<i16> {
411        match self {
412            Self::Absent => None,
413            Self::Materialized(values) => values.first().copied(),
414            Self::Uniform { value, count } => (count > 0).then_some(value),
415        }
416    }
417
418    #[cfg(feature = "arrow")]
419    pub(crate) fn value_at(self, idx: usize) -> Option<i16> {
420        match self {
421            Self::Absent => None,
422            Self::Materialized(values) => values.get(idx).copied(),
423            Self::Uniform { value, count } => (idx < count).then_some(value),
424        }
425    }
426
427    pub(crate) fn slice(self, offset: usize, len: usize) -> Self {
428        match self {
429            Self::Absent => Self::Absent,
430            Self::Materialized(values) => Self::Materialized(&values[offset..offset + len]),
431            Self::Uniform { value, .. } => Self::Uniform { value, count: len },
432        }
433    }
434
435    /// Count of positions in this slice that represent an actual value
436    /// (definition level equal to `max_def`). `Absent` means the column has
437    /// `max_def == 0` and every position is a value, so the implicit count
438    /// is the caller-supplied `total`.
439    pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize {
440        match self {
441            Self::Absent => total,
442            Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(),
443            Self::Uniform { value, count } => {
444                if value == max_def {
445                    count
446                } else {
447                    0
448                }
449            }
450        }
451    }
452}
453
454/// Typed column writer for a primitive column.
455pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl<T>>;
456
457/// Generic column writer for a primitive Parquet column
458pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> {
459    // Column writer properties
460    descr: ColumnDescPtr,
461    props: WriterPropertiesPtr,
462    /// Per-column settings for [`Self::descr`], resolved once here so that the
463    /// per-batch and per-page write paths never search the per-column override
464    /// map in `props` again.
465    column_props: ResolvedColumnProperties,
466
467    page_writer: Box<dyn PageWriter + 'a>,
468    codec: Compression,
469    compressor: Option<Box<dyn Codec>>,
470    encoder: E,
471
472    page_metrics: PageMetrics,
473    // Metrics per column writer
474    column_metrics: ColumnMetrics<E::T>,
475
476    /// Pre-computed distinct count to write into column chunk statistics.
477    /// When set, takes precedence over `column_metrics.column_distinct_count`.
478    distinct_count_override: Option<u64>,
479
480    /// The order of encodings within the generated metadata does not impact its meaning,
481    /// but we use a BTreeSet so that the output is deterministic
482    encodings: BTreeSet<Encoding>,
483    encoding_stats: Vec<PageEncodingStats>,
484    // Streaming level encoders for definition/repetition levels.
485    def_levels_encoder: LevelEncoder,
486    rep_levels_encoder: LevelEncoder,
487    data_pages: VecDeque<CompressedPage>,
488    // column index and offset index
489    column_index_builder: ColumnIndexBuilder,
490    offset_index_builder: Option<OffsetIndexBuilder>,
491
492    // Below fields used to incrementally check boundary order across data pages.
493    // We assume they are ascending/descending until proven wrong.
494    data_page_boundary_ascending: bool,
495    data_page_boundary_descending: bool,
496    /// (min, max)
497    last_non_null_data_page_min_max: Option<(E::T, E::T)>,
498}
499
500impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
501    /// Returns a new instance of [`GenericColumnWriter`].
502    pub fn new(
503        descr: ColumnDescPtr,
504        props: WriterPropertiesPtr,
505        page_writer: Box<dyn PageWriter + 'a>,
506    ) -> Self {
507        let column_props = props.resolve_column_properties(descr.path());
508        let codec = column_props.compression;
509        let codec_options = CodecOptionsBuilder::default().build();
510        let compressor = create_codec(codec, &codec_options).unwrap();
511        let encoder = E::try_new(&descr, props.as_ref(), &column_props).unwrap();
512
513        let statistics_enabled = column_props.statistics_enabled;
514
515        let mut encodings = BTreeSet::new();
516        // Used for level information
517        encodings.insert(Encoding::RLE);
518
519        let mut page_metrics = PageMetrics::new();
520        let mut column_metrics = ColumnMetrics::<E::T>::new();
521
522        // Initialize level histograms if collecting page or chunk statistics
523        if statistics_enabled != EnabledStatistics::None {
524            page_metrics = page_metrics
525                .with_repetition_level_histogram(descr.max_rep_level())
526                .with_definition_level_histogram(descr.max_def_level());
527            column_metrics = column_metrics
528                .with_repetition_level_histogram(descr.max_rep_level())
529                .with_definition_level_histogram(descr.max_def_level())
530        }
531
532        // Disable column_index_builder if not collecting page statistics.
533        let mut column_index_builder = ColumnIndexBuilder::new(descr.physical_type());
534        if statistics_enabled != EnabledStatistics::Page {
535            column_index_builder.to_invalid()
536        }
537
538        // Disable offset_index_builder if requested by user.
539        let offset_index_builder = match props.offset_index_disabled() {
540            false => Some(OffsetIndexBuilder::new()),
541            _ => None,
542        };
543
544        Self {
545            def_levels_encoder: Self::create_level_encoder(descr.max_def_level(), &props),
546            rep_levels_encoder: Self::create_level_encoder(descr.max_rep_level(), &props),
547            descr,
548            props,
549            column_props,
550            page_writer,
551            codec,
552            compressor,
553            encoder,
554            data_pages: VecDeque::new(),
555            page_metrics,
556            column_metrics,
557            distinct_count_override: None,
558            column_index_builder,
559            offset_index_builder,
560            encodings,
561            encoding_stats: vec![],
562            data_page_boundary_ascending: true,
563            data_page_boundary_descending: true,
564            last_non_null_data_page_min_max: None,
565        }
566    }
567
568    /// Sets a pre-computed distinct count to write into column chunk statistics.
569    ///
570    /// When set, this value is written as `distinct_count` in the row group statistics
571    /// footer. It takes precedence over any `distinct_count` passed through
572    /// [`Self::write_batch_with_statistics`].
573    #[cfg(feature = "arrow")]
574    pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
575        self.distinct_count_override = Some(count);
576    }
577
578    #[expect(clippy::too_many_arguments)]
579    pub(crate) fn write_batch_internal(
580        &mut self,
581        values: &E::Values,
582        value_indices: Option<&[usize]>,
583        def_levels: LevelDataRef<'_>,
584        rep_levels: LevelDataRef<'_>,
585        min: Option<&E::T>,
586        max: Option<&E::T>,
587        distinct_count: Option<u64>,
588    ) -> Result<usize> {
589        // Check if number of definition levels is the same as number of repetition levels.
590        if def_levels.len() != 0 && rep_levels.len() != 0 && def_levels.len() != rep_levels.len() {
591            return Err(general_err!(
592                "Inconsistent length of definition and repetition levels: {} != {}",
593                def_levels.len(),
594                rep_levels.len()
595            ));
596        }
597
598        // We check for DataPage limits only after we have inserted the values. If a user
599        // writes a large number of values, the DataPage size can be well above the limit.
600        //
601        // The purpose of this chunking is to bound this. Even if a user writes large
602        // number of values, the chunking will ensure that we add data page at a
603        // reasonable pagesize limit.
604
605        // TODO: find out why we don't account for size of levels when we estimate page
606        // size.
607        let num_levels = def_levels.len().max(rep_levels.len());
608        let num_levels = if num_levels > 0 {
609            num_levels
610        } else {
611            value_indices.map_or_else(|| values.len(), |i| i.len())
612        };
613
614        if let Some(min) = min {
615            update_min(&self.descr, min, &mut self.column_metrics.min_column_value);
616        }
617        if let Some(max) = max {
618            update_max(&self.descr, max, &mut self.column_metrics.max_column_value);
619        }
620
621        // Encoder counts reset per page; row metrics retain column-wide history.
622        let has_prior_data = self.column_metrics.total_rows_written != 0
623            || self.page_metrics.num_buffered_values != 0;
624        self.column_metrics.column_distinct_count =
625            if has_prior_data { None } else { distinct_count };
626
627        let mut values_offset = 0;
628        let mut levels_offset = 0;
629        let both_levels_compact = !matches!(def_levels, LevelDataRef::Materialized(_))
630            && !matches!(rep_levels, LevelDataRef::Materialized(_));
631        let has_levels = !matches!(def_levels, LevelDataRef::Absent)
632            || !matches!(rep_levels, LevelDataRef::Absent);
633
634        // When both level vectors are compact (Uniform or Absent), there is no
635        // materialized slice to split and the per-mini-batch work is O(1), so we
636        // can safely use a much larger batch size.
637        let base_batch_size = if both_levels_compact && has_levels {
638            self.column_props.data_page_row_count_limit
639        } else {
640            self.props.write_batch_size()
641        };
642        debug_assert!(base_batch_size > 0);
643
644        let chunker = ByteBudgetChunker::new(&self.descr, &self.column_props, base_batch_size);
645        while levels_offset < num_levels {
646            let mut end_offset = num_levels.min(levels_offset + base_batch_size);
647
648            // Split at record boundary
649            if let LevelDataRef::Materialized(levels) = rep_levels {
650                while end_offset < levels.len() && levels[end_offset] != 0 {
651                    end_offset += 1;
652                }
653            }
654
655            let chunk_size = end_offset - levels_offset;
656            let chunk_def = def_levels.slice(levels_offset, chunk_size);
657            let chunk_rep = rep_levels.slice(levels_offset, chunk_size);
658
659            // Key decision point: can we write this whole chunk as one
660            // mini-batch (the common case — small or fixed-width values, no
661            // further page-size accounting needed), or must we fall back to
662            // byte-budget-aware sub-batching to keep a page from overshooting
663            // `data_page_size_limit`? `pick_sub_batch` returns `None` for
664            // the former, and otherwise how wide one mini-batch may be.
665            let sub_batch = chunker.pick_sub_batch(
666                &self.encoder,
667                values,
668                value_indices,
669                chunk_def,
670                values_offset,
671                chunk_size,
672            );
673
674            match sub_batch {
675                None => {
676                    values_offset += self.write_mini_batch(
677                        values,
678                        values_offset,
679                        value_indices,
680                        chunk_size,
681                        chunk_def,
682                        chunk_rep,
683                    )?;
684                }
685                Some(sub_batch) => {
686                    values_offset += self.write_granular_chunk(
687                        values,
688                        values_offset,
689                        value_indices,
690                        chunk_size,
691                        chunk_def,
692                        chunk_rep,
693                        sub_batch,
694                    )?;
695                }
696            }
697            levels_offset = end_offset;
698        }
699
700        // Return total number of values processed.
701        Ok(values_offset)
702    }
703
704    /// Writes batch of values, definition levels and repetition levels.
705    /// Returns number of values processed (written).
706    ///
707    /// If definition and repetition levels are provided, we write fully those levels and
708    /// select how many values to write (this number will be returned), since number of
709    /// actual written values may be smaller than provided values.
710    ///
711    /// If only values are provided, then all values are written and the length of
712    /// of the values buffer is returned.
713    ///
714    /// Definition and/or repetition levels can be omitted, if values are
715    /// non-nullable and/or non-repeated.
716    pub fn write_batch(
717        &mut self,
718        values: &E::Values,
719        def_levels: Option<&[i16]>,
720        rep_levels: Option<&[i16]>,
721    ) -> Result<usize> {
722        self.write_batch_internal(
723            values,
724            None,
725            LevelDataRef::from(def_levels),
726            LevelDataRef::from(rep_levels),
727            None,
728            None,
729            None,
730        )
731    }
732
733    /// Writer may optionally provide pre-calculated statistics for use when computing
734    /// chunk-level statistics
735    ///
736    /// NB: [`WriterProperties::statistics_enabled`] must be set to [`EnabledStatistics::Chunk`]
737    /// for these statistics to take effect. If [`EnabledStatistics::None`] they will be ignored,
738    /// and if [`EnabledStatistics::Page`] the chunk statistics will instead be computed from the
739    /// computed page statistics
740    pub fn write_batch_with_statistics(
741        &mut self,
742        values: &E::Values,
743        def_levels: Option<&[i16]>,
744        rep_levels: Option<&[i16]>,
745        min: Option<&E::T>,
746        max: Option<&E::T>,
747        distinct_count: Option<u64>,
748    ) -> Result<usize> {
749        self.write_batch_internal(
750            values,
751            None,
752            LevelDataRef::from(def_levels),
753            LevelDataRef::from(rep_levels),
754            min,
755            max,
756            distinct_count,
757        )
758    }
759
760    /// Returns the estimated total memory usage.
761    ///
762    /// Unlike [`Self::get_estimated_total_bytes`] this is an estimate
763    /// of the current memory usage and not the final anticipated encoded size.
764    #[cfg(feature = "arrow")]
765    pub(crate) fn memory_size(&self) -> usize {
766        // In-flight encoder buffers, plus any completed pages still held on the
767        // heap: the dictionary-column data pages buffered here (column-at-a-time
768        // path), plus whatever the page writer keeps resident. A page writer
769        // that spills completed pages off-heap reports far less than the bytes
770        // it was handed, so this tracks real memory rather than bytes written.
771        self.encoder.estimated_memory_size()
772            + self
773                .data_pages
774                .iter()
775                .map(|page| page.memory_usage())
776                .sum::<usize>()
777            + self.page_writer.buffered_memory_size()
778    }
779
780    /// Returns total number of bytes written by this column writer so far.
781    /// This value is also returned when column writer is closed.
782    ///
783    /// Note: this value does not include any buffered data that has not
784    /// yet been flushed to a page.
785    pub fn get_total_bytes_written(&self) -> u64 {
786        self.column_metrics.total_bytes_written
787    }
788
789    /// Returns the estimated total encoded bytes for this column writer.
790    ///
791    /// Unlike [`Self::get_total_bytes_written`] this includes an estimate
792    /// of any data that has not yet been flushed to a page, based on it's
793    /// anticipated encoded size.
794    #[cfg(feature = "arrow")]
795    pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
796        self.data_pages
797            .iter()
798            .map(|page| page.data().len() as u64)
799            .sum::<u64>()
800            + self.column_metrics.total_bytes_written
801            + self.encoder.estimated_data_page_size() as u64
802            + self.encoder.estimated_dict_page_size().unwrap_or_default() as u64
803    }
804
805    /// Returns total number of rows written by this column writer so far.
806    /// This value is also returned when column writer is closed.
807    pub fn get_total_rows_written(&self) -> u64 {
808        self.column_metrics.total_rows_written
809    }
810
811    /// Returns a reference to a [`ColumnDescPtr`]
812    pub fn get_descriptor(&self) -> &ColumnDescPtr {
813        &self.descr
814    }
815
816    /// Finalizes writes and closes the column writer.
817    /// Returns total bytes written, total rows written and column chunk metadata.
818    pub fn close(mut self) -> Result<ColumnCloseResult> {
819        if self.page_metrics.num_buffered_values > 0 {
820            self.add_data_page()?;
821        }
822        if self.encoder.has_dictionary() {
823            self.write_dictionary_page()?;
824        }
825        self.flush_data_pages()?;
826        let metadata = self.build_column_metadata()?;
827        self.page_writer.close()?;
828
829        let write_bloom_filter = self.props.bloom_filter_for_dictionary_encoded_chunks()
830            || self.has_non_dictionary_data_page();
831        let bloom_filter = self
832            .encoder
833            .flush_bloom_filter()
834            .filter(|_| write_bloom_filter);
835
836        let boundary_order = match (
837            self.data_page_boundary_ascending,
838            self.data_page_boundary_descending,
839        ) {
840            // If the lists are composed of equal elements then will be marked as ascending
841            // (Also the case if all pages are null pages)
842            (true, _) => BoundaryOrder::ASCENDING,
843            (false, true) => BoundaryOrder::DESCENDING,
844            (false, false) => BoundaryOrder::UNORDERED,
845        };
846        self.column_index_builder.set_boundary_order(boundary_order);
847
848        let column_index = match self.column_index_builder.valid() {
849            true => Some(self.column_index_builder.build()?),
850            false => None,
851        };
852
853        let offset_index = self.offset_index_builder.map(|b| b.build());
854
855        Ok(ColumnCloseResult {
856            bytes_written: self.column_metrics.total_bytes_written,
857            rows_written: self.column_metrics.total_rows_written,
858            bloom_filter,
859            metadata,
860            column_index,
861            offset_index,
862        })
863    }
864
865    /// Writes a chunk in sub-batches sized by `sub_batch`, checking the page
866    /// byte limit after each. This keeps the page size close to
867    /// `data_page_size_limit` instead of overshooting it by a whole chunk.
868    ///
869    /// [`SubBatchStrategy::Values`] windows are cut on value boundaries,
870    /// walking definition levels to find where the value budget is used up;
871    /// [`SubBatchStrategy::Levels`] windows are a fixed level count. See
872    /// [`SubBatchStrategy`] for which budget gets which and why.
873    ///
874    /// For repeated/nested columns sub-batches then extend to the next
875    /// `rep == 0` boundary so a record never spans data pages, matching the
876    /// parquet format rule. A record holding several over-limit values
877    /// therefore still exceeds the budget; that is inherent to the format.
878    ///
879    /// Returns the total number of values consumed across all sub-batches.
880    ///
881    /// `#[inline(never)]` keeps this slow path — only reached for
882    /// variable-width columns whose values need page splitting — out of
883    /// the hot `write_batch_internal` loop.
884    #[expect(clippy::too_many_arguments)]
885    #[inline(never)]
886    fn write_granular_chunk(
887        &mut self,
888        values: &E::Values,
889        values_offset: usize,
890        value_indices: Option<&[usize]>,
891        chunk_size: usize,
892        chunk_def: LevelDataRef<'_>,
893        chunk_rep: LevelDataRef<'_>,
894        sub_batch: SubBatchStrategy,
895    ) -> Result<usize> {
896        // The chunker always sizes a sub-batch to at least one value or one
897        // level, and a value-exact window spans at least one level, so each
898        // iteration below makes progress (`sub_end > sub_start`).
899        debug_assert!(
900            matches!(sub_batch, SubBatchStrategy::Values(n) | SubBatchStrategy::Levels(n) if n >= 1),
901            "chunker must size at least one value or level"
902        );
903        let max_def_level = self.descr.max_def_level();
904        let mut values_consumed = 0;
905        let mut sub_start = 0;
906        while sub_start < chunk_size {
907            let window_end = match sub_batch {
908                SubBatchStrategy::Values(n) => {
909                    Self::window_end_for_values(chunk_def, chunk_size, max_def_level, sub_start, n)
910                }
911                SubBatchStrategy::Levels(n) => (sub_start + n).min(chunk_size),
912            };
913            let sub_end = match chunk_rep {
914                LevelDataRef::Materialized(levels) => {
915                    // Extend the window to the next record boundary
916                    // (rep == 0) so a record never spans data pages. Packing
917                    // whole records rather than stepping one record at a time
918                    // avoids calling `write_mini_batch` per record: records
919                    // average only a handful of levels, so a record-at-a-time
920                    // step would issue many more mini-batches than necessary.
921                    let mut e = window_end;
922                    while e < chunk_size && levels[e] != 0 {
923                        e += 1;
924                    }
925                    e
926                }
927                _ => window_end,
928            };
929            let sub_len = sub_end - sub_start;
930            let written = self.write_mini_batch(
931                values,
932                values_offset + values_consumed,
933                value_indices,
934                sub_len,
935                chunk_def.slice(sub_start, sub_len),
936                chunk_rep.slice(sub_start, sub_len),
937            )?;
938            values_consumed += written;
939            sub_start = sub_end;
940        }
941        Ok(values_consumed)
942    }
943
944    fn has_non_dictionary_data_page(&self) -> bool {
945        self.encoding_stats.iter().any(|stats| {
946            matches!(
947                stats.page_type,
948                PageType::DATA_PAGE | PageType::DATA_PAGE_V2
949            ) && !matches!(
950                stats.encoding,
951                Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY
952            )
953        })
954    }
955
956    /// Index one past the last level of a sub-batch window that starts at
957    /// `start` and covers at most `max_values` values, clamped to
958    /// `chunk_size`.
959    ///
960    /// Nulls trailing the last value are left to the next window, so a window
961    /// ends immediately after the value that exhausts its budget. The window
962    /// always spans at least one level, so callers make progress: `start` is
963    /// below `chunk_size` and `max_values` is at least one.
964    ///
965    /// Only nullable and nested columns pay the level walk. It runs solely in
966    /// the granular path, whose values are by definition large enough to
967    /// overflow a page budget, so touching each of the chunk's levels once is
968    /// noise next to writing those values — and that path already makes full
969    /// def-level and rep-level passes.
970    fn window_end_for_values(
971        chunk_def: LevelDataRef<'_>,
972        chunk_size: usize,
973        max_def_level: i16,
974        start: usize,
975        max_values: usize,
976    ) -> usize {
977        match chunk_def {
978            // `max_def_level == 0`: every level is a value.
979            LevelDataRef::Absent => (start + max_values).min(chunk_size),
980            LevelDataRef::Uniform { value, .. } => {
981                if value == max_def_level {
982                    (start + max_values).min(chunk_size)
983                } else {
984                    // Uniformly below max def: the chunk holds no values at
985                    // all, so no window boundary can help. (The chunker
986                    // returns `None` for such a chunk, so this is defensive.)
987                    chunk_size
988                }
989            }
990            LevelDataRef::Materialized(levels) => {
991                let mut seen = 0;
992                let mut end = start;
993                while end < chunk_size && seen < max_values {
994                    if levels[end] == max_def_level {
995                        seen += 1;
996                    }
997                    end += 1;
998                }
999                end
1000            }
1001        }
1002    }
1003
1004    /// Creates a new streaming level encoder appropriate for the writer version.
1005    fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder {
1006        match props.writer_version() {
1007            WriterVersion::PARQUET_1_0 => LevelEncoder::v1_streaming(max_level),
1008            WriterVersion::PARQUET_2_0 => LevelEncoder::v2_streaming(max_level),
1009        }
1010    }
1011
1012    /// Writes mini batch of values, definition and repetition levels.
1013    /// This allows fine-grained processing of values and maintaining a reasonable
1014    /// page size.
1015    fn write_mini_batch(
1016        &mut self,
1017        values: &E::Values,
1018        values_offset: usize,
1019        value_indices: Option<&[usize]>,
1020        num_levels: usize,
1021        def_levels: LevelDataRef<'_>,
1022        rep_levels: LevelDataRef<'_>,
1023    ) -> Result<usize> {
1024        // Process definition levels and determine how many values to write.
1025        let values_to_write = if self.descr.max_def_level() > 0 {
1026            let max_def = self.descr.max_def_level();
1027            match def_levels {
1028                LevelDataRef::Absent => {
1029                    return Err(general_err!(
1030                        "Definition levels are required, because max definition level = {}",
1031                        self.descr.max_def_level()
1032                    ));
1033                }
1034                LevelDataRef::Materialized(levels) => {
1035                    // General path for caller-provided or already-materialized
1036                    // level buffers.
1037                    let mut values_to_write = 0usize;
1038                    let encoder = &mut self.def_levels_encoder;
1039                    match self.page_metrics.definition_level_histogram.as_mut() {
1040                        Some(histogram) => encoder.put_with_observer(levels, |level, count| {
1041                            values_to_write += count * (level == max_def) as usize;
1042                            histogram.increment_by(level, count as i64);
1043                        }),
1044                        None => encoder.put_with_observer(levels, |level, count| {
1045                            values_to_write += count * (level == max_def) as usize;
1046                        }),
1047                    };
1048                    self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64;
1049                    values_to_write
1050                }
1051                LevelDataRef::Uniform { value, count } => {
1052                    // Fast path for all-null, all-valid, or otherwise uniform
1053                    // definition levels without materializing a level buffer.
1054                    let encoder = &mut self.def_levels_encoder;
1055                    match self.page_metrics.definition_level_histogram.as_mut() {
1056                        Some(histogram) => {
1057                            encoder.put_n_with_observer(value, count, |level, run_len| {
1058                                histogram.increment_by(level, run_len as i64);
1059                            })
1060                        }
1061                        None => encoder.put_n_with_observer(value, count, |_, _| {}),
1062                    }
1063                    let values_to_write = count * (value == max_def) as usize;
1064                    self.page_metrics.num_page_nulls += (count - values_to_write) as u64;
1065                    values_to_write
1066                }
1067            }
1068        } else {
1069            num_levels
1070        };
1071
1072        // Process repetition levels and determine how many rows we are about to process.
1073        if self.descr.max_rep_level() > 0 {
1074            // A row could contain more than one value.
1075            let first_level = rep_levels.first().ok_or_else(|| {
1076                general_err!(
1077                    "Repetition levels are required, because max repetition level = {}",
1078                    self.descr.max_rep_level()
1079                )
1080            })?;
1081
1082            if first_level != 0 {
1083                return Err(general_err!(
1084                    "Write must start at a record boundary, got non-zero repetition level of {}",
1085                    first_level
1086                ));
1087            }
1088
1089            let mut new_rows = 0u32;
1090            match rep_levels {
1091                LevelDataRef::Absent => unreachable!(),
1092                LevelDataRef::Materialized(levels) => {
1093                    let encoder = &mut self.rep_levels_encoder;
1094                    match self.page_metrics.repetition_level_histogram.as_mut() {
1095                        Some(histogram) => encoder.put_with_observer(levels, |level, count| {
1096                            new_rows += (count as u32) * (level == 0) as u32;
1097                            histogram.increment_by(level, count as i64);
1098                        }),
1099                        None => encoder.put_with_observer(levels, |level, count| {
1100                            new_rows += (count as u32) * (level == 0) as u32;
1101                        }),
1102                    };
1103                }
1104                LevelDataRef::Uniform { value, count } => {
1105                    let encoder = &mut self.rep_levels_encoder;
1106                    match self.page_metrics.repetition_level_histogram.as_mut() {
1107                        Some(histogram) => {
1108                            encoder.put_n_with_observer(value, count, |level, run_len| {
1109                                new_rows += (run_len as u32) * (level == 0) as u32;
1110                                histogram.increment_by(level, run_len as i64);
1111                            })
1112                        }
1113                        None => encoder.put_n_with_observer(value, count, |level, run_len| {
1114                            new_rows += (run_len as u32) * (level == 0) as u32;
1115                        }),
1116                    }
1117                }
1118            }
1119            self.page_metrics.num_buffered_rows += new_rows;
1120        } else {
1121            // Each value is exactly one row.
1122            // Equals to the number of values, we count nulls as well.
1123            self.page_metrics.num_buffered_rows += num_levels as u32;
1124        }
1125
1126        match value_indices {
1127            Some(indices) => {
1128                let indices = &indices[values_offset..values_offset + values_to_write];
1129                self.encoder.write_gather(values, indices)?;
1130            }
1131            None => self.encoder.write(values, values_offset, values_to_write)?,
1132        }
1133
1134        let page_was_empty = self.page_metrics.num_buffered_values == 0;
1135        self.page_metrics.num_buffered_values += num_levels as u32;
1136
1137        if page_was_empty && values_to_write == 1 {
1138            self.set_page_size_exemption();
1139        }
1140
1141        if self.should_add_data_page() {
1142            self.add_data_page()?;
1143        }
1144
1145        if self.should_dict_fallback() {
1146            self.dict_fallback()?;
1147        }
1148
1149        Ok(values_to_write)
1150    }
1151
1152    /// Returns true if we need to fall back to non-dictionary encoding.
1153    ///
1154    /// We can only fall back if dictionary encoder is set and we have exceeded dictionary
1155    /// size.
1156    #[inline]
1157    fn should_dict_fallback(&self) -> bool {
1158        match self.encoder.estimated_dict_page_size() {
1159            Some(size) => size >= self.column_props.dictionary_page_size_limit,
1160            None => false,
1161        }
1162    }
1163
1164    /// Exempt a page's mandatory first value from the data page byte limit,
1165    /// when that value alone already exceeds it.
1166    ///
1167    /// Parquet requires every data page to hold at least one value, so such a
1168    /// value cannot be split out no matter how the limit is set. Counting it
1169    /// against the limit makes the limit unsatisfiable, and
1170    /// [`Self::should_add_data_page`] then cuts a page after every single
1171    /// value.
1172    ///
1173    /// For `DELTA_BYTE_ARRAY` that costs more than the extra pages. A value is
1174    /// stored as a suffix of the value before it, and a page boundary resets
1175    /// what "the value before it" refers to, so one value per page means every
1176    /// value is stored in full: a column of large values sharing long prefixes
1177    /// writes exactly the bytes `PLAIN` would
1178    /// ([#10489](https://github.com/apache/arrow-rs/issues/10489)).
1179    ///
1180    /// Only encodings that compress against the preceding value opt in, so
1181    /// `PLAIN` and `DELTA_LENGTH_BYTE_ARRAY` keep their tighter one-value page
1182    /// bound.
1183    ///
1184    /// The caller's trigger keys on a page-opening mini-batch holding exactly
1185    /// one value. `write_granular_chunk` cuts windows after an exact value
1186    /// count, so an over-limit value gets a single-value mini-batch whether
1187    /// or not the chunk contains nulls; see
1188    /// `test_column_writer_delta_byte_array_nullable_shared_prefix_dedup`.
1189    /// The exception is a repeated column, where a record holding several
1190    /// over-limit values cannot be split across pages at all.
1191    #[cold]
1192    fn set_page_size_exemption(&mut self) {
1193        if !self.encoder.compresses_against_previous_value() {
1194            return;
1195        }
1196        let size = self.encoder.estimated_data_page_size();
1197        if size >= self.column_props.data_page_size_limit {
1198            self.page_metrics.page_size_exemption = size;
1199        }
1200    }
1201
1202    /// Returns true if there is enough data for a data page, false otherwise.
1203    #[inline]
1204    fn should_add_data_page(&self) -> bool {
1205        // This is necessary in the event of a much larger dictionary size than page size
1206        //
1207        // In such a scenario the dictionary decoder may return an estimated encoded
1208        // size in excess of the page size limit, even when there are no buffered values
1209        if self.page_metrics.num_buffered_values == 0 {
1210            return false;
1211        }
1212
1213        self.page_metrics.num_buffered_rows as usize >= self.column_props.data_page_row_count_limit
1214            || self
1215                .encoder
1216                .estimated_data_page_size()
1217                .saturating_sub(self.page_metrics.page_size_exemption)
1218                >= self.column_props.data_page_size_limit
1219    }
1220
1221    /// Performs dictionary fallback.
1222    /// Prepares and writes dictionary and all data pages into page writer.
1223    fn dict_fallback(&mut self) -> Result<()> {
1224        // At this point we know that we need to fall back.
1225        if self.page_metrics.num_buffered_values > 0 {
1226            self.add_data_page()?;
1227        }
1228        self.write_dictionary_page()?;
1229        self.flush_data_pages()?;
1230        Ok(())
1231    }
1232
1233    // For float columns, always provide Some(n), even if n is 0
1234    // For non-float columns, always provide None
1235    fn get_nan_count<T: ParquetValueType>(&self) -> Option<i64> {
1236        let nan_count = || {
1237            let nan_count = self.page_metrics.num_page_nans.unwrap_or(0);
1238            match i64::try_from(nan_count) {
1239                Ok(count) => Some(count),
1240                _ => Some(i64::MAX),
1241            }
1242        };
1243        match T::PHYSICAL_TYPE {
1244            Type::FLOAT | Type::DOUBLE => nan_count(),
1245            Type::FIXED_LEN_BYTE_ARRAY
1246                if matches!(self.descr.logical_type_ref(), Some(LogicalType::Float16)) =>
1247            {
1248                nan_count()
1249            }
1250            _ => None,
1251        }
1252    }
1253
1254    /// Update the column index and offset index when adding the data page
1255    fn update_column_offset_index(
1256        &mut self,
1257        page_statistics: Option<&ValueStatistics<E::T>>,
1258        page_variable_length_bytes: Option<i64>,
1259    ) {
1260        // update the column index
1261        let null_page =
1262            (self.page_metrics.num_buffered_rows as u64) == self.page_metrics.num_page_nulls;
1263        // a page contains only null values,
1264        // and writers have to set the corresponding entries in min_values and max_values to byte[0]
1265        if null_page && self.column_index_builder.valid() {
1266            self.column_index_builder.append(
1267                null_page,
1268                vec![],
1269                vec![],
1270                self.page_metrics.num_page_nulls as i64,
1271                self.get_nan_count::<E::T>(),
1272            );
1273        } else if self.column_index_builder.valid() {
1274            // from page statistics
1275            // If can't get the page statistics, ignore this column/offset index for this column chunk
1276            match &page_statistics {
1277                None => {
1278                    self.column_index_builder.to_invalid();
1279                }
1280                Some(stat) => {
1281                    // Check if min/max are still ascending/descending across pages
1282                    let new_min = stat.min_opt().unwrap();
1283                    let new_max = stat.max_opt().unwrap();
1284                    if let Some((last_min, last_max)) = &self.last_non_null_data_page_min_max {
1285                        let basic_info = self.descr.get_basic_info();
1286                        if self.data_page_boundary_ascending {
1287                            // If last min/max are greater than new min/max then not ascending anymore
1288                            let not_ascending = compare_greater(basic_info, last_min, new_min)
1289                                || compare_greater(basic_info, last_max, new_max);
1290                            if not_ascending {
1291                                self.data_page_boundary_ascending = false;
1292                            }
1293                        }
1294
1295                        if self.data_page_boundary_descending {
1296                            // If new min/max are greater than last min/max then not descending anymore
1297                            let not_descending = compare_greater(basic_info, new_min, last_min)
1298                                || compare_greater(basic_info, new_max, last_max);
1299                            if not_descending {
1300                                self.data_page_boundary_descending = false;
1301                            }
1302                        }
1303                    }
1304                    self.last_non_null_data_page_min_max = Some((new_min.clone(), new_max.clone()));
1305
1306                    if self.can_truncate_value() {
1307                        self.column_index_builder.append(
1308                            null_page,
1309                            self.truncate_min_value(
1310                                self.props.column_index_truncate_length(),
1311                                stat.min_bytes_opt().unwrap(),
1312                            )
1313                            .0,
1314                            self.truncate_max_value(
1315                                self.props.column_index_truncate_length(),
1316                                stat.max_bytes_opt().unwrap(),
1317                            )
1318                            .0,
1319                            self.page_metrics.num_page_nulls as i64,
1320                            self.get_nan_count::<E::T>(),
1321                        );
1322                    } else {
1323                        self.column_index_builder.append(
1324                            null_page,
1325                            stat.min_bytes_opt().unwrap().to_vec(),
1326                            stat.max_bytes_opt().unwrap().to_vec(),
1327                            self.page_metrics.num_page_nulls as i64,
1328                            self.get_nan_count::<E::T>(),
1329                        );
1330                    }
1331                }
1332            }
1333        }
1334
1335        // Append page histograms to the `ColumnIndex` histograms
1336        self.column_index_builder.append_histograms(
1337            &self.page_metrics.repetition_level_histogram,
1338            &self.page_metrics.definition_level_histogram,
1339        );
1340
1341        // Update the offset index
1342        if let Some(builder) = self.offset_index_builder.as_mut() {
1343            builder.append_row_count(self.page_metrics.num_buffered_rows as i64);
1344            builder.append_unencoded_byte_array_data_bytes(page_variable_length_bytes);
1345        }
1346    }
1347
1348    /// Determine if we should allow truncating min/max values for this column's statistics
1349    fn can_truncate_value(&self) -> bool {
1350        match self.descr.physical_type() {
1351            // Don't truncate for Float16 and Decimal because their sort order is different
1352            // from that of FIXED_LEN_BYTE_ARRAY sort order.
1353            // So truncation of those types could lead to inaccurate min/max statistics
1354            Type::FIXED_LEN_BYTE_ARRAY
1355                if !is_decimal_descr(self.descr.get_basic_info())
1356                    && !matches!(self.descr.logical_type_ref(), Some(&LogicalType::Float16)) =>
1357            {
1358                true
1359            }
1360            // As with FIXED_LEN_BYTE_ARRAY, do not truncate Decimal values
1361            Type::BYTE_ARRAY if !is_decimal_descr(self.descr.get_basic_info()) => true,
1362            // Truncation only applies for fba/binary physical types
1363            _ => false,
1364        }
1365    }
1366
1367    /// Returns `true` if this column's logical type is a UTF-8 string.
1368    fn is_utf8(&self) -> bool {
1369        self.get_descriptor().logical_type_ref() == Some(&LogicalType::String)
1370            || self.get_descriptor().converted_type() == ConvertedType::UTF8
1371    }
1372
1373    /// Truncates a binary statistic to at most `truncation_length` bytes.
1374    ///
1375    /// If truncation is not possible, returns `data`.
1376    ///
1377    /// The `bool` in the returned tuple indicates whether truncation occurred or not.
1378    ///
1379    /// UTF-8 Note:
1380    /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will
1381    /// also remain valid UTF-8, but may be less than `truncation_length` bytes to avoid splitting
1382    /// on non-character boundaries.
1383    fn truncate_min_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
1384        truncation_length
1385            .filter(|l| data.len() > *l)
1386            .and_then(|l|
1387                // don't do extra work if this column isn't UTF-8
1388                if self.is_utf8() {
1389                    match str::from_utf8(data) {
1390                        Ok(str_data) => truncate_utf8(str_data, l),
1391                        Err(_) => Some(data[..l].to_vec()),
1392                    }
1393                } else {
1394                    Some(data[..l].to_vec())
1395                }
1396            )
1397            .map(|truncated| (truncated, true))
1398            .unwrap_or_else(|| (data.to_vec(), false))
1399    }
1400
1401    /// Truncates a binary statistic to at most `truncation_length` bytes, and then increment the
1402    /// final byte(s) to yield a valid upper bound. This may result in a result of less than
1403    /// `truncation_length` bytes if the last byte(s) overflows.
1404    ///
1405    /// If truncation is not possible, returns `data`.
1406    ///
1407    /// The `bool` in the returned tuple indicates whether truncation occurred or not.
1408    ///
1409    /// UTF-8 Note:
1410    /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will
1411    /// also remain valid UTF-8 (but again may be less than `truncation_length` bytes). If `data`
1412    /// does not contain valid UTF-8, then truncation will occur as if the column is non-string
1413    /// binary.
1414    fn truncate_max_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
1415        truncation_length
1416            .filter(|l| data.len() > *l)
1417            .and_then(|l|
1418                // don't do extra work if this column isn't UTF-8
1419                if self.is_utf8() {
1420                    match str::from_utf8(data) {
1421                        Ok(str_data) => truncate_and_increment_utf8(str_data, l),
1422                        Err(_) => increment(data[..l].to_vec()),
1423                    }
1424                } else {
1425                    increment(data[..l].to_vec())
1426                }
1427            )
1428            .map(|truncated| (truncated, true))
1429            .unwrap_or_else(|| (data.to_vec(), false))
1430    }
1431
1432    /// Truncate the min and max values that will be written to a data page
1433    /// header or column chunk Statistics
1434    fn truncate_statistics(&self, statistics: Statistics) -> Statistics {
1435        let backwards_compatible_min_max = self.descr.sort_order().is_signed();
1436        match statistics {
1437            Statistics::ByteArray(stats)
1438                if (stats._internal_has_min_max_set() && self.can_truncate_value()) =>
1439            {
1440                let (min, did_truncate_min) = self.truncate_min_value(
1441                    self.props.statistics_truncate_length(),
1442                    stats.min_bytes_opt().unwrap(),
1443                );
1444                let (max, did_truncate_max) = self.truncate_max_value(
1445                    self.props.statistics_truncate_length(),
1446                    stats.max_bytes_opt().unwrap(),
1447                );
1448                Statistics::ByteArray(
1449                    ValueStatistics::new(
1450                        Some(min.into()),
1451                        Some(max.into()),
1452                        stats.distinct_count(),
1453                        stats.null_count_opt(),
1454                        backwards_compatible_min_max,
1455                    )
1456                    .with_max_is_exact(!did_truncate_max)
1457                    .with_min_is_exact(!did_truncate_min),
1458                )
1459            }
1460            Statistics::FixedLenByteArray(stats)
1461                if (stats._internal_has_min_max_set() && self.can_truncate_value()) =>
1462            {
1463                let (min, did_truncate_min) = self.truncate_min_value(
1464                    self.props.statistics_truncate_length(),
1465                    stats.min_bytes_opt().unwrap(),
1466                );
1467                let (max, did_truncate_max) = self.truncate_max_value(
1468                    self.props.statistics_truncate_length(),
1469                    stats.max_bytes_opt().unwrap(),
1470                );
1471                Statistics::FixedLenByteArray(
1472                    ValueStatistics::new(
1473                        Some(min.into()),
1474                        Some(max.into()),
1475                        stats.distinct_count(),
1476                        stats.null_count_opt(),
1477                        backwards_compatible_min_max,
1478                    )
1479                    .with_max_is_exact(!did_truncate_max)
1480                    .with_min_is_exact(!did_truncate_min),
1481                )
1482            }
1483            stats => stats,
1484        }
1485    }
1486
1487    /// Adds data page.
1488    /// Data page is either buffered in case of dictionary encoding or written directly.
1489    pub(crate) fn add_data_page(&mut self) -> Result<()> {
1490        // Extract encoded values
1491        let values_data = self.encoder.flush_data_page()?;
1492
1493        let max_def_level = self.descr.max_def_level();
1494        let max_rep_level = self.descr.max_rep_level();
1495
1496        self.column_metrics.num_column_nulls += self.page_metrics.num_page_nulls;
1497
1498        if let Some(nan_count) = values_data.nan_count {
1499            *self.column_metrics.num_column_nans.get_or_insert(0) += nan_count;
1500            self.page_metrics.num_page_nans = Some(nan_count);
1501        }
1502
1503        let page_statistics = match (values_data.min_value, values_data.max_value) {
1504            (Some(min), Some(max)) => {
1505                // Update chunk level statistics
1506                update_min(&self.descr, &min, &mut self.column_metrics.min_column_value);
1507                update_max(&self.descr, &max, &mut self.column_metrics.max_column_value);
1508
1509                (self.column_props.statistics_enabled == EnabledStatistics::Page).then_some(
1510                    ValueStatistics::new(
1511                        Some(min),
1512                        Some(max),
1513                        None,
1514                        Some(self.page_metrics.num_page_nulls),
1515                        false,
1516                    )
1517                    .with_nan_count(values_data.nan_count),
1518                )
1519            }
1520            _ => None,
1521        };
1522
1523        // update column and offset index
1524        self.update_column_offset_index(
1525            page_statistics.as_ref(),
1526            values_data.variable_length_bytes,
1527        );
1528
1529        // Update histograms and variable_length_bytes in column_metrics
1530        self.column_metrics
1531            .update_from_page_metrics(&self.page_metrics);
1532        self.column_metrics
1533            .update_variable_length_bytes(values_data.variable_length_bytes);
1534
1535        // From here on, we only need page statistics if they will be written to the page header.
1536        let page_statistics = page_statistics
1537            .filter(|_| self.column_props.write_page_header_statistics)
1538            .map(|stats| self.truncate_statistics(Statistics::from(stats)));
1539
1540        let compressed_page = match self.props.writer_version() {
1541            WriterVersion::PARQUET_1_0 => {
1542                let mut buffer = vec![];
1543
1544                if max_rep_level > 0 {
1545                    self.rep_levels_encoder
1546                        .flush_to(|data| buffer.extend_from_slice(data));
1547                }
1548
1549                if max_def_level > 0 {
1550                    self.def_levels_encoder
1551                        .flush_to(|data| buffer.extend_from_slice(data));
1552                }
1553
1554                buffer.extend_from_slice(&values_data.buf);
1555                let uncompressed_size = buffer.len();
1556
1557                if let Some(ref mut cmpr) = self.compressor {
1558                    let mut compressed_buf = Vec::with_capacity(uncompressed_size);
1559                    cmpr.compress(&buffer[..], &mut compressed_buf)?;
1560                    compressed_buf.shrink_to_fit();
1561                    buffer = compressed_buf;
1562                }
1563
1564                let data_page = Page::DataPage {
1565                    buf: buffer.into(),
1566                    num_values: self.page_metrics.num_buffered_values,
1567                    encoding: values_data.encoding,
1568                    def_level_encoding: Encoding::RLE,
1569                    rep_level_encoding: Encoding::RLE,
1570                    statistics: page_statistics,
1571                };
1572
1573                CompressedPage::new(data_page, uncompressed_size)
1574            }
1575            WriterVersion::PARQUET_2_0 => {
1576                let mut rep_levels_byte_len = 0;
1577                let mut def_levels_byte_len = 0;
1578                let mut buffer = vec![];
1579
1580                if max_rep_level > 0 {
1581                    self.rep_levels_encoder
1582                        .flush_to(|data| buffer.extend_from_slice(data));
1583                    rep_levels_byte_len = buffer.len();
1584                }
1585
1586                if max_def_level > 0 {
1587                    self.def_levels_encoder
1588                        .flush_to(|data| buffer.extend_from_slice(data));
1589                    def_levels_byte_len = buffer.len() - rep_levels_byte_len;
1590                }
1591
1592                let uncompressed_size =
1593                    rep_levels_byte_len + def_levels_byte_len + values_data.buf.len();
1594
1595                // Data Page v2 compresses values only.
1596                let is_compressed = match self.compressor {
1597                    Some(ref mut cmpr) => {
1598                        let buffer_len = buffer.len();
1599                        cmpr.compress(&values_data.buf, &mut buffer)?;
1600                        let compressed_values_size = buffer.len() - buffer_len;
1601                        let threshold = self.column_props.data_page_v2_compression_ratio_threshold;
1602                        if (compressed_values_size as f64) >= (uncompressed_size as f64) * threshold
1603                        {
1604                            buffer.truncate(buffer_len);
1605                            buffer.extend_from_slice(&values_data.buf);
1606                            false
1607                        } else {
1608                            true
1609                        }
1610                    }
1611                    None => {
1612                        buffer.extend_from_slice(&values_data.buf);
1613                        false
1614                    }
1615                };
1616
1617                let data_page = Page::DataPageV2 {
1618                    buf: buffer.into(),
1619                    num_values: self.page_metrics.num_buffered_values,
1620                    encoding: values_data.encoding,
1621                    num_nulls: self.page_metrics.num_page_nulls as u32,
1622                    num_rows: self.page_metrics.num_buffered_rows,
1623                    def_levels_byte_len: def_levels_byte_len as u32,
1624                    rep_levels_byte_len: rep_levels_byte_len as u32,
1625                    is_compressed,
1626                    statistics: page_statistics,
1627                };
1628
1629                CompressedPage::new(data_page, uncompressed_size)
1630            }
1631        };
1632
1633        // Check if we need to buffer data page or flush it to the sink directly.
1634        //
1635        // For dictionary-encoded columns the dictionary page must be written
1636        // first, but it is not final until all values are seen, so completed
1637        // data pages are normally buffered here until `close`. A page writer
1638        // that defers final layout (the Arrow path) instead orders pages itself
1639        // at flush, so we stream the data pages straight through and never let
1640        // them accumulate in memory.
1641        if self.encoder.has_dictionary() && !self.page_writer.defers_dictionary_ordering() {
1642            self.data_pages.push_back(compressed_page);
1643        } else {
1644            self.write_data_page(compressed_page)?;
1645        }
1646
1647        // Update total number of rows.
1648        self.column_metrics.total_rows_written += self.page_metrics.num_buffered_rows as u64;
1649        self.page_metrics.new_page();
1650
1651        Ok(())
1652    }
1653
1654    /// Finalises any outstanding data pages and flushes buffered data pages from
1655    /// dictionary encoding into underlying sink.
1656    #[inline]
1657    fn flush_data_pages(&mut self) -> Result<()> {
1658        // Write all outstanding data to a new page.
1659        if self.page_metrics.num_buffered_values > 0 {
1660            self.add_data_page()?;
1661        }
1662
1663        while let Some(page) = self.data_pages.pop_front() {
1664            self.write_data_page(page)?;
1665        }
1666
1667        Ok(())
1668    }
1669
1670    /// Assembles column chunk metadata.
1671    fn build_column_metadata(&mut self) -> Result<ColumnChunkMetaData> {
1672        let total_compressed_size = self.column_metrics.total_compressed_size as i64;
1673        let total_uncompressed_size = self.column_metrics.total_uncompressed_size as i64;
1674        let num_values = self.column_metrics.total_num_values as i64;
1675        let dict_page_offset = self.column_metrics.dictionary_page_offset.map(|v| v as i64);
1676        // If data page offset is not set, then no pages have been written
1677        let data_page_offset = self.column_metrics.data_page_offset.unwrap_or(0) as i64;
1678
1679        let mut builder = ColumnChunkMetaData::builder(self.descr.clone())
1680            .set_compression(self.codec)
1681            .set_encodings_mask(EncodingMask::new_from_encodings(self.encodings.iter()))
1682            .set_page_encoding_stats(self.encoding_stats.clone())
1683            .set_total_compressed_size(total_compressed_size)
1684            .set_total_uncompressed_size(total_uncompressed_size)
1685            .set_num_values(num_values)
1686            .set_data_page_offset(data_page_offset)
1687            .set_dictionary_page_offset(dict_page_offset);
1688
1689        if self.column_props.statistics_enabled != EnabledStatistics::None {
1690            let backwards_compatible_min_max = self.descr.sort_order().is_signed();
1691
1692            let distinct_count = self
1693                .distinct_count_override
1694                .or(self.column_metrics.column_distinct_count);
1695            let statistics = ValueStatistics::<E::T>::new(
1696                self.column_metrics.min_column_value.clone(),
1697                self.column_metrics.max_column_value.clone(),
1698                distinct_count,
1699                Some(self.column_metrics.num_column_nulls),
1700                false,
1701            )
1702            .with_nan_count(self.column_metrics.num_column_nans)
1703            .with_backwards_compatible_min_max(backwards_compatible_min_max)
1704            .into();
1705
1706            let statistics = self.truncate_statistics(statistics);
1707
1708            builder = builder
1709                .set_statistics(statistics)
1710                .set_unencoded_byte_array_data_bytes(self.column_metrics.variable_length_bytes)
1711                .set_repetition_level_histogram(
1712                    self.column_metrics.repetition_level_histogram.take(),
1713                )
1714                .set_definition_level_histogram(
1715                    self.column_metrics.definition_level_histogram.take(),
1716                );
1717
1718            if let Some(geo_stats) = self.encoder.flush_geospatial_statistics() {
1719                builder = builder.set_geo_statistics(geo_stats);
1720            }
1721        }
1722
1723        builder = self.set_column_chunk_encryption_properties(builder);
1724
1725        let metadata = builder.build()?;
1726        Ok(metadata)
1727    }
1728
1729    /// Writes compressed data page into underlying sink and updates global metrics.
1730    #[inline]
1731    fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {
1732        self.encodings.insert(page.encoding());
1733        match self.encoding_stats.last_mut() {
1734            Some(encoding_stats)
1735                if encoding_stats.page_type == page.page_type()
1736                    && encoding_stats.encoding == page.encoding() =>
1737            {
1738                encoding_stats.count += 1;
1739            }
1740            _ => {
1741                // data page type does not change inside a file
1742                // encoding can currently only change from dictionary to non-dictionary once
1743                self.encoding_stats.push(PageEncodingStats {
1744                    page_type: page.page_type(),
1745                    encoding: page.encoding(),
1746                    count: 1,
1747                });
1748            }
1749        }
1750        let page_spec = self.page_writer.write_page(page)?;
1751        // update offset index
1752        // compressed_size = header_size + compressed_data_size
1753        if let Some(builder) = self.offset_index_builder.as_mut() {
1754            builder
1755                .append_offset_and_size(page_spec.offset as i64, page_spec.compressed_size as i32)
1756        }
1757        self.update_metrics_for_page(page_spec);
1758        Ok(())
1759    }
1760
1761    /// Writes dictionary page into underlying sink.
1762    #[inline]
1763    fn write_dictionary_page(&mut self) -> Result<()> {
1764        let compressed_page = {
1765            let mut page = self
1766                .encoder
1767                .flush_dict_page()?
1768                .ok_or_else(|| general_err!("Dictionary encoder is not set"))?;
1769
1770            let uncompressed_size = page.buf.len();
1771
1772            if let Some(ref mut cmpr) = self.compressor {
1773                let mut output_buf = Vec::with_capacity(uncompressed_size);
1774                cmpr.compress(&page.buf, &mut output_buf)?;
1775                page.buf = Bytes::from(output_buf);
1776            }
1777
1778            let dict_page = Page::DictionaryPage {
1779                buf: page.buf,
1780                num_values: page.num_values as u32,
1781                encoding: self.props.dictionary_page_encoding(),
1782                is_sorted: page.is_sorted,
1783            };
1784            CompressedPage::new(dict_page, uncompressed_size)
1785        };
1786
1787        self.encodings.insert(compressed_page.encoding());
1788        // The dictionary page is always laid out first in the column chunk,
1789        // even when a page writer that defers dictionary ordering receives it
1790        // after the data pages. Keep `encoding_stats` in file order, as some
1791        // readers (e.g. Arrow C++) expect the dictionary page entry first.
1792        self.encoding_stats.insert(
1793            0,
1794            PageEncodingStats {
1795                page_type: PageType::DICTIONARY_PAGE,
1796                encoding: compressed_page.encoding(),
1797                count: 1,
1798            },
1799        );
1800        let page_spec = self.page_writer.write_page(compressed_page)?;
1801        self.update_metrics_for_page(page_spec);
1802        // For the directory page, don't need to update column/offset index.
1803        Ok(())
1804    }
1805
1806    /// Updates column writer metrics with each page metadata.
1807    #[inline]
1808    fn update_metrics_for_page(&mut self, page_spec: PageWriteSpec) {
1809        self.column_metrics.total_uncompressed_size += page_spec.uncompressed_size as u64;
1810        self.column_metrics.total_compressed_size += page_spec.compressed_size as u64;
1811        self.column_metrics.total_bytes_written += page_spec.bytes_written;
1812
1813        match page_spec.page_type {
1814            PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
1815                self.column_metrics.total_num_values += page_spec.num_values as u64;
1816                if self.column_metrics.data_page_offset.is_none() {
1817                    self.column_metrics.data_page_offset = Some(page_spec.offset);
1818                }
1819            }
1820            PageType::DICTIONARY_PAGE => {
1821                assert!(
1822                    self.column_metrics.dictionary_page_offset.is_none(),
1823                    "Dictionary offset is already set"
1824                );
1825                self.column_metrics.dictionary_page_offset = Some(page_spec.offset);
1826            }
1827            PageType::INDEX_PAGE => {}
1828        }
1829    }
1830
1831    #[inline]
1832    #[cfg(feature = "encryption")]
1833    fn set_column_chunk_encryption_properties(
1834        &self,
1835        builder: ColumnChunkMetaDataBuilder,
1836    ) -> ColumnChunkMetaDataBuilder {
1837        if let Some(encryption_properties) = self.props.file_encryption_properties.as_ref() {
1838            builder.set_column_crypto_metadata(get_column_crypto_metadata(
1839                encryption_properties,
1840                &self.descr,
1841            ))
1842        } else {
1843            builder
1844        }
1845    }
1846
1847    #[inline]
1848    #[cfg(not(feature = "encryption"))]
1849    fn set_column_chunk_encryption_properties(
1850        &self,
1851        builder: ColumnChunkMetaDataBuilder,
1852    ) -> ColumnChunkMetaDataBuilder {
1853        builder
1854    }
1855}
1856
1857fn update_min<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, min: &mut Option<T>) {
1858    match min {
1859        None => *min = Some(val.clone()),
1860        Some(min) => {
1861            let basic_type_info = descr.get_basic_info();
1862            let is_min_nan = is_nan(basic_type_info, min);
1863            let is_val_nan = is_nan(basic_type_info, val);
1864            match (is_min_nan, is_val_nan) {
1865                // current min is not NaN, but incoming is NaN: skip
1866                (false, true) => {}
1867                // current min is NaN, but incoming is not: assign val to min
1868                (true, false) => *min = val.clone(),
1869                // both NaN or non-NaN, safe to call update_stat()
1870                _ => {
1871                    update_stat::<T, _>(val, min, |cur| compare_greater(basic_type_info, cur, val))
1872                }
1873            }
1874        }
1875    }
1876}
1877
1878fn update_max<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, max: &mut Option<T>) {
1879    match max {
1880        None => *max = Some(val.clone()),
1881        Some(max) => {
1882            let basic_type_info = descr.get_basic_info();
1883            let is_max_nan = is_nan(basic_type_info, max);
1884            let is_val_nan = is_nan(basic_type_info, val);
1885            match (is_max_nan, is_val_nan) {
1886                // current max is not NaN, but incoming is NaN: skip
1887                (false, true) => {}
1888                // current max is NaN, but incoming is not: assign val to max
1889                (true, false) => *max = val.clone(),
1890                // both NaN or non-NaN, safe to call update_stat()
1891                _ => {
1892                    update_stat::<T, _>(val, max, |cur| compare_greater(basic_type_info, val, cur))
1893                }
1894            }
1895        }
1896    }
1897}
1898
1899#[inline]
1900#[expect(clippy::eq_op)]
1901fn is_nan<T: ParquetValueType>(basic_type_info: &BasicTypeInfo, val: &T) -> bool {
1902    match T::PHYSICAL_TYPE {
1903        Type::FLOAT | Type::DOUBLE => val != val,
1904        Type::FIXED_LEN_BYTE_ARRAY
1905            if matches!(basic_type_info.sort_order(), SortOrder::TOTAL_ORDER) =>
1906        {
1907            // taken from f16 impl, but skips creating f16. just compare the bits as u16.
1908            let val = val.as_bytes();
1909            // Float16 is stored little endian
1910            let uval = ((val[1] as u16) << 8) | val[0] as u16;
1911            uval & 0x7FFFu16 > 0x7C00u16
1912        }
1913        _ => false,
1914    }
1915}
1916
1917/// Perform a conditional update of `cur`
1918///
1919/// Calls `should_update` with the value of `cur`, and updates `cur` to `Some(val)` if it
1920/// returns `true`. `cur` must not be `None` or this will panic.
1921fn update_stat<T: ParquetValueType, F>(val: &T, cur: &mut T, should_update: F)
1922where
1923    F: Fn(&T) -> bool,
1924{
1925    if should_update(cur) {
1926        *cur = val.clone();
1927    }
1928}
1929
1930/// Evaluate `a > b` according to underlying logical type.
1931fn compare_greater<T: ParquetValueType>(basic_type_info: &BasicTypeInfo, a: &T, b: &T) -> bool {
1932    match T::PHYSICAL_TYPE {
1933        Type::FLOAT => {
1934            let a = f32::from_le_bytes(a.as_bytes().try_into().unwrap());
1935            let b = f32::from_le_bytes(b.as_bytes().try_into().unwrap());
1936            return a.total_cmp(&b) == Ordering::Greater;
1937        }
1938        Type::DOUBLE => {
1939            let a = f64::from_le_bytes(a.as_bytes().try_into().unwrap());
1940            let b = f64::from_le_bytes(b.as_bytes().try_into().unwrap());
1941            return a.total_cmp(&b) == Ordering::Greater;
1942        }
1943        Type::INT32 | Type::INT64
1944            if matches!(basic_type_info.sort_order(), SortOrder::UNSIGNED) =>
1945        {
1946            return compare_greater_unsigned_int(a, b);
1947        }
1948        Type::FIXED_LEN_BYTE_ARRAY
1949            if matches!(basic_type_info.sort_order(), SortOrder::TOTAL_ORDER) =>
1950        {
1951            return compare_greater_f16(a.as_bytes(), b.as_bytes());
1952        }
1953        Type::FIXED_LEN_BYTE_ARRAY | Type::BYTE_ARRAY if is_decimal_descr(basic_type_info) => {
1954            return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes());
1955        }
1956
1957        _ => {}
1958    }
1959
1960    // compare independent of logical / converted type
1961    a > b
1962}
1963
1964// ----------------------------------------------------------------------
1965// Encoding support for column writer.
1966// This mirrors parquet-mr default encodings for writes. See:
1967// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV1ValuesWriterFactory.java
1968// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV2ValuesWriterFactory.java
1969
1970/// Returns encoding for a column when no other encoding is provided in writer properties.
1971fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding {
1972    match (kind, props.writer_version()) {
1973        (Type::BOOLEAN, WriterVersion::PARQUET_2_0) => Encoding::RLE,
1974        (Type::INT32, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1975        (Type::INT64, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1976        (Type::BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1977        (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1978        _ => Encoding::PLAIN,
1979    }
1980}
1981
1982/// Returns true if dictionary is supported for column writer, false otherwise.
1983fn has_dictionary_support(kind: Type) -> bool {
1984    match kind {
1985        // Booleans do not support dict encoding and should use a fallback encoding.
1986        Type::BOOLEAN => false,
1987        _ => true,
1988    }
1989}
1990
1991#[inline]
1992fn compare_greater_unsigned_int<T: ParquetValueType>(a: &T, b: &T) -> bool {
1993    a.as_u64().unwrap() > b.as_u64().unwrap()
1994}
1995
1996#[inline]
1997fn compare_greater_f16(a: &[u8], b: &[u8]) -> bool {
1998    let a = f16::from_le_bytes(a.try_into().unwrap());
1999    let b = f16::from_le_bytes(b.try_into().unwrap());
2000    a.total_cmp(&b) == Ordering::Greater
2001}
2002
2003/// Returns `true` if the column described by `basic_type_info` is a Decimal column
2004#[inline]
2005pub(crate) fn is_decimal_descr(basic_type_info: &BasicTypeInfo) -> bool {
2006    basic_type_info.converted_type() == ConvertedType::DECIMAL
2007        || matches!(
2008            basic_type_info.logical_type_ref(),
2009            Some(LogicalType::Decimal(_))
2010        )
2011}
2012
2013/// Signed comparison of bytes arrays
2014pub(crate) fn compare_greater_byte_array_decimals(a: &[u8], b: &[u8]) -> bool {
2015    let a_length = a.len();
2016    let b_length = b.len();
2017
2018    if a_length == 0 || b_length == 0 {
2019        return a_length > 0;
2020    }
2021
2022    let first_a: u8 = a[0];
2023    let first_b: u8 = b[0];
2024
2025    // We can short circuit for different signed numbers or
2026    // for equal length bytes arrays that have different first bytes.
2027    // The equality requirement is necessary for sign extension cases.
2028    // 0xFF10 should be equal to 0x10 (due to big endian sign extension).
2029    if (0x80 & first_a) != (0x80 & first_b) || (a_length == b_length && first_a != first_b) {
2030        return (first_a as i8) > (first_b as i8);
2031    }
2032
2033    // When the lengths are unequal and the numbers are of the same sign,
2034    // sign-extend the shorter value: if any of the longer value's extra
2035    // leading bytes differs from the sign-extension byte it has the larger
2036    // magnitude, and otherwise those bytes are redundant and the aligned
2037    // equal-length tails decide via unsigned lexicographical comparison.
2038
2039    let extension: u8 = if (first_a as i8) < 0 { 0xFF } else { 0 };
2040
2041    if a_length != b_length {
2042        let not_equal = if a_length > b_length {
2043            let lead_length = a_length - b_length;
2044            a[0..lead_length].iter().any(|&x| x != extension)
2045        } else {
2046            let lead_length = b_length - a_length;
2047            b[0..lead_length].iter().any(|&x| x != extension)
2048        };
2049
2050        if not_equal {
2051            let negative_values: bool = (first_a as i8) < 0;
2052            let a_longer: bool = a_length > b_length;
2053            return if negative_values { !a_longer } else { a_longer };
2054        }
2055    }
2056
2057    let tail_length = a_length.min(b_length);
2058    (a[a_length - tail_length..]) > (b[b_length - tail_length..])
2059}
2060
2061/// Truncate a UTF-8 slice to the longest prefix that is still a valid UTF-8 string,
2062/// while being less than `length` bytes and non-empty. Returns `None` if truncation
2063/// is not possible within those constraints.
2064///
2065/// The caller guarantees that data.len() > length.
2066fn truncate_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
2067    let split = (1..=length).rfind(|x| data.is_char_boundary(*x))?;
2068    Some(data.as_bytes()[..split].to_vec())
2069}
2070
2071/// Truncate a UTF-8 slice and increment it's final character. The returned value is the
2072/// longest such slice that is still a valid UTF-8 string while being less than `length`
2073/// bytes and non-empty. Returns `None` if no such transformation is possible.
2074///
2075/// The caller guarantees that data.len() > length.
2076fn truncate_and_increment_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
2077    // UTF-8 is max 4 bytes, so start search 3 back from desired length
2078    let lower_bound = length.saturating_sub(3);
2079    let split = (lower_bound..=length).rfind(|x| data.is_char_boundary(*x))?;
2080    increment_utf8(data.get(..split)?)
2081}
2082
2083/// Increment the final character in a UTF-8 string in such a way that the returned result
2084/// is still a valid UTF-8 string. The returned string may be shorter than the input if the
2085/// last character(s) cannot be incremented (due to overflow or producing invalid code points).
2086/// Returns `None` if the string cannot be incremented.
2087///
2088/// Note that this implementation will not promote an N-byte code point to (N+1) bytes.
2089fn increment_utf8(data: &str) -> Option<Vec<u8>> {
2090    for (idx, original_char) in data.char_indices().rev() {
2091        let original_len = original_char.len_utf8();
2092        if let Some(next_char) = char::from_u32(original_char as u32 + 1) {
2093            // do not allow increasing byte width of incremented char
2094            if next_char.len_utf8() == original_len {
2095                let mut result = data.as_bytes()[..idx + original_len].to_vec();
2096                next_char.encode_utf8(&mut result[idx..]);
2097                return Some(result);
2098            }
2099        }
2100    }
2101
2102    None
2103}
2104
2105/// Try and increment the bytes from right to left.
2106///
2107/// Returns `None` if all bytes are set to `u8::MAX`.
2108fn increment(mut data: Vec<u8>) -> Option<Vec<u8>> {
2109    for byte in data.iter_mut().rev() {
2110        let (incremented, overflow) = byte.overflowing_add(1);
2111        *byte = incremented;
2112
2113        if !overflow {
2114            return Some(data);
2115        }
2116    }
2117
2118    None
2119}
2120
2121#[cfg(test)]
2122mod tests {
2123    use crate::{
2124        file::{properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, writer::SerializedFileWriter},
2125        schema::parser::parse_message_type,
2126    };
2127    use core::str;
2128    use rand::distr::uniform::SampleUniform;
2129    use std::{fs::File, sync::Arc};
2130
2131    use crate::column::{
2132        page::PageReader,
2133        reader::{ColumnReaderImpl, get_column_reader, get_typed_column_reader},
2134    };
2135    use crate::file::writer::TrackedWrite;
2136    use crate::file::{
2137        properties::ReaderProperties, reader::SerializedPageReader, writer::SerializedPageWriter,
2138    };
2139    use crate::schema::types::{ColumnPath, Type as SchemaType};
2140    use crate::util::test_common::rand_gen::random_numbers_range;
2141
2142    use super::*;
2143
2144    #[test]
2145    fn test_column_writer_inconsistent_def_rep_length() {
2146        let page_writer = get_test_page_writer();
2147        let props = Default::default();
2148        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 1, props);
2149        let res = writer.write_batch(&[1, 2, 3, 4], Some(&[1, 1, 1]), Some(&[0, 0]));
2150        assert!(res.is_err());
2151        if let Err(err) = res {
2152            assert_eq!(
2153                format!("{err}"),
2154                "Parquet error: Inconsistent length of definition and repetition levels: 3 != 2"
2155            );
2156        }
2157    }
2158
2159    #[test]
2160    fn test_column_writer_invalid_def_levels() {
2161        let page_writer = get_test_page_writer();
2162        let props = Default::default();
2163        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2164        let res = writer.write_batch(&[1, 2, 3, 4], None, None);
2165        assert!(res.is_err());
2166        if let Err(err) = res {
2167            assert_eq!(
2168                format!("{err}"),
2169                "Parquet error: Definition levels are required, because max definition level = 1"
2170            );
2171        }
2172    }
2173
2174    #[test]
2175    fn test_column_writer_invalid_rep_levels() {
2176        let page_writer = get_test_page_writer();
2177        let props = Default::default();
2178        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 1, props);
2179        let res = writer.write_batch(&[1, 2, 3, 4], None, None);
2180        assert!(res.is_err());
2181        if let Err(err) = res {
2182            assert_eq!(
2183                format!("{err}"),
2184                "Parquet error: Repetition levels are required, because max repetition level = 1"
2185            );
2186        }
2187    }
2188
2189    #[test]
2190    fn test_column_writer_not_enough_values_to_write() {
2191        let page_writer = get_test_page_writer();
2192        let props = Default::default();
2193        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2194        let res = writer.write_batch(&[1, 2], Some(&[1, 1, 1, 1]), None);
2195        assert!(res.is_err());
2196        if let Err(err) = res {
2197            assert_eq!(
2198                format!("{err}"),
2199                "Parquet error: Expected to write 4 values, but have only 2"
2200            );
2201        }
2202    }
2203
2204    #[test]
2205    fn test_column_writer_write_only_one_dictionary_page() {
2206        let page_writer = get_test_page_writer();
2207        let props = Default::default();
2208        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2209        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2210        // First page should be correctly written.
2211        writer.add_data_page().unwrap();
2212        writer.write_dictionary_page().unwrap();
2213        let err = writer.write_dictionary_page().unwrap_err().to_string();
2214        assert_eq!(err, "Parquet error: Dictionary encoder is not set");
2215    }
2216
2217    #[test]
2218    fn test_column_writer_error_when_writing_disabled_dictionary() {
2219        let page_writer = get_test_page_writer();
2220        let props = Arc::new(
2221            WriterProperties::builder()
2222                .set_dictionary_enabled(false)
2223                .build(),
2224        );
2225        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2226        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2227        let err = writer.write_dictionary_page().unwrap_err().to_string();
2228        assert_eq!(err, "Parquet error: Dictionary encoder is not set");
2229    }
2230
2231    #[test]
2232    fn test_column_writer_boolean_type_does_not_support_dictionary() {
2233        let page_writer = get_test_page_writer();
2234        let props = Arc::new(
2235            WriterProperties::builder()
2236                .set_dictionary_enabled(true)
2237                .build(),
2238        );
2239        let mut writer = get_test_column_writer::<BoolType>(page_writer, 0, 0, props);
2240        writer
2241            .write_batch(&[true, false, true, false], None, None)
2242            .unwrap();
2243
2244        let r = writer.close().unwrap();
2245        // PlainEncoder uses bit writer to write boolean values, which all fit into 1
2246        // byte.
2247        assert_eq!(r.bytes_written, 1);
2248        assert_eq!(r.rows_written, 4);
2249
2250        let metadata = r.metadata;
2251        assert_eq!(
2252            metadata.encodings().collect::<Vec<_>>(),
2253            vec![Encoding::PLAIN, Encoding::RLE]
2254        );
2255        assert_eq!(metadata.num_values(), 4); // just values
2256        assert_eq!(metadata.dictionary_page_offset(), None);
2257    }
2258
2259    #[test]
2260    fn test_bloom_filter_for_dictionary_encoded_chunks() {
2261        fn bloom_filter_written(dictionary_enabled: bool, for_dictionary_chunks: bool) -> bool {
2262            let props = Arc::new(
2263                WriterProperties::builder()
2264                    .set_dictionary_enabled(dictionary_enabled)
2265                    .set_bloom_filter_enabled(true)
2266                    .set_bloom_filter_for_dictionary_encoded_chunks(for_dictionary_chunks)
2267                    .build(),
2268            );
2269            let mut writer =
2270                get_test_column_writer::<Int32Type>(get_test_page_writer(), 0, 0, props);
2271            writer.write_batch(&[1, 2, 1, 2], None, None).unwrap();
2272            writer.close().unwrap().bloom_filter.is_some()
2273        }
2274
2275        assert!(bloom_filter_written(true, true));
2276        assert!(bloom_filter_written(false, true));
2277        assert!(bloom_filter_written(false, false));
2278        assert!(!bloom_filter_written(true, false));
2279    }
2280
2281    #[test]
2282    fn test_column_writer_default_encoding_support_bool() {
2283        check_encoding_write_support::<BoolType>(
2284            WriterVersion::PARQUET_1_0,
2285            true,
2286            &[true, false],
2287            None,
2288            &[Encoding::PLAIN, Encoding::RLE],
2289            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2290        );
2291        check_encoding_write_support::<BoolType>(
2292            WriterVersion::PARQUET_1_0,
2293            false,
2294            &[true, false],
2295            None,
2296            &[Encoding::PLAIN, Encoding::RLE],
2297            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2298        );
2299        check_encoding_write_support::<BoolType>(
2300            WriterVersion::PARQUET_2_0,
2301            true,
2302            &[true, false],
2303            None,
2304            &[Encoding::RLE],
2305            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
2306        );
2307        check_encoding_write_support::<BoolType>(
2308            WriterVersion::PARQUET_2_0,
2309            false,
2310            &[true, false],
2311            None,
2312            &[Encoding::RLE],
2313            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
2314        );
2315    }
2316
2317    #[test]
2318    fn test_column_writer_default_encoding_support_int32() {
2319        check_encoding_write_support::<Int32Type>(
2320            WriterVersion::PARQUET_1_0,
2321            true,
2322            &[1, 2],
2323            Some(0),
2324            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2325            &[
2326                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2327                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2328            ],
2329        );
2330        check_encoding_write_support::<Int32Type>(
2331            WriterVersion::PARQUET_1_0,
2332            false,
2333            &[1, 2],
2334            None,
2335            &[Encoding::PLAIN, Encoding::RLE],
2336            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2337        );
2338        check_encoding_write_support::<Int32Type>(
2339            WriterVersion::PARQUET_2_0,
2340            true,
2341            &[1, 2],
2342            Some(0),
2343            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2344            &[
2345                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2346                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2347            ],
2348        );
2349        check_encoding_write_support::<Int32Type>(
2350            WriterVersion::PARQUET_2_0,
2351            false,
2352            &[1, 2],
2353            None,
2354            &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
2355            &[encoding_stats(
2356                PageType::DATA_PAGE_V2,
2357                Encoding::DELTA_BINARY_PACKED,
2358                1,
2359            )],
2360        );
2361    }
2362
2363    #[test]
2364    fn test_column_writer_default_encoding_support_int64() {
2365        check_encoding_write_support::<Int64Type>(
2366            WriterVersion::PARQUET_1_0,
2367            true,
2368            &[1, 2],
2369            Some(0),
2370            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2371            &[
2372                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2373                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2374            ],
2375        );
2376        check_encoding_write_support::<Int64Type>(
2377            WriterVersion::PARQUET_1_0,
2378            false,
2379            &[1, 2],
2380            None,
2381            &[Encoding::PLAIN, Encoding::RLE],
2382            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2383        );
2384        check_encoding_write_support::<Int64Type>(
2385            WriterVersion::PARQUET_2_0,
2386            true,
2387            &[1, 2],
2388            Some(0),
2389            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2390            &[
2391                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2392                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2393            ],
2394        );
2395        check_encoding_write_support::<Int64Type>(
2396            WriterVersion::PARQUET_2_0,
2397            false,
2398            &[1, 2],
2399            None,
2400            &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
2401            &[encoding_stats(
2402                PageType::DATA_PAGE_V2,
2403                Encoding::DELTA_BINARY_PACKED,
2404                1,
2405            )],
2406        );
2407    }
2408
2409    #[test]
2410    fn test_column_writer_default_encoding_support_int96() {
2411        check_encoding_write_support::<Int96Type>(
2412            WriterVersion::PARQUET_1_0,
2413            true,
2414            &[Int96::from(vec![1, 2, 3])],
2415            Some(0),
2416            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2417            &[
2418                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2419                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2420            ],
2421        );
2422        check_encoding_write_support::<Int96Type>(
2423            WriterVersion::PARQUET_1_0,
2424            false,
2425            &[Int96::from(vec![1, 2, 3])],
2426            None,
2427            &[Encoding::PLAIN, Encoding::RLE],
2428            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2429        );
2430        check_encoding_write_support::<Int96Type>(
2431            WriterVersion::PARQUET_2_0,
2432            true,
2433            &[Int96::from(vec![1, 2, 3])],
2434            Some(0),
2435            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2436            &[
2437                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2438                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2439            ],
2440        );
2441        check_encoding_write_support::<Int96Type>(
2442            WriterVersion::PARQUET_2_0,
2443            false,
2444            &[Int96::from(vec![1, 2, 3])],
2445            None,
2446            &[Encoding::PLAIN, Encoding::RLE],
2447            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2448        );
2449    }
2450
2451    #[test]
2452    fn test_column_writer_default_encoding_support_float() {
2453        check_encoding_write_support::<FloatType>(
2454            WriterVersion::PARQUET_1_0,
2455            true,
2456            &[1.0, 2.0],
2457            Some(0),
2458            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2459            &[
2460                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2461                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2462            ],
2463        );
2464        check_encoding_write_support::<FloatType>(
2465            WriterVersion::PARQUET_1_0,
2466            false,
2467            &[1.0, 2.0],
2468            None,
2469            &[Encoding::PLAIN, Encoding::RLE],
2470            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2471        );
2472        check_encoding_write_support::<FloatType>(
2473            WriterVersion::PARQUET_2_0,
2474            true,
2475            &[1.0, 2.0],
2476            Some(0),
2477            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2478            &[
2479                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2480                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2481            ],
2482        );
2483        check_encoding_write_support::<FloatType>(
2484            WriterVersion::PARQUET_2_0,
2485            false,
2486            &[1.0, 2.0],
2487            None,
2488            &[Encoding::PLAIN, Encoding::RLE],
2489            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2490        );
2491    }
2492
2493    #[test]
2494    fn test_column_writer_default_encoding_support_double() {
2495        check_encoding_write_support::<DoubleType>(
2496            WriterVersion::PARQUET_1_0,
2497            true,
2498            &[1.0, 2.0],
2499            Some(0),
2500            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2501            &[
2502                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2503                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2504            ],
2505        );
2506        check_encoding_write_support::<DoubleType>(
2507            WriterVersion::PARQUET_1_0,
2508            false,
2509            &[1.0, 2.0],
2510            None,
2511            &[Encoding::PLAIN, Encoding::RLE],
2512            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2513        );
2514        check_encoding_write_support::<DoubleType>(
2515            WriterVersion::PARQUET_2_0,
2516            true,
2517            &[1.0, 2.0],
2518            Some(0),
2519            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2520            &[
2521                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2522                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2523            ],
2524        );
2525        check_encoding_write_support::<DoubleType>(
2526            WriterVersion::PARQUET_2_0,
2527            false,
2528            &[1.0, 2.0],
2529            None,
2530            &[Encoding::PLAIN, Encoding::RLE],
2531            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2532        );
2533    }
2534
2535    #[test]
2536    fn test_column_writer_default_encoding_support_byte_array() {
2537        check_encoding_write_support::<ByteArrayType>(
2538            WriterVersion::PARQUET_1_0,
2539            true,
2540            &[ByteArray::from(vec![1u8])],
2541            Some(0),
2542            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2543            &[
2544                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2545                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2546            ],
2547        );
2548        check_encoding_write_support::<ByteArrayType>(
2549            WriterVersion::PARQUET_1_0,
2550            false,
2551            &[ByteArray::from(vec![1u8])],
2552            None,
2553            &[Encoding::PLAIN, Encoding::RLE],
2554            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2555        );
2556        check_encoding_write_support::<ByteArrayType>(
2557            WriterVersion::PARQUET_2_0,
2558            true,
2559            &[ByteArray::from(vec![1u8])],
2560            Some(0),
2561            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2562            &[
2563                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2564                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2565            ],
2566        );
2567        check_encoding_write_support::<ByteArrayType>(
2568            WriterVersion::PARQUET_2_0,
2569            false,
2570            &[ByteArray::from(vec![1u8])],
2571            None,
2572            &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2573            &[encoding_stats(
2574                PageType::DATA_PAGE_V2,
2575                Encoding::DELTA_BYTE_ARRAY,
2576                1,
2577            )],
2578        );
2579    }
2580
2581    #[test]
2582    fn test_column_writer_default_encoding_support_fixed_len_byte_array() {
2583        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
2584            let default_props = WriterProperties::builder()
2585                .set_writer_version(version)
2586                .build();
2587            let meta = column_write_and_get_metadata::<FixedLenByteArrayType>(
2588                default_props,
2589                &[ByteArray::from(vec![1u8]).into()],
2590            );
2591            assert_eq!(meta.dictionary_page_offset(), Some(0));
2592        }
2593
2594        check_encoding_write_support::<FixedLenByteArrayType>(
2595            WriterVersion::PARQUET_1_0,
2596            true,
2597            &[ByteArray::from(vec![1u8]).into()],
2598            Some(0),
2599            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2600            &[
2601                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2602                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2603            ],
2604        );
2605        let column_props = WriterProperties::builder()
2606            .set_writer_version(WriterVersion::PARQUET_1_0)
2607            .set_dictionary_enabled(false)
2608            .set_column_dictionary_enabled(ColumnPath::from("col"), true)
2609            .build();
2610        let meta = column_write_and_get_metadata::<FixedLenByteArrayType>(
2611            column_props,
2612            &[ByteArray::from(vec![1u8]).into()],
2613        );
2614        assert_eq!(meta.dictionary_page_offset(), Some(0));
2615        check_encoding_write_support::<FixedLenByteArrayType>(
2616            WriterVersion::PARQUET_1_0,
2617            false,
2618            &[ByteArray::from(vec![1u8]).into()],
2619            None,
2620            &[Encoding::PLAIN, Encoding::RLE],
2621            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2622        );
2623        check_encoding_write_support::<FixedLenByteArrayType>(
2624            WriterVersion::PARQUET_2_0,
2625            true,
2626            &[ByteArray::from(vec![1u8]).into()],
2627            Some(0),
2628            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2629            &[
2630                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2631                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2632            ],
2633        );
2634        check_encoding_write_support::<FixedLenByteArrayType>(
2635            WriterVersion::PARQUET_2_0,
2636            false,
2637            &[ByteArray::from(vec![1u8]).into()],
2638            None,
2639            &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2640            &[encoding_stats(
2641                PageType::DATA_PAGE_V2,
2642                Encoding::DELTA_BYTE_ARRAY,
2643                1,
2644            )],
2645        );
2646    }
2647
2648    #[test]
2649    fn test_column_writer_check_metadata() {
2650        let page_writer = get_test_page_writer();
2651        let props = Default::default();
2652        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2653        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2654
2655        let r = writer.close().unwrap();
2656        assert_eq!(r.bytes_written, 20);
2657        assert_eq!(r.rows_written, 4);
2658
2659        let metadata = r.metadata;
2660        assert_eq!(
2661            metadata.encodings().collect::<Vec<_>>(),
2662            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2663        );
2664        assert_eq!(metadata.num_values(), 4);
2665        assert_eq!(metadata.compressed_size(), 20);
2666        assert_eq!(metadata.uncompressed_size(), 20);
2667        assert_eq!(metadata.data_page_offset(), 0);
2668        assert_eq!(metadata.dictionary_page_offset(), Some(0));
2669        if let Some(stats) = metadata.statistics() {
2670            assert_eq!(stats.null_count_opt(), Some(0));
2671            assert_eq!(stats.distinct_count_opt(), None);
2672            if let Statistics::Int32(stats) = stats {
2673                assert_eq!(stats.min_opt().unwrap(), &1);
2674                assert_eq!(stats.max_opt().unwrap(), &4);
2675            } else {
2676                panic!("expecting Statistics::Int32");
2677            }
2678        } else {
2679            panic!("metadata missing statistics");
2680        }
2681    }
2682
2683    #[test]
2684    fn test_column_writer_check_byte_array_min_max() {
2685        let page_writer = get_test_page_writer();
2686        let props = Default::default();
2687        let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
2688        writer
2689            .write_batch(
2690                &[
2691                    ByteArray::from(vec![
2692                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2693                        35u8, 231u8, 90u8, 0u8, 0u8,
2694                    ]),
2695                    ByteArray::from(vec![
2696                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 228u8, 62u8, 146u8,
2697                        152u8, 177u8, 56u8, 0u8, 0u8,
2698                    ]),
2699                    ByteArray::from(vec![
2700                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8,
2701                        0u8,
2702                    ]),
2703                    ByteArray::from(vec![
2704                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2705                        44u8, 0u8, 0u8,
2706                    ]),
2707                ],
2708                None,
2709                None,
2710            )
2711            .unwrap();
2712        let metadata = writer.close().unwrap().metadata;
2713        if let Some(stats) = metadata.statistics() {
2714            if let Statistics::ByteArray(stats) = stats {
2715                assert_eq!(
2716                    stats.min_opt().unwrap(),
2717                    &ByteArray::from(vec![
2718                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2719                        35u8, 231u8, 90u8, 0u8, 0u8,
2720                    ])
2721                );
2722                assert_eq!(
2723                    stats.max_opt().unwrap(),
2724                    &ByteArray::from(vec![
2725                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2726                        44u8, 0u8, 0u8,
2727                    ])
2728                );
2729            } else {
2730                panic!("expecting Statistics::ByteArray");
2731            }
2732        } else {
2733            panic!("metadata missing statistics");
2734        }
2735    }
2736
2737    #[test]
2738    fn test_column_writer_byte_array_min_max_unequal_lengths() {
2739        // Byte-array decimal min/max with values of different encoded lengths
2740        // https://github.com/apache/arrow-rs/issues/10860
2741        let page_writer = get_test_page_writer();
2742        let props = Default::default();
2743        let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
2744        writer
2745            .write_batch(
2746                &[
2747                    ByteArray::from(vec![0u8, 255u8]),      // 255
2748                    ByteArray::from(vec![0u8, 128u8, 0u8]), // 32768
2749                    ByteArray::from(vec![255u8, 127u8]),    // -129
2750                    ByteArray::from(vec![128u8]),           // -128
2751                ],
2752                None,
2753                None,
2754            )
2755            .unwrap();
2756        let metadata = writer.close().unwrap().metadata;
2757        let stats = metadata.statistics().expect("metadata missing statistics");
2758        let Statistics::ByteArray(stats) = stats else {
2759            panic!("expecting Statistics::ByteArray");
2760        };
2761        // -129
2762        assert_eq!(
2763            stats.min_opt().unwrap(),
2764            &ByteArray::from(vec![255u8, 127u8])
2765        );
2766        // 32768
2767        assert_eq!(
2768            stats.max_opt().unwrap(),
2769            &ByteArray::from(vec![0u8, 128u8, 0u8])
2770        );
2771    }
2772
2773    #[test]
2774    fn test_column_writer_uint32_converted_type_min_max() {
2775        let page_writer = get_test_page_writer();
2776        let props = Default::default();
2777        let mut writer = get_test_unsigned_int_given_as_converted_column_writer::<Int32Type>(
2778            page_writer,
2779            0,
2780            0,
2781            props,
2782        );
2783        writer.write_batch(&[0, 1, 2, 3, 4, 5], None, None).unwrap();
2784        let metadata = writer.close().unwrap().metadata;
2785        if let Some(stats) = metadata.statistics() {
2786            if let Statistics::Int32(stats) = stats {
2787                assert_eq!(stats.min_opt().unwrap(), &0,);
2788                assert_eq!(stats.max_opt().unwrap(), &5,);
2789            } else {
2790                panic!("expecting Statistics::Int32");
2791            }
2792        } else {
2793            panic!("metadata missing statistics");
2794        }
2795    }
2796
2797    #[test]
2798    fn test_column_writer_precalculated_statistics() {
2799        let page_writer = get_test_page_writer();
2800        let props = Arc::new(
2801            WriterProperties::builder()
2802                .set_statistics_enabled(EnabledStatistics::Chunk)
2803                .build(),
2804        );
2805        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2806        writer
2807            .write_batch_with_statistics(
2808                &[1, 2, 3, 4],
2809                None,
2810                None,
2811                Some(&-17),
2812                Some(&9000),
2813                Some(55),
2814            )
2815            .unwrap();
2816
2817        let r = writer.close().unwrap();
2818        assert_eq!(r.bytes_written, 20);
2819        assert_eq!(r.rows_written, 4);
2820
2821        let metadata = r.metadata;
2822        assert_eq!(
2823            metadata.encodings().collect::<Vec<_>>(),
2824            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2825        );
2826        assert_eq!(metadata.num_values(), 4);
2827        assert_eq!(metadata.compressed_size(), 20);
2828        assert_eq!(metadata.uncompressed_size(), 20);
2829        assert_eq!(metadata.data_page_offset(), 0);
2830        assert_eq!(metadata.dictionary_page_offset(), Some(0));
2831        if let Some(stats) = metadata.statistics() {
2832            assert_eq!(stats.null_count_opt(), Some(0));
2833            assert_eq!(stats.distinct_count_opt().unwrap_or(0), 55);
2834            if let Statistics::Int32(stats) = stats {
2835                assert_eq!(stats.min_opt().unwrap(), &-17);
2836                assert_eq!(stats.max_opt().unwrap(), &9000);
2837            } else {
2838                panic!("expecting Statistics::Int32");
2839            }
2840        } else {
2841            panic!("metadata missing statistics");
2842        }
2843    }
2844
2845    #[test]
2846    fn test_mixed_precomputed_statistics() {
2847        let mut buf = Vec::with_capacity(100);
2848        let mut write = TrackedWrite::new(&mut buf);
2849        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2850        let props = Arc::new(
2851            WriterProperties::builder()
2852                .set_write_page_header_statistics(true)
2853                .set_data_page_row_count_limit(4)
2854                .build(),
2855        );
2856        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2857
2858        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2859        writer
2860            .write_batch_with_statistics(&[5, 6, 7], None, None, Some(&5), Some(&7), Some(3))
2861            .unwrap();
2862
2863        let r = writer.close().unwrap();
2864
2865        let stats = r.metadata.statistics().unwrap();
2866        assert_eq!(stats.min_bytes_opt().unwrap(), 1_i32.to_le_bytes());
2867        assert_eq!(stats.max_bytes_opt().unwrap(), 7_i32.to_le_bytes());
2868        assert_eq!(stats.null_count_opt(), Some(0));
2869        assert!(stats.distinct_count_opt().is_none());
2870
2871        drop(write);
2872
2873        let props = ReaderProperties::builder()
2874            .set_backward_compatible_lz4(false)
2875            .set_read_page_statistics(true)
2876            .build();
2877        let reader = SerializedPageReader::new_with_properties(
2878            Arc::new(Bytes::from(buf)),
2879            &r.metadata,
2880            r.rows_written as usize,
2881            None,
2882            Arc::new(props),
2883        )
2884        .unwrap();
2885
2886        let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2887        assert_eq!(pages.len(), 3);
2888
2889        assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2890        assert_eq!(pages[1].page_type(), PageType::DATA_PAGE);
2891        assert_eq!(pages[2].page_type(), PageType::DATA_PAGE);
2892        for (page, min, max) in [(&pages[1], 1_i32, 4_i32), (&pages[2], 5_i32, 7_i32)] {
2893            let stats = page.statistics().unwrap();
2894            assert_eq!(stats.min_bytes_opt().unwrap(), min.to_le_bytes());
2895            assert_eq!(stats.max_bytes_opt().unwrap(), max.to_le_bytes());
2896            assert_eq!(stats.null_count_opt(), Some(0));
2897            assert!(stats.distinct_count_opt().is_none());
2898        }
2899    }
2900
2901    #[test]
2902    fn test_disabled_statistics() {
2903        let mut buf = Vec::with_capacity(100);
2904        let mut write = TrackedWrite::new(&mut buf);
2905        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2906        let props = WriterProperties::builder()
2907            .set_statistics_enabled(EnabledStatistics::None)
2908            .set_writer_version(WriterVersion::PARQUET_2_0)
2909            .build();
2910        let props = Arc::new(props);
2911
2912        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2913        writer
2914            .write_batch(&[1, 2, 3, 4], Some(&[1, 0, 0, 1, 1, 1]), None)
2915            .unwrap();
2916
2917        let r = writer.close().unwrap();
2918        assert!(r.metadata.statistics().is_none());
2919
2920        drop(write);
2921
2922        let props = ReaderProperties::builder()
2923            .set_backward_compatible_lz4(false)
2924            .build();
2925        let reader = SerializedPageReader::new_with_properties(
2926            Arc::new(Bytes::from(buf)),
2927            &r.metadata,
2928            r.rows_written as usize,
2929            None,
2930            Arc::new(props),
2931        )
2932        .unwrap();
2933
2934        let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2935        assert_eq!(pages.len(), 2);
2936
2937        assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2938        assert_eq!(pages[1].page_type(), PageType::DATA_PAGE_V2);
2939
2940        match &pages[1] {
2941            Page::DataPageV2 {
2942                num_values,
2943                num_nulls,
2944                num_rows,
2945                statistics,
2946                ..
2947            } => {
2948                assert_eq!(*num_values, 6);
2949                assert_eq!(*num_nulls, 2);
2950                assert_eq!(*num_rows, 6);
2951                assert!(statistics.is_none());
2952            }
2953            _ => unreachable!(),
2954        }
2955    }
2956
2957    #[test]
2958    fn test_column_writer_empty_column_roundtrip() {
2959        let props = Default::default();
2960        column_roundtrip::<Int32Type>(props, &[], None, None);
2961    }
2962
2963    #[test]
2964    fn test_column_writer_non_nullable_values_roundtrip() {
2965        let props = Default::default();
2966        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 0, 0);
2967    }
2968
2969    #[test]
2970    fn test_column_writer_nullable_non_repeated_values_roundtrip() {
2971        let props = Default::default();
2972        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 0);
2973    }
2974
2975    #[test]
2976    fn test_column_writer_nullable_repeated_values_roundtrip() {
2977        let props = Default::default();
2978        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2979    }
2980
2981    #[test]
2982    fn test_column_writer_dictionary_fallback_small_data_page() {
2983        let props = WriterProperties::builder()
2984            .set_dictionary_page_size_limit(32)
2985            .set_data_page_size_limit(32)
2986            .build();
2987        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2988    }
2989
2990    #[test]
2991    #[cfg_attr(miri, ignore)] // Takes too long
2992    fn test_column_writer_small_write_batch_size() {
2993        for i in &[1usize, 2, 5, 10, 11, 1023] {
2994            let props = WriterProperties::builder().set_write_batch_size(*i).build();
2995
2996            column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2997        }
2998    }
2999
3000    #[test]
3001    fn test_column_writer_dictionary_disabled_v1() {
3002        let props = WriterProperties::builder()
3003            .set_writer_version(WriterVersion::PARQUET_1_0)
3004            .set_dictionary_enabled(false)
3005            .build();
3006        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
3007    }
3008
3009    #[test]
3010    fn test_column_writer_dictionary_disabled_v2() {
3011        let props = WriterProperties::builder()
3012            .set_writer_version(WriterVersion::PARQUET_2_0)
3013            .set_dictionary_enabled(false)
3014            .build();
3015        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
3016    }
3017
3018    #[test]
3019    fn test_column_writer_compression_v1() {
3020        let props = WriterProperties::builder()
3021            .set_writer_version(WriterVersion::PARQUET_1_0)
3022            .set_compression(Compression::SNAPPY)
3023            .build();
3024        column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
3025    }
3026
3027    #[test]
3028    fn test_column_writer_compression_v2() {
3029        let props = WriterProperties::builder()
3030            .set_writer_version(WriterVersion::PARQUET_2_0)
3031            .set_compression(Compression::SNAPPY)
3032            .build();
3033        column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
3034    }
3035
3036    #[test]
3037    fn test_column_writer_v2_compression_ratio_threshold() {
3038        fn write_v2_page(threshold: f64) -> bool {
3039            let mut buf = Vec::with_capacity(4096);
3040            let mut write = TrackedWrite::new(&mut buf);
3041            let page_writer = Box::new(SerializedPageWriter::new(&mut write));
3042            let props = Arc::new(
3043                WriterProperties::builder()
3044                    .set_writer_version(WriterVersion::PARQUET_2_0)
3045                    .set_compression(Compression::SNAPPY)
3046                    .set_dictionary_enabled(false)
3047                    .set_data_page_v2_compression_ratio_threshold(threshold)
3048                    .build(),
3049            );
3050
3051            let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3052            let values: Vec<i32> = vec![42; 4096];
3053            writer.write_batch(&values, None, None).unwrap();
3054            let r = writer.close().unwrap();
3055            drop(write);
3056
3057            let reader_props = ReaderProperties::builder()
3058                .set_backward_compatible_lz4(false)
3059                .build();
3060            let reader = SerializedPageReader::new_with_properties(
3061                Arc::new(Bytes::from(buf)),
3062                &r.metadata,
3063                r.rows_written as usize,
3064                None,
3065                Arc::new(reader_props),
3066            )
3067            .unwrap();
3068            let pages = reader.collect::<Result<Vec<_>>>().unwrap();
3069            let data_page = pages
3070                .iter()
3071                .find(|p| p.page_type() == PageType::DATA_PAGE_V2)
3072                .expect("expected a v2 data page");
3073            match data_page {
3074                Page::DataPageV2 { is_compressed, .. } => *is_compressed,
3075                _ => unreachable!(),
3076            }
3077        }
3078
3079        // Default threshold keeps the compressed buffer for constant data.
3080        assert!(write_v2_page(1.0));
3081        // A strict threshold (require >1000x reduction) discards it.
3082        assert!(!write_v2_page(0.001));
3083    }
3084
3085    #[test]
3086    fn test_column_writer_add_data_pages_with_dict() {
3087        // ARROW-5129: Test verifies that we add data page in case of dictionary encoding
3088        // and no fallback occurred so far.
3089        let mut file = tempfile::tempfile().unwrap();
3090        let mut write = TrackedWrite::new(&mut file);
3091        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
3092        let props = Arc::new(
3093            WriterProperties::builder()
3094                .set_data_page_size_limit(10)
3095                .set_write_batch_size(3) // write 3 values at a time
3096                .build(),
3097        );
3098        let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
3099        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3100        writer.write_batch(data, None, None).unwrap();
3101        let r = writer.close().unwrap();
3102
3103        drop(write);
3104
3105        // Read pages and check the sequence
3106        let props = ReaderProperties::builder()
3107            .set_backward_compatible_lz4(false)
3108            .build();
3109        let mut page_reader = Box::new(
3110            SerializedPageReader::new_with_properties(
3111                Arc::new(file),
3112                &r.metadata,
3113                r.rows_written as usize,
3114                None,
3115                Arc::new(props),
3116            )
3117            .unwrap(),
3118        );
3119        let mut res = Vec::new();
3120        while let Some(page) = page_reader.get_next_page().unwrap() {
3121            res.push((page.page_type(), page.num_values(), page.buffer().len()));
3122        }
3123        assert_eq!(
3124            res,
3125            vec![
3126                (PageType::DICTIONARY_PAGE, 10, 40),
3127                (PageType::DATA_PAGE, 9, 10),
3128                (PageType::DATA_PAGE, 1, 3),
3129            ]
3130        );
3131        assert_eq!(
3132            r.metadata.page_encoding_stats(),
3133            Some(&vec![
3134                PageEncodingStats {
3135                    page_type: PageType::DICTIONARY_PAGE,
3136                    encoding: Encoding::PLAIN,
3137                    count: 1
3138                },
3139                PageEncodingStats {
3140                    page_type: PageType::DATA_PAGE,
3141                    encoding: Encoding::RLE_DICTIONARY,
3142                    count: 2,
3143                }
3144            ])
3145        );
3146    }
3147
3148    #[test]
3149    fn test_column_writer_column_data_page_size_limit() {
3150        let props = Arc::new(
3151            WriterProperties::builder()
3152                .set_writer_version(WriterVersion::PARQUET_1_0)
3153                .set_dictionary_enabled(false)
3154                .set_data_page_size_limit(1000)
3155                .set_column_data_page_size_limit(ColumnPath::from("col"), 10)
3156                .set_write_batch_size(3)
3157                .build(),
3158        );
3159        let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
3160
3161        let col_values =
3162            write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data);
3163        let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data);
3164
3165        assert_eq!(col_values, vec![3, 3, 3, 1]);
3166        assert_eq!(other_values, vec![10]);
3167    }
3168
3169    #[test]
3170    fn test_column_writer_column_data_page_row_count_limit() {
3171        let props = Arc::new(
3172            WriterProperties::builder()
3173                .set_writer_version(WriterVersion::PARQUET_1_0)
3174                .set_dictionary_enabled(false)
3175                .set_data_page_row_count_limit(100)
3176                .set_column_data_page_row_count_limit(ColumnPath::from("col"), 3)
3177                .set_write_batch_size(1)
3178                .build(),
3179        );
3180        let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
3181
3182        let col_values =
3183            write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data);
3184        let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data);
3185
3186        assert_eq!(col_values, vec![3, 3, 3, 1]);
3187        assert_eq!(other_values, vec![10]);
3188    }
3189
3190    #[test]
3191    fn test_column_writer_caps_page_size_for_large_byte_array_values() {
3192        // Regression: the post-write data page byte limit check only fires
3193        // at mini-batch boundaries, so a 1024-row mini-batch of multi-MiB
3194        // BYTE_ARRAY values used to buffer multiple GiB into a single page
3195        // before the limit was even consulted. With the threshold-based
3196        // granular mode this batch should split into ~one page per value.
3197        let value_size = 64 * 1024; // 64 KiB per value
3198        let page_byte_limit = 16 * 1024; // 16 KiB page limit
3199        let num_rows = 64;
3200
3201        let props = WriterProperties::builder()
3202            .set_writer_version(WriterVersion::PARQUET_1_0)
3203            .set_dictionary_enabled(false)
3204            .set_encoding(Encoding::PLAIN)
3205            .set_data_page_size_limit(page_byte_limit)
3206            // Default write_batch_size (1024) — without the fix this
3207            // buffers the entire input into a single ~4 MiB page.
3208            .build();
3209
3210        let data: Vec<_> = (0..num_rows)
3211            .map(|i| ByteArray::from(vec![i as u8; value_size]))
3212            .collect();
3213        let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3214
3215        // Every value must end up somewhere.
3216        let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3217        assert_eq!(total_values as usize, num_rows);
3218        // Without the fix this assertion fired with one ~4 MiB page; the
3219        // threshold splits the input so that no page holds more than a
3220        // single oversized value's worth of bytes.
3221        assert!(
3222            pages.data_pages.len() >= num_rows / 2,
3223            "expected pages to be cut close to one per value, got {:?}",
3224            pages.data_pages,
3225        );
3226        // Each page must be bounded by roughly one value's worth of bytes;
3227        // parquet allows a single oversized value to occupy a page by
3228        // itself but never lets us pile many of them together.
3229        for (size, _) in &pages.data_pages {
3230            assert!(
3231                *size <= value_size + 64,
3232                "page size {size} exceeds one-value bound ({}B) — pages {:?}",
3233                value_size + 64,
3234                pages.data_pages,
3235            );
3236        }
3237    }
3238
3239    #[test]
3240    fn test_column_writer_delta_byte_array_dedups_large_shared_prefix_values() {
3241        // Regression for https://github.com/apache/arrow-rs/issues/10489.
3242        // 16 identical 64 KiB values against a 16 KiB page limit: every value
3243        // is over the limit on its own, and `DELTA_BYTE_ARRAY` should still
3244        // dedup them down to about one value's worth of bytes in total.
3245        let value_size = 64 * 1024; // 64 KiB per value, > the page limit
3246        let page_byte_limit = 16 * 1024;
3247        let num_rows = 16;
3248
3249        let props = WriterProperties::builder()
3250            .set_writer_version(WriterVersion::PARQUET_1_0)
3251            .set_dictionary_enabled(false)
3252            .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3253            .set_data_page_size_limit(page_byte_limit)
3254            .set_statistics_enabled(EnabledStatistics::None)
3255            .build();
3256
3257        // Identical values: one full value plus `num_rows - 1` zero-length
3258        // suffixes is all this column should cost.
3259        let data: Vec<_> = (0..num_rows)
3260            .map(|_| ByteArray::from(vec![b'a'; value_size]))
3261            .collect();
3262        let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3263
3264        // Every value must still end up somewhere.
3265        let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3266        assert_eq!(total_values as usize, num_rows);
3267
3268        // Before the fix this was `num_rows * value_size` — byte for byte
3269        // what PLAIN produces, i.e. the encoding doing no work at all.
3270        let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum();
3271        assert!(
3272            total_bytes < 2 * value_size,
3273            "expected under 2x a single value ({}B) for {num_rows} identical \
3274             values, got {total_bytes}B across pages {:?}",
3275            2 * value_size,
3276            pages.data_pages,
3277        );
3278    }
3279
3280    #[test]
3281    fn test_column_writer_delta_byte_array_bounds_pages_without_shared_prefix() {
3282        // Companion to the test above: same shape, but the values share no
3283        // prefix, so there is nothing to dedup and pages must stay bounded
3284        // by the value size. This is why the exemption covers one value
3285        // rather than dropping the byte budget altogether.
3286        let value_size = 64 * 1024;
3287        let page_byte_limit = 16 * 1024;
3288        let num_rows = 16;
3289
3290        let props = WriterProperties::builder()
3291            .set_writer_version(WriterVersion::PARQUET_1_0)
3292            .set_dictionary_enabled(false)
3293            .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3294            .set_data_page_size_limit(page_byte_limit)
3295            .set_statistics_enabled(EnabledStatistics::None)
3296            .build();
3297
3298        // No two values share a prefix: they differ at the first byte.
3299        let data: Vec<_> = (0..num_rows)
3300            .map(|i| ByteArray::from(vec![i as u8; value_size]))
3301            .collect();
3302        let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3303
3304        let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3305        assert_eq!(total_values as usize, num_rows);
3306
3307        // Expect at most two values per page: the exempted first value plus
3308        // one more that trips the budget.
3309        let upper_bound = 2 * value_size + 64;
3310        for (size, n_values) in &pages.data_pages {
3311            assert!(
3312                *size <= upper_bound,
3313                "page size {size} exceeds two-value bound ({upper_bound}B); pages {:?}",
3314                pages.data_pages,
3315            );
3316            assert!(
3317                *n_values <= 2,
3318                "page holds {n_values} values, expected at most 2; pages {:?}",
3319                pages.data_pages,
3320            );
3321        }
3322    }
3323
3324    #[test]
3325    fn test_column_writer_caps_page_size_with_sparse_nulls() {
3326        // `PLAIN` keeps ratio-scaled windows, so a sparsely-null chunk puts
3327        // *two* over-limit values on a page rather than one. That is the
3328        // deliberate choice: value-exact windows would cut this to one, but
3329        // `PLAIN` stores a value identically wherever it lands, so the output
3330        // is byte for byte the same either way while the page count doubles.
3331        //
3332        // What has to hold is that the bound stays a small constant and does
3333        // not scale with `write_batch_size` — the failure #9972 fixed, where
3334        // a page took a whole mini-batch. A window spans
3335        // `ceil(values * levels / values_in_chunk)` levels, which covers at
3336        // most two values however sparse the nulls are, so two is the whole
3337        // exposure. Assert it exactly: this test fails if the encoding gate
3338        // on value-exact windows is dropped (pages would hold one value) as
3339        // well as if the bound is lost (they would hold many).
3340        let value_size = 64 * 1024;
3341        let page_byte_limit = 16 * 1024;
3342        let num_values = 16;
3343
3344        let props = WriterProperties::builder()
3345            .set_dictionary_enabled(false)
3346            .set_encoding(Encoding::PLAIN)
3347            .set_data_page_size_limit(page_byte_limit)
3348            .set_statistics_enabled(EnabledStatistics::None)
3349            .build();
3350
3351        let data: Vec<_> = (0..num_values)
3352            .map(|_| ByteArray::from(vec![b'a'; value_size]))
3353            .collect();
3354        // 17 levels: a null at index 8, values everywhere else.
3355        let def_levels: Vec<i16> = (0..num_values as i16 + 1)
3356            .map(|i| i16::from(i != 8))
3357            .collect();
3358        let pages =
3359            write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3360
3361        // At most two values' payload on any page, and never a whole
3362        // mini-batch's worth.
3363        let upper_bound = 2 * value_size + 64;
3364        for (size, n_levels) in &pages.data_pages {
3365            assert!(
3366                *size <= upper_bound,
3367                "page size {size} exceeds two-value bound ({upper_bound}B); pages {:?}",
3368                pages.data_pages,
3369            );
3370            assert!(
3371                *n_levels <= 3,
3372                "page holds {n_levels} levels, expected at most 3 (two values + a null); \
3373                 pages {:?}",
3374                pages.data_pages,
3375            );
3376        }
3377        // Two-level windows over 17 levels: eight pages carrying two levels
3378        // and a ninth holding the remainder.
3379        let num_levels: usize = num_values + 1;
3380        assert_eq!(pages.data_pages.len(), num_levels.div_ceil(2));
3381    }
3382
3383    #[test]
3384    fn test_column_writer_delta_byte_array_nullable_shared_prefix_dedup() {
3385        // A null in the chunk must not cost dedup. The first-value exemption
3386        // fires when a page's opening mini-batch holds exactly one value, and
3387        // `write_granular_chunk` cuts windows after an exact value count, so
3388        // the over-limit value that opens the page gets a single-value
3389        // mini-batch regardless of where nulls fall.
3390        //
3391        // This pinned `[2, 2, 2, 2, 9]` before #10538: the chunker scaled the
3392        // one-value budget by the chunk's 17:16 level:value ratio and rounded
3393        // up to two-level windows, so most pages opened with two values,
3394        // missed the exemption, and stored their first value in full — ~5
3395        // full values against ~1 here (`PLAIN` stores all 16).
3396        let value_size = 64 * 1024;
3397        let page_byte_limit = 16 * 1024;
3398        let num_values = 16;
3399
3400        let props = WriterProperties::builder()
3401            .set_writer_version(WriterVersion::PARQUET_1_0)
3402            .set_dictionary_enabled(false)
3403            .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3404            .set_data_page_size_limit(page_byte_limit)
3405            .set_statistics_enabled(EnabledStatistics::None)
3406            .build();
3407
3408        let data: Vec<_> = (0..num_values)
3409            .map(|_| ByteArray::from(vec![b'a'; value_size]))
3410            .collect();
3411        // 17 levels: a null at index 8, values everywhere else.
3412        let def_levels: Vec<i16> = (0..num_values as i16 + 1)
3413            .map(|i| i16::from(i != 8))
3414            .collect();
3415        let pages =
3416            write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3417
3418        // One page holding all 17 levels: the first value is exempt and every
3419        // later value dedups down to a length pair.
3420        let per_page_values: Vec<u32> = pages.data_pages.iter().map(|(_, n)| *n).collect();
3421        assert_eq!(per_page_values, vec![num_values as u32 + 1]);
3422
3423        let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum();
3424        assert!(
3425            total_bytes < 2 * value_size,
3426            "expected ~one full value's worth of bytes (full dedup), \
3427             got {total_bytes}B across pages {:?}",
3428            pages.data_pages,
3429        );
3430    }
3431
3432    #[test]
3433    fn test_column_writer_caps_page_size_for_large_values_in_list() {
3434        // Coverage for the Materialized-rep branch of
3435        // `write_granular_chunk`. The flat-column regression test
3436        // exercises the per-level step; this exercises the
3437        // record-by-record step used when rep levels are present.
3438        //
3439        // Column is `list<required binary>` (max_def = 1, max_rep = 1)
3440        // with 3 records of 3 large blobs each. The page byte limit is
3441        // smaller than a single blob, so granular mode kicks in, and the
3442        // Materialized-rep arm of `write_granular_chunk` steps from one
3443        // `rep == 0` boundary to the next so a record never spans pages.
3444        let value_size = 32 * 1024;
3445        let page_byte_limit = 16 * 1024;
3446        let values_per_record = 3;
3447        let num_records = 3;
3448        let num_values = values_per_record * num_records;
3449
3450        // rep levels: 0, 1, 1, 0, 1, 1, 0, 1, 1
3451        let mut rep_levels = Vec::with_capacity(num_values);
3452        for _ in 0..num_records {
3453            rep_levels.push(0i16);
3454            rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1));
3455        }
3456        let def_levels = vec![1i16; num_values];
3457
3458        let props = WriterProperties::builder()
3459            .set_writer_version(WriterVersion::PARQUET_1_0)
3460            .set_dictionary_enabled(false)
3461            .set_encoding(Encoding::PLAIN)
3462            .set_data_page_size_limit(page_byte_limit)
3463            .build();
3464
3465        let data: Vec<_> = (0..num_values)
3466            .map(|i| ByteArray::from(vec![i as u8; value_size]))
3467            .collect();
3468        let pages = write_and_collect_pages::<ByteArrayType>(
3469            props,
3470            1,
3471            1,
3472            &data,
3473            Some(&def_levels),
3474            Some(&rep_levels),
3475        );
3476        let data_pages = pages.data_pages;
3477
3478        // The Materialized-rep arm groups levels by record, and each
3479        // record's bytes blow the page byte limit on its own, so we get
3480        // exactly one page per record.
3481        assert_eq!(
3482            data_pages.len(),
3483            num_records,
3484            "expected one data page per record, got {data_pages:?}"
3485        );
3486        for (bytes, n_values) in &data_pages {
3487            assert_eq!(
3488                *n_values as usize, values_per_record,
3489                "each page must hold a whole record's leaves, got {data_pages:?}"
3490            );
3491            // Each page is one full record (its leaves cannot be split),
3492            // so allow up to `values_per_record` blobs of payload plus a
3493            // small fudge for level encoding overhead.
3494            let upper_bound = values_per_record * (value_size + 16);
3495            assert!(
3496                *bytes <= upper_bound,
3497                "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}"
3498            );
3499        }
3500    }
3501
3502    #[test]
3503    fn test_column_writer_caps_page_size_with_nullable_large_values() {
3504        // Coverage for `LevelDataRef::value_count` on Materialized def
3505        // levels: a nullable column with mixed nulls and large values.
3506        // `value_count` must return the actual non-null count so the
3507        // byte estimate reflects bytes that will actually be written,
3508        // not the level count.
3509        let value_size = 32 * 1024;
3510        let page_byte_limit = 16 * 1024;
3511        let num_levels = 32;
3512
3513        // Alternating null / non-null: 16 nulls and 16 values.
3514        let def_levels: Vec<i16> = (0..num_levels as i16).map(|i| i % 2).collect();
3515        let num_values = def_levels.iter().filter(|&&d| d == 1).count();
3516
3517        let props = WriterProperties::builder()
3518            .set_writer_version(WriterVersion::PARQUET_1_0)
3519            .set_dictionary_enabled(false)
3520            .set_encoding(Encoding::PLAIN)
3521            .set_data_page_size_limit(page_byte_limit)
3522            .build();
3523
3524        let data: Vec<_> = (0..num_values)
3525            .map(|i| ByteArray::from(vec![i as u8; value_size]))
3526            .collect();
3527        let pages =
3528            write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3529        let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3530
3531        // With 16 actual values of 32 KiB each and a 16 KiB page limit,
3532        // every non-null value should get its own page (plus possibly
3533        // adjacent nulls). At minimum, the number of pages must be
3534        // roughly the value count, not 1 (which is what `main` produced).
3535        assert!(
3536            data_pages.len() >= num_values / 2,
3537            "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}",
3538            num_values / 2,
3539            data_pages.len(),
3540        );
3541        // No page contains more than ~one value's worth of payload bytes.
3542        for size in &data_pages {
3543            assert!(
3544                *size <= value_size + 64,
3545                "page size {size} exceeds one-value bound; pages {data_pages:?}"
3546            );
3547        }
3548    }
3549
3550    #[test]
3551    fn test_column_writer_dict_enabled_large_values_post_spill() {
3552        // While dictionary encoding is active, `has_dictionary()` short-
3553        // circuits `estimated_value_bytes` — the byte estimate is plain-
3554        // encoded size but dict-encoded pages only store small RLE
3555        // indices, so we'd otherwise shrink pages spuriously. Once the
3556        // dictionary spills (each value is large + unique), plain
3557        // encoding takes over and the byte-budget sub-batch kicks in.
3558        //
3559        // This test makes sure the writer survives that transition and
3560        // produces bounded pages thereafter.
3561        let value_size = 64 * 1024;
3562        let page_byte_limit = 16 * 1024;
3563        let num_rows = 32;
3564
3565        let props = WriterProperties::builder()
3566            .set_writer_version(WriterVersion::PARQUET_1_0)
3567            .set_dictionary_enabled(true)
3568            // Force a small dict so it spills quickly even though
3569            // each value here is unique.
3570            .set_dictionary_page_size_limit(1024)
3571            .set_data_page_size_limit(page_byte_limit)
3572            // Small mini-batches so dict fallback happens part-way
3573            // through the input, leaving subsequent mini-batches to
3574            // exercise the post-spill plain-encoding path that the
3575            // page-size fix actually targets.
3576            .set_write_batch_size(4)
3577            .build();
3578
3579        let data: Vec<_> = (0..num_rows)
3580            .map(|i| ByteArray::from(vec![i as u8; value_size]))
3581            .collect();
3582        let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3583        let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3584
3585        // After spill, plain encoding writes one ~64 KiB value per page.
3586        // Without the fix, post-spill writes still buffered all 32
3587        // values into a single ~2 MiB page.
3588        assert!(
3589            data_pages.len() >= num_rows / 2,
3590            "expected >= {} data pages after dict spill, got {} ({data_pages:?})",
3591            num_rows / 2,
3592            data_pages.len(),
3593        );
3594        for size in &data_pages {
3595            assert!(
3596                *size <= value_size + 64,
3597                "page size {size} exceeds one-value bound; pages {data_pages:?}"
3598            );
3599        }
3600    }
3601
3602    #[test]
3603    fn test_column_writer_caps_dictionary_page_size() {
3604        // A column of large *distinct* values with dictionary encoding on:
3605        // the dictionary page accumulates the values themselves, and its
3606        // spill check runs only once per mini-batch. Without bounding the
3607        // dictionary-encoding mini-batch, one `write_batch_size` mini-batch
3608        // would intern `write_batch_size * value_size` bytes into the
3609        // dictionary page before the check fires (~16 MiB here). The chunker
3610        // must sub-batch the dictionary-encoding phase too.
3611        let value_size = 8 * 1024;
3612        let dict_page_limit = 64 * 1024;
3613        let num_rows = 2048;
3614
3615        let props = WriterProperties::builder()
3616            .set_writer_version(WriterVersion::PARQUET_1_0)
3617            .set_dictionary_enabled(true)
3618            .set_dictionary_page_size_limit(dict_page_limit)
3619            .build();
3620
3621        let data: Vec<_> = (0..num_rows)
3622            .map(|i| {
3623                // each value distinct, so the dictionary cannot dedup them
3624                let mut v = vec![0u8; value_size];
3625                v[..8].copy_from_slice(&(i as u64).to_le_bytes());
3626                ByteArray::from(v)
3627            })
3628            .collect();
3629        let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3630        let dict_page_size = pages.dict_page_size;
3631
3632        assert!(
3633            dict_page_size > 0,
3634            "expected the column to dictionary-encode"
3635        );
3636        // Bounded near the limit (~2x from the post-mini-batch check). Before
3637        // the fix the dictionary page reached num_rows * value_size (~16 MiB,
3638        // 256x the limit).
3639        assert!(
3640            dict_page_size <= 3 * dict_page_limit,
3641            "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit",
3642        );
3643    }
3644
3645    #[test]
3646    fn test_column_writer_caps_page_size_for_fixed_len_byte_array() {
3647        // Coverage for `ParquetValueType::byte_size` override on
3648        // `FixedLenByteArray`. With `type_length = 1`, each plain-encoded
3649        // value is one byte, so a 4-byte page byte limit forces the
3650        // sub-batch sizer to write ~4 values per page rather than one
3651        // page for the whole batch.
3652        let page_byte_limit = 4;
3653        let num_values = 128;
3654
3655        let props = WriterProperties::builder()
3656            .set_writer_version(WriterVersion::PARQUET_1_0)
3657            .set_dictionary_enabled(false)
3658            .set_encoding(Encoding::PLAIN)
3659            .set_data_page_size_limit(page_byte_limit)
3660            .build();
3661
3662        let data: Vec<_> = (0..num_values)
3663            .map(|i| {
3664                let mut fla = FixedLenByteArray::default();
3665                fla.set_data(Bytes::from(vec![i as u8]));
3666                fla
3667            })
3668            .collect();
3669        let pages =
3670            write_and_collect_pages::<FixedLenByteArrayType>(props, 0, 0, &data, None, None);
3671        let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3672
3673        // Without the fix this is a single 128-byte page; with the fix
3674        // the byte budget caps each page at ~`page_byte_limit` bytes.
3675        assert!(
3676            data_pages.len() >= num_values / 8,
3677            "expected pages capped by byte budget, got {data_pages:?}"
3678        );
3679        for size in &data_pages {
3680            assert!(
3681                *size <= page_byte_limit * 4,
3682                "page size {size} larger than expected; pages {data_pages:?}"
3683            );
3684        }
3685    }
3686
3687    #[test]
3688    fn test_bool_statistics() {
3689        let stats = statistics_roundtrip::<BoolType>(&[true, false, false, true]);
3690        // Booleans have an unsigned sort order and so are not compatible
3691        // with the deprecated `min` and `max` statistics
3692        assert!(!stats.is_min_max_backwards_compatible());
3693        if let Statistics::Boolean(stats) = stats {
3694            assert_eq!(stats.min_opt().unwrap(), &false);
3695            assert_eq!(stats.max_opt().unwrap(), &true);
3696        } else {
3697            panic!("expecting Statistics::Boolean, got {stats:?}");
3698        }
3699    }
3700
3701    #[test]
3702    fn test_int32_statistics() {
3703        let stats = statistics_roundtrip::<Int32Type>(&[-1, 3, -2, 2]);
3704        assert!(stats.is_min_max_backwards_compatible());
3705        if let Statistics::Int32(stats) = stats {
3706            assert_eq!(stats.min_opt().unwrap(), &-2);
3707            assert_eq!(stats.max_opt().unwrap(), &3);
3708        } else {
3709            panic!("expecting Statistics::Int32, got {stats:?}");
3710        }
3711    }
3712
3713    #[test]
3714    fn test_int64_statistics() {
3715        let stats = statistics_roundtrip::<Int64Type>(&[-1, 3, -2, 2]);
3716        assert!(stats.is_min_max_backwards_compatible());
3717        if let Statistics::Int64(stats) = stats {
3718            assert_eq!(stats.min_opt().unwrap(), &-2);
3719            assert_eq!(stats.max_opt().unwrap(), &3);
3720        } else {
3721            panic!("expecting Statistics::Int64, got {stats:?}");
3722        }
3723    }
3724
3725    #[test]
3726    fn test_int96_statistics() {
3727        let input = vec![
3728            Int96::from(vec![1, 20, 30]),
3729            Int96::from(vec![3, 20, 10]),
3730            Int96::from(vec![0, 20, 30]),
3731            Int96::from(vec![2, 20, 30]),
3732        ]
3733        .into_iter()
3734        .collect::<Vec<Int96>>();
3735
3736        let stats = statistics_roundtrip::<Int96Type>(&input);
3737        assert!(!stats.is_min_max_backwards_compatible());
3738        if let Statistics::Int96(stats) = stats {
3739            assert_eq!(stats.min_opt().unwrap(), &Int96::from(vec![3, 20, 10]));
3740            assert_eq!(stats.max_opt().unwrap(), &Int96::from(vec![2, 20, 30]));
3741        } else {
3742            panic!("expecting Statistics::Int96, got {stats:?}");
3743        }
3744    }
3745
3746    #[test]
3747    fn test_float_statistics() {
3748        let stats = statistics_roundtrip::<FloatType>(&[-1.0, 3.0, -2.0, 2.0]);
3749        assert!(!stats.is_min_max_backwards_compatible());
3750        if let Statistics::Float(stats) = stats {
3751            assert_eq!(stats.min_opt().unwrap(), &-2.0);
3752            assert_eq!(stats.max_opt().unwrap(), &3.0);
3753        } else {
3754            panic!("expecting Statistics::Float, got {stats:?}");
3755        }
3756    }
3757
3758    #[test]
3759    fn test_double_statistics() {
3760        let stats = statistics_roundtrip::<DoubleType>(&[-1.0, 3.0, -2.0, 2.0]);
3761        assert!(!stats.is_min_max_backwards_compatible());
3762        if let Statistics::Double(stats) = stats {
3763            assert_eq!(stats.min_opt().unwrap(), &-2.0);
3764            assert_eq!(stats.max_opt().unwrap(), &3.0);
3765        } else {
3766            panic!("expecting Statistics::Double, got {stats:?}");
3767        }
3768    }
3769
3770    #[test]
3771    fn test_byte_array_statistics() {
3772        let input = ["aawaa", "zz", "aaw", "m", "qrs"]
3773            .iter()
3774            .map(|&s| s.into())
3775            .collect::<Vec<_>>();
3776
3777        let stats = statistics_roundtrip::<ByteArrayType>(&input);
3778        assert!(!stats.is_min_max_backwards_compatible());
3779        if let Statistics::ByteArray(stats) = stats {
3780            assert_eq!(stats.min_opt().unwrap(), &ByteArray::from("aaw"));
3781            assert_eq!(stats.max_opt().unwrap(), &ByteArray::from("zz"));
3782        } else {
3783            panic!("expecting Statistics::ByteArray, got {stats:?}");
3784        }
3785    }
3786
3787    #[test]
3788    fn test_fixed_len_byte_array_statistics() {
3789        let input = ["aawaa", "zz   ", "aaw  ", "m    ", "qrs  "]
3790            .iter()
3791            .map(|&s| ByteArray::from(s).into())
3792            .collect::<Vec<_>>();
3793
3794        let stats = statistics_roundtrip::<FixedLenByteArrayType>(&input);
3795        assert!(!stats.is_min_max_backwards_compatible());
3796        if let Statistics::FixedLenByteArray(stats) = stats {
3797            let expected_min: FixedLenByteArray = ByteArray::from("aaw  ").into();
3798            assert_eq!(stats.min_opt().unwrap(), &expected_min);
3799            let expected_max: FixedLenByteArray = ByteArray::from("zz   ").into();
3800            assert_eq!(stats.max_opt().unwrap(), &expected_max);
3801        } else {
3802            panic!("expecting Statistics::FixedLenByteArray, got {stats:?}");
3803        }
3804    }
3805
3806    #[test]
3807    fn test_ieee754_total_order_float() {
3808        // Test IEEE 754 total order for f32
3809        // Order should be: -NaN < -Inf < -1.0 < -0.0 < +0.0 < 1.0 < +Inf < +NaN
3810        let neg_nan = f32::from_bits(0xffc00000); // a NaN with the sign bit set
3811        let neg_inf = f32::NEG_INFINITY;
3812        let neg_one = -1.0_f32;
3813        let neg_zero = -0.0_f32;
3814        let pos_zero = 0.0_f32;
3815        let pos_one = 1.0_f32;
3816        let pos_inf = f32::INFINITY;
3817        let pos_nan = f32::from_bits(0x7fc00000); // a NaN with the sign bit unset
3818
3819        let values = vec![
3820            pos_nan, neg_zero, pos_inf, neg_one, neg_nan, pos_one, neg_inf, pos_zero,
3821        ];
3822
3823        let stats = statistics_roundtrip::<FloatType>(&values);
3824        if let Statistics::Float(stats) = stats {
3825            // With IEEE 754 total order, min should be -NaN, max should be +NaN
3826            // But since we filter out NaN values, min should be -Inf, max should be +Inf
3827            assert_eq!(stats.min_opt().unwrap(), &neg_inf);
3828            assert_eq!(stats.max_opt().unwrap(), &pos_inf);
3829            assert_eq!(stats.nan_count_opt(), Some(2)); // neg_nan and pos_nan
3830        } else {
3831            panic!("Expected float statistics");
3832        }
3833    }
3834
3835    #[test]
3836    fn test_ieee754_total_order_float_only_nan() {
3837        // Test IEEE 754 total order for various NaN representations
3838        // They should be ordered by the significand
3839        let neg_nan1 = f32::from_bits(0xffc00000); // sign bit set, significand x400000
3840        let neg_nan2 = f32::from_bits(0xffc00001); // sign bit set, significand x400001
3841        let neg_nan3 = f32::from_bits(0xffc00002); // sign bit set, significand x400002
3842        let pos_nan1 = f32::from_bits(0x7fc00000); // sign bit unset, significand x400000
3843        let pos_nan2 = f32::from_bits(0x7fc00001); // sign bit unset, significand x400001
3844        let pos_nan3 = f32::from_bits(0x7fc00002); // sign bit unset, significand x400002
3845
3846        let values = vec![neg_nan1, neg_nan2, neg_nan3, pos_nan1, pos_nan2, pos_nan3];
3847
3848        let stats = statistics_roundtrip::<FloatType>(&values);
3849        if let Statistics::Float(stats) = stats {
3850            // With IEEE 754 total order, min should be `neg_nan3`, max `pos_nan3`
3851            assert_eq!(
3852                stats.min_opt().unwrap().total_cmp(&neg_nan3),
3853                Ordering::Equal
3854            );
3855            assert_eq!(
3856                stats.max_opt().unwrap().total_cmp(&pos_nan3),
3857                Ordering::Equal
3858            );
3859            assert_eq!(stats.nan_count_opt(), Some(6));
3860        } else {
3861            panic!("Expected float statistics");
3862        }
3863    }
3864
3865    #[test]
3866    fn test_ieee754_total_order_double() {
3867        // Test IEEE 754 total order for f64
3868        let neg_nan = f64::from_bits(0xfff8000000000000);
3869        let neg_inf = f64::NEG_INFINITY;
3870        let neg_one = -1.0_f64;
3871        let neg_zero = -0.0_f64;
3872        let pos_zero = 0.0_f64;
3873        let pos_one = 1.0_f64;
3874        let pos_inf = f64::INFINITY;
3875        let pos_nan = f64::from_bits(0x7ff8000000000000);
3876
3877        let values = vec![
3878            pos_nan, neg_zero, pos_inf, neg_one, neg_nan, pos_one, neg_inf, pos_zero,
3879        ];
3880
3881        let stats = statistics_roundtrip::<DoubleType>(&values);
3882        if let Statistics::Double(stats) = stats {
3883            // With IEEE 754 total order, and NaN filtering
3884            assert_eq!(stats.min_opt().unwrap(), &neg_inf);
3885            assert_eq!(stats.max_opt().unwrap(), &pos_inf);
3886            assert_eq!(stats.nan_count_opt(), Some(2));
3887        } else {
3888            panic!("Expected double statistics");
3889        }
3890    }
3891
3892    #[test]
3893    fn test_ieee754_total_order_double_only_nan() {
3894        // Test IEEE 754 total order for various NaN representations
3895        // They should be ordered by the significand
3896        let neg_nan1 = f64::from_bits(0xfff8000000000000);
3897        let neg_nan2 = f64::from_bits(0xfff8000000000001);
3898        let neg_nan3 = f64::from_bits(0xfff8000000000002);
3899        let pos_nan1 = f64::from_bits(0x7ff8000000000000);
3900        let pos_nan2 = f64::from_bits(0x7ff8000000000001);
3901        let pos_nan3 = f64::from_bits(0x7ff8000000000002);
3902
3903        let values = vec![neg_nan1, neg_nan2, neg_nan3, pos_nan1, pos_nan2, pos_nan3];
3904
3905        let stats = statistics_roundtrip::<DoubleType>(&values);
3906        if let Statistics::Double(stats) = stats {
3907            // With IEEE 754 total order, min should be `neg_nan3`, max `pos_nan3`
3908            assert_eq!(
3909                stats.min_opt().unwrap().total_cmp(&neg_nan3),
3910                Ordering::Equal
3911            );
3912            assert_eq!(
3913                stats.max_opt().unwrap().total_cmp(&pos_nan3),
3914                Ordering::Equal
3915            );
3916            assert_eq!(stats.nan_count_opt(), Some(6));
3917        } else {
3918            panic!("Expected float statistics");
3919        }
3920    }
3921
3922    #[test]
3923    fn test_ieee754_total_order_zeros() {
3924        // Test that -0.0 and +0.0 are handled correctly
3925        let values = vec![-0.0_f32, 0.0_f32, -0.0_f32, 0.0_f32];
3926
3927        let stats = statistics_roundtrip::<FloatType>(&values);
3928        if let Statistics::Float(stats) = stats {
3929            // With IEEE 754 total order, -0.0 < +0.0
3930            assert_eq!(stats.min_opt().unwrap().to_bits(), (-0.0_f32).to_bits());
3931            assert_eq!(stats.max_opt().unwrap().to_bits(), 0.0_f32.to_bits());
3932        } else {
3933            panic!("Expected float statistics");
3934        }
3935    }
3936
3937    #[test]
3938    #[cfg_attr(miri, ignore)] // inline assembly is not supported
3939    fn test_column_writer_check_float16_min_max() {
3940        let input = [
3941            -f16::ONE,
3942            f16::from_f32(3.0),
3943            -f16::from_f32(2.0),
3944            f16::from_f32(2.0),
3945        ]
3946        .into_iter()
3947        .map(|s| ByteArray::from(s).into())
3948        .collect::<Vec<_>>();
3949
3950        let stats = float16_statistics_roundtrip(&input);
3951        assert!(!stats.is_min_max_backwards_compatible());
3952        assert_eq!(
3953            stats.min_opt().unwrap(),
3954            &ByteArray::from(-f16::from_f32(2.0))
3955        );
3956        assert_eq!(
3957            stats.max_opt().unwrap(),
3958            &ByteArray::from(f16::from_f32(3.0))
3959        );
3960    }
3961
3962    #[test]
3963    #[cfg_attr(miri, ignore)] // inline assembly is not supported
3964    fn test_column_writer_check_float16_nan_middle() {
3965        let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
3966            .into_iter()
3967            .map(|s| ByteArray::from(s).into())
3968            .collect::<Vec<_>>();
3969
3970        let stats = float16_statistics_roundtrip(&input);
3971        assert!(!stats.is_min_max_backwards_compatible());
3972        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
3973        assert_eq!(
3974            stats.max_opt().unwrap(),
3975            &ByteArray::from(f16::ONE + f16::ONE)
3976        );
3977        assert_eq!(stats.nan_count_opt(), Some(1));
3978    }
3979
3980    #[test]
3981    #[cfg_attr(miri, ignore)] // inline assembly is not supported
3982    fn test_float16_statistics_nan_middle() {
3983        let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
3984            .into_iter()
3985            .map(|s| ByteArray::from(s).into())
3986            .collect::<Vec<_>>();
3987
3988        let stats = float16_statistics_roundtrip(&input);
3989        assert!(!stats.is_min_max_backwards_compatible());
3990        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
3991        assert_eq!(
3992            stats.max_opt().unwrap(),
3993            &ByteArray::from(f16::ONE + f16::ONE)
3994        );
3995        assert_eq!(stats.nan_count_opt(), Some(1));
3996    }
3997
3998    #[test]
3999    #[cfg_attr(miri, ignore)] // inline assembly is not supported
4000    fn test_float16_statistics_nan_start() {
4001        let input = [f16::NAN, f16::ONE, f16::ONE + f16::ONE]
4002            .into_iter()
4003            .map(|s| ByteArray::from(s).into())
4004            .collect::<Vec<_>>();
4005
4006        let stats = float16_statistics_roundtrip(&input);
4007        assert!(!stats.is_min_max_backwards_compatible());
4008        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
4009        assert_eq!(
4010            stats.max_opt().unwrap(),
4011            &ByteArray::from(f16::ONE + f16::ONE)
4012        );
4013        assert_eq!(stats.nan_count_opt(), Some(1));
4014    }
4015
4016    #[test]
4017    fn test_float16_statistics_nan_only() {
4018        let input = [f16::NAN, f16::NAN]
4019            .into_iter()
4020            .map(|s| ByteArray::from(s).into())
4021            .collect::<Vec<_>>();
4022
4023        let stats = float16_statistics_roundtrip(&input);
4024        assert_eq!(
4025            stats.min_bytes_opt(),
4026            Some(ByteArray::from(f16::NAN).as_bytes())
4027        );
4028        assert_eq!(
4029            stats.max_bytes_opt(),
4030            Some(ByteArray::from(f16::NAN).as_bytes())
4031        );
4032        assert!(!stats.is_min_max_backwards_compatible());
4033        assert_eq!(stats.nan_count_opt(), Some(2));
4034    }
4035
4036    #[test]
4037    fn test_float16_statistics_zero_only() {
4038        let input = std::iter::once(f16::ZERO)
4039            .map(|s| ByteArray::from(s).into())
4040            .collect::<Vec<_>>();
4041
4042        let stats = float16_statistics_roundtrip(&input);
4043        assert!(!stats.is_min_max_backwards_compatible());
4044        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ZERO));
4045        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO));
4046    }
4047
4048    #[test]
4049    fn test_float16_statistics_neg_zero_only() {
4050        let input = std::iter::once(f16::NEG_ZERO)
4051            .map(|s| ByteArray::from(s).into())
4052            .collect::<Vec<_>>();
4053
4054        let stats = float16_statistics_roundtrip(&input);
4055        assert!(!stats.is_min_max_backwards_compatible());
4056        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4057        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4058    }
4059
4060    #[test]
4061    fn test_float16_statistics_zero_min() {
4062        let input = [f16::ZERO, f16::ONE, f16::NAN, f16::PI]
4063            .into_iter()
4064            .map(|s| ByteArray::from(s).into())
4065            .collect::<Vec<_>>();
4066
4067        let stats = float16_statistics_roundtrip(&input);
4068        assert!(!stats.is_min_max_backwards_compatible());
4069        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ZERO));
4070        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::PI));
4071    }
4072
4073    #[test]
4074    fn test_float16_statistics_neg_zero_max() {
4075        let input = [f16::NEG_ZERO, f16::NEG_ONE, f16::NAN, -f16::PI]
4076            .into_iter()
4077            .map(|s| ByteArray::from(s).into())
4078            .collect::<Vec<_>>();
4079
4080        let stats = float16_statistics_roundtrip(&input);
4081        assert!(!stats.is_min_max_backwards_compatible());
4082        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(-f16::PI));
4083        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4084    }
4085
4086    #[test]
4087    fn test_float_statistics_nan_middle() {
4088        let stats = statistics_roundtrip::<FloatType>(&[1.0, f32::NAN, 2.0]);
4089        assert!(!stats.is_min_max_backwards_compatible());
4090        if let Statistics::Float(stats) = stats {
4091            assert_eq!(stats.min_opt().unwrap(), &1.0);
4092            assert_eq!(stats.max_opt().unwrap(), &2.0);
4093            assert_eq!(stats.nan_count_opt(), Some(1))
4094        } else {
4095            panic!("expecting Statistics::Float");
4096        }
4097    }
4098
4099    #[test]
4100    fn test_float_statistics_nan_start() {
4101        let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, 1.0, 2.0]);
4102        assert!(!stats.is_min_max_backwards_compatible());
4103        if let Statistics::Float(stats) = stats {
4104            assert_eq!(stats.min_opt().unwrap(), &1.0);
4105            assert_eq!(stats.max_opt().unwrap(), &2.0);
4106            assert_eq!(stats.nan_count_opt(), Some(1))
4107        } else {
4108            panic!("expecting Statistics::Float");
4109        }
4110    }
4111
4112    #[test]
4113    fn test_float_statistics_nan_only() {
4114        let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, f32::NAN]);
4115        assert_eq!(stats.min_bytes_opt(), Some(f32::NAN.as_bytes()));
4116        assert_eq!(stats.max_bytes_opt(), Some(f32::NAN.as_bytes()));
4117        assert_eq!(stats.nan_count_opt(), Some(2));
4118        assert!(!stats.is_min_max_backwards_compatible());
4119        assert!(matches!(stats, Statistics::Float(_)));
4120    }
4121
4122    #[test]
4123    fn test_float_statistics_zero_only() {
4124        let stats = statistics_roundtrip::<FloatType>(&[0.0]);
4125        assert!(!stats.is_min_max_backwards_compatible());
4126        if let Statistics::Float(stats) = stats {
4127            assert_eq!(stats.min_opt().unwrap(), &0.0);
4128            assert!(stats.min_opt().unwrap().is_sign_positive());
4129            assert_eq!(stats.max_opt().unwrap(), &0.0);
4130            assert!(stats.max_opt().unwrap().is_sign_positive());
4131        } else {
4132            panic!("expecting Statistics::Float");
4133        }
4134    }
4135
4136    #[test]
4137    fn test_float_statistics_neg_zero_only() {
4138        let stats = statistics_roundtrip::<FloatType>(&[-0.0]);
4139        assert!(!stats.is_min_max_backwards_compatible());
4140        if let Statistics::Float(stats) = stats {
4141            assert_eq!(stats.min_opt().unwrap(), &-0.0);
4142            assert!(stats.min_opt().unwrap().is_sign_negative());
4143            assert_eq!(stats.max_opt().unwrap(), &-0.0);
4144            assert!(stats.max_opt().unwrap().is_sign_negative());
4145        } else {
4146            panic!("expecting Statistics::Float");
4147        }
4148    }
4149
4150    #[test]
4151    fn test_float_statistics_zero_min() {
4152        let stats = statistics_roundtrip::<FloatType>(&[0.0, 1.0, f32::NAN, 2.0]);
4153        assert!(!stats.is_min_max_backwards_compatible());
4154        if let Statistics::Float(stats) = stats {
4155            assert_eq!(stats.min_opt().unwrap(), &0.0);
4156            assert!(stats.min_opt().unwrap().is_sign_positive());
4157            assert_eq!(stats.max_opt().unwrap(), &2.0);
4158        } else {
4159            panic!("expecting Statistics::Float");
4160        }
4161    }
4162
4163    #[test]
4164    fn test_float_statistics_neg_zero_max() {
4165        let stats = statistics_roundtrip::<FloatType>(&[-0.0, -1.0, f32::NAN, -2.0]);
4166        assert!(!stats.is_min_max_backwards_compatible());
4167        if let Statistics::Float(stats) = stats {
4168            assert_eq!(stats.min_opt().unwrap(), &-2.0);
4169            assert_eq!(stats.max_opt().unwrap(), &-0.0);
4170            assert!(stats.max_opt().unwrap().is_sign_negative());
4171        } else {
4172            panic!("expecting Statistics::Float");
4173        }
4174    }
4175
4176    #[test]
4177    fn test_double_statistics_nan_middle() {
4178        let stats = statistics_roundtrip::<DoubleType>(&[1.0, f64::NAN, 2.0]);
4179        assert!(!stats.is_min_max_backwards_compatible());
4180        if let Statistics::Double(stats) = stats {
4181            assert_eq!(stats.min_opt().unwrap(), &1.0);
4182            assert_eq!(stats.max_opt().unwrap(), &2.0);
4183            assert_eq!(stats.nan_count_opt(), Some(1))
4184        } else {
4185            panic!("expecting Statistics::Double");
4186        }
4187    }
4188
4189    #[test]
4190    fn test_double_statistics_nan_start() {
4191        let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, 1.0, 2.0]);
4192        assert!(!stats.is_min_max_backwards_compatible());
4193        if let Statistics::Double(stats) = stats {
4194            assert_eq!(stats.min_opt().unwrap(), &1.0);
4195            assert_eq!(stats.max_opt().unwrap(), &2.0);
4196            assert_eq!(stats.nan_count_opt(), Some(1))
4197        } else {
4198            panic!("expecting Statistics::Double");
4199        }
4200    }
4201
4202    #[test]
4203    fn test_double_statistics_nan_only() {
4204        let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, f64::NAN]);
4205        assert_eq!(stats.min_bytes_opt(), Some(f64::NAN.as_bytes()));
4206        assert_eq!(stats.max_bytes_opt(), Some(f64::NAN.as_bytes()));
4207        assert_eq!(stats.nan_count_opt(), Some(2));
4208        assert!(matches!(stats, Statistics::Double(_)));
4209        assert!(!stats.is_min_max_backwards_compatible());
4210    }
4211
4212    #[test]
4213    fn test_double_statistics_zero_only() {
4214        let stats = statistics_roundtrip::<DoubleType>(&[0.0]);
4215        assert!(!stats.is_min_max_backwards_compatible());
4216        if let Statistics::Double(stats) = stats {
4217            assert_eq!(stats.min_opt().unwrap(), &0.0);
4218            assert!(stats.min_opt().unwrap().is_sign_positive());
4219            assert_eq!(stats.max_opt().unwrap(), &0.0);
4220            assert!(stats.max_opt().unwrap().is_sign_positive());
4221        } else {
4222            panic!("expecting Statistics::Double");
4223        }
4224    }
4225
4226    #[test]
4227    fn test_double_statistics_neg_zero_only() {
4228        let stats = statistics_roundtrip::<DoubleType>(&[-0.0]);
4229        assert!(!stats.is_min_max_backwards_compatible());
4230        if let Statistics::Double(stats) = stats {
4231            assert_eq!(stats.min_opt().unwrap(), &-0.0);
4232            assert!(stats.min_opt().unwrap().is_sign_negative());
4233            assert_eq!(stats.max_opt().unwrap(), &-0.0);
4234            assert!(stats.max_opt().unwrap().is_sign_negative());
4235        } else {
4236            panic!("expecting Statistics::Double");
4237        }
4238    }
4239
4240    #[test]
4241    fn test_double_statistics_zero_min() {
4242        let stats = statistics_roundtrip::<DoubleType>(&[0.0, 1.0, f64::NAN, 2.0]);
4243        assert!(!stats.is_min_max_backwards_compatible());
4244        if let Statistics::Double(stats) = stats {
4245            assert_eq!(stats.min_opt().unwrap(), &0.0);
4246            assert!(stats.min_opt().unwrap().is_sign_positive());
4247            assert_eq!(stats.max_opt().unwrap(), &2.0);
4248        } else {
4249            panic!("expecting Statistics::Double");
4250        }
4251    }
4252
4253    #[test]
4254    fn test_double_statistics_neg_zero_max() {
4255        let stats = statistics_roundtrip::<DoubleType>(&[-0.0, -1.0, f64::NAN, -2.0]);
4256        assert!(!stats.is_min_max_backwards_compatible());
4257        if let Statistics::Double(stats) = stats {
4258            assert_eq!(stats.min_opt().unwrap(), &-2.0);
4259            assert_eq!(stats.max_opt().unwrap(), &-0.0);
4260            assert!(stats.max_opt().unwrap().is_sign_negative());
4261        } else {
4262            panic!("expecting Statistics::Double");
4263        }
4264    }
4265
4266    #[test]
4267    fn test_compare_greater_byte_array_decimals() {
4268        assert!(!compare_greater_byte_array_decimals(&[], &[],),);
4269        assert!(compare_greater_byte_array_decimals(&[1u8,], &[],),);
4270        assert!(!compare_greater_byte_array_decimals(&[], &[1u8,],),);
4271        assert!(compare_greater_byte_array_decimals(&[1u8,], &[0u8,],),);
4272        assert!(!compare_greater_byte_array_decimals(&[1u8,], &[1u8,],),);
4273        assert!(compare_greater_byte_array_decimals(&[1u8, 0u8,], &[0u8,],),);
4274        assert!(!compare_greater_byte_array_decimals(
4275            &[0u8, 1u8,],
4276            &[1u8, 0u8,],
4277        ),);
4278        assert!(!compare_greater_byte_array_decimals(
4279            &[255u8, 35u8, 0u8, 0u8,],
4280            &[0u8,],
4281        ),);
4282        assert!(compare_greater_byte_array_decimals(
4283            &[0u8,],
4284            &[255u8, 35u8, 0u8, 0u8,],
4285        ),);
4286
4287        // Unequal lengths where the longer value's extra leading bytes are all
4288        // sign extension, so the aligned tails decide.
4289        // https://github.com/apache/arrow-rs/issues/10860
4290
4291        // 32768 > 255
4292        assert!(compare_greater_byte_array_decimals(
4293            &[0u8, 128u8, 0u8,],
4294            &[0u8, 255u8,],
4295        ),);
4296        assert!(!compare_greater_byte_array_decimals(
4297            &[0u8, 255u8,],
4298            &[0u8, 128u8, 0u8,],
4299        ),);
4300        // -128 > -129
4301        assert!(compare_greater_byte_array_decimals(
4302            &[128u8,],
4303            &[255u8, 127u8,],
4304        ),);
4305        assert!(!compare_greater_byte_array_decimals(
4306            &[255u8, 127u8,],
4307            &[128u8,],
4308        ),);
4309        // -128 > -256
4310        assert!(compare_greater_byte_array_decimals(
4311            &[128u8,],
4312            &[255u8, 0u8,],
4313        ),);
4314        // 10 (with a redundant leading zero) > 5
4315        assert!(compare_greater_byte_array_decimals(&[0u8, 10u8,], &[5u8,],),);
4316        assert!(compare_greater_byte_array_decimals(&[10u8,], &[0u8, 5u8,],),);
4317        // equal values of different lengths are not greater in either direction
4318        assert!(!compare_greater_byte_array_decimals(
4319            &[255u8, 128u8,],
4320            &[128u8,],
4321        ),);
4322        assert!(!compare_greater_byte_array_decimals(
4323            &[128u8,],
4324            &[255u8, 128u8,],
4325        ),);
4326    }
4327
4328    #[test]
4329    fn test_column_index_with_null_pages() {
4330        // write a single page of all nulls
4331        let page_writer = get_test_page_writer();
4332        let props = Default::default();
4333        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
4334        writer.write_batch(&[], Some(&[0, 0, 0, 0]), None).unwrap();
4335
4336        let r = writer.close().unwrap();
4337        assert!(r.column_index.is_some());
4338        let col_idx = r.column_index.unwrap();
4339        let ColumnIndexMetaData::INT32(col_idx) = col_idx else {
4340            panic!("wrong stats type")
4341        };
4342        // null_pages should be true for page 0
4343        assert!(col_idx.is_null_page(0));
4344        // min and max should be empty byte arrays
4345        assert!(col_idx.min_value(0).is_none());
4346        assert!(col_idx.max_value(0).is_none());
4347        // null_counts should be defined and be 4 for page 0
4348        assert!(col_idx.null_count(0).is_some());
4349        assert_eq!(col_idx.null_count(0), Some(4));
4350        // there is no repetition so rep histogram should be absent
4351        assert!(col_idx.repetition_level_histogram(0).is_none());
4352        // definition_level_histogram should be present and should be 0:4, 1:0
4353        assert!(col_idx.definition_level_histogram(0).is_some());
4354        assert_eq!(col_idx.definition_level_histogram(0).unwrap(), &[4, 0]);
4355    }
4356
4357    #[test]
4358    fn test_column_offset_index_metadata() {
4359        // write data
4360        // and check the offset index and column index
4361        let page_writer = get_test_page_writer();
4362        let props = Default::default();
4363        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
4364        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
4365        // first page
4366        writer.flush_data_pages().unwrap();
4367        // second page
4368        writer.write_batch(&[4, 8, 2, -5], None, None).unwrap();
4369
4370        let r = writer.close().unwrap();
4371        let column_index = r.column_index.unwrap();
4372        let offset_index = r.offset_index.unwrap();
4373
4374        assert_eq!(8, r.rows_written);
4375
4376        // column index
4377        let ColumnIndexMetaData::INT32(column_index) = column_index else {
4378            panic!("wrong stats type")
4379        };
4380        assert_eq!(2, column_index.num_pages());
4381        assert_eq!(2, offset_index.page_locations.len());
4382        assert_eq!(BoundaryOrder::UNORDERED, column_index.boundary_order);
4383        for idx in 0..2 {
4384            assert!(!column_index.is_null_page(idx));
4385            assert_eq!(0, column_index.null_count(0).unwrap());
4386        }
4387
4388        if let Some(stats) = r.metadata.statistics() {
4389            assert_eq!(stats.null_count_opt(), Some(0));
4390            assert_eq!(stats.distinct_count_opt(), None);
4391            if let Statistics::Int32(stats) = stats {
4392                // first page is [1,2,3,4]
4393                // second page is [-5,2,4,8]
4394                // note that we don't increment here, as this is a non BinaryArray type.
4395                assert_eq!(stats.min_opt(), column_index.min_value(1));
4396                assert_eq!(stats.max_opt(), column_index.max_value(1));
4397            } else {
4398                panic!("expecting Statistics::Int32");
4399            }
4400        } else {
4401            panic!("metadata missing statistics");
4402        }
4403
4404        // page location
4405        assert_eq!(0, offset_index.page_locations[0].first_row_index);
4406        assert_eq!(4, offset_index.page_locations[1].first_row_index);
4407    }
4408
4409    /// Verify min/max value truncation in the column index works as expected
4410    #[test]
4411    fn test_column_offset_index_metadata_truncating() {
4412        // write data
4413        // and check the offset index and column index
4414        let page_writer = get_test_page_writer();
4415        let props = WriterProperties::builder()
4416            .set_statistics_truncate_length(None) // disable column index truncation
4417            .build()
4418            .into();
4419        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4420
4421        let mut data = vec![FixedLenByteArray::default(); 3];
4422        // This is the expected min value - "aaa..."
4423        data[0].set_data(Bytes::from(vec![97_u8; 200]));
4424        // This is the expected max value - "ZZZ..."
4425        data[1].set_data(Bytes::from(vec![112_u8; 200]));
4426        data[2].set_data(Bytes::from(vec![98_u8; 200]));
4427
4428        writer.write_batch(&data, None, None).unwrap();
4429
4430        writer.flush_data_pages().unwrap();
4431
4432        let r = writer.close().unwrap();
4433        let column_index = r.column_index.unwrap();
4434        let offset_index = r.offset_index.unwrap();
4435
4436        let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4437            panic!("wrong stats type")
4438        };
4439
4440        assert_eq!(3, r.rows_written);
4441
4442        // column index
4443        assert_eq!(1, column_index.num_pages());
4444        assert_eq!(1, offset_index.page_locations.len());
4445        assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
4446        assert!(!column_index.is_null_page(0));
4447        assert_eq!(Some(0), column_index.null_count(0));
4448
4449        if let Some(stats) = r.metadata.statistics() {
4450            assert_eq!(stats.null_count_opt(), Some(0));
4451            assert_eq!(stats.distinct_count_opt(), None);
4452            if let Statistics::FixedLenByteArray(stats) = stats {
4453                let column_index_min_value = column_index.min_value(0).unwrap();
4454                let column_index_max_value = column_index.max_value(0).unwrap();
4455
4456                // Column index stats are truncated, while the column chunk's aren't.
4457                assert_ne!(stats.min_bytes_opt().unwrap(), column_index_min_value);
4458                assert_ne!(stats.max_bytes_opt().unwrap(), column_index_max_value);
4459
4460                assert_eq!(
4461                    column_index_min_value.len(),
4462                    DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
4463                );
4464                assert_eq!(column_index_min_value, &[97_u8; 64]);
4465                assert_eq!(
4466                    column_index_max_value.len(),
4467                    DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
4468                );
4469
4470                // We expect the last byte to be incremented
4471                assert_eq!(
4472                    *column_index_max_value.last().unwrap(),
4473                    *column_index_max_value.first().unwrap() + 1
4474                );
4475            } else {
4476                panic!("expecting Statistics::FixedLenByteArray");
4477            }
4478        } else {
4479            panic!("metadata missing statistics");
4480        }
4481    }
4482
4483    #[test]
4484    fn test_column_offset_index_truncating_spec_example() {
4485        // write data
4486        // and check the offset index and column index
4487        let page_writer = get_test_page_writer();
4488
4489        // Truncate values at 1 byte
4490        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4491        let props = Arc::new(builder.build());
4492        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4493
4494        let mut data = vec![FixedLenByteArray::default(); 1];
4495        // This is the expected min value
4496        data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
4497
4498        writer.write_batch(&data, None, None).unwrap();
4499
4500        writer.flush_data_pages().unwrap();
4501
4502        let r = writer.close().unwrap();
4503        let column_index = r.column_index.unwrap();
4504        let offset_index = r.offset_index.unwrap();
4505
4506        let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4507            panic!("wrong stats type")
4508        };
4509
4510        assert_eq!(1, r.rows_written);
4511
4512        // column index
4513        assert_eq!(1, column_index.num_pages());
4514        assert_eq!(1, offset_index.page_locations.len());
4515        assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
4516        assert!(!column_index.is_null_page(0));
4517        assert_eq!(Some(0), column_index.null_count(0));
4518
4519        if let Some(stats) = r.metadata.statistics() {
4520            assert_eq!(stats.null_count_opt(), Some(0));
4521            assert_eq!(stats.distinct_count_opt(), None);
4522            if let Statistics::FixedLenByteArray(_stats) = stats {
4523                let column_index_min_value = column_index.min_value(0).unwrap();
4524                let column_index_max_value = column_index.max_value(0).unwrap();
4525
4526                assert_eq!(column_index_min_value.len(), 1);
4527                assert_eq!(column_index_max_value.len(), 1);
4528
4529                assert_eq!(b"B", column_index_min_value);
4530                assert_eq!(b"C", column_index_max_value);
4531
4532                assert_ne!(column_index_min_value, stats.min_bytes_opt().unwrap());
4533                assert_ne!(column_index_max_value, stats.max_bytes_opt().unwrap());
4534            } else {
4535                panic!("expecting Statistics::FixedLenByteArray");
4536            }
4537        } else {
4538            panic!("metadata missing statistics");
4539        }
4540    }
4541
4542    #[test]
4543    fn test_float16_min_max_no_truncation() {
4544        // Even if we set truncation to occur at 1 byte, we should not truncate for Float16
4545        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4546        let props = Arc::new(builder.build());
4547        let page_writer = get_test_page_writer();
4548        let mut writer = get_test_float16_column_writer(page_writer, props);
4549
4550        let expected_value = f16::PI.to_le_bytes().to_vec();
4551        let data = vec![ByteArray::from(expected_value.clone()).into()];
4552        writer.write_batch(&data, None, None).unwrap();
4553        writer.flush_data_pages().unwrap();
4554
4555        let r = writer.close().unwrap();
4556
4557        // stats should still be written
4558        // ensure bytes weren't truncated for column index
4559        let column_index = r.column_index.unwrap();
4560        let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4561            panic!("wrong stats type")
4562        };
4563        let column_index_min_bytes = column_index.min_value(0).unwrap();
4564        let column_index_max_bytes = column_index.max_value(0).unwrap();
4565        assert_eq!(expected_value, column_index_min_bytes);
4566        assert_eq!(expected_value, column_index_max_bytes);
4567
4568        // ensure bytes weren't truncated for statistics
4569        let stats = r.metadata.statistics().unwrap();
4570        if let Statistics::FixedLenByteArray(stats) = stats {
4571            let stats_min_bytes = stats.min_bytes_opt().unwrap();
4572            let stats_max_bytes = stats.max_bytes_opt().unwrap();
4573            assert_eq!(expected_value, stats_min_bytes);
4574            assert_eq!(expected_value, stats_max_bytes);
4575        } else {
4576            panic!("expecting Statistics::FixedLenByteArray");
4577        }
4578    }
4579
4580    #[test]
4581    fn test_decimal_min_max_no_truncation() {
4582        // Even if we set truncation to occur at 1 byte, we should not truncate for Decimal
4583        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4584        let props = Arc::new(builder.build());
4585        let page_writer = get_test_page_writer();
4586        let mut writer =
4587            get_test_decimals_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4588
4589        let expected_value = vec![
4590            255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, 35u8,
4591            231u8, 90u8, 0u8, 0u8,
4592        ];
4593        let data = vec![ByteArray::from(expected_value.clone()).into()];
4594        writer.write_batch(&data, None, None).unwrap();
4595        writer.flush_data_pages().unwrap();
4596
4597        let r = writer.close().unwrap();
4598
4599        // stats should still be written
4600        // ensure bytes weren't truncated for column index
4601        let column_index = r.column_index.unwrap();
4602        let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4603            panic!("wrong stats type")
4604        };
4605        let column_index_min_bytes = column_index.min_value(0).unwrap();
4606        let column_index_max_bytes = column_index.max_value(0).unwrap();
4607        assert_eq!(expected_value, column_index_min_bytes);
4608        assert_eq!(expected_value, column_index_max_bytes);
4609
4610        // ensure bytes weren't truncated for statistics
4611        let stats = r.metadata.statistics().unwrap();
4612        if let Statistics::FixedLenByteArray(stats) = stats {
4613            let stats_min_bytes = stats.min_bytes_opt().unwrap();
4614            let stats_max_bytes = stats.max_bytes_opt().unwrap();
4615            assert_eq!(expected_value, stats_min_bytes);
4616            assert_eq!(expected_value, stats_max_bytes);
4617        } else {
4618            panic!("expecting Statistics::FixedLenByteArray");
4619        }
4620    }
4621
4622    #[test]
4623    fn test_decimal_byte_array_min_max_no_statistics_truncation() {
4624        // See https://github.com/apache/arrow-rs/issues/11073
4625        let page_writer = get_test_page_writer();
4626
4627        // Truncate at 1 byte -- far shorter than either encoded value below.
4628        let builder = WriterProperties::builder().set_statistics_truncate_length(Some(1));
4629        let props = Arc::new(builder.build());
4630        let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
4631
4632        // Two's-complement, big-endian encodings of 200 and -200 (2 bytes each).
4633        // Naive unsigned byte-wise truncation to 1 byte followed by an unsigned
4634        // increment of the max value would corrupt these:
4635        //   max: truncate([0x00, 0xC8], 1) -> [0x00], increment -> [0x01],
4636        //        which decodes to 1 -- far less than the true max of 200.
4637        let positive_200 = vec![0x00u8, 0xC8u8];
4638        let negative_200 = vec![0xFFu8, 0x38u8];
4639
4640        let data = vec![
4641            ByteArray::from(positive_200.clone()),
4642            ByteArray::from(negative_200.clone()),
4643        ];
4644        writer.write_batch(&data, None, None).unwrap();
4645        writer.flush_data_pages().unwrap();
4646
4647        let r = writer.close().unwrap();
4648        assert_eq!(2, r.rows_written);
4649
4650        let stats = r.metadata.statistics().expect("statistics");
4651        if let Statistics::ByteArray(stats) = stats {
4652            // Decimal columns must never be truncated, regardless of
4653            // `statistics_truncate_length`.
4654            assert!(stats.min_is_exact());
4655            assert!(stats.max_is_exact());
4656
4657            let min_value = stats.min_bytes_opt().unwrap();
4658            let max_value = stats.max_bytes_opt().unwrap();
4659
4660            // Correct signed (two's-complement) ordering: -200 < 200.
4661            assert_eq!(negative_200, min_value);
4662            assert_eq!(positive_200, max_value);
4663        } else {
4664            panic!("expecting Statistics::ByteArray");
4665        }
4666    }
4667
4668    #[test]
4669    fn test_statistics_truncating_byte_array_default() {
4670        let page_writer = get_test_page_writer();
4671
4672        // The default truncate length is 64 bytes
4673        let props = WriterProperties::builder().build().into();
4674        let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
4675
4676        let mut data = vec![ByteArray::default(); 1];
4677        data[0].set_data(Bytes::from(String::from(
4678            "This string is longer than 64 bytes, so it will almost certainly be truncated.",
4679        )));
4680        writer.write_batch(&data, None, None).unwrap();
4681        writer.flush_data_pages().unwrap();
4682
4683        let r = writer.close().unwrap();
4684
4685        assert_eq!(1, r.rows_written);
4686
4687        let stats = r.metadata.statistics().expect("statistics");
4688        if let Statistics::ByteArray(stats) = stats {
4689            let min_value = stats.min_opt().unwrap();
4690            let max_value = stats.max_opt().unwrap();
4691
4692            assert!(!stats.min_is_exact());
4693            assert!(!stats.max_is_exact());
4694
4695            let expected_len = 64;
4696            assert_eq!(min_value.len(), expected_len);
4697            assert_eq!(max_value.len(), expected_len);
4698
4699            let expected_min =
4700                "This string is longer than 64 bytes, so it will almost certainly".as_bytes();
4701            assert_eq!(expected_min, min_value.as_bytes());
4702            // note the max value is different from the min value: the last byte is incremented
4703            let expected_max =
4704                "This string is longer than 64 bytes, so it will almost certainlz".as_bytes();
4705            assert_eq!(expected_max, max_value.as_bytes());
4706        } else {
4707            panic!("expecting Statistics::ByteArray");
4708        }
4709    }
4710
4711    #[test]
4712    fn test_statistics_truncating_byte_array() {
4713        let page_writer = get_test_page_writer();
4714
4715        const TEST_TRUNCATE_LENGTH: usize = 1;
4716
4717        // Truncate values at 1 byte
4718        let builder =
4719            WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
4720        let props = Arc::new(builder.build());
4721        let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
4722
4723        let mut data = vec![ByteArray::default(); 1];
4724        // This is the expected min value
4725        data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
4726
4727        writer.write_batch(&data, None, None).unwrap();
4728
4729        writer.flush_data_pages().unwrap();
4730
4731        let r = writer.close().unwrap();
4732
4733        assert_eq!(1, r.rows_written);
4734
4735        let stats = r.metadata.statistics().expect("statistics");
4736        assert_eq!(stats.null_count_opt(), Some(0));
4737        assert_eq!(stats.distinct_count_opt(), None);
4738        if let Statistics::ByteArray(stats) = stats {
4739            let min_value = stats.min_opt().unwrap();
4740            let max_value = stats.max_opt().unwrap();
4741
4742            assert!(!stats.min_is_exact());
4743            assert!(!stats.max_is_exact());
4744
4745            assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
4746            assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
4747
4748            assert_eq!(b"B", min_value.as_bytes());
4749            assert_eq!(b"C", max_value.as_bytes());
4750        } else {
4751            panic!("expecting Statistics::ByteArray");
4752        }
4753    }
4754
4755    #[test]
4756    fn test_statistics_truncating_fixed_len_byte_array() {
4757        let page_writer = get_test_page_writer();
4758
4759        const TEST_TRUNCATE_LENGTH: usize = 1;
4760
4761        // Truncate values at 1 byte
4762        let builder =
4763            WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
4764        let props = Arc::new(builder.build());
4765        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4766
4767        let mut data = vec![FixedLenByteArray::default(); 1];
4768
4769        const PSEUDO_DECIMAL_VALUE: i128 = 6541894651216648486512564456564654;
4770        const PSEUDO_DECIMAL_BYTES: [u8; 16] = PSEUDO_DECIMAL_VALUE.to_be_bytes();
4771
4772        const EXPECTED_MIN: [u8; TEST_TRUNCATE_LENGTH] = [PSEUDO_DECIMAL_BYTES[0]]; // parquet specifies big-endian order for decimals
4773        const EXPECTED_MAX: [u8; TEST_TRUNCATE_LENGTH] =
4774            [PSEUDO_DECIMAL_BYTES[0].overflowing_add(1).0];
4775
4776        // This is the expected min value
4777        data[0].set_data(Bytes::from(PSEUDO_DECIMAL_BYTES.as_slice()));
4778
4779        writer.write_batch(&data, None, None).unwrap();
4780
4781        writer.flush_data_pages().unwrap();
4782
4783        let r = writer.close().unwrap();
4784
4785        assert_eq!(1, r.rows_written);
4786
4787        let stats = r.metadata.statistics().expect("statistics");
4788        assert_eq!(stats.null_count_opt(), Some(0));
4789        assert_eq!(stats.distinct_count_opt(), None);
4790        if let Statistics::FixedLenByteArray(stats) = stats {
4791            let min_value = stats.min_opt().unwrap();
4792            let max_value = stats.max_opt().unwrap();
4793
4794            assert!(!stats.min_is_exact());
4795            assert!(!stats.max_is_exact());
4796
4797            assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
4798            assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
4799
4800            assert_eq!(EXPECTED_MIN.as_slice(), min_value.as_bytes());
4801            assert_eq!(EXPECTED_MAX.as_slice(), max_value.as_bytes());
4802
4803            let reconstructed_min = i128::from_be_bytes([
4804                min_value.as_bytes()[0],
4805                0,
4806                0,
4807                0,
4808                0,
4809                0,
4810                0,
4811                0,
4812                0,
4813                0,
4814                0,
4815                0,
4816                0,
4817                0,
4818                0,
4819                0,
4820            ]);
4821
4822            let reconstructed_max = i128::from_be_bytes([
4823                max_value.as_bytes()[0],
4824                0,
4825                0,
4826                0,
4827                0,
4828                0,
4829                0,
4830                0,
4831                0,
4832                0,
4833                0,
4834                0,
4835                0,
4836                0,
4837                0,
4838                0,
4839            ]);
4840
4841            // check that the inner value is correctly bounded by the min/max
4842            println!("min: {reconstructed_min} {PSEUDO_DECIMAL_VALUE}");
4843            assert!(reconstructed_min <= PSEUDO_DECIMAL_VALUE);
4844            println!("max {reconstructed_max} {PSEUDO_DECIMAL_VALUE}");
4845            assert!(reconstructed_max >= PSEUDO_DECIMAL_VALUE);
4846        } else {
4847            panic!("expecting Statistics::FixedLenByteArray");
4848        }
4849    }
4850
4851    #[test]
4852    fn test_send() {
4853        fn test<T: Send>() {}
4854        test::<ColumnWriterImpl<Int32Type>>();
4855    }
4856
4857    #[test]
4858    fn test_increment() {
4859        let v = increment(vec![0, 0, 0]).unwrap();
4860        assert_eq!(&v, &[0, 0, 1]);
4861
4862        // Handle overflow
4863        let v = increment(vec![0, 255, 255]).unwrap();
4864        assert_eq!(&v, &[1, 0, 0]);
4865
4866        // Return `None` if all bytes are u8::MAX
4867        let v = increment(vec![255, 255, 255]);
4868        assert!(v.is_none());
4869    }
4870
4871    #[test]
4872    fn test_increment_utf8() {
4873        let test_inc = |o: &str, expected: &str| {
4874            if let Ok(v) = String::from_utf8(increment_utf8(o).unwrap()) {
4875                // Got the expected result...
4876                assert_eq!(v, expected);
4877                // and it's greater than the original string
4878                assert!(*v > *o);
4879                // Also show that BinaryArray level comparison works here
4880                let mut greater = ByteArray::new();
4881                greater.set_data(Bytes::from(v));
4882                let mut original = ByteArray::new();
4883                original.set_data(Bytes::from(o.as_bytes().to_vec()));
4884                assert!(greater > original);
4885            } else {
4886                panic!("Expected incremented UTF8 string to also be valid.");
4887            }
4888        };
4889
4890        // Basic ASCII case
4891        test_inc("hello", "hellp");
4892
4893        // 1-byte ending in max 1-byte
4894        test_inc("a\u{7f}", "b");
4895
4896        // 1-byte max should not truncate as it would need 2-byte code points
4897        assert!(increment_utf8("\u{7f}\u{7f}").is_none());
4898
4899        // UTF8 string
4900        test_inc("❤️🧡💛💚💙💜", "❤️🧡💛💚💙💝");
4901
4902        // 2-byte without overflow
4903        test_inc("éééé", "éééê");
4904
4905        // 2-byte that overflows lowest byte
4906        test_inc("\u{ff}\u{ff}", "\u{ff}\u{100}");
4907
4908        // 2-byte ending in max 2-byte
4909        test_inc("a\u{7ff}", "b");
4910
4911        // Max 2-byte should not truncate as it would need 3-byte code points
4912        assert!(increment_utf8("\u{7ff}\u{7ff}").is_none());
4913
4914        // 3-byte without overflow [U+800, U+800] -> [U+800, U+801] (note that these
4915        // characters should render right to left).
4916        test_inc("ࠀࠀ", "ࠀࠁ");
4917
4918        // 3-byte ending in max 3-byte
4919        test_inc("a\u{ffff}", "b");
4920
4921        // Max 3-byte should not truncate as it would need 4-byte code points
4922        assert!(increment_utf8("\u{ffff}\u{ffff}").is_none());
4923
4924        // 4-byte without overflow
4925        test_inc("𐀀𐀀", "𐀀𐀁");
4926
4927        // 4-byte ending in max unicode
4928        test_inc("a\u{10ffff}", "b");
4929
4930        // Max 4-byte should not truncate
4931        assert!(increment_utf8("\u{10ffff}\u{10ffff}").is_none());
4932
4933        // Skip over surrogate pair range (0xD800..=0xDFFF)
4934        //test_inc("a\u{D7FF}", "a\u{e000}");
4935        test_inc("a\u{D7FF}", "b");
4936    }
4937
4938    #[test]
4939    fn test_truncate_utf8() {
4940        // No-op
4941        let data = "❤️🧡💛💚💙💜";
4942        let r = truncate_utf8(data, data.len()).unwrap();
4943        assert_eq!(r.len(), data.len());
4944        assert_eq!(&r, data.as_bytes());
4945
4946        // We slice it away from the UTF8 boundary
4947        let r = truncate_utf8(data, 13).unwrap();
4948        assert_eq!(r.len(), 10);
4949        assert_eq!(&r, "❤️🧡".as_bytes());
4950
4951        // One multi-byte code point, and a length shorter than it, so we can't slice it
4952        let r = truncate_utf8("\u{0836}", 1);
4953        assert!(r.is_none());
4954
4955        // Test truncate and increment for max bounds on UTF-8 statistics
4956        // 7-bit (i.e. ASCII)
4957        let r = truncate_and_increment_utf8("yyyyyyyyy", 8).unwrap();
4958        assert_eq!(&r, b"yyyyyyyz");
4959
4960        // 2-byte without overflow
4961        let r = truncate_and_increment_utf8("ééééé", 7).unwrap();
4962        assert_eq!(&r, "ééê".as_bytes());
4963
4964        // 2-byte that overflows lowest byte
4965        let r = truncate_and_increment_utf8("\u{ff}\u{ff}\u{ff}\u{ff}\u{ff}", 8).unwrap();
4966        assert_eq!(&r, "\u{ff}\u{ff}\u{ff}\u{100}".as_bytes());
4967
4968        // max 2-byte should not truncate as it would need 3-byte code points
4969        let r = truncate_and_increment_utf8("߿߿߿߿߿", 8);
4970        assert!(r.is_none());
4971
4972        // 3-byte without overflow [U+800, U+800, U+800] -> [U+800, U+801] (note that these
4973        // characters should render right to left).
4974        let r = truncate_and_increment_utf8("ࠀࠀࠀࠀ", 8).unwrap();
4975        assert_eq!(&r, "ࠀࠁ".as_bytes());
4976
4977        // max 3-byte should not truncate as it would need 4-byte code points
4978        let r = truncate_and_increment_utf8("\u{ffff}\u{ffff}\u{ffff}", 8);
4979        assert!(r.is_none());
4980
4981        // 4-byte without overflow
4982        let r = truncate_and_increment_utf8("𐀀𐀀𐀀𐀀", 9).unwrap();
4983        assert_eq!(&r, "𐀀𐀁".as_bytes());
4984
4985        // max 4-byte should not truncate
4986        let r = truncate_and_increment_utf8("\u{10ffff}\u{10ffff}", 8);
4987        assert!(r.is_none());
4988    }
4989
4990    #[test]
4991    // Check fallback truncation of statistics that should be UTF-8, but aren't
4992    // (see https://github.com/apache/arrow-rs/pull/6870).
4993    fn test_byte_array_truncate_invalid_utf8_statistics() {
4994        let message_type = "
4995            message test_schema {
4996                OPTIONAL BYTE_ARRAY a (UTF8);
4997            }
4998        ";
4999        let schema = Arc::new(parse_message_type(message_type).unwrap());
5000
5001        // Create Vec<ByteArray> containing non-UTF8 bytes
5002        let data = vec![ByteArray::from(vec![128u8; 32]); 7];
5003        let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1];
5004        let file: File = tempfile::tempfile().unwrap();
5005        let props = Arc::new(
5006            WriterProperties::builder()
5007                .set_statistics_enabled(EnabledStatistics::Chunk)
5008                .set_statistics_truncate_length(Some(8))
5009                .build(),
5010        );
5011
5012        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
5013        let mut row_group_writer = writer.next_row_group().unwrap();
5014
5015        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
5016        col_writer
5017            .typed::<ByteArrayType>()
5018            .write_batch(&data, Some(&def_levels), None)
5019            .unwrap();
5020        col_writer.close().unwrap();
5021        row_group_writer.close().unwrap();
5022        let file_metadata = writer.close().unwrap();
5023        let stats = file_metadata.row_group(0).column(0).statistics().unwrap();
5024        assert!(!stats.max_is_exact());
5025        // Truncation of invalid UTF-8 should fall back to binary truncation, so last byte should
5026        // be incremented by 1.
5027        assert_eq!(
5028            stats.max_bytes_opt().map(|v| v.to_vec()),
5029            Some([128, 128, 128, 128, 128, 128, 128, 129].to_vec())
5030        );
5031    }
5032
5033    #[test]
5034    fn test_increment_max_binary_chars() {
5035        let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]);
5036        assert_eq!(&r.unwrap(), &[0xFF, 0xFE, 0xFE, 0x00, 0x00]);
5037
5038        let incremented = increment(vec![0xFF, 0xFF, 0xFF]);
5039        assert!(incremented.is_none())
5040    }
5041
5042    #[test]
5043    fn test_no_column_index_when_stats_disabled() {
5044        // https://github.com/apache/arrow-rs/issues/6010
5045        // Test that column index is not created/written for all-nulls column when page
5046        // statistics are disabled.
5047        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
5048        let props = Arc::new(
5049            WriterProperties::builder()
5050                .set_statistics_enabled(EnabledStatistics::None)
5051                .build(),
5052        );
5053        let column_writer = get_column_writer(descr, props, get_test_page_writer());
5054        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
5055
5056        let data = Vec::new();
5057        let def_levels = vec![0; 10];
5058        writer.write_batch(&data, Some(&def_levels), None).unwrap();
5059        writer.flush_data_pages().unwrap();
5060
5061        let column_close_result = writer.close().unwrap();
5062        assert!(column_close_result.offset_index.is_some());
5063        assert!(column_close_result.column_index.is_none());
5064    }
5065
5066    #[test]
5067    fn test_no_offset_index_when_disabled() {
5068        // Test that offset indexes can be disabled
5069        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
5070        let props = Arc::new(
5071            WriterProperties::builder()
5072                .set_statistics_enabled(EnabledStatistics::None)
5073                .set_offset_index_disabled(true)
5074                .build(),
5075        );
5076        let column_writer = get_column_writer(descr, props, get_test_page_writer());
5077        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
5078
5079        let data = Vec::new();
5080        let def_levels = vec![0; 10];
5081        writer.write_batch(&data, Some(&def_levels), None).unwrap();
5082        writer.flush_data_pages().unwrap();
5083
5084        let column_close_result = writer.close().unwrap();
5085        assert!(column_close_result.offset_index.is_none());
5086        assert!(column_close_result.column_index.is_none());
5087    }
5088
5089    #[test]
5090    fn test_offset_index_overridden() {
5091        // Test that offset indexes are not disabled when gathering page statistics
5092        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
5093        let props = Arc::new(
5094            WriterProperties::builder()
5095                .set_statistics_enabled(EnabledStatistics::Page)
5096                .set_offset_index_disabled(true)
5097                .build(),
5098        );
5099        let column_writer = get_column_writer(descr, props, get_test_page_writer());
5100        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
5101
5102        let data = Vec::new();
5103        let def_levels = vec![0; 10];
5104        writer.write_batch(&data, Some(&def_levels), None).unwrap();
5105        writer.flush_data_pages().unwrap();
5106
5107        let column_close_result = writer.close().unwrap();
5108        assert!(column_close_result.offset_index.is_some());
5109        assert!(column_close_result.column_index.is_some());
5110    }
5111
5112    #[test]
5113    fn test_boundary_order() -> Result<()> {
5114        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
5115        // min max both ascending
5116        let column_close_result = write_multiple_pages::<Int32Type>(
5117            &descr,
5118            &[
5119                &[Some(-10), Some(10)],
5120                &[Some(-5), Some(11)],
5121                &[None],
5122                &[Some(-5), Some(11)],
5123            ],
5124        )?;
5125        let boundary_order = column_close_result
5126            .column_index
5127            .unwrap()
5128            .get_boundary_order();
5129        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5130
5131        // min max both descending
5132        let column_close_result = write_multiple_pages::<Int32Type>(
5133            &descr,
5134            &[
5135                &[Some(10), Some(11)],
5136                &[Some(5), Some(11)],
5137                &[None],
5138                &[Some(-5), Some(0)],
5139            ],
5140        )?;
5141        let boundary_order = column_close_result
5142            .column_index
5143            .unwrap()
5144            .get_boundary_order();
5145        assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
5146
5147        // min max both equal
5148        let column_close_result = write_multiple_pages::<Int32Type>(
5149            &descr,
5150            &[&[Some(10), Some(11)], &[None], &[Some(10), Some(11)]],
5151        )?;
5152        let boundary_order = column_close_result
5153            .column_index
5154            .unwrap()
5155            .get_boundary_order();
5156        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5157
5158        // only nulls
5159        let column_close_result =
5160            write_multiple_pages::<Int32Type>(&descr, &[&[None], &[None], &[None]])?;
5161        let boundary_order = column_close_result
5162            .column_index
5163            .unwrap()
5164            .get_boundary_order();
5165        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5166
5167        // one page
5168        let column_close_result =
5169            write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)]])?;
5170        let boundary_order = column_close_result
5171            .column_index
5172            .unwrap()
5173            .get_boundary_order();
5174        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5175
5176        // one non-null page
5177        let column_close_result =
5178            write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)], &[None]])?;
5179        let boundary_order = column_close_result
5180            .column_index
5181            .unwrap()
5182            .get_boundary_order();
5183        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5184
5185        // min max both unordered
5186        let column_close_result = write_multiple_pages::<Int32Type>(
5187            &descr,
5188            &[
5189                &[Some(10), Some(11)],
5190                &[Some(11), Some(16)],
5191                &[None],
5192                &[Some(-5), Some(0)],
5193            ],
5194        )?;
5195        let boundary_order = column_close_result
5196            .column_index
5197            .unwrap()
5198            .get_boundary_order();
5199        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5200
5201        // min max both ordered in different orders
5202        let column_close_result = write_multiple_pages::<Int32Type>(
5203            &descr,
5204            &[
5205                &[Some(1), Some(9)],
5206                &[Some(2), Some(8)],
5207                &[None],
5208                &[Some(3), Some(7)],
5209            ],
5210        )?;
5211        let boundary_order = column_close_result
5212            .column_index
5213            .unwrap()
5214            .get_boundary_order();
5215        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5216
5217        Ok(())
5218    }
5219
5220    #[test]
5221    fn test_boundary_order_logical_type() -> Result<()> {
5222        // ensure that logical types account for different sort order than underlying
5223        // physical type representation
5224        let f16_descr = Arc::new(get_test_float16_column_descr(1, 0));
5225        let fba_descr = {
5226            let type_ = SchemaType::primitive_type_builder(
5227                "col",
5228                FixedLenByteArrayType::get_physical_type(),
5229            )
5230            .with_length(2)
5231            .build()?;
5232            Arc::new(ColumnDescriptor::new(
5233                Arc::new(type_),
5234                1,
5235                0,
5236                ColumnPath::from("col"),
5237            ))
5238        };
5239
5240        let values: &[&[Option<FixedLenByteArray>]] = &[
5241            &[Some(FixedLenByteArray::from(ByteArray::from(f16::ONE)))],
5242            &[Some(FixedLenByteArray::from(ByteArray::from(f16::ZERO)))],
5243            &[Some(FixedLenByteArray::from(ByteArray::from(
5244                f16::NEG_ZERO,
5245            )))],
5246            &[Some(FixedLenByteArray::from(ByteArray::from(f16::NEG_ONE)))],
5247        ];
5248
5249        // f16 descending
5250        let column_close_result =
5251            write_multiple_pages::<FixedLenByteArrayType>(&f16_descr, values)?;
5252        let boundary_order = column_close_result
5253            .column_index
5254            .unwrap()
5255            .get_boundary_order();
5256        assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
5257
5258        // same bytes, but fba unordered
5259        let column_close_result =
5260            write_multiple_pages::<FixedLenByteArrayType>(&fba_descr, values)?;
5261        let boundary_order = column_close_result
5262            .column_index
5263            .unwrap()
5264            .get_boundary_order();
5265        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5266
5267        Ok(())
5268    }
5269
5270    #[test]
5271    fn test_interval_stats_should_not_have_min_max() {
5272        let input = [
5273            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
5274            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
5275            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2],
5276        ]
5277        .into_iter()
5278        .map(|s| ByteArray::from(s).into())
5279        .collect::<Vec<_>>();
5280
5281        let page_writer = get_test_page_writer();
5282        let mut writer = get_test_interval_column_writer(page_writer);
5283        writer.write_batch(&input, None, None).unwrap();
5284
5285        let metadata = writer.close().unwrap().metadata;
5286        let stats = if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
5287            stats.clone()
5288        } else {
5289            panic!("metadata missing statistics");
5290        };
5291        assert!(stats.min_bytes_opt().is_none());
5292        assert!(stats.max_bytes_opt().is_none());
5293    }
5294
5295    #[test]
5296    #[cfg(feature = "arrow")]
5297    fn test_column_writer_get_estimated_total_bytes() {
5298        let page_writer = get_test_page_writer();
5299        let props = Default::default();
5300        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
5301        assert_eq!(writer.get_estimated_total_bytes(), 0);
5302
5303        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
5304        writer.add_data_page().unwrap();
5305        let size_with_one_page = writer.get_estimated_total_bytes();
5306        assert_eq!(size_with_one_page, 20);
5307
5308        writer.write_batch(&[5, 6, 7, 8], None, None).unwrap();
5309        writer.add_data_page().unwrap();
5310        let size_with_two_pages = writer.get_estimated_total_bytes();
5311        // different pages have different compressed lengths
5312        assert_eq!(size_with_two_pages, 20 + 21);
5313    }
5314
5315    fn write_multiple_pages<T: DataType>(
5316        column_descr: &Arc<ColumnDescriptor>,
5317        pages: &[&[Option<T::T>]],
5318    ) -> Result<ColumnCloseResult> {
5319        let column_writer = get_column_writer(
5320            column_descr.clone(),
5321            Default::default(),
5322            get_test_page_writer(),
5323        );
5324        let mut writer = get_typed_column_writer::<T>(column_writer);
5325
5326        for &page in pages {
5327            let values = page.iter().filter_map(Clone::clone).collect::<Vec<_>>();
5328            let def_levels = page
5329                .iter()
5330                .map(|maybe_value| i16::from(maybe_value.is_some()))
5331                .collect::<Vec<_>>();
5332            writer.write_batch(&values, Some(&def_levels), None)?;
5333            writer.flush_data_pages()?;
5334        }
5335
5336        writer.close()
5337    }
5338
5339    /// Performs write-read roundtrip with randomly generated values and levels.
5340    /// `max_size` is maximum number of values or levels (if `max_def_level` > 0) to write
5341    /// for a column.
5342    fn column_roundtrip_random<T: DataType>(
5343        props: WriterProperties,
5344        max_size: usize,
5345        min_value: T::T,
5346        max_value: T::T,
5347        max_def_level: i16,
5348        max_rep_level: i16,
5349    ) where
5350        T::T: PartialOrd + SampleUniform + Copy,
5351    {
5352        let mut num_values: usize = 0;
5353
5354        let mut buf: Vec<i16> = Vec::new();
5355        let def_levels = if max_def_level > 0 {
5356            random_numbers_range(max_size, 0, max_def_level + 1, &mut buf);
5357            for &dl in &buf[..] {
5358                if dl == max_def_level {
5359                    num_values += 1;
5360                }
5361            }
5362            Some(&buf[..])
5363        } else {
5364            num_values = max_size;
5365            None
5366        };
5367
5368        let mut buf: Vec<i16> = Vec::new();
5369        let rep_levels = if max_rep_level > 0 {
5370            random_numbers_range(max_size, 0, max_rep_level + 1, &mut buf);
5371            buf[0] = 0; // Must start on record boundary
5372            Some(&buf[..])
5373        } else {
5374            None
5375        };
5376
5377        let mut values: Vec<T::T> = Vec::new();
5378        random_numbers_range(num_values, min_value, max_value, &mut values);
5379
5380        column_roundtrip::<T>(props, &values[..], def_levels, rep_levels);
5381    }
5382
5383    /// Performs write-read roundtrip and asserts written values and levels.
5384    fn column_roundtrip<T: DataType>(
5385        props: WriterProperties,
5386        values: &[T::T],
5387        def_levels: Option<&[i16]>,
5388        rep_levels: Option<&[i16]>,
5389    ) {
5390        let mut file = tempfile::tempfile().unwrap();
5391        let mut write = TrackedWrite::new(&mut file);
5392        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5393
5394        let max_def_level = match def_levels {
5395            Some(buf) => *buf.iter().max().unwrap_or(&0i16),
5396            None => 0i16,
5397        };
5398
5399        let max_rep_level = match rep_levels {
5400            Some(buf) => *buf.iter().max().unwrap_or(&0i16),
5401            None => 0i16,
5402        };
5403
5404        let mut max_batch_size = values.len();
5405        if let Some(levels) = def_levels {
5406            max_batch_size = max_batch_size.max(levels.len());
5407        }
5408        if let Some(levels) = rep_levels {
5409            max_batch_size = max_batch_size.max(levels.len());
5410        }
5411
5412        let mut writer =
5413            get_test_column_writer::<T>(page_writer, max_def_level, max_rep_level, Arc::new(props));
5414
5415        let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap();
5416        assert_eq!(values_written, values.len());
5417        let result = writer.close().unwrap();
5418
5419        drop(write);
5420
5421        let props = ReaderProperties::builder()
5422            .set_backward_compatible_lz4(false)
5423            .build();
5424        let page_reader = Box::new(
5425            SerializedPageReader::new_with_properties(
5426                Arc::new(file),
5427                &result.metadata,
5428                result.rows_written as usize,
5429                None,
5430                Arc::new(props),
5431            )
5432            .unwrap(),
5433        );
5434        let mut reader = get_test_column_reader::<T>(page_reader, max_def_level, max_rep_level);
5435
5436        let mut actual_values = Vec::with_capacity(max_batch_size);
5437        let mut actual_def_levels = def_levels.map(|_| Vec::with_capacity(max_batch_size));
5438        let mut actual_rep_levels = rep_levels.map(|_| Vec::with_capacity(max_batch_size));
5439
5440        let (_, values_read, levels_read) = reader
5441            .read_records(
5442                max_batch_size,
5443                actual_def_levels.as_mut(),
5444                actual_rep_levels.as_mut(),
5445                &mut actual_values,
5446            )
5447            .unwrap();
5448
5449        // Assert values, definition and repetition levels.
5450
5451        assert_eq!(&actual_values[..values_read], values);
5452        match actual_def_levels {
5453            Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), def_levels),
5454            None => assert_eq!(None, def_levels),
5455        }
5456        match actual_rep_levels {
5457            Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), rep_levels),
5458            None => assert_eq!(None, rep_levels),
5459        }
5460
5461        // Assert written rows.
5462
5463        if let Some(levels) = actual_rep_levels {
5464            let mut actual_rows_written = 0;
5465            for l in levels {
5466                if l == 0 {
5467                    actual_rows_written += 1;
5468                }
5469            }
5470            assert_eq!(actual_rows_written, result.rows_written);
5471        } else if actual_def_levels.is_some() {
5472            assert_eq!(levels_read as u64, result.rows_written);
5473        } else {
5474            assert_eq!(values_read as u64, result.rows_written);
5475        }
5476    }
5477
5478    /// Performs write of provided values and returns column metadata of those values.
5479    /// Used to test encoding support for column writer.
5480    fn column_write_and_get_metadata<T: DataType>(
5481        props: WriterProperties,
5482        values: &[T::T],
5483    ) -> ColumnChunkMetaData {
5484        let page_writer = get_test_page_writer();
5485        let props = Arc::new(props);
5486        let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
5487        writer.write_batch(values, None, None).unwrap();
5488        writer.close().unwrap().metadata
5489    }
5490
5491    // Helper function to more compactly create a PageEncodingStats struct.
5492    fn encoding_stats(page_type: PageType, encoding: Encoding, count: i32) -> PageEncodingStats {
5493        PageEncodingStats {
5494            page_type,
5495            encoding,
5496            count,
5497        }
5498    }
5499
5500    // Function to use in tests for EncodingWriteSupport. This checks that dictionary
5501    // offset and encodings to make sure that column writer uses provided by trait
5502    // encodings.
5503    fn check_encoding_write_support<T: DataType>(
5504        version: WriterVersion,
5505        dict_enabled: bool,
5506        data: &[T::T],
5507        dictionary_page_offset: Option<i64>,
5508        encodings: &[Encoding],
5509        page_encoding_stats: &[PageEncodingStats],
5510    ) {
5511        let props = WriterProperties::builder()
5512            .set_writer_version(version)
5513            .set_dictionary_enabled(dict_enabled)
5514            .build();
5515        let meta = column_write_and_get_metadata::<T>(props, data);
5516        assert_eq!(meta.dictionary_page_offset(), dictionary_page_offset);
5517        assert_eq!(meta.encodings().collect::<Vec<_>>(), encodings);
5518        assert_eq!(meta.page_encoding_stats().unwrap(), page_encoding_stats);
5519    }
5520
5521    /// Returns column writer.
5522    fn get_test_column_writer<'a, T: DataType>(
5523        page_writer: Box<dyn PageWriter + 'a>,
5524        max_def_level: i16,
5525        max_rep_level: i16,
5526        props: WriterPropertiesPtr,
5527    ) -> ColumnWriterImpl<'a, T> {
5528        let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
5529        let column_writer = get_column_writer(descr, props, page_writer);
5530        get_typed_column_writer::<T>(column_writer)
5531    }
5532
5533    fn get_test_column_writer_with_path<'a, T: DataType>(
5534        page_writer: Box<dyn PageWriter + 'a>,
5535        max_def_level: i16,
5536        max_rep_level: i16,
5537        props: WriterPropertiesPtr,
5538        path: ColumnPath,
5539    ) -> ColumnWriterImpl<'a, T> {
5540        let descr = Arc::new(get_test_column_descr_with_path::<T>(
5541            max_def_level,
5542            max_rep_level,
5543            path,
5544        ));
5545        let column_writer = get_column_writer(descr, props, page_writer);
5546        get_typed_column_writer::<T>(column_writer)
5547    }
5548
5549    /// Pages collected by [`write_and_collect_pages`].
5550    struct CollectedPages {
5551        /// `(compressed byte size, value count)` for every data page, in order.
5552        data_pages: Vec<(usize, u32)>,
5553        /// Largest dictionary page seen, or 0 if the column wasn't dict-encoded.
5554        dict_page_size: usize,
5555    }
5556
5557    /// Writes `data` (with optional def/rep levels) through a raw
5558    /// `ColumnWriterImpl` configured by `props`, then re-reads the file and
5559    /// returns its page layout. Shared by the page-size regression tests so
5560    /// each only has to express its props, input, and assertions.
5561    fn write_and_collect_pages<T: DataType>(
5562        props: WriterProperties,
5563        max_def_level: i16,
5564        max_rep_level: i16,
5565        data: &[T::T],
5566        def_levels: Option<&[i16]>,
5567        rep_levels: Option<&[i16]>,
5568    ) -> CollectedPages {
5569        let mut file = tempfile::tempfile().unwrap();
5570        let mut write = TrackedWrite::new(&mut file);
5571        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5572        let mut writer =
5573            get_test_column_writer::<T>(page_writer, max_def_level, max_rep_level, Arc::new(props));
5574        writer.write_batch(data, def_levels, rep_levels).unwrap();
5575        let r = writer.close().unwrap();
5576        drop(write);
5577
5578        let read_props = ReaderProperties::builder()
5579            .set_backward_compatible_lz4(false)
5580            .build();
5581        let mut page_reader = Box::new(
5582            SerializedPageReader::new_with_properties(
5583                Arc::new(file),
5584                &r.metadata,
5585                r.rows_written as usize,
5586                None,
5587                Arc::new(read_props),
5588            )
5589            .unwrap(),
5590        );
5591
5592        let mut collected = CollectedPages {
5593            data_pages: Vec::new(),
5594            dict_page_size: 0,
5595        };
5596        while let Some(page) = page_reader.get_next_page().unwrap() {
5597            match page.page_type() {
5598                PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
5599                    collected
5600                        .data_pages
5601                        .push((page.buffer().len(), page.num_values()));
5602                }
5603                PageType::DICTIONARY_PAGE => {
5604                    collected.dict_page_size = collected.dict_page_size.max(page.buffer().len());
5605                }
5606                PageType::INDEX_PAGE => {}
5607            }
5608        }
5609        collected
5610    }
5611
5612    /// Returns column reader.
5613    fn get_test_column_reader<T: DataType>(
5614        page_reader: Box<dyn PageReader>,
5615        max_def_level: i16,
5616        max_rep_level: i16,
5617    ) -> ColumnReaderImpl<T> {
5618        let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
5619        let column_reader = get_column_reader(descr, page_reader);
5620        get_typed_column_reader::<T>(column_reader)
5621    }
5622
5623    /// Returns descriptor for primitive column.
5624    fn get_test_column_descr<T: DataType>(
5625        max_def_level: i16,
5626        max_rep_level: i16,
5627    ) -> ColumnDescriptor {
5628        let path = ColumnPath::from("col");
5629        let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5630            // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type,
5631            // it should be no-op for other types
5632            .with_length(1)
5633            .build()
5634            .unwrap();
5635        ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5636    }
5637
5638    fn get_test_column_descr_with_path<T: DataType>(
5639        max_def_level: i16,
5640        max_rep_level: i16,
5641        path: ColumnPath,
5642    ) -> ColumnDescriptor {
5643        let name = path.string();
5644        let type_ = SchemaType::primitive_type_builder(&name, T::get_physical_type())
5645            // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type,
5646            // it should be no-op for other types
5647            .with_length(1)
5648            .build()
5649            .unwrap();
5650        ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5651    }
5652
5653    fn write_and_collect_page_values(
5654        path: ColumnPath,
5655        props: WriterPropertiesPtr,
5656        data: &[i32],
5657    ) -> Vec<u32> {
5658        let mut file = tempfile::tempfile().unwrap();
5659        let mut write = TrackedWrite::new(&mut file);
5660        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5661        let mut writer =
5662            get_test_column_writer_with_path::<Int32Type>(page_writer, 0, 0, props, path);
5663        writer.write_batch(data, None, None).unwrap();
5664        let r = writer.close().unwrap();
5665
5666        drop(write);
5667
5668        let props = ReaderProperties::builder()
5669            .set_backward_compatible_lz4(false)
5670            .build();
5671        let mut page_reader = Box::new(
5672            SerializedPageReader::new_with_properties(
5673                Arc::new(file),
5674                &r.metadata,
5675                r.rows_written as usize,
5676                None,
5677                Arc::new(props),
5678            )
5679            .unwrap(),
5680        );
5681
5682        let mut values_per_page = Vec::new();
5683        while let Some(page) = page_reader.get_next_page().unwrap() {
5684            assert_eq!(page.page_type(), PageType::DATA_PAGE);
5685            values_per_page.push(page.num_values());
5686        }
5687
5688        values_per_page
5689    }
5690
5691    /// Returns page writer that collects pages without serializing them.
5692    fn get_test_page_writer() -> Box<dyn PageWriter> {
5693        Box::new(TestPageWriter {})
5694    }
5695
5696    struct TestPageWriter {}
5697
5698    impl PageWriter for TestPageWriter {
5699        fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
5700            let mut res = PageWriteSpec::new();
5701            res.page_type = page.page_type();
5702            res.uncompressed_size = page.uncompressed_size();
5703            res.compressed_size = page.compressed_size();
5704            res.num_values = page.num_values();
5705            res.offset = 0;
5706            res.bytes_written = page.data().len() as u64;
5707            Ok(res)
5708        }
5709
5710        fn close(&mut self) -> Result<()> {
5711            Ok(())
5712        }
5713    }
5714
5715    /// Write data into parquet using [`get_test_page_writer`] and [`get_test_column_writer`] and returns generated statistics.
5716    fn statistics_roundtrip<T: DataType>(values: &[<T as DataType>::T]) -> Statistics {
5717        let page_writer = get_test_page_writer();
5718        let props = Default::default();
5719        let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
5720        writer.write_batch(values, None, None).unwrap();
5721
5722        let metadata = writer.close().unwrap().metadata;
5723        if let Some(stats) = metadata.statistics() {
5724            stats.clone()
5725        } else {
5726            panic!("metadata missing statistics");
5727        }
5728    }
5729
5730    /// Returns Decimals column writer.
5731    fn get_test_decimals_column_writer<T: DataType>(
5732        page_writer: Box<dyn PageWriter>,
5733        max_def_level: i16,
5734        max_rep_level: i16,
5735        props: WriterPropertiesPtr,
5736    ) -> ColumnWriterImpl<'static, T> {
5737        let descr = Arc::new(get_test_decimals_column_descr::<T>(
5738            max_def_level,
5739            max_rep_level,
5740        ));
5741        let column_writer = get_column_writer(descr, props, page_writer);
5742        get_typed_column_writer::<T>(column_writer)
5743    }
5744
5745    /// Returns descriptor for Decimal type with primitive column.
5746    fn get_test_decimals_column_descr<T: DataType>(
5747        max_def_level: i16,
5748        max_rep_level: i16,
5749    ) -> ColumnDescriptor {
5750        let path = ColumnPath::from("col");
5751        let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5752            .with_length(16)
5753            .with_logical_type(Some(LogicalType::decimal(2, 3)))
5754            .with_scale(2)
5755            .with_precision(3)
5756            .build()
5757            .unwrap();
5758        ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5759    }
5760
5761    fn float16_statistics_roundtrip(
5762        values: &[FixedLenByteArray],
5763    ) -> ValueStatistics<FixedLenByteArray> {
5764        let page_writer = get_test_page_writer();
5765        let mut writer = get_test_float16_column_writer(page_writer, Default::default());
5766        writer.write_batch(values, None, None).unwrap();
5767
5768        let metadata = writer.close().unwrap().metadata;
5769        if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
5770            stats.clone()
5771        } else {
5772            panic!("metadata missing statistics");
5773        }
5774    }
5775
5776    fn get_test_float16_column_writer(
5777        page_writer: Box<dyn PageWriter>,
5778        props: WriterPropertiesPtr,
5779    ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
5780        let descr = Arc::new(get_test_float16_column_descr(0, 0));
5781        let column_writer = get_column_writer(descr, props, page_writer);
5782        get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
5783    }
5784
5785    fn get_test_float16_column_descr(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor {
5786        let path = ColumnPath::from("col");
5787        let type_ =
5788            SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
5789                .with_length(2)
5790                .with_logical_type(Some(LogicalType::Float16))
5791                .build()
5792                .unwrap();
5793        ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5794    }
5795
5796    fn get_test_interval_column_writer(
5797        page_writer: Box<dyn PageWriter>,
5798    ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
5799        let descr = Arc::new(get_test_interval_column_descr());
5800        let column_writer = get_column_writer(descr, Default::default(), page_writer);
5801        get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
5802    }
5803
5804    fn get_test_interval_column_descr() -> ColumnDescriptor {
5805        let path = ColumnPath::from("col");
5806        let type_ =
5807            SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
5808                .with_length(12)
5809                .with_converted_type(ConvertedType::INTERVAL)
5810                .build()
5811                .unwrap();
5812        ColumnDescriptor::new(Arc::new(type_), 0, 0, path)
5813    }
5814
5815    /// Returns column writer for UINT32 Column provided as ConvertedType only
5816    fn get_test_unsigned_int_given_as_converted_column_writer<'a, T: DataType>(
5817        page_writer: Box<dyn PageWriter + 'a>,
5818        max_def_level: i16,
5819        max_rep_level: i16,
5820        props: WriterPropertiesPtr,
5821    ) -> ColumnWriterImpl<'a, T> {
5822        let descr = Arc::new(get_test_converted_type_unsigned_integer_column_descr::<T>(
5823            max_def_level,
5824            max_rep_level,
5825        ));
5826        let column_writer = get_column_writer(descr, props, page_writer);
5827        get_typed_column_writer::<T>(column_writer)
5828    }
5829
5830    /// Returns column descriptor for UINT32 Column provided as ConvertedType only
5831    fn get_test_converted_type_unsigned_integer_column_descr<T: DataType>(
5832        max_def_level: i16,
5833        max_rep_level: i16,
5834    ) -> ColumnDescriptor {
5835        let path = ColumnPath::from("col");
5836        let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5837            .with_converted_type(ConvertedType::UINT_32)
5838            .build()
5839            .unwrap();
5840        ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5841    }
5842
5843    #[test]
5844    fn test_page_v2_snappy_compression_fallback() {
5845        // Test that PageV2 sets is_compressed to false when Snappy compression increases data size
5846        let page_writer = TestPageWriter {};
5847
5848        // Create WriterProperties with PageV2 and Snappy compression
5849        let props = WriterProperties::builder()
5850            .set_writer_version(WriterVersion::PARQUET_2_0)
5851            // Disable dictionary to ensure data is written directly
5852            .set_dictionary_enabled(false)
5853            .set_compression(Compression::SNAPPY)
5854            .build();
5855
5856        let mut column_writer =
5857            get_test_column_writer::<ByteArrayType>(Box::new(page_writer), 0, 0, Arc::new(props));
5858
5859        // Create small, simple data that Snappy compression will likely increase in size
5860        // due to compression overhead for very small data
5861        let values = vec![ByteArray::from("a")];
5862
5863        column_writer.write_batch(&values, None, None).unwrap();
5864
5865        let result = column_writer.close().unwrap();
5866        assert_eq!(
5867            result.metadata.uncompressed_size(),
5868            result.metadata.compressed_size()
5869        );
5870    }
5871
5872    struct ColumnRoundTripUniform<'a, T: DataType> {
5873        props: WriterProperties,
5874        values: &'a [T::T],
5875        def_levels: LevelDataRef<'a>,
5876        rep_levels: LevelDataRef<'a>,
5877        max_def_level: i16,
5878        max_rep_level: i16,
5879        expected_values: &'a [T::T],
5880        expected_def_levels: Option<&'a [i16]>,
5881        expected_rep_levels: Option<&'a [i16]>,
5882    }
5883
5884    impl<'a, T: DataType> ColumnRoundTripUniform<'a, T>
5885    where
5886        T::T: PartialEq + std::fmt::Debug,
5887    {
5888        fn new() -> Self {
5889            Self {
5890                props: Default::default(),
5891                values: &[],
5892                def_levels: LevelDataRef::Absent,
5893                rep_levels: LevelDataRef::Absent,
5894                max_def_level: 0,
5895                max_rep_level: 0,
5896                expected_values: &[],
5897                expected_def_levels: None,
5898                expected_rep_levels: None,
5899            }
5900        }
5901
5902        fn with_props(mut self, props: WriterProperties) -> Self {
5903            self.props = props;
5904            self
5905        }
5906
5907        fn with_values(mut self, values: &'a [T::T]) -> Self {
5908            self.values = values;
5909            self
5910        }
5911
5912        fn with_def_levels(mut self, def_levels: LevelDataRef<'a>) -> Self {
5913            self.def_levels = def_levels;
5914            self
5915        }
5916
5917        fn with_rep_levels(mut self, rep_levels: LevelDataRef<'a>) -> Self {
5918            self.rep_levels = rep_levels;
5919            self
5920        }
5921
5922        fn with_max_def_level(mut self, max_def_level: i16) -> Self {
5923            self.max_def_level = max_def_level;
5924            self
5925        }
5926
5927        fn with_max_rep_level(mut self, max_rep_level: i16) -> Self {
5928            self.max_rep_level = max_rep_level;
5929            self
5930        }
5931
5932        fn with_expected_values(mut self, expected_values: &'a [T::T]) -> Self {
5933            self.expected_values = expected_values;
5934            self
5935        }
5936
5937        fn with_expected_def_levels(mut self, expected_def_levels: &'a [i16]) -> Self {
5938            self.expected_def_levels = Some(expected_def_levels);
5939            self
5940        }
5941
5942        fn with_expected_rep_levels(mut self, expected_rep_levels: &'a [i16]) -> Self {
5943            self.expected_rep_levels = Some(expected_rep_levels);
5944            self
5945        }
5946
5947        /// Write-then-read roundtrip using `write_batch_internal` with the given
5948        /// [`LevelDataRef`] variants, and assert the read-back matches `expected_*`.
5949        fn run(self) {
5950            let mut file = tempfile::tempfile().unwrap();
5951            let mut write = TrackedWrite::new(&mut file);
5952            let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5953            let mut writer = get_test_column_writer::<T>(
5954                page_writer,
5955                self.max_def_level,
5956                self.max_rep_level,
5957                Arc::new(self.props),
5958            );
5959
5960            writer
5961                .write_batch_internal(
5962                    self.values,
5963                    None,
5964                    self.def_levels,
5965                    self.rep_levels,
5966                    None,
5967                    None,
5968                    None,
5969                )
5970                .unwrap();
5971            let result = writer.close().unwrap();
5972            drop(write);
5973
5974            let props = ReaderProperties::builder()
5975                .set_backward_compatible_lz4(false)
5976                .build();
5977            let page_reader = Box::new(
5978                SerializedPageReader::new_with_properties(
5979                    Arc::new(file),
5980                    &result.metadata,
5981                    result.rows_written as usize,
5982                    None,
5983                    Arc::new(props),
5984                )
5985                .unwrap(),
5986            );
5987            let mut reader =
5988                get_test_column_reader::<T>(page_reader, self.max_def_level, self.max_rep_level);
5989
5990            let batch_size = self
5991                .expected_def_levels
5992                .map_or(self.expected_values.len(), |l| l.len());
5993            let mut actual_values = Vec::with_capacity(batch_size);
5994            let mut actual_def = self
5995                .expected_def_levels
5996                .map(|_| Vec::with_capacity(batch_size));
5997            let mut actual_rep = self
5998                .expected_rep_levels
5999                .map(|_| Vec::with_capacity(batch_size));
6000
6001            let (_, values_read, levels_read) = reader
6002                .read_records(
6003                    batch_size,
6004                    actual_def.as_mut(),
6005                    actual_rep.as_mut(),
6006                    &mut actual_values,
6007                )
6008                .unwrap();
6009
6010            assert_eq!(&actual_values[..values_read], self.expected_values);
6011            if let Some(ref v) = actual_def {
6012                assert_eq!(&v[..levels_read], self.expected_def_levels.unwrap());
6013            }
6014            if let Some(ref v) = actual_rep {
6015                assert_eq!(&v[..levels_read], self.expected_rep_levels.unwrap());
6016            }
6017        }
6018    }
6019
6020    #[test]
6021    fn test_level_data_ref_value_count() {
6022        // `value_count` is what the byte-budget chunker uses to convert a
6023        // chunk's level span into a leaf-value count. It must work for any
6024        // column shape — flat, nullable, or nested — because the leaf
6025        // values array is decoupled from the rep/def level stream.
6026        let max_def = 2;
6027        // Non-nullable / unrepeated: no def levels materialized — every
6028        // level is a value.
6029        assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64);
6030        // Uniform run of present values, and of nulls.
6031        assert_eq!(
6032            LevelDataRef::Uniform {
6033                value: max_def,
6034                count: 40
6035            }
6036            .value_count(40, max_def),
6037            40
6038        );
6039        assert_eq!(
6040            LevelDataRef::Uniform {
6041                value: max_def - 1,
6042                count: 40
6043            }
6044            .value_count(40, max_def),
6045            0
6046        );
6047        // Materialized def levels (nullable / nested): only levels equal to
6048        // `max_def` are values; empty-list / null levels are not.
6049        let levels = [2i16, 0, 2, 1, 2, 2, 0];
6050        assert_eq!(
6051            LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def),
6052            4
6053        );
6054    }
6055
6056    #[test]
6057    fn test_uniform_def_levels_all_null() {
6058        // All-null column: def_level=0 (null) for every slot, no values written.
6059        let max_def_level = 1;
6060        let count = 100;
6061        let expected_def_levels = vec![0i16; count];
6062        ColumnRoundTripUniform::<Int32Type>::new()
6063            .with_def_levels(LevelDataRef::Uniform { value: 0, count })
6064            .with_max_def_level(max_def_level)
6065            .with_expected_def_levels(&expected_def_levels)
6066            .run();
6067    }
6068
6069    #[test]
6070    fn test_uniform_def_levels_all_valid() {
6071        // All-valid column: def_level=max for every slot, all values written.
6072        let max_def_level = 1;
6073        let values: Vec<i32> = (0..50).collect();
6074        let expected_def_levels = vec![max_def_level; values.len()];
6075        ColumnRoundTripUniform::<Int32Type>::new()
6076            .with_values(&values)
6077            .with_def_levels(LevelDataRef::Uniform {
6078                value: max_def_level,
6079                count: values.len(),
6080            })
6081            .with_max_def_level(max_def_level)
6082            .with_expected_values(&values)
6083            .with_expected_def_levels(&expected_def_levels)
6084            .run();
6085    }
6086
6087    #[test]
6088    fn test_uniform_def_and_rep_levels() {
6089        // Simulates a list column where every row is null:
6090        // def=0, rep=0 for each row (one row = one entry with no child values).
6091        let max_def_level = 2;
6092        let max_rep_level = 1;
6093        let count = 200;
6094        let expected_def_levels = vec![0i16; count];
6095        let expected_rep_levels = vec![0i16; count];
6096        ColumnRoundTripUniform::<Int32Type>::new()
6097            .with_def_levels(LevelDataRef::Uniform { value: 0, count })
6098            .with_rep_levels(LevelDataRef::Uniform { value: 0, count })
6099            .with_max_def_level(max_def_level)
6100            .with_max_rep_level(max_rep_level)
6101            .with_expected_def_levels(&expected_def_levels)
6102            .with_expected_rep_levels(&expected_rep_levels)
6103            .run();
6104    }
6105
6106    #[test]
6107    fn test_uniform_levels_v1_and_v2() {
6108        // Verify uniform levels work identically for both Parquet writer versions.
6109        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
6110            let props = WriterProperties::builder()
6111                .set_writer_version(version)
6112                .build();
6113            let max_def = 1;
6114            let count = 100;
6115            let expected_def_levels = vec![0i16; count];
6116            ColumnRoundTripUniform::<Int32Type>::new()
6117                .with_props(props)
6118                .with_def_levels(LevelDataRef::Uniform { value: 0, count })
6119                .with_max_def_level(max_def)
6120                .with_expected_def_levels(&expected_def_levels)
6121                .run();
6122        }
6123    }
6124}