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