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