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