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::collections::{BTreeSet, VecDeque};
27use std::str;
28
29use crate::basic::{
30    BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, LogicalType, PageType, Type,
31};
32use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter};
33use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues};
34use crate::compression::{Codec, CodecOptionsBuilder, create_codec};
35use crate::data_type::private::ParquetValueType;
36use crate::data_type::*;
37use crate::encodings::levels::LevelEncoder;
38#[cfg(feature = "encryption")]
39use crate::encryption::encrypt::get_column_crypto_metadata;
40use crate::errors::{ParquetError, Result};
41use crate::file::metadata::{
42    ColumnChunkMetaData, ColumnChunkMetaDataBuilder, ColumnIndexBuilder, LevelHistogram,
43    OffsetIndexBuilder, PageEncodingStats,
44};
45use crate::file::properties::{
46    EnabledStatistics, WriterProperties, WriterPropertiesPtr, WriterVersion,
47};
48use crate::file::statistics::{Statistics, ValueStatistics};
49use crate::schema::types::{ColumnDescPtr, ColumnDescriptor};
50
51pub(crate) mod encoder;
52
53macro_rules! downcast_writer {
54    ($e:expr, $i:ident, $b:expr) => {
55        match $e {
56            Self::BoolColumnWriter($i) => $b,
57            Self::Int32ColumnWriter($i) => $b,
58            Self::Int64ColumnWriter($i) => $b,
59            Self::Int96ColumnWriter($i) => $b,
60            Self::FloatColumnWriter($i) => $b,
61            Self::DoubleColumnWriter($i) => $b,
62            Self::ByteArrayColumnWriter($i) => $b,
63            Self::FixedLenByteArrayColumnWriter($i) => $b,
64        }
65    };
66}
67
68/// Column writer for a Parquet type.
69///
70/// See [`get_column_writer`] to create instances of this type
71pub enum ColumnWriter<'a> {
72    /// Column writer for boolean type
73    BoolColumnWriter(ColumnWriterImpl<'a, BoolType>),
74    /// Column writer for int32 type
75    Int32ColumnWriter(ColumnWriterImpl<'a, Int32Type>),
76    /// Column writer for int64 type
77    Int64ColumnWriter(ColumnWriterImpl<'a, Int64Type>),
78    /// Column writer for int96 (timestamp) type
79    Int96ColumnWriter(ColumnWriterImpl<'a, Int96Type>),
80    /// Column writer for float type
81    FloatColumnWriter(ColumnWriterImpl<'a, FloatType>),
82    /// Column writer for double type
83    DoubleColumnWriter(ColumnWriterImpl<'a, DoubleType>),
84    /// Column writer for byte array type
85    ByteArrayColumnWriter(ColumnWriterImpl<'a, ByteArrayType>),
86    /// Column writer for fixed length byte array type
87    FixedLenByteArrayColumnWriter(ColumnWriterImpl<'a, FixedLenByteArrayType>),
88}
89
90impl ColumnWriter<'_> {
91    /// Returns the estimated total memory usage
92    #[cfg(feature = "arrow")]
93    pub(crate) fn memory_size(&self) -> usize {
94        downcast_writer!(self, typed, typed.memory_size())
95    }
96
97    /// Returns the estimated total encoded bytes for this column writer
98    #[cfg(feature = "arrow")]
99    pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
100        downcast_writer!(self, typed, typed.get_estimated_total_bytes())
101    }
102
103    /// Finalize the currently buffered values as a data page.
104    ///
105    /// This is used by content-defined chunking to force a page boundary at
106    /// content-determined positions.
107    #[cfg(feature = "arrow")]
108    pub(crate) fn add_data_page(&mut self) -> Result<()> {
109        downcast_writer!(self, typed, typed.add_data_page())
110    }
111
112    /// Close this [`ColumnWriter`], returning the metadata for the column chunk.
113    pub fn close(self) -> Result<ColumnCloseResult> {
114        downcast_writer!(self, typed, typed.close())
115    }
116}
117
118/// Create a specific column writer corresponding to column descriptor `descr`.
119pub fn get_column_writer<'a>(
120    descr: ColumnDescPtr,
121    props: WriterPropertiesPtr,
122    page_writer: Box<dyn PageWriter + 'a>,
123) -> ColumnWriter<'a> {
124    match descr.physical_type() {
125        Type::BOOLEAN => {
126            ColumnWriter::BoolColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
127        }
128        Type::INT32 => {
129            ColumnWriter::Int32ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
130        }
131        Type::INT64 => {
132            ColumnWriter::Int64ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
133        }
134        Type::INT96 => {
135            ColumnWriter::Int96ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
136        }
137        Type::FLOAT => {
138            ColumnWriter::FloatColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
139        }
140        Type::DOUBLE => {
141            ColumnWriter::DoubleColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
142        }
143        Type::BYTE_ARRAY => {
144            ColumnWriter::ByteArrayColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
145        }
146        Type::FIXED_LEN_BYTE_ARRAY => ColumnWriter::FixedLenByteArrayColumnWriter(
147            ColumnWriterImpl::new(descr, props, page_writer),
148        ),
149    }
150}
151
152/// Gets a typed column writer for the specific type `T`, by "up-casting" `col_writer` of
153/// non-generic type to a generic column writer type `ColumnWriterImpl`.
154///
155/// Panics if actual enum value for `col_writer` does not match the type `T`.
156pub fn get_typed_column_writer<T: DataType>(col_writer: ColumnWriter) -> ColumnWriterImpl<T> {
157    T::get_column_writer(col_writer).unwrap_or_else(|| {
158        panic!(
159            "Failed to convert column writer into a typed column writer for `{}` type",
160            T::get_physical_type()
161        )
162    })
163}
164
165/// Similar to `get_typed_column_writer` but returns a reference.
166pub fn get_typed_column_writer_ref<'a, 'b: 'a, T: DataType>(
167    col_writer: &'b ColumnWriter<'a>,
168) -> &'b ColumnWriterImpl<'a, T> {
169    T::get_column_writer_ref(col_writer).unwrap_or_else(|| {
170        panic!(
171            "Failed to convert column writer into a typed column writer for `{}` type",
172            T::get_physical_type()
173        )
174    })
175}
176
177/// Similar to `get_typed_column_writer` but returns a reference.
178pub fn get_typed_column_writer_mut<'a, 'b: 'a, T: DataType>(
179    col_writer: &'a mut ColumnWriter<'b>,
180) -> &'a mut ColumnWriterImpl<'b, T> {
181    T::get_column_writer_mut(col_writer).unwrap_or_else(|| {
182        panic!(
183            "Failed to convert column writer into a typed column writer for `{}` type",
184            T::get_physical_type()
185        )
186    })
187}
188
189/// Metadata for a column chunk of a Parquet file.
190///
191/// Note this structure is returned by [`ColumnWriter::close`].
192#[derive(Debug, Clone)]
193pub struct ColumnCloseResult {
194    /// The total number of bytes written
195    pub bytes_written: u64,
196    /// The total number of rows written
197    pub rows_written: u64,
198    /// Metadata for this column chunk
199    pub metadata: ColumnChunkMetaData,
200    /// Optional bloom filter for this column
201    pub bloom_filter: Option<Sbbf>,
202    /// Optional column index, for filtering
203    pub column_index: Option<ColumnIndexMetaData>,
204    /// Optional offset index, identifying page locations
205    pub offset_index: Option<OffsetIndexMetaData>,
206}
207
208// Metrics per page
209#[derive(Default)]
210struct PageMetrics {
211    num_buffered_values: u32,
212    num_buffered_rows: u32,
213    num_page_nulls: u64,
214    repetition_level_histogram: Option<LevelHistogram>,
215    definition_level_histogram: Option<LevelHistogram>,
216}
217
218impl PageMetrics {
219    fn new() -> Self {
220        Default::default()
221    }
222
223    /// Initialize the repetition level histogram
224    fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
225        self.repetition_level_histogram = LevelHistogram::try_new(max_level);
226        self
227    }
228
229    /// Initialize the definition level histogram
230    fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
231        self.definition_level_histogram = LevelHistogram::try_new(max_level);
232        self
233    }
234
235    /// Resets the state of this `PageMetrics` to the initial state.
236    /// If histograms have been initialized their contents will be reset to zero.
237    fn new_page(&mut self) {
238        self.num_buffered_values = 0;
239        self.num_buffered_rows = 0;
240        self.num_page_nulls = 0;
241        self.repetition_level_histogram
242            .as_mut()
243            .map(LevelHistogram::reset);
244        self.definition_level_histogram
245            .as_mut()
246            .map(LevelHistogram::reset);
247    }
248}
249
250// Metrics per column writer
251#[derive(Default)]
252struct ColumnMetrics<T: Default> {
253    total_bytes_written: u64,
254    total_rows_written: u64,
255    total_uncompressed_size: u64,
256    total_compressed_size: u64,
257    total_num_values: u64,
258    dictionary_page_offset: Option<u64>,
259    data_page_offset: Option<u64>,
260    min_column_value: Option<T>,
261    max_column_value: Option<T>,
262    num_column_nulls: u64,
263    column_distinct_count: Option<u64>,
264    variable_length_bytes: Option<i64>,
265    repetition_level_histogram: Option<LevelHistogram>,
266    definition_level_histogram: Option<LevelHistogram>,
267}
268
269impl<T: Default> ColumnMetrics<T> {
270    fn new() -> Self {
271        Default::default()
272    }
273
274    /// Initialize the repetition level histogram
275    fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
276        self.repetition_level_histogram = LevelHistogram::try_new(max_level);
277        self
278    }
279
280    /// Initialize the definition level histogram
281    fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
282        self.definition_level_histogram = LevelHistogram::try_new(max_level);
283        self
284    }
285
286    /// Sum `page_histogram` into `chunk_histogram`
287    fn update_histogram(
288        chunk_histogram: &mut Option<LevelHistogram>,
289        page_histogram: &Option<LevelHistogram>,
290    ) {
291        if let (Some(page_hist), Some(chunk_hist)) = (page_histogram, chunk_histogram) {
292            chunk_hist.add(page_hist);
293        }
294    }
295
296    /// Sum the provided PageMetrics histograms into the chunk histograms. Does nothing if
297    /// page histograms are not initialized.
298    fn update_from_page_metrics(&mut self, page_metrics: &PageMetrics) {
299        ColumnMetrics::<T>::update_histogram(
300            &mut self.definition_level_histogram,
301            &page_metrics.definition_level_histogram,
302        );
303        ColumnMetrics::<T>::update_histogram(
304            &mut self.repetition_level_histogram,
305            &page_metrics.repetition_level_histogram,
306        );
307    }
308
309    /// Sum the provided page variable_length_bytes into the chunk variable_length_bytes
310    fn update_variable_length_bytes(&mut self, variable_length_bytes: Option<i64>) {
311        if let Some(var_bytes) = variable_length_bytes {
312            *self.variable_length_bytes.get_or_insert(0) += var_bytes;
313        }
314    }
315}
316
317/// Typed column writer for a primitive column.
318pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl<T>>;
319
320/// Generic column writer for a primitive Parquet column
321pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> {
322    // Column writer properties
323    descr: ColumnDescPtr,
324    props: WriterPropertiesPtr,
325    statistics_enabled: EnabledStatistics,
326
327    page_writer: Box<dyn PageWriter + 'a>,
328    codec: Compression,
329    compressor: Option<Box<dyn Codec>>,
330    encoder: E,
331
332    page_metrics: PageMetrics,
333    // Metrics per column writer
334    column_metrics: ColumnMetrics<E::T>,
335
336    /// The order of encodings within the generated metadata does not impact its meaning,
337    /// but we use a BTreeSet so that the output is deterministic
338    encodings: BTreeSet<Encoding>,
339    encoding_stats: Vec<PageEncodingStats>,
340    // Streaming level encoders for definition/repetition levels.
341    def_levels_encoder: LevelEncoder,
342    rep_levels_encoder: LevelEncoder,
343    data_pages: VecDeque<CompressedPage>,
344    // column index and offset index
345    column_index_builder: ColumnIndexBuilder,
346    offset_index_builder: Option<OffsetIndexBuilder>,
347
348    // Below fields used to incrementally check boundary order across data pages.
349    // We assume they are ascending/descending until proven wrong.
350    data_page_boundary_ascending: bool,
351    data_page_boundary_descending: bool,
352    /// (min, max)
353    last_non_null_data_page_min_max: Option<(E::T, E::T)>,
354}
355
356impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
357    /// Returns a new instance of [`GenericColumnWriter`].
358    pub fn new(
359        descr: ColumnDescPtr,
360        props: WriterPropertiesPtr,
361        page_writer: Box<dyn PageWriter + 'a>,
362    ) -> Self {
363        let codec = props.compression(descr.path());
364        let codec_options = CodecOptionsBuilder::default().build();
365        let compressor = create_codec(codec, &codec_options).unwrap();
366        let encoder = E::try_new(&descr, props.as_ref()).unwrap();
367
368        let statistics_enabled = props.statistics_enabled(descr.path());
369
370        let mut encodings = BTreeSet::new();
371        // Used for level information
372        encodings.insert(Encoding::RLE);
373
374        let mut page_metrics = PageMetrics::new();
375        let mut column_metrics = ColumnMetrics::<E::T>::new();
376
377        // Initialize level histograms if collecting page or chunk statistics
378        if statistics_enabled != EnabledStatistics::None {
379            page_metrics = page_metrics
380                .with_repetition_level_histogram(descr.max_rep_level())
381                .with_definition_level_histogram(descr.max_def_level());
382            column_metrics = column_metrics
383                .with_repetition_level_histogram(descr.max_rep_level())
384                .with_definition_level_histogram(descr.max_def_level())
385        }
386
387        // Disable column_index_builder if not collecting page statistics.
388        let mut column_index_builder = ColumnIndexBuilder::new(descr.physical_type());
389        if statistics_enabled != EnabledStatistics::Page {
390            column_index_builder.to_invalid()
391        }
392
393        // Disable offset_index_builder if requested by user.
394        let offset_index_builder = match props.offset_index_disabled() {
395            false => Some(OffsetIndexBuilder::new()),
396            _ => None,
397        };
398
399        Self {
400            def_levels_encoder: Self::create_level_encoder(descr.max_def_level(), &props),
401            rep_levels_encoder: Self::create_level_encoder(descr.max_rep_level(), &props),
402            descr,
403            props,
404            statistics_enabled,
405            page_writer,
406            codec,
407            compressor,
408            encoder,
409            data_pages: VecDeque::new(),
410            page_metrics,
411            column_metrics,
412            column_index_builder,
413            offset_index_builder,
414            encodings,
415            encoding_stats: vec![],
416            data_page_boundary_ascending: true,
417            data_page_boundary_descending: true,
418            last_non_null_data_page_min_max: None,
419        }
420    }
421
422    #[allow(clippy::too_many_arguments)]
423    pub(crate) fn write_batch_internal(
424        &mut self,
425        values: &E::Values,
426        value_indices: Option<&[usize]>,
427        def_levels: Option<&[i16]>,
428        rep_levels: Option<&[i16]>,
429        min: Option<&E::T>,
430        max: Option<&E::T>,
431        distinct_count: Option<u64>,
432    ) -> Result<usize> {
433        // Check if number of definition levels is the same as number of repetition levels.
434        if let (Some(def), Some(rep)) = (def_levels, rep_levels) {
435            if def.len() != rep.len() {
436                return Err(general_err!(
437                    "Inconsistent length of definition and repetition levels: {} != {}",
438                    def.len(),
439                    rep.len()
440                ));
441            }
442        }
443
444        // We check for DataPage limits only after we have inserted the values. If a user
445        // writes a large number of values, the DataPage size can be well above the limit.
446        //
447        // The purpose of this chunking is to bound this. Even if a user writes large
448        // number of values, the chunking will ensure that we add data page at a
449        // reasonable pagesize limit.
450
451        // TODO: find out why we don't account for size of levels when we estimate page
452        // size.
453
454        let num_levels = match def_levels {
455            Some(def_levels) => def_levels.len(),
456            None => values.len(),
457        };
458
459        if let Some(min) = min {
460            update_min(&self.descr, min, &mut self.column_metrics.min_column_value);
461        }
462        if let Some(max) = max {
463            update_max(&self.descr, max, &mut self.column_metrics.max_column_value);
464        }
465
466        // We can only set the distinct count if there are no other writes
467        if self.encoder.num_values() == 0 {
468            self.column_metrics.column_distinct_count = distinct_count;
469        } else {
470            self.column_metrics.column_distinct_count = None;
471        }
472
473        let mut values_offset = 0;
474        let mut levels_offset = 0;
475        let base_batch_size = self.props.write_batch_size();
476        while levels_offset < num_levels {
477            let mut end_offset = num_levels.min(levels_offset + base_batch_size);
478
479            // Split at record boundary
480            if let Some(r) = rep_levels {
481                while end_offset < r.len() && r[end_offset] != 0 {
482                    end_offset += 1;
483                }
484            }
485
486            values_offset += self.write_mini_batch(
487                values,
488                values_offset,
489                value_indices,
490                end_offset - levels_offset,
491                def_levels.map(|lv| &lv[levels_offset..end_offset]),
492                rep_levels.map(|lv| &lv[levels_offset..end_offset]),
493            )?;
494            levels_offset = end_offset;
495        }
496
497        // Return total number of values processed.
498        Ok(values_offset)
499    }
500
501    /// Writes batch of values, definition levels and repetition levels.
502    /// Returns number of values processed (written).
503    ///
504    /// If definition and repetition levels are provided, we write fully those levels and
505    /// select how many values to write (this number will be returned), since number of
506    /// actual written values may be smaller than provided values.
507    ///
508    /// If only values are provided, then all values are written and the length of
509    /// of the values buffer is returned.
510    ///
511    /// Definition and/or repetition levels can be omitted, if values are
512    /// non-nullable and/or non-repeated.
513    pub fn write_batch(
514        &mut self,
515        values: &E::Values,
516        def_levels: Option<&[i16]>,
517        rep_levels: Option<&[i16]>,
518    ) -> Result<usize> {
519        self.write_batch_internal(values, None, def_levels, rep_levels, None, None, None)
520    }
521
522    /// Writer may optionally provide pre-calculated statistics for use when computing
523    /// chunk-level statistics
524    ///
525    /// NB: [`WriterProperties::statistics_enabled`] must be set to [`EnabledStatistics::Chunk`]
526    /// for these statistics to take effect. If [`EnabledStatistics::None`] they will be ignored,
527    /// and if [`EnabledStatistics::Page`] the chunk statistics will instead be computed from the
528    /// computed page statistics
529    pub fn write_batch_with_statistics(
530        &mut self,
531        values: &E::Values,
532        def_levels: Option<&[i16]>,
533        rep_levels: Option<&[i16]>,
534        min: Option<&E::T>,
535        max: Option<&E::T>,
536        distinct_count: Option<u64>,
537    ) -> Result<usize> {
538        self.write_batch_internal(
539            values,
540            None,
541            def_levels,
542            rep_levels,
543            min,
544            max,
545            distinct_count,
546        )
547    }
548
549    /// Returns the estimated total memory usage.
550    ///
551    /// Unlike [`Self::get_estimated_total_bytes`] this is an estimate
552    /// of the current memory usage and not the final anticipated encoded size.
553    #[cfg(feature = "arrow")]
554    pub(crate) fn memory_size(&self) -> usize {
555        self.column_metrics.total_bytes_written as usize + self.encoder.estimated_memory_size()
556    }
557
558    /// Returns total number of bytes written by this column writer so far.
559    /// This value is also returned when column writer is closed.
560    ///
561    /// Note: this value does not include any buffered data that has not
562    /// yet been flushed to a page.
563    pub fn get_total_bytes_written(&self) -> u64 {
564        self.column_metrics.total_bytes_written
565    }
566
567    /// Returns the estimated total encoded bytes for this column writer.
568    ///
569    /// Unlike [`Self::get_total_bytes_written`] this includes an estimate
570    /// of any data that has not yet been flushed to a page, based on it's
571    /// anticipated encoded size.
572    #[cfg(feature = "arrow")]
573    pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
574        self.data_pages
575            .iter()
576            .map(|page| page.data().len() as u64)
577            .sum::<u64>()
578            + self.column_metrics.total_bytes_written
579            + self.encoder.estimated_data_page_size() as u64
580            + self.encoder.estimated_dict_page_size().unwrap_or_default() as u64
581    }
582
583    /// Returns total number of rows written by this column writer so far.
584    /// This value is also returned when column writer is closed.
585    pub fn get_total_rows_written(&self) -> u64 {
586        self.column_metrics.total_rows_written
587    }
588
589    /// Returns a reference to a [`ColumnDescPtr`]
590    pub fn get_descriptor(&self) -> &ColumnDescPtr {
591        &self.descr
592    }
593
594    /// Finalizes writes and closes the column writer.
595    /// Returns total bytes written, total rows written and column chunk metadata.
596    pub fn close(mut self) -> Result<ColumnCloseResult> {
597        if self.page_metrics.num_buffered_values > 0 {
598            self.add_data_page()?;
599        }
600        if self.encoder.has_dictionary() {
601            self.write_dictionary_page()?;
602        }
603        self.flush_data_pages()?;
604        let metadata = self.build_column_metadata()?;
605        self.page_writer.close()?;
606
607        let boundary_order = match (
608            self.data_page_boundary_ascending,
609            self.data_page_boundary_descending,
610        ) {
611            // If the lists are composed of equal elements then will be marked as ascending
612            // (Also the case if all pages are null pages)
613            (true, _) => BoundaryOrder::ASCENDING,
614            (false, true) => BoundaryOrder::DESCENDING,
615            (false, false) => BoundaryOrder::UNORDERED,
616        };
617        self.column_index_builder.set_boundary_order(boundary_order);
618
619        let column_index = match self.column_index_builder.valid() {
620            true => Some(self.column_index_builder.build()?),
621            false => None,
622        };
623
624        let offset_index = self.offset_index_builder.map(|b| b.build());
625
626        Ok(ColumnCloseResult {
627            bytes_written: self.column_metrics.total_bytes_written,
628            rows_written: self.column_metrics.total_rows_written,
629            bloom_filter: self.encoder.flush_bloom_filter(),
630            metadata,
631            column_index,
632            offset_index,
633        })
634    }
635
636    /// Creates a new streaming level encoder appropriate for the writer version.
637    fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder {
638        match props.writer_version() {
639            WriterVersion::PARQUET_1_0 => LevelEncoder::v1_streaming(max_level),
640            WriterVersion::PARQUET_2_0 => LevelEncoder::v2_streaming(max_level),
641        }
642    }
643
644    /// Writes mini batch of values, definition and repetition levels.
645    /// This allows fine-grained processing of values and maintaining a reasonable
646    /// page size.
647    fn write_mini_batch(
648        &mut self,
649        values: &E::Values,
650        values_offset: usize,
651        value_indices: Option<&[usize]>,
652        num_levels: usize,
653        def_levels: Option<&[i16]>,
654        rep_levels: Option<&[i16]>,
655    ) -> Result<usize> {
656        // Process definition levels and determine how many values to write.
657        let values_to_write = if self.descr.max_def_level() > 0 {
658            let levels = def_levels.ok_or_else(|| {
659                general_err!(
660                    "Definition levels are required, because max definition level = {}",
661                    self.descr.max_def_level()
662                )
663            })?;
664
665            let mut values_to_write = 0usize;
666            let max_def = self.descr.max_def_level();
667            let encoder = &mut self.def_levels_encoder;
668            match self.page_metrics.definition_level_histogram.as_mut() {
669                Some(histogram) => encoder.put_with_observer(levels, |level, count| {
670                    values_to_write += count * (level == max_def) as usize;
671                    histogram.increment_by(level, count as i64);
672                }),
673                None => encoder.put_with_observer(levels, |level, count| {
674                    values_to_write += count * (level == max_def) as usize;
675                }),
676            };
677            self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64;
678            values_to_write
679        } else {
680            num_levels
681        };
682
683        // Process repetition levels and determine how many rows we are about to process.
684        if self.descr.max_rep_level() > 0 {
685            // A row could contain more than one value.
686            let levels = rep_levels.ok_or_else(|| {
687                general_err!(
688                    "Repetition levels are required, because max repetition level = {}",
689                    self.descr.max_rep_level()
690                )
691            })?;
692
693            if !levels.is_empty() && levels[0] != 0 {
694                return Err(general_err!(
695                    "Write must start at a record boundary, got non-zero repetition level of {}",
696                    levels[0]
697                ));
698            }
699
700            let mut new_rows = 0u32;
701            let encoder = &mut self.rep_levels_encoder;
702            match self.page_metrics.repetition_level_histogram.as_mut() {
703                Some(histogram) => encoder.put_with_observer(levels, |level, count| {
704                    new_rows += (count as u32) * (level == 0) as u32;
705                    histogram.increment_by(level, count as i64);
706                }),
707                None => encoder.put_with_observer(levels, |level, count| {
708                    new_rows += (count as u32) * (level == 0) as u32;
709                }),
710            };
711            self.page_metrics.num_buffered_rows += new_rows;
712        } else {
713            // Each value is exactly one row.
714            // Equals to the number of values, we count nulls as well.
715            self.page_metrics.num_buffered_rows += num_levels as u32;
716        }
717
718        match value_indices {
719            Some(indices) => {
720                let indices = &indices[values_offset..values_offset + values_to_write];
721                self.encoder.write_gather(values, indices)?;
722            }
723            None => self.encoder.write(values, values_offset, values_to_write)?,
724        }
725
726        self.page_metrics.num_buffered_values += num_levels as u32;
727
728        if self.should_add_data_page() {
729            self.add_data_page()?;
730        }
731
732        if self.should_dict_fallback() {
733            self.dict_fallback()?;
734        }
735
736        Ok(values_to_write)
737    }
738
739    /// Returns true if we need to fall back to non-dictionary encoding.
740    ///
741    /// We can only fall back if dictionary encoder is set and we have exceeded dictionary
742    /// size.
743    #[inline]
744    fn should_dict_fallback(&self) -> bool {
745        match self.encoder.estimated_dict_page_size() {
746            Some(size) => {
747                size >= self
748                    .props
749                    .column_dictionary_page_size_limit(self.descr.path())
750            }
751            None => false,
752        }
753    }
754
755    /// Returns true if there is enough data for a data page, false otherwise.
756    #[inline]
757    fn should_add_data_page(&self) -> bool {
758        // This is necessary in the event of a much larger dictionary size than page size
759        //
760        // In such a scenario the dictionary decoder may return an estimated encoded
761        // size in excess of the page size limit, even when there are no buffered values
762        if self.page_metrics.num_buffered_values == 0 {
763            return false;
764        }
765
766        self.page_metrics.num_buffered_rows as usize >= self.props.data_page_row_count_limit()
767            || self.encoder.estimated_data_page_size()
768                >= self.props.column_data_page_size_limit(self.descr.path())
769    }
770
771    /// Performs dictionary fallback.
772    /// Prepares and writes dictionary and all data pages into page writer.
773    fn dict_fallback(&mut self) -> Result<()> {
774        // At this point we know that we need to fall back.
775        if self.page_metrics.num_buffered_values > 0 {
776            self.add_data_page()?;
777        }
778        self.write_dictionary_page()?;
779        self.flush_data_pages()?;
780        Ok(())
781    }
782
783    /// Update the column index and offset index when adding the data page
784    fn update_column_offset_index(
785        &mut self,
786        page_statistics: Option<&ValueStatistics<E::T>>,
787        page_variable_length_bytes: Option<i64>,
788    ) {
789        // update the column index
790        let null_page =
791            (self.page_metrics.num_buffered_rows as u64) == self.page_metrics.num_page_nulls;
792        // a page contains only null values,
793        // and writers have to set the corresponding entries in min_values and max_values to byte[0]
794        if null_page && self.column_index_builder.valid() {
795            self.column_index_builder.append(
796                null_page,
797                vec![],
798                vec![],
799                self.page_metrics.num_page_nulls as i64,
800            );
801        } else if self.column_index_builder.valid() {
802            // from page statistics
803            // If can't get the page statistics, ignore this column/offset index for this column chunk
804            match &page_statistics {
805                None => {
806                    self.column_index_builder.to_invalid();
807                }
808                Some(stat) => {
809                    // Check if min/max are still ascending/descending across pages
810                    let new_min = stat.min_opt().unwrap();
811                    let new_max = stat.max_opt().unwrap();
812                    if let Some((last_min, last_max)) = &self.last_non_null_data_page_min_max {
813                        if self.data_page_boundary_ascending {
814                            // If last min/max are greater than new min/max then not ascending anymore
815                            let not_ascending = compare_greater(&self.descr, last_min, new_min)
816                                || compare_greater(&self.descr, last_max, new_max);
817                            if not_ascending {
818                                self.data_page_boundary_ascending = false;
819                            }
820                        }
821
822                        if self.data_page_boundary_descending {
823                            // If new min/max are greater than last min/max then not descending anymore
824                            let not_descending = compare_greater(&self.descr, new_min, last_min)
825                                || compare_greater(&self.descr, new_max, last_max);
826                            if not_descending {
827                                self.data_page_boundary_descending = false;
828                            }
829                        }
830                    }
831                    self.last_non_null_data_page_min_max = Some((new_min.clone(), new_max.clone()));
832
833                    if self.can_truncate_value() {
834                        self.column_index_builder.append(
835                            null_page,
836                            self.truncate_min_value(
837                                self.props.column_index_truncate_length(),
838                                stat.min_bytes_opt().unwrap(),
839                            )
840                            .0,
841                            self.truncate_max_value(
842                                self.props.column_index_truncate_length(),
843                                stat.max_bytes_opt().unwrap(),
844                            )
845                            .0,
846                            self.page_metrics.num_page_nulls as i64,
847                        );
848                    } else {
849                        self.column_index_builder.append(
850                            null_page,
851                            stat.min_bytes_opt().unwrap().to_vec(),
852                            stat.max_bytes_opt().unwrap().to_vec(),
853                            self.page_metrics.num_page_nulls as i64,
854                        );
855                    }
856                }
857            }
858        }
859
860        // Append page histograms to the `ColumnIndex` histograms
861        self.column_index_builder.append_histograms(
862            &self.page_metrics.repetition_level_histogram,
863            &self.page_metrics.definition_level_histogram,
864        );
865
866        // Update the offset index
867        if let Some(builder) = self.offset_index_builder.as_mut() {
868            builder.append_row_count(self.page_metrics.num_buffered_rows as i64);
869            builder.append_unencoded_byte_array_data_bytes(page_variable_length_bytes);
870        }
871    }
872
873    /// Determine if we should allow truncating min/max values for this column's statistics
874    fn can_truncate_value(&self) -> bool {
875        match self.descr.physical_type() {
876            // Don't truncate for Float16 and Decimal because their sort order is different
877            // from that of FIXED_LEN_BYTE_ARRAY sort order.
878            // So truncation of those types could lead to inaccurate min/max statistics
879            Type::FIXED_LEN_BYTE_ARRAY
880                if !matches!(
881                    self.descr.logical_type_ref(),
882                    Some(&LogicalType::Decimal { .. }) | Some(&LogicalType::Float16)
883                ) =>
884            {
885                true
886            }
887            Type::BYTE_ARRAY => true,
888            // Truncation only applies for fba/binary physical types
889            _ => false,
890        }
891    }
892
893    /// Returns `true` if this column's logical type is a UTF-8 string.
894    fn is_utf8(&self) -> bool {
895        self.get_descriptor().logical_type_ref() == Some(&LogicalType::String)
896            || self.get_descriptor().converted_type() == ConvertedType::UTF8
897    }
898
899    /// Truncates a binary statistic to at most `truncation_length` bytes.
900    ///
901    /// If truncation is not possible, returns `data`.
902    ///
903    /// The `bool` in the returned tuple indicates whether truncation occurred or not.
904    ///
905    /// UTF-8 Note:
906    /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will
907    /// also remain valid UTF-8, but may be less tnan `truncation_length` bytes to avoid splitting
908    /// on non-character boundaries.
909    fn truncate_min_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
910        truncation_length
911            .filter(|l| data.len() > *l)
912            .and_then(|l|
913                // don't do extra work if this column isn't UTF-8
914                if self.is_utf8() {
915                    match str::from_utf8(data) {
916                        Ok(str_data) => truncate_utf8(str_data, l),
917                        Err(_) => Some(data[..l].to_vec()),
918                    }
919                } else {
920                    Some(data[..l].to_vec())
921                }
922            )
923            .map(|truncated| (truncated, true))
924            .unwrap_or_else(|| (data.to_vec(), false))
925    }
926
927    /// Truncates a binary statistic to at most `truncation_length` bytes, and then increment the
928    /// final byte(s) to yield a valid upper bound. This may result in a result of less than
929    /// `truncation_length` bytes if the last byte(s) overflows.
930    ///
931    /// If truncation is not possible, returns `data`.
932    ///
933    /// The `bool` in the returned tuple indicates whether truncation occurred or not.
934    ///
935    /// UTF-8 Note:
936    /// If the column type indicates UTF-8, and `data` contains valid UTF-8, then the result will
937    /// also remain valid UTF-8 (but again may be less than `truncation_length` bytes). If `data`
938    /// does not contain valid UTF-8, then truncation will occur as if the column is non-string
939    /// binary.
940    fn truncate_max_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
941        truncation_length
942            .filter(|l| data.len() > *l)
943            .and_then(|l|
944                // don't do extra work if this column isn't UTF-8
945                if self.is_utf8() {
946                    match str::from_utf8(data) {
947                        Ok(str_data) => truncate_and_increment_utf8(str_data, l),
948                        Err(_) => increment(data[..l].to_vec()),
949                    }
950                } else {
951                    increment(data[..l].to_vec())
952                }
953            )
954            .map(|truncated| (truncated, true))
955            .unwrap_or_else(|| (data.to_vec(), false))
956    }
957
958    /// Truncate the min and max values that will be written to a data page
959    /// header or column chunk Statistics
960    fn truncate_statistics(&self, statistics: Statistics) -> Statistics {
961        let backwards_compatible_min_max = self.descr.sort_order().is_signed();
962        match statistics {
963            Statistics::ByteArray(stats) if stats._internal_has_min_max_set() => {
964                let (min, did_truncate_min) = self.truncate_min_value(
965                    self.props.statistics_truncate_length(),
966                    stats.min_bytes_opt().unwrap(),
967                );
968                let (max, did_truncate_max) = self.truncate_max_value(
969                    self.props.statistics_truncate_length(),
970                    stats.max_bytes_opt().unwrap(),
971                );
972                Statistics::ByteArray(
973                    ValueStatistics::new(
974                        Some(min.into()),
975                        Some(max.into()),
976                        stats.distinct_count(),
977                        stats.null_count_opt(),
978                        backwards_compatible_min_max,
979                    )
980                    .with_max_is_exact(!did_truncate_max)
981                    .with_min_is_exact(!did_truncate_min),
982                )
983            }
984            Statistics::FixedLenByteArray(stats)
985                if (stats._internal_has_min_max_set() && self.can_truncate_value()) =>
986            {
987                let (min, did_truncate_min) = self.truncate_min_value(
988                    self.props.statistics_truncate_length(),
989                    stats.min_bytes_opt().unwrap(),
990                );
991                let (max, did_truncate_max) = self.truncate_max_value(
992                    self.props.statistics_truncate_length(),
993                    stats.max_bytes_opt().unwrap(),
994                );
995                Statistics::FixedLenByteArray(
996                    ValueStatistics::new(
997                        Some(min.into()),
998                        Some(max.into()),
999                        stats.distinct_count(),
1000                        stats.null_count_opt(),
1001                        backwards_compatible_min_max,
1002                    )
1003                    .with_max_is_exact(!did_truncate_max)
1004                    .with_min_is_exact(!did_truncate_min),
1005                )
1006            }
1007            stats => stats,
1008        }
1009    }
1010
1011    /// Adds data page.
1012    /// Data page is either buffered in case of dictionary encoding or written directly.
1013    pub(crate) fn add_data_page(&mut self) -> Result<()> {
1014        // Extract encoded values
1015        let values_data = self.encoder.flush_data_page()?;
1016
1017        let max_def_level = self.descr.max_def_level();
1018        let max_rep_level = self.descr.max_rep_level();
1019
1020        self.column_metrics.num_column_nulls += self.page_metrics.num_page_nulls;
1021
1022        let page_statistics = match (values_data.min_value, values_data.max_value) {
1023            (Some(min), Some(max)) => {
1024                // Update chunk level statistics
1025                update_min(&self.descr, &min, &mut self.column_metrics.min_column_value);
1026                update_max(&self.descr, &max, &mut self.column_metrics.max_column_value);
1027
1028                (self.statistics_enabled == EnabledStatistics::Page).then_some(
1029                    ValueStatistics::new(
1030                        Some(min),
1031                        Some(max),
1032                        None,
1033                        Some(self.page_metrics.num_page_nulls),
1034                        false,
1035                    ),
1036                )
1037            }
1038            _ => None,
1039        };
1040
1041        // update column and offset index
1042        self.update_column_offset_index(
1043            page_statistics.as_ref(),
1044            values_data.variable_length_bytes,
1045        );
1046
1047        // Update histograms and variable_length_bytes in column_metrics
1048        self.column_metrics
1049            .update_from_page_metrics(&self.page_metrics);
1050        self.column_metrics
1051            .update_variable_length_bytes(values_data.variable_length_bytes);
1052
1053        // From here on, we only need page statistics if they will be written to the page header.
1054        let page_statistics = page_statistics
1055            .filter(|_| self.props.write_page_header_statistics(self.descr.path()))
1056            .map(|stats| self.truncate_statistics(Statistics::from(stats)));
1057
1058        let compressed_page = match self.props.writer_version() {
1059            WriterVersion::PARQUET_1_0 => {
1060                let mut buffer = vec![];
1061
1062                if max_rep_level > 0 {
1063                    self.rep_levels_encoder
1064                        .flush_to(|data| buffer.extend_from_slice(data));
1065                }
1066
1067                if max_def_level > 0 {
1068                    self.def_levels_encoder
1069                        .flush_to(|data| buffer.extend_from_slice(data));
1070                }
1071
1072                buffer.extend_from_slice(&values_data.buf);
1073                let uncompressed_size = buffer.len();
1074
1075                if let Some(ref mut cmpr) = self.compressor {
1076                    let mut compressed_buf = Vec::with_capacity(uncompressed_size);
1077                    cmpr.compress(&buffer[..], &mut compressed_buf)?;
1078                    compressed_buf.shrink_to_fit();
1079                    buffer = compressed_buf;
1080                }
1081
1082                let data_page = Page::DataPage {
1083                    buf: buffer.into(),
1084                    num_values: self.page_metrics.num_buffered_values,
1085                    encoding: values_data.encoding,
1086                    def_level_encoding: Encoding::RLE,
1087                    rep_level_encoding: Encoding::RLE,
1088                    statistics: page_statistics,
1089                };
1090
1091                CompressedPage::new(data_page, uncompressed_size)
1092            }
1093            WriterVersion::PARQUET_2_0 => {
1094                let mut rep_levels_byte_len = 0;
1095                let mut def_levels_byte_len = 0;
1096                let mut buffer = vec![];
1097
1098                if max_rep_level > 0 {
1099                    self.rep_levels_encoder
1100                        .flush_to(|data| buffer.extend_from_slice(data));
1101                    rep_levels_byte_len = buffer.len();
1102                }
1103
1104                if max_def_level > 0 {
1105                    self.def_levels_encoder
1106                        .flush_to(|data| buffer.extend_from_slice(data));
1107                    def_levels_byte_len = buffer.len() - rep_levels_byte_len;
1108                }
1109
1110                let uncompressed_size =
1111                    rep_levels_byte_len + def_levels_byte_len + values_data.buf.len();
1112
1113                // Data Page v2 compresses values only.
1114                let is_compressed = match self.compressor {
1115                    Some(ref mut cmpr) => {
1116                        let buffer_len = buffer.len();
1117                        cmpr.compress(&values_data.buf, &mut buffer)?;
1118                        let compressed_values_size = buffer.len() - buffer_len;
1119                        let threshold = self
1120                            .props
1121                            .column_data_page_v2_compression_ratio_threshold(self.descr.path());
1122                        if (compressed_values_size as f64) >= (uncompressed_size as f64) * threshold
1123                        {
1124                            buffer.truncate(buffer_len);
1125                            buffer.extend_from_slice(&values_data.buf);
1126                            false
1127                        } else {
1128                            true
1129                        }
1130                    }
1131                    None => {
1132                        buffer.extend_from_slice(&values_data.buf);
1133                        false
1134                    }
1135                };
1136
1137                let data_page = Page::DataPageV2 {
1138                    buf: buffer.into(),
1139                    num_values: self.page_metrics.num_buffered_values,
1140                    encoding: values_data.encoding,
1141                    num_nulls: self.page_metrics.num_page_nulls as u32,
1142                    num_rows: self.page_metrics.num_buffered_rows,
1143                    def_levels_byte_len: def_levels_byte_len as u32,
1144                    rep_levels_byte_len: rep_levels_byte_len as u32,
1145                    is_compressed,
1146                    statistics: page_statistics,
1147                };
1148
1149                CompressedPage::new(data_page, uncompressed_size)
1150            }
1151        };
1152
1153        // Check if we need to buffer data page or flush it to the sink directly.
1154        if self.encoder.has_dictionary() {
1155            self.data_pages.push_back(compressed_page);
1156        } else {
1157            self.write_data_page(compressed_page)?;
1158        }
1159
1160        // Update total number of rows.
1161        self.column_metrics.total_rows_written += self.page_metrics.num_buffered_rows as u64;
1162        self.page_metrics.new_page();
1163
1164        Ok(())
1165    }
1166
1167    /// Finalises any outstanding data pages and flushes buffered data pages from
1168    /// dictionary encoding into underlying sink.
1169    #[inline]
1170    fn flush_data_pages(&mut self) -> Result<()> {
1171        // Write all outstanding data to a new page.
1172        if self.page_metrics.num_buffered_values > 0 {
1173            self.add_data_page()?;
1174        }
1175
1176        while let Some(page) = self.data_pages.pop_front() {
1177            self.write_data_page(page)?;
1178        }
1179
1180        Ok(())
1181    }
1182
1183    /// Assembles column chunk metadata.
1184    fn build_column_metadata(&mut self) -> Result<ColumnChunkMetaData> {
1185        let total_compressed_size = self.column_metrics.total_compressed_size as i64;
1186        let total_uncompressed_size = self.column_metrics.total_uncompressed_size as i64;
1187        let num_values = self.column_metrics.total_num_values as i64;
1188        let dict_page_offset = self.column_metrics.dictionary_page_offset.map(|v| v as i64);
1189        // If data page offset is not set, then no pages have been written
1190        let data_page_offset = self.column_metrics.data_page_offset.unwrap_or(0) as i64;
1191
1192        let mut builder = ColumnChunkMetaData::builder(self.descr.clone())
1193            .set_compression(self.codec)
1194            .set_encodings_mask(EncodingMask::new_from_encodings(self.encodings.iter()))
1195            .set_page_encoding_stats(self.encoding_stats.clone())
1196            .set_total_compressed_size(total_compressed_size)
1197            .set_total_uncompressed_size(total_uncompressed_size)
1198            .set_num_values(num_values)
1199            .set_data_page_offset(data_page_offset)
1200            .set_dictionary_page_offset(dict_page_offset);
1201
1202        if self.statistics_enabled != EnabledStatistics::None {
1203            let backwards_compatible_min_max = self.descr.sort_order().is_signed();
1204
1205            let statistics = ValueStatistics::<E::T>::new(
1206                self.column_metrics.min_column_value.clone(),
1207                self.column_metrics.max_column_value.clone(),
1208                self.column_metrics.column_distinct_count,
1209                Some(self.column_metrics.num_column_nulls),
1210                false,
1211            )
1212            .with_backwards_compatible_min_max(backwards_compatible_min_max)
1213            .into();
1214
1215            let statistics = self.truncate_statistics(statistics);
1216
1217            builder = builder
1218                .set_statistics(statistics)
1219                .set_unencoded_byte_array_data_bytes(self.column_metrics.variable_length_bytes)
1220                .set_repetition_level_histogram(
1221                    self.column_metrics.repetition_level_histogram.take(),
1222                )
1223                .set_definition_level_histogram(
1224                    self.column_metrics.definition_level_histogram.take(),
1225                );
1226
1227            if let Some(geo_stats) = self.encoder.flush_geospatial_statistics() {
1228                builder = builder.set_geo_statistics(geo_stats);
1229            }
1230        }
1231
1232        builder = self.set_column_chunk_encryption_properties(builder);
1233
1234        let metadata = builder.build()?;
1235        Ok(metadata)
1236    }
1237
1238    /// Writes compressed data page into underlying sink and updates global metrics.
1239    #[inline]
1240    fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {
1241        self.encodings.insert(page.encoding());
1242        match self.encoding_stats.last_mut() {
1243            Some(encoding_stats)
1244                if encoding_stats.page_type == page.page_type()
1245                    && encoding_stats.encoding == page.encoding() =>
1246            {
1247                encoding_stats.count += 1;
1248            }
1249            _ => {
1250                // data page type does not change inside a file
1251                // encoding can currently only change from dictionary to non-dictionary once
1252                self.encoding_stats.push(PageEncodingStats {
1253                    page_type: page.page_type(),
1254                    encoding: page.encoding(),
1255                    count: 1,
1256                });
1257            }
1258        }
1259        let page_spec = self.page_writer.write_page(page)?;
1260        // update offset index
1261        // compressed_size = header_size + compressed_data_size
1262        if let Some(builder) = self.offset_index_builder.as_mut() {
1263            builder
1264                .append_offset_and_size(page_spec.offset as i64, page_spec.compressed_size as i32)
1265        }
1266        self.update_metrics_for_page(page_spec);
1267        Ok(())
1268    }
1269
1270    /// Writes dictionary page into underlying sink.
1271    #[inline]
1272    fn write_dictionary_page(&mut self) -> Result<()> {
1273        let compressed_page = {
1274            let mut page = self
1275                .encoder
1276                .flush_dict_page()?
1277                .ok_or_else(|| general_err!("Dictionary encoder is not set"))?;
1278
1279            let uncompressed_size = page.buf.len();
1280
1281            if let Some(ref mut cmpr) = self.compressor {
1282                let mut output_buf = Vec::with_capacity(uncompressed_size);
1283                cmpr.compress(&page.buf, &mut output_buf)?;
1284                page.buf = Bytes::from(output_buf);
1285            }
1286
1287            let dict_page = Page::DictionaryPage {
1288                buf: page.buf,
1289                num_values: page.num_values as u32,
1290                encoding: self.props.dictionary_page_encoding(),
1291                is_sorted: page.is_sorted,
1292            };
1293            CompressedPage::new(dict_page, uncompressed_size)
1294        };
1295
1296        self.encodings.insert(compressed_page.encoding());
1297        self.encoding_stats.push(PageEncodingStats {
1298            page_type: PageType::DICTIONARY_PAGE,
1299            encoding: compressed_page.encoding(),
1300            count: 1,
1301        });
1302        let page_spec = self.page_writer.write_page(compressed_page)?;
1303        self.update_metrics_for_page(page_spec);
1304        // For the directory page, don't need to update column/offset index.
1305        Ok(())
1306    }
1307
1308    /// Updates column writer metrics with each page metadata.
1309    #[inline]
1310    fn update_metrics_for_page(&mut self, page_spec: PageWriteSpec) {
1311        self.column_metrics.total_uncompressed_size += page_spec.uncompressed_size as u64;
1312        self.column_metrics.total_compressed_size += page_spec.compressed_size as u64;
1313        self.column_metrics.total_bytes_written += page_spec.bytes_written;
1314
1315        match page_spec.page_type {
1316            PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
1317                self.column_metrics.total_num_values += page_spec.num_values as u64;
1318                if self.column_metrics.data_page_offset.is_none() {
1319                    self.column_metrics.data_page_offset = Some(page_spec.offset);
1320                }
1321            }
1322            PageType::DICTIONARY_PAGE => {
1323                assert!(
1324                    self.column_metrics.dictionary_page_offset.is_none(),
1325                    "Dictionary offset is already set"
1326                );
1327                self.column_metrics.dictionary_page_offset = Some(page_spec.offset);
1328            }
1329            _ => {}
1330        }
1331    }
1332
1333    #[inline]
1334    #[cfg(feature = "encryption")]
1335    fn set_column_chunk_encryption_properties(
1336        &self,
1337        builder: ColumnChunkMetaDataBuilder,
1338    ) -> ColumnChunkMetaDataBuilder {
1339        if let Some(encryption_properties) = self.props.file_encryption_properties.as_ref() {
1340            builder.set_column_crypto_metadata(get_column_crypto_metadata(
1341                encryption_properties,
1342                &self.descr,
1343            ))
1344        } else {
1345            builder
1346        }
1347    }
1348
1349    #[inline]
1350    #[cfg(not(feature = "encryption"))]
1351    fn set_column_chunk_encryption_properties(
1352        &self,
1353        builder: ColumnChunkMetaDataBuilder,
1354    ) -> ColumnChunkMetaDataBuilder {
1355        builder
1356    }
1357}
1358
1359fn update_min<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, min: &mut Option<T>) {
1360    update_stat::<T, _>(descr, val, min, |cur| compare_greater(descr, cur, val))
1361}
1362
1363fn update_max<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, max: &mut Option<T>) {
1364    update_stat::<T, _>(descr, val, max, |cur| compare_greater(descr, val, cur))
1365}
1366
1367#[inline]
1368#[allow(clippy::eq_op)]
1369fn is_nan<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T) -> bool {
1370    match T::PHYSICAL_TYPE {
1371        Type::FLOAT | Type::DOUBLE => val != val,
1372        Type::FIXED_LEN_BYTE_ARRAY if descr.logical_type_ref() == Some(&LogicalType::Float16) => {
1373            let val = val.as_bytes();
1374            let val = f16::from_le_bytes([val[0], val[1]]);
1375            val.is_nan()
1376        }
1377        _ => false,
1378    }
1379}
1380
1381/// Perform a conditional update of `cur`, skipping any NaN values
1382///
1383/// If `cur` is `None`, sets `cur` to `Some(val)`, otherwise calls `should_update` with
1384/// the value of `cur`, and updates `cur` to `Some(val)` if it returns `true`
1385fn update_stat<T: ParquetValueType, F>(
1386    descr: &ColumnDescriptor,
1387    val: &T,
1388    cur: &mut Option<T>,
1389    should_update: F,
1390) where
1391    F: Fn(&T) -> bool,
1392{
1393    if is_nan(descr, val) {
1394        return;
1395    }
1396
1397    if cur.as_ref().is_none_or(should_update) {
1398        *cur = Some(val.clone());
1399    }
1400}
1401
1402/// Evaluate `a > b` according to underlying logical type.
1403fn compare_greater<T: ParquetValueType>(descr: &ColumnDescriptor, a: &T, b: &T) -> bool {
1404    match T::PHYSICAL_TYPE {
1405        Type::INT32 | Type::INT64 => {
1406            if let Some(LogicalType::Integer {
1407                is_signed: false, ..
1408            }) = descr.logical_type_ref()
1409            {
1410                // need to compare unsigned
1411                return compare_greater_unsigned_int(a, b);
1412            }
1413
1414            match descr.converted_type() {
1415                ConvertedType::UINT_8
1416                | ConvertedType::UINT_16
1417                | ConvertedType::UINT_32
1418                | ConvertedType::UINT_64 => {
1419                    return compare_greater_unsigned_int(a, b);
1420                }
1421                _ => {}
1422            };
1423        }
1424        Type::FIXED_LEN_BYTE_ARRAY | Type::BYTE_ARRAY => {
1425            if let Some(LogicalType::Decimal { .. }) = descr.logical_type_ref() {
1426                return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes());
1427            }
1428            if let ConvertedType::DECIMAL = descr.converted_type() {
1429                return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes());
1430            }
1431            if let Some(LogicalType::Float16) = descr.logical_type_ref() {
1432                return compare_greater_f16(a.as_bytes(), b.as_bytes());
1433            }
1434        }
1435
1436        _ => {}
1437    }
1438
1439    // compare independent of logical / converted type
1440    a > b
1441}
1442
1443// ----------------------------------------------------------------------
1444// Encoding support for column writer.
1445// This mirrors parquet-mr default encodings for writes. See:
1446// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV1ValuesWriterFactory.java
1447// https://github.com/apache/parquet-mr/blob/master/parquet-column/src/main/java/org/apache/parquet/column/values/factory/DefaultV2ValuesWriterFactory.java
1448
1449/// Returns encoding for a column when no other encoding is provided in writer properties.
1450fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding {
1451    match (kind, props.writer_version()) {
1452        (Type::BOOLEAN, WriterVersion::PARQUET_2_0) => Encoding::RLE,
1453        (Type::INT32, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1454        (Type::INT64, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1455        (Type::BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1456        (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1457        _ => Encoding::PLAIN,
1458    }
1459}
1460
1461/// Returns true if dictionary is supported for column writer, false otherwise.
1462fn has_dictionary_support(kind: Type, props: &WriterProperties) -> bool {
1463    match (kind, props.writer_version()) {
1464        // Booleans do not support dict encoding and should use a fallback encoding.
1465        (Type::BOOLEAN, _) => false,
1466        // Dictionary encoding was not enabled in PARQUET 1.0
1467        (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_1_0) => false,
1468        (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => true,
1469        _ => true,
1470    }
1471}
1472
1473#[inline]
1474fn compare_greater_unsigned_int<T: ParquetValueType>(a: &T, b: &T) -> bool {
1475    a.as_u64().unwrap() > b.as_u64().unwrap()
1476}
1477
1478#[inline]
1479fn compare_greater_f16(a: &[u8], b: &[u8]) -> bool {
1480    let a = f16::from_le_bytes(a.try_into().unwrap());
1481    let b = f16::from_le_bytes(b.try_into().unwrap());
1482    a > b
1483}
1484
1485/// Signed comparison of bytes arrays
1486fn compare_greater_byte_array_decimals(a: &[u8], b: &[u8]) -> bool {
1487    let a_length = a.len();
1488    let b_length = b.len();
1489
1490    if a_length == 0 || b_length == 0 {
1491        return a_length > 0;
1492    }
1493
1494    let first_a: u8 = a[0];
1495    let first_b: u8 = b[0];
1496
1497    // We can short circuit for different signed numbers or
1498    // for equal length bytes arrays that have different first bytes.
1499    // The equality requirement is necessary for sign extension cases.
1500    // 0xFF10 should be equal to 0x10 (due to big endian sign extension).
1501    if (0x80 & first_a) != (0x80 & first_b) || (a_length == b_length && first_a != first_b) {
1502        return (first_a as i8) > (first_b as i8);
1503    }
1504
1505    // When the lengths are unequal and the numbers are of the same
1506    // sign we need to do comparison by sign extending the shorter
1507    // value first, and once we get to equal sized arrays, lexicographical
1508    // unsigned comparison of everything but the first byte is sufficient.
1509
1510    let extension: u8 = if (first_a as i8) < 0 { 0xFF } else { 0 };
1511
1512    if a_length != b_length {
1513        let not_equal = if a_length > b_length {
1514            let lead_length = a_length - b_length;
1515            a[0..lead_length].iter().any(|&x| x != extension)
1516        } else {
1517            let lead_length = b_length - a_length;
1518            b[0..lead_length].iter().any(|&x| x != extension)
1519        };
1520
1521        if not_equal {
1522            let negative_values: bool = (first_a as i8) < 0;
1523            let a_longer: bool = a_length > b_length;
1524            return if negative_values { !a_longer } else { a_longer };
1525        }
1526    }
1527
1528    (a[1..]) > (b[1..])
1529}
1530
1531/// Truncate a UTF-8 slice to the longest prefix that is still a valid UTF-8 string,
1532/// while being less than `length` bytes and non-empty. Returns `None` if truncation
1533/// is not possible within those constraints.
1534///
1535/// The caller guarantees that data.len() > length.
1536fn truncate_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
1537    let split = (1..=length).rfind(|x| data.is_char_boundary(*x))?;
1538    Some(data.as_bytes()[..split].to_vec())
1539}
1540
1541/// Truncate a UTF-8 slice and increment it's final character. The returned value is the
1542/// longest such slice that is still a valid UTF-8 string while being less than `length`
1543/// bytes and non-empty. Returns `None` if no such transformation is possible.
1544///
1545/// The caller guarantees that data.len() > length.
1546fn truncate_and_increment_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
1547    // UTF-8 is max 4 bytes, so start search 3 back from desired length
1548    let lower_bound = length.saturating_sub(3);
1549    let split = (lower_bound..=length).rfind(|x| data.is_char_boundary(*x))?;
1550    increment_utf8(data.get(..split)?)
1551}
1552
1553/// Increment the final character in a UTF-8 string in such a way that the returned result
1554/// is still a valid UTF-8 string. The returned string may be shorter than the input if the
1555/// last character(s) cannot be incremented (due to overflow or producing invalid code points).
1556/// Returns `None` if the string cannot be incremented.
1557///
1558/// Note that this implementation will not promote an N-byte code point to (N+1) bytes.
1559fn increment_utf8(data: &str) -> Option<Vec<u8>> {
1560    for (idx, original_char) in data.char_indices().rev() {
1561        let original_len = original_char.len_utf8();
1562        if let Some(next_char) = char::from_u32(original_char as u32 + 1) {
1563            // do not allow increasing byte width of incremented char
1564            if next_char.len_utf8() == original_len {
1565                let mut result = data.as_bytes()[..idx + original_len].to_vec();
1566                next_char.encode_utf8(&mut result[idx..]);
1567                return Some(result);
1568            }
1569        }
1570    }
1571
1572    None
1573}
1574
1575/// Try and increment the bytes from right to left.
1576///
1577/// Returns `None` if all bytes are set to `u8::MAX`.
1578fn increment(mut data: Vec<u8>) -> Option<Vec<u8>> {
1579    for byte in data.iter_mut().rev() {
1580        let (incremented, overflow) = byte.overflowing_add(1);
1581        *byte = incremented;
1582
1583        if !overflow {
1584            return Some(data);
1585        }
1586    }
1587
1588    None
1589}
1590
1591#[cfg(test)]
1592mod tests {
1593    use crate::{
1594        file::{properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, writer::SerializedFileWriter},
1595        schema::parser::parse_message_type,
1596    };
1597    use core::str;
1598    use rand::distr::uniform::SampleUniform;
1599    use std::{fs::File, sync::Arc};
1600
1601    use crate::column::{
1602        page::PageReader,
1603        reader::{ColumnReaderImpl, get_column_reader, get_typed_column_reader},
1604    };
1605    use crate::file::writer::TrackedWrite;
1606    use crate::file::{
1607        properties::ReaderProperties, reader::SerializedPageReader, writer::SerializedPageWriter,
1608    };
1609    use crate::schema::types::{ColumnPath, Type as SchemaType};
1610    use crate::util::test_common::rand_gen::random_numbers_range;
1611
1612    use super::*;
1613
1614    #[test]
1615    fn test_column_writer_inconsistent_def_rep_length() {
1616        let page_writer = get_test_page_writer();
1617        let props = Default::default();
1618        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 1, props);
1619        let res = writer.write_batch(&[1, 2, 3, 4], Some(&[1, 1, 1]), Some(&[0, 0]));
1620        assert!(res.is_err());
1621        if let Err(err) = res {
1622            assert_eq!(
1623                format!("{err}"),
1624                "Parquet error: Inconsistent length of definition and repetition levels: 3 != 2"
1625            );
1626        }
1627    }
1628
1629    #[test]
1630    fn test_column_writer_invalid_def_levels() {
1631        let page_writer = get_test_page_writer();
1632        let props = Default::default();
1633        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
1634        let res = writer.write_batch(&[1, 2, 3, 4], None, None);
1635        assert!(res.is_err());
1636        if let Err(err) = res {
1637            assert_eq!(
1638                format!("{err}"),
1639                "Parquet error: Definition levels are required, because max definition level = 1"
1640            );
1641        }
1642    }
1643
1644    #[test]
1645    fn test_column_writer_invalid_rep_levels() {
1646        let page_writer = get_test_page_writer();
1647        let props = Default::default();
1648        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 1, props);
1649        let res = writer.write_batch(&[1, 2, 3, 4], None, None);
1650        assert!(res.is_err());
1651        if let Err(err) = res {
1652            assert_eq!(
1653                format!("{err}"),
1654                "Parquet error: Repetition levels are required, because max repetition level = 1"
1655            );
1656        }
1657    }
1658
1659    #[test]
1660    fn test_column_writer_not_enough_values_to_write() {
1661        let page_writer = get_test_page_writer();
1662        let props = Default::default();
1663        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
1664        let res = writer.write_batch(&[1, 2], Some(&[1, 1, 1, 1]), None);
1665        assert!(res.is_err());
1666        if let Err(err) = res {
1667            assert_eq!(
1668                format!("{err}"),
1669                "Parquet error: Expected to write 4 values, but have only 2"
1670            );
1671        }
1672    }
1673
1674    #[test]
1675    fn test_column_writer_write_only_one_dictionary_page() {
1676        let page_writer = get_test_page_writer();
1677        let props = Default::default();
1678        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
1679        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
1680        // First page should be correctly written.
1681        writer.add_data_page().unwrap();
1682        writer.write_dictionary_page().unwrap();
1683        let err = writer.write_dictionary_page().unwrap_err().to_string();
1684        assert_eq!(err, "Parquet error: Dictionary encoder is not set");
1685    }
1686
1687    #[test]
1688    fn test_column_writer_error_when_writing_disabled_dictionary() {
1689        let page_writer = get_test_page_writer();
1690        let props = Arc::new(
1691            WriterProperties::builder()
1692                .set_dictionary_enabled(false)
1693                .build(),
1694        );
1695        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
1696        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
1697        let err = writer.write_dictionary_page().unwrap_err().to_string();
1698        assert_eq!(err, "Parquet error: Dictionary encoder is not set");
1699    }
1700
1701    #[test]
1702    fn test_column_writer_boolean_type_does_not_support_dictionary() {
1703        let page_writer = get_test_page_writer();
1704        let props = Arc::new(
1705            WriterProperties::builder()
1706                .set_dictionary_enabled(true)
1707                .build(),
1708        );
1709        let mut writer = get_test_column_writer::<BoolType>(page_writer, 0, 0, props);
1710        writer
1711            .write_batch(&[true, false, true, false], None, None)
1712            .unwrap();
1713
1714        let r = writer.close().unwrap();
1715        // PlainEncoder uses bit writer to write boolean values, which all fit into 1
1716        // byte.
1717        assert_eq!(r.bytes_written, 1);
1718        assert_eq!(r.rows_written, 4);
1719
1720        let metadata = r.metadata;
1721        assert_eq!(
1722            metadata.encodings().collect::<Vec<_>>(),
1723            vec![Encoding::PLAIN, Encoding::RLE]
1724        );
1725        assert_eq!(metadata.num_values(), 4); // just values
1726        assert_eq!(metadata.dictionary_page_offset(), None);
1727    }
1728
1729    #[test]
1730    fn test_column_writer_default_encoding_support_bool() {
1731        check_encoding_write_support::<BoolType>(
1732            WriterVersion::PARQUET_1_0,
1733            true,
1734            &[true, false],
1735            None,
1736            &[Encoding::PLAIN, Encoding::RLE],
1737            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1738        );
1739        check_encoding_write_support::<BoolType>(
1740            WriterVersion::PARQUET_1_0,
1741            false,
1742            &[true, false],
1743            None,
1744            &[Encoding::PLAIN, Encoding::RLE],
1745            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1746        );
1747        check_encoding_write_support::<BoolType>(
1748            WriterVersion::PARQUET_2_0,
1749            true,
1750            &[true, false],
1751            None,
1752            &[Encoding::RLE],
1753            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
1754        );
1755        check_encoding_write_support::<BoolType>(
1756            WriterVersion::PARQUET_2_0,
1757            false,
1758            &[true, false],
1759            None,
1760            &[Encoding::RLE],
1761            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
1762        );
1763    }
1764
1765    #[test]
1766    fn test_column_writer_default_encoding_support_int32() {
1767        check_encoding_write_support::<Int32Type>(
1768            WriterVersion::PARQUET_1_0,
1769            true,
1770            &[1, 2],
1771            Some(0),
1772            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1773            &[
1774                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1775                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1776            ],
1777        );
1778        check_encoding_write_support::<Int32Type>(
1779            WriterVersion::PARQUET_1_0,
1780            false,
1781            &[1, 2],
1782            None,
1783            &[Encoding::PLAIN, Encoding::RLE],
1784            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1785        );
1786        check_encoding_write_support::<Int32Type>(
1787            WriterVersion::PARQUET_2_0,
1788            true,
1789            &[1, 2],
1790            Some(0),
1791            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1792            &[
1793                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1794                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
1795            ],
1796        );
1797        check_encoding_write_support::<Int32Type>(
1798            WriterVersion::PARQUET_2_0,
1799            false,
1800            &[1, 2],
1801            None,
1802            &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
1803            &[encoding_stats(
1804                PageType::DATA_PAGE_V2,
1805                Encoding::DELTA_BINARY_PACKED,
1806                1,
1807            )],
1808        );
1809    }
1810
1811    #[test]
1812    fn test_column_writer_default_encoding_support_int64() {
1813        check_encoding_write_support::<Int64Type>(
1814            WriterVersion::PARQUET_1_0,
1815            true,
1816            &[1, 2],
1817            Some(0),
1818            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1819            &[
1820                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1821                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1822            ],
1823        );
1824        check_encoding_write_support::<Int64Type>(
1825            WriterVersion::PARQUET_1_0,
1826            false,
1827            &[1, 2],
1828            None,
1829            &[Encoding::PLAIN, Encoding::RLE],
1830            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1831        );
1832        check_encoding_write_support::<Int64Type>(
1833            WriterVersion::PARQUET_2_0,
1834            true,
1835            &[1, 2],
1836            Some(0),
1837            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1838            &[
1839                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1840                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
1841            ],
1842        );
1843        check_encoding_write_support::<Int64Type>(
1844            WriterVersion::PARQUET_2_0,
1845            false,
1846            &[1, 2],
1847            None,
1848            &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
1849            &[encoding_stats(
1850                PageType::DATA_PAGE_V2,
1851                Encoding::DELTA_BINARY_PACKED,
1852                1,
1853            )],
1854        );
1855    }
1856
1857    #[test]
1858    fn test_column_writer_default_encoding_support_int96() {
1859        check_encoding_write_support::<Int96Type>(
1860            WriterVersion::PARQUET_1_0,
1861            true,
1862            &[Int96::from(vec![1, 2, 3])],
1863            Some(0),
1864            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1865            &[
1866                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1867                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1868            ],
1869        );
1870        check_encoding_write_support::<Int96Type>(
1871            WriterVersion::PARQUET_1_0,
1872            false,
1873            &[Int96::from(vec![1, 2, 3])],
1874            None,
1875            &[Encoding::PLAIN, Encoding::RLE],
1876            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1877        );
1878        check_encoding_write_support::<Int96Type>(
1879            WriterVersion::PARQUET_2_0,
1880            true,
1881            &[Int96::from(vec![1, 2, 3])],
1882            Some(0),
1883            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1884            &[
1885                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1886                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
1887            ],
1888        );
1889        check_encoding_write_support::<Int96Type>(
1890            WriterVersion::PARQUET_2_0,
1891            false,
1892            &[Int96::from(vec![1, 2, 3])],
1893            None,
1894            &[Encoding::PLAIN, Encoding::RLE],
1895            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
1896        );
1897    }
1898
1899    #[test]
1900    fn test_column_writer_default_encoding_support_float() {
1901        check_encoding_write_support::<FloatType>(
1902            WriterVersion::PARQUET_1_0,
1903            true,
1904            &[1.0, 2.0],
1905            Some(0),
1906            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1907            &[
1908                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1909                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1910            ],
1911        );
1912        check_encoding_write_support::<FloatType>(
1913            WriterVersion::PARQUET_1_0,
1914            false,
1915            &[1.0, 2.0],
1916            None,
1917            &[Encoding::PLAIN, Encoding::RLE],
1918            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1919        );
1920        check_encoding_write_support::<FloatType>(
1921            WriterVersion::PARQUET_2_0,
1922            true,
1923            &[1.0, 2.0],
1924            Some(0),
1925            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1926            &[
1927                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1928                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
1929            ],
1930        );
1931        check_encoding_write_support::<FloatType>(
1932            WriterVersion::PARQUET_2_0,
1933            false,
1934            &[1.0, 2.0],
1935            None,
1936            &[Encoding::PLAIN, Encoding::RLE],
1937            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
1938        );
1939    }
1940
1941    #[test]
1942    fn test_column_writer_default_encoding_support_double() {
1943        check_encoding_write_support::<DoubleType>(
1944            WriterVersion::PARQUET_1_0,
1945            true,
1946            &[1.0, 2.0],
1947            Some(0),
1948            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1949            &[
1950                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1951                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1952            ],
1953        );
1954        check_encoding_write_support::<DoubleType>(
1955            WriterVersion::PARQUET_1_0,
1956            false,
1957            &[1.0, 2.0],
1958            None,
1959            &[Encoding::PLAIN, Encoding::RLE],
1960            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
1961        );
1962        check_encoding_write_support::<DoubleType>(
1963            WriterVersion::PARQUET_2_0,
1964            true,
1965            &[1.0, 2.0],
1966            Some(0),
1967            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1968            &[
1969                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1970                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
1971            ],
1972        );
1973        check_encoding_write_support::<DoubleType>(
1974            WriterVersion::PARQUET_2_0,
1975            false,
1976            &[1.0, 2.0],
1977            None,
1978            &[Encoding::PLAIN, Encoding::RLE],
1979            &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
1980        );
1981    }
1982
1983    #[test]
1984    fn test_column_writer_default_encoding_support_byte_array() {
1985        check_encoding_write_support::<ByteArrayType>(
1986            WriterVersion::PARQUET_1_0,
1987            true,
1988            &[ByteArray::from(vec![1u8])],
1989            Some(0),
1990            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1991            &[
1992                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
1993                encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
1994            ],
1995        );
1996        check_encoding_write_support::<ByteArrayType>(
1997            WriterVersion::PARQUET_1_0,
1998            false,
1999            &[ByteArray::from(vec![1u8])],
2000            None,
2001            &[Encoding::PLAIN, Encoding::RLE],
2002            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2003        );
2004        check_encoding_write_support::<ByteArrayType>(
2005            WriterVersion::PARQUET_2_0,
2006            true,
2007            &[ByteArray::from(vec![1u8])],
2008            Some(0),
2009            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2010            &[
2011                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2012                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2013            ],
2014        );
2015        check_encoding_write_support::<ByteArrayType>(
2016            WriterVersion::PARQUET_2_0,
2017            false,
2018            &[ByteArray::from(vec![1u8])],
2019            None,
2020            &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2021            &[encoding_stats(
2022                PageType::DATA_PAGE_V2,
2023                Encoding::DELTA_BYTE_ARRAY,
2024                1,
2025            )],
2026        );
2027    }
2028
2029    #[test]
2030    fn test_column_writer_default_encoding_support_fixed_len_byte_array() {
2031        check_encoding_write_support::<FixedLenByteArrayType>(
2032            WriterVersion::PARQUET_1_0,
2033            true,
2034            &[ByteArray::from(vec![1u8]).into()],
2035            None,
2036            &[Encoding::PLAIN, Encoding::RLE],
2037            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2038        );
2039        check_encoding_write_support::<FixedLenByteArrayType>(
2040            WriterVersion::PARQUET_1_0,
2041            false,
2042            &[ByteArray::from(vec![1u8]).into()],
2043            None,
2044            &[Encoding::PLAIN, Encoding::RLE],
2045            &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2046        );
2047        check_encoding_write_support::<FixedLenByteArrayType>(
2048            WriterVersion::PARQUET_2_0,
2049            true,
2050            &[ByteArray::from(vec![1u8]).into()],
2051            Some(0),
2052            &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2053            &[
2054                encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2055                encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2056            ],
2057        );
2058        check_encoding_write_support::<FixedLenByteArrayType>(
2059            WriterVersion::PARQUET_2_0,
2060            false,
2061            &[ByteArray::from(vec![1u8]).into()],
2062            None,
2063            &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2064            &[encoding_stats(
2065                PageType::DATA_PAGE_V2,
2066                Encoding::DELTA_BYTE_ARRAY,
2067                1,
2068            )],
2069        );
2070    }
2071
2072    #[test]
2073    fn test_column_writer_check_metadata() {
2074        let page_writer = get_test_page_writer();
2075        let props = Default::default();
2076        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2077        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2078
2079        let r = writer.close().unwrap();
2080        assert_eq!(r.bytes_written, 20);
2081        assert_eq!(r.rows_written, 4);
2082
2083        let metadata = r.metadata;
2084        assert_eq!(
2085            metadata.encodings().collect::<Vec<_>>(),
2086            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2087        );
2088        assert_eq!(metadata.num_values(), 4);
2089        assert_eq!(metadata.compressed_size(), 20);
2090        assert_eq!(metadata.uncompressed_size(), 20);
2091        assert_eq!(metadata.data_page_offset(), 0);
2092        assert_eq!(metadata.dictionary_page_offset(), Some(0));
2093        if let Some(stats) = metadata.statistics() {
2094            assert_eq!(stats.null_count_opt(), Some(0));
2095            assert_eq!(stats.distinct_count_opt(), None);
2096            if let Statistics::Int32(stats) = stats {
2097                assert_eq!(stats.min_opt().unwrap(), &1);
2098                assert_eq!(stats.max_opt().unwrap(), &4);
2099            } else {
2100                panic!("expecting Statistics::Int32");
2101            }
2102        } else {
2103            panic!("metadata missing statistics");
2104        }
2105    }
2106
2107    #[test]
2108    fn test_column_writer_check_byte_array_min_max() {
2109        let page_writer = get_test_page_writer();
2110        let props = Default::default();
2111        let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
2112        writer
2113            .write_batch(
2114                &[
2115                    ByteArray::from(vec![
2116                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2117                        35u8, 231u8, 90u8, 0u8, 0u8,
2118                    ]),
2119                    ByteArray::from(vec![
2120                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 228u8, 62u8, 146u8,
2121                        152u8, 177u8, 56u8, 0u8, 0u8,
2122                    ]),
2123                    ByteArray::from(vec![
2124                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8,
2125                        0u8,
2126                    ]),
2127                    ByteArray::from(vec![
2128                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2129                        44u8, 0u8, 0u8,
2130                    ]),
2131                ],
2132                None,
2133                None,
2134            )
2135            .unwrap();
2136        let metadata = writer.close().unwrap().metadata;
2137        if let Some(stats) = metadata.statistics() {
2138            if let Statistics::ByteArray(stats) = stats {
2139                assert_eq!(
2140                    stats.min_opt().unwrap(),
2141                    &ByteArray::from(vec![
2142                        255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2143                        35u8, 231u8, 90u8, 0u8, 0u8,
2144                    ])
2145                );
2146                assert_eq!(
2147                    stats.max_opt().unwrap(),
2148                    &ByteArray::from(vec![
2149                        0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2150                        44u8, 0u8, 0u8,
2151                    ])
2152                );
2153            } else {
2154                panic!("expecting Statistics::ByteArray");
2155            }
2156        } else {
2157            panic!("metadata missing statistics");
2158        }
2159    }
2160
2161    #[test]
2162    fn test_column_writer_uint32_converted_type_min_max() {
2163        let page_writer = get_test_page_writer();
2164        let props = Default::default();
2165        let mut writer = get_test_unsigned_int_given_as_converted_column_writer::<Int32Type>(
2166            page_writer,
2167            0,
2168            0,
2169            props,
2170        );
2171        writer.write_batch(&[0, 1, 2, 3, 4, 5], None, None).unwrap();
2172        let metadata = writer.close().unwrap().metadata;
2173        if let Some(stats) = metadata.statistics() {
2174            if let Statistics::Int32(stats) = stats {
2175                assert_eq!(stats.min_opt().unwrap(), &0,);
2176                assert_eq!(stats.max_opt().unwrap(), &5,);
2177            } else {
2178                panic!("expecting Statistics::Int32");
2179            }
2180        } else {
2181            panic!("metadata missing statistics");
2182        }
2183    }
2184
2185    #[test]
2186    fn test_column_writer_precalculated_statistics() {
2187        let page_writer = get_test_page_writer();
2188        let props = Arc::new(
2189            WriterProperties::builder()
2190                .set_statistics_enabled(EnabledStatistics::Chunk)
2191                .build(),
2192        );
2193        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2194        writer
2195            .write_batch_with_statistics(
2196                &[1, 2, 3, 4],
2197                None,
2198                None,
2199                Some(&-17),
2200                Some(&9000),
2201                Some(55),
2202            )
2203            .unwrap();
2204
2205        let r = writer.close().unwrap();
2206        assert_eq!(r.bytes_written, 20);
2207        assert_eq!(r.rows_written, 4);
2208
2209        let metadata = r.metadata;
2210        assert_eq!(
2211            metadata.encodings().collect::<Vec<_>>(),
2212            vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2213        );
2214        assert_eq!(metadata.num_values(), 4);
2215        assert_eq!(metadata.compressed_size(), 20);
2216        assert_eq!(metadata.uncompressed_size(), 20);
2217        assert_eq!(metadata.data_page_offset(), 0);
2218        assert_eq!(metadata.dictionary_page_offset(), Some(0));
2219        if let Some(stats) = metadata.statistics() {
2220            assert_eq!(stats.null_count_opt(), Some(0));
2221            assert_eq!(stats.distinct_count_opt().unwrap_or(0), 55);
2222            if let Statistics::Int32(stats) = stats {
2223                assert_eq!(stats.min_opt().unwrap(), &-17);
2224                assert_eq!(stats.max_opt().unwrap(), &9000);
2225            } else {
2226                panic!("expecting Statistics::Int32");
2227            }
2228        } else {
2229            panic!("metadata missing statistics");
2230        }
2231    }
2232
2233    #[test]
2234    fn test_mixed_precomputed_statistics() {
2235        let mut buf = Vec::with_capacity(100);
2236        let mut write = TrackedWrite::new(&mut buf);
2237        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2238        let props = Arc::new(
2239            WriterProperties::builder()
2240                .set_write_page_header_statistics(true)
2241                .build(),
2242        );
2243        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2244
2245        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2246        writer
2247            .write_batch_with_statistics(&[5, 6, 7], None, None, Some(&5), Some(&7), Some(3))
2248            .unwrap();
2249
2250        let r = writer.close().unwrap();
2251
2252        let stats = r.metadata.statistics().unwrap();
2253        assert_eq!(stats.min_bytes_opt().unwrap(), 1_i32.to_le_bytes());
2254        assert_eq!(stats.max_bytes_opt().unwrap(), 7_i32.to_le_bytes());
2255        assert_eq!(stats.null_count_opt(), Some(0));
2256        assert!(stats.distinct_count_opt().is_none());
2257
2258        drop(write);
2259
2260        let props = ReaderProperties::builder()
2261            .set_backward_compatible_lz4(false)
2262            .set_read_page_statistics(true)
2263            .build();
2264        let reader = SerializedPageReader::new_with_properties(
2265            Arc::new(Bytes::from(buf)),
2266            &r.metadata,
2267            r.rows_written as usize,
2268            None,
2269            Arc::new(props),
2270        )
2271        .unwrap();
2272
2273        let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2274        assert_eq!(pages.len(), 2);
2275
2276        assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2277        assert_eq!(pages[1].page_type(), PageType::DATA_PAGE);
2278
2279        let page_statistics = pages[1].statistics().unwrap();
2280        assert_eq!(
2281            page_statistics.min_bytes_opt().unwrap(),
2282            1_i32.to_le_bytes()
2283        );
2284        assert_eq!(
2285            page_statistics.max_bytes_opt().unwrap(),
2286            7_i32.to_le_bytes()
2287        );
2288        assert_eq!(page_statistics.null_count_opt(), Some(0));
2289        assert!(page_statistics.distinct_count_opt().is_none());
2290    }
2291
2292    #[test]
2293    fn test_disabled_statistics() {
2294        let mut buf = Vec::with_capacity(100);
2295        let mut write = TrackedWrite::new(&mut buf);
2296        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2297        let props = WriterProperties::builder()
2298            .set_statistics_enabled(EnabledStatistics::None)
2299            .set_writer_version(WriterVersion::PARQUET_2_0)
2300            .build();
2301        let props = Arc::new(props);
2302
2303        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2304        writer
2305            .write_batch(&[1, 2, 3, 4], Some(&[1, 0, 0, 1, 1, 1]), None)
2306            .unwrap();
2307
2308        let r = writer.close().unwrap();
2309        assert!(r.metadata.statistics().is_none());
2310
2311        drop(write);
2312
2313        let props = ReaderProperties::builder()
2314            .set_backward_compatible_lz4(false)
2315            .build();
2316        let reader = SerializedPageReader::new_with_properties(
2317            Arc::new(Bytes::from(buf)),
2318            &r.metadata,
2319            r.rows_written as usize,
2320            None,
2321            Arc::new(props),
2322        )
2323        .unwrap();
2324
2325        let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2326        assert_eq!(pages.len(), 2);
2327
2328        assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2329        assert_eq!(pages[1].page_type(), PageType::DATA_PAGE_V2);
2330
2331        match &pages[1] {
2332            Page::DataPageV2 {
2333                num_values,
2334                num_nulls,
2335                num_rows,
2336                statistics,
2337                ..
2338            } => {
2339                assert_eq!(*num_values, 6);
2340                assert_eq!(*num_nulls, 2);
2341                assert_eq!(*num_rows, 6);
2342                assert!(statistics.is_none());
2343            }
2344            _ => unreachable!(),
2345        }
2346    }
2347
2348    #[test]
2349    fn test_column_writer_empty_column_roundtrip() {
2350        let props = Default::default();
2351        column_roundtrip::<Int32Type>(props, &[], None, None);
2352    }
2353
2354    #[test]
2355    fn test_column_writer_non_nullable_values_roundtrip() {
2356        let props = Default::default();
2357        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 0, 0);
2358    }
2359
2360    #[test]
2361    fn test_column_writer_nullable_non_repeated_values_roundtrip() {
2362        let props = Default::default();
2363        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 0);
2364    }
2365
2366    #[test]
2367    fn test_column_writer_nullable_repeated_values_roundtrip() {
2368        let props = Default::default();
2369        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2370    }
2371
2372    #[test]
2373    fn test_column_writer_dictionary_fallback_small_data_page() {
2374        let props = WriterProperties::builder()
2375            .set_dictionary_page_size_limit(32)
2376            .set_data_page_size_limit(32)
2377            .build();
2378        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2379    }
2380
2381    #[test]
2382    fn test_column_writer_small_write_batch_size() {
2383        for i in &[1usize, 2, 5, 10, 11, 1023] {
2384            let props = WriterProperties::builder().set_write_batch_size(*i).build();
2385
2386            column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2387        }
2388    }
2389
2390    #[test]
2391    fn test_column_writer_dictionary_disabled_v1() {
2392        let props = WriterProperties::builder()
2393            .set_writer_version(WriterVersion::PARQUET_1_0)
2394            .set_dictionary_enabled(false)
2395            .build();
2396        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2397    }
2398
2399    #[test]
2400    fn test_column_writer_dictionary_disabled_v2() {
2401        let props = WriterProperties::builder()
2402            .set_writer_version(WriterVersion::PARQUET_2_0)
2403            .set_dictionary_enabled(false)
2404            .build();
2405        column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2406    }
2407
2408    #[test]
2409    fn test_column_writer_compression_v1() {
2410        let props = WriterProperties::builder()
2411            .set_writer_version(WriterVersion::PARQUET_1_0)
2412            .set_compression(Compression::SNAPPY)
2413            .build();
2414        column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
2415    }
2416
2417    #[test]
2418    fn test_column_writer_compression_v2() {
2419        let props = WriterProperties::builder()
2420            .set_writer_version(WriterVersion::PARQUET_2_0)
2421            .set_compression(Compression::SNAPPY)
2422            .build();
2423        column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
2424    }
2425
2426    #[test]
2427    fn test_column_writer_v2_compression_ratio_threshold() {
2428        fn write_v2_page(threshold: f64) -> bool {
2429            let mut buf = Vec::with_capacity(4096);
2430            let mut write = TrackedWrite::new(&mut buf);
2431            let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2432            let props = Arc::new(
2433                WriterProperties::builder()
2434                    .set_writer_version(WriterVersion::PARQUET_2_0)
2435                    .set_compression(Compression::SNAPPY)
2436                    .set_dictionary_enabled(false)
2437                    .set_data_page_v2_compression_ratio_threshold(threshold)
2438                    .build(),
2439            );
2440
2441            let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2442            let values: Vec<i32> = vec![42; 4096];
2443            writer.write_batch(&values, None, None).unwrap();
2444            let r = writer.close().unwrap();
2445            drop(write);
2446
2447            let reader_props = ReaderProperties::builder()
2448                .set_backward_compatible_lz4(false)
2449                .build();
2450            let reader = SerializedPageReader::new_with_properties(
2451                Arc::new(Bytes::from(buf)),
2452                &r.metadata,
2453                r.rows_written as usize,
2454                None,
2455                Arc::new(reader_props),
2456            )
2457            .unwrap();
2458            let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2459            let data_page = pages
2460                .iter()
2461                .find(|p| p.page_type() == PageType::DATA_PAGE_V2)
2462                .expect("expected a v2 data page");
2463            match data_page {
2464                Page::DataPageV2 { is_compressed, .. } => *is_compressed,
2465                _ => unreachable!(),
2466            }
2467        }
2468
2469        // Default threshold keeps the compressed buffer for constant data.
2470        assert!(write_v2_page(1.0));
2471        // A strict threshold (require >1000x reduction) discards it.
2472        assert!(!write_v2_page(0.001));
2473    }
2474
2475    #[test]
2476    fn test_column_writer_add_data_pages_with_dict() {
2477        // ARROW-5129: Test verifies that we add data page in case of dictionary encoding
2478        // and no fallback occurred so far.
2479        let mut file = tempfile::tempfile().unwrap();
2480        let mut write = TrackedWrite::new(&mut file);
2481        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2482        let props = Arc::new(
2483            WriterProperties::builder()
2484                .set_data_page_size_limit(10)
2485                .set_write_batch_size(3) // write 3 values at a time
2486                .build(),
2487        );
2488        let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
2489        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2490        writer.write_batch(data, None, None).unwrap();
2491        let r = writer.close().unwrap();
2492
2493        drop(write);
2494
2495        // Read pages and check the sequence
2496        let props = ReaderProperties::builder()
2497            .set_backward_compatible_lz4(false)
2498            .build();
2499        let mut page_reader = Box::new(
2500            SerializedPageReader::new_with_properties(
2501                Arc::new(file),
2502                &r.metadata,
2503                r.rows_written as usize,
2504                None,
2505                Arc::new(props),
2506            )
2507            .unwrap(),
2508        );
2509        let mut res = Vec::new();
2510        while let Some(page) = page_reader.get_next_page().unwrap() {
2511            res.push((page.page_type(), page.num_values(), page.buffer().len()));
2512        }
2513        assert_eq!(
2514            res,
2515            vec![
2516                (PageType::DICTIONARY_PAGE, 10, 40),
2517                (PageType::DATA_PAGE, 9, 10),
2518                (PageType::DATA_PAGE, 1, 3),
2519            ]
2520        );
2521        assert_eq!(
2522            r.metadata.page_encoding_stats(),
2523            Some(&vec![
2524                PageEncodingStats {
2525                    page_type: PageType::DICTIONARY_PAGE,
2526                    encoding: Encoding::PLAIN,
2527                    count: 1
2528                },
2529                PageEncodingStats {
2530                    page_type: PageType::DATA_PAGE,
2531                    encoding: Encoding::RLE_DICTIONARY,
2532                    count: 2,
2533                }
2534            ])
2535        );
2536    }
2537
2538    #[test]
2539    fn test_column_writer_column_data_page_size_limit() {
2540        let props = Arc::new(
2541            WriterProperties::builder()
2542                .set_writer_version(WriterVersion::PARQUET_1_0)
2543                .set_dictionary_enabled(false)
2544                .set_data_page_size_limit(1000)
2545                .set_column_data_page_size_limit(ColumnPath::from("col"), 10)
2546                .set_write_batch_size(3)
2547                .build(),
2548        );
2549        let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
2550
2551        let col_values =
2552            write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data);
2553        let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data);
2554
2555        assert_eq!(col_values, vec![3, 3, 3, 1]);
2556        assert_eq!(other_values, vec![10]);
2557    }
2558
2559    #[test]
2560    fn test_bool_statistics() {
2561        let stats = statistics_roundtrip::<BoolType>(&[true, false, false, true]);
2562        // Booleans have an unsigned sort order and so are not compatible
2563        // with the deprecated `min` and `max` statistics
2564        assert!(!stats.is_min_max_backwards_compatible());
2565        if let Statistics::Boolean(stats) = stats {
2566            assert_eq!(stats.min_opt().unwrap(), &false);
2567            assert_eq!(stats.max_opt().unwrap(), &true);
2568        } else {
2569            panic!("expecting Statistics::Boolean, got {stats:?}");
2570        }
2571    }
2572
2573    #[test]
2574    fn test_int32_statistics() {
2575        let stats = statistics_roundtrip::<Int32Type>(&[-1, 3, -2, 2]);
2576        assert!(stats.is_min_max_backwards_compatible());
2577        if let Statistics::Int32(stats) = stats {
2578            assert_eq!(stats.min_opt().unwrap(), &-2);
2579            assert_eq!(stats.max_opt().unwrap(), &3);
2580        } else {
2581            panic!("expecting Statistics::Int32, got {stats:?}");
2582        }
2583    }
2584
2585    #[test]
2586    fn test_int64_statistics() {
2587        let stats = statistics_roundtrip::<Int64Type>(&[-1, 3, -2, 2]);
2588        assert!(stats.is_min_max_backwards_compatible());
2589        if let Statistics::Int64(stats) = stats {
2590            assert_eq!(stats.min_opt().unwrap(), &-2);
2591            assert_eq!(stats.max_opt().unwrap(), &3);
2592        } else {
2593            panic!("expecting Statistics::Int64, got {stats:?}");
2594        }
2595    }
2596
2597    #[test]
2598    fn test_int96_statistics() {
2599        let input = vec![
2600            Int96::from(vec![1, 20, 30]),
2601            Int96::from(vec![3, 20, 10]),
2602            Int96::from(vec![0, 20, 30]),
2603            Int96::from(vec![2, 20, 30]),
2604        ]
2605        .into_iter()
2606        .collect::<Vec<Int96>>();
2607
2608        let stats = statistics_roundtrip::<Int96Type>(&input);
2609        assert!(!stats.is_min_max_backwards_compatible());
2610        if let Statistics::Int96(stats) = stats {
2611            assert_eq!(stats.min_opt().unwrap(), &Int96::from(vec![3, 20, 10]));
2612            assert_eq!(stats.max_opt().unwrap(), &Int96::from(vec![2, 20, 30]));
2613        } else {
2614            panic!("expecting Statistics::Int96, got {stats:?}");
2615        }
2616    }
2617
2618    #[test]
2619    fn test_float_statistics() {
2620        let stats = statistics_roundtrip::<FloatType>(&[-1.0, 3.0, -2.0, 2.0]);
2621        assert!(stats.is_min_max_backwards_compatible());
2622        if let Statistics::Float(stats) = stats {
2623            assert_eq!(stats.min_opt().unwrap(), &-2.0);
2624            assert_eq!(stats.max_opt().unwrap(), &3.0);
2625        } else {
2626            panic!("expecting Statistics::Float, got {stats:?}");
2627        }
2628    }
2629
2630    #[test]
2631    fn test_double_statistics() {
2632        let stats = statistics_roundtrip::<DoubleType>(&[-1.0, 3.0, -2.0, 2.0]);
2633        assert!(stats.is_min_max_backwards_compatible());
2634        if let Statistics::Double(stats) = stats {
2635            assert_eq!(stats.min_opt().unwrap(), &-2.0);
2636            assert_eq!(stats.max_opt().unwrap(), &3.0);
2637        } else {
2638            panic!("expecting Statistics::Double, got {stats:?}");
2639        }
2640    }
2641
2642    #[test]
2643    fn test_byte_array_statistics() {
2644        let input = ["aawaa", "zz", "aaw", "m", "qrs"]
2645            .iter()
2646            .map(|&s| s.into())
2647            .collect::<Vec<_>>();
2648
2649        let stats = statistics_roundtrip::<ByteArrayType>(&input);
2650        assert!(!stats.is_min_max_backwards_compatible());
2651        if let Statistics::ByteArray(stats) = stats {
2652            assert_eq!(stats.min_opt().unwrap(), &ByteArray::from("aaw"));
2653            assert_eq!(stats.max_opt().unwrap(), &ByteArray::from("zz"));
2654        } else {
2655            panic!("expecting Statistics::ByteArray, got {stats:?}");
2656        }
2657    }
2658
2659    #[test]
2660    fn test_fixed_len_byte_array_statistics() {
2661        let input = ["aawaa", "zz   ", "aaw  ", "m    ", "qrs  "]
2662            .iter()
2663            .map(|&s| ByteArray::from(s).into())
2664            .collect::<Vec<_>>();
2665
2666        let stats = statistics_roundtrip::<FixedLenByteArrayType>(&input);
2667        assert!(!stats.is_min_max_backwards_compatible());
2668        if let Statistics::FixedLenByteArray(stats) = stats {
2669            let expected_min: FixedLenByteArray = ByteArray::from("aaw  ").into();
2670            assert_eq!(stats.min_opt().unwrap(), &expected_min);
2671            let expected_max: FixedLenByteArray = ByteArray::from("zz   ").into();
2672            assert_eq!(stats.max_opt().unwrap(), &expected_max);
2673        } else {
2674            panic!("expecting Statistics::FixedLenByteArray, got {stats:?}");
2675        }
2676    }
2677
2678    #[test]
2679    fn test_column_writer_check_float16_min_max() {
2680        let input = [
2681            -f16::ONE,
2682            f16::from_f32(3.0),
2683            -f16::from_f32(2.0),
2684            f16::from_f32(2.0),
2685        ]
2686        .into_iter()
2687        .map(|s| ByteArray::from(s).into())
2688        .collect::<Vec<_>>();
2689
2690        let stats = float16_statistics_roundtrip(&input);
2691        assert!(stats.is_min_max_backwards_compatible());
2692        assert_eq!(
2693            stats.min_opt().unwrap(),
2694            &ByteArray::from(-f16::from_f32(2.0))
2695        );
2696        assert_eq!(
2697            stats.max_opt().unwrap(),
2698            &ByteArray::from(f16::from_f32(3.0))
2699        );
2700    }
2701
2702    #[test]
2703    fn test_column_writer_check_float16_nan_middle() {
2704        let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
2705            .into_iter()
2706            .map(|s| ByteArray::from(s).into())
2707            .collect::<Vec<_>>();
2708
2709        let stats = float16_statistics_roundtrip(&input);
2710        assert!(stats.is_min_max_backwards_compatible());
2711        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
2712        assert_eq!(
2713            stats.max_opt().unwrap(),
2714            &ByteArray::from(f16::ONE + f16::ONE)
2715        );
2716    }
2717
2718    #[test]
2719    fn test_float16_statistics_nan_middle() {
2720        let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
2721            .into_iter()
2722            .map(|s| ByteArray::from(s).into())
2723            .collect::<Vec<_>>();
2724
2725        let stats = float16_statistics_roundtrip(&input);
2726        assert!(stats.is_min_max_backwards_compatible());
2727        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
2728        assert_eq!(
2729            stats.max_opt().unwrap(),
2730            &ByteArray::from(f16::ONE + f16::ONE)
2731        );
2732    }
2733
2734    #[test]
2735    fn test_float16_statistics_nan_start() {
2736        let input = [f16::NAN, f16::ONE, f16::ONE + f16::ONE]
2737            .into_iter()
2738            .map(|s| ByteArray::from(s).into())
2739            .collect::<Vec<_>>();
2740
2741        let stats = float16_statistics_roundtrip(&input);
2742        assert!(stats.is_min_max_backwards_compatible());
2743        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
2744        assert_eq!(
2745            stats.max_opt().unwrap(),
2746            &ByteArray::from(f16::ONE + f16::ONE)
2747        );
2748    }
2749
2750    #[test]
2751    fn test_float16_statistics_nan_only() {
2752        let input = [f16::NAN, f16::NAN]
2753            .into_iter()
2754            .map(|s| ByteArray::from(s).into())
2755            .collect::<Vec<_>>();
2756
2757        let stats = float16_statistics_roundtrip(&input);
2758        assert!(stats.min_bytes_opt().is_none());
2759        assert!(stats.max_bytes_opt().is_none());
2760        assert!(stats.is_min_max_backwards_compatible());
2761    }
2762
2763    #[test]
2764    fn test_float16_statistics_zero_only() {
2765        let input = [f16::ZERO]
2766            .into_iter()
2767            .map(|s| ByteArray::from(s).into())
2768            .collect::<Vec<_>>();
2769
2770        let stats = float16_statistics_roundtrip(&input);
2771        assert!(stats.is_min_max_backwards_compatible());
2772        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
2773        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO));
2774    }
2775
2776    #[test]
2777    fn test_float16_statistics_neg_zero_only() {
2778        let input = [f16::NEG_ZERO]
2779            .into_iter()
2780            .map(|s| ByteArray::from(s).into())
2781            .collect::<Vec<_>>();
2782
2783        let stats = float16_statistics_roundtrip(&input);
2784        assert!(stats.is_min_max_backwards_compatible());
2785        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
2786        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO));
2787    }
2788
2789    #[test]
2790    fn test_float16_statistics_zero_min() {
2791        let input = [f16::ZERO, f16::ONE, f16::NAN, f16::PI]
2792            .into_iter()
2793            .map(|s| ByteArray::from(s).into())
2794            .collect::<Vec<_>>();
2795
2796        let stats = float16_statistics_roundtrip(&input);
2797        assert!(stats.is_min_max_backwards_compatible());
2798        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
2799        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::PI));
2800    }
2801
2802    #[test]
2803    fn test_float16_statistics_neg_zero_max() {
2804        let input = [f16::NEG_ZERO, f16::NEG_ONE, f16::NAN, -f16::PI]
2805            .into_iter()
2806            .map(|s| ByteArray::from(s).into())
2807            .collect::<Vec<_>>();
2808
2809        let stats = float16_statistics_roundtrip(&input);
2810        assert!(stats.is_min_max_backwards_compatible());
2811        assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(-f16::PI));
2812        assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO));
2813    }
2814
2815    #[test]
2816    fn test_float_statistics_nan_middle() {
2817        let stats = statistics_roundtrip::<FloatType>(&[1.0, f32::NAN, 2.0]);
2818        assert!(stats.is_min_max_backwards_compatible());
2819        if let Statistics::Float(stats) = stats {
2820            assert_eq!(stats.min_opt().unwrap(), &1.0);
2821            assert_eq!(stats.max_opt().unwrap(), &2.0);
2822        } else {
2823            panic!("expecting Statistics::Float");
2824        }
2825    }
2826
2827    #[test]
2828    fn test_float_statistics_nan_start() {
2829        let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, 1.0, 2.0]);
2830        assert!(stats.is_min_max_backwards_compatible());
2831        if let Statistics::Float(stats) = stats {
2832            assert_eq!(stats.min_opt().unwrap(), &1.0);
2833            assert_eq!(stats.max_opt().unwrap(), &2.0);
2834        } else {
2835            panic!("expecting Statistics::Float");
2836        }
2837    }
2838
2839    #[test]
2840    fn test_float_statistics_nan_only() {
2841        let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, f32::NAN]);
2842        assert!(stats.min_bytes_opt().is_none());
2843        assert!(stats.max_bytes_opt().is_none());
2844        assert!(stats.is_min_max_backwards_compatible());
2845        assert!(matches!(stats, Statistics::Float(_)));
2846    }
2847
2848    #[test]
2849    fn test_float_statistics_zero_only() {
2850        let stats = statistics_roundtrip::<FloatType>(&[0.0]);
2851        assert!(stats.is_min_max_backwards_compatible());
2852        if let Statistics::Float(stats) = stats {
2853            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2854            assert!(stats.min_opt().unwrap().is_sign_negative());
2855            assert_eq!(stats.max_opt().unwrap(), &0.0);
2856            assert!(stats.max_opt().unwrap().is_sign_positive());
2857        } else {
2858            panic!("expecting Statistics::Float");
2859        }
2860    }
2861
2862    #[test]
2863    fn test_float_statistics_neg_zero_only() {
2864        let stats = statistics_roundtrip::<FloatType>(&[-0.0]);
2865        assert!(stats.is_min_max_backwards_compatible());
2866        if let Statistics::Float(stats) = stats {
2867            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2868            assert!(stats.min_opt().unwrap().is_sign_negative());
2869            assert_eq!(stats.max_opt().unwrap(), &0.0);
2870            assert!(stats.max_opt().unwrap().is_sign_positive());
2871        } else {
2872            panic!("expecting Statistics::Float");
2873        }
2874    }
2875
2876    #[test]
2877    fn test_float_statistics_zero_min() {
2878        let stats = statistics_roundtrip::<FloatType>(&[0.0, 1.0, f32::NAN, 2.0]);
2879        assert!(stats.is_min_max_backwards_compatible());
2880        if let Statistics::Float(stats) = stats {
2881            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2882            assert!(stats.min_opt().unwrap().is_sign_negative());
2883            assert_eq!(stats.max_opt().unwrap(), &2.0);
2884        } else {
2885            panic!("expecting Statistics::Float");
2886        }
2887    }
2888
2889    #[test]
2890    fn test_float_statistics_neg_zero_max() {
2891        let stats = statistics_roundtrip::<FloatType>(&[-0.0, -1.0, f32::NAN, -2.0]);
2892        assert!(stats.is_min_max_backwards_compatible());
2893        if let Statistics::Float(stats) = stats {
2894            assert_eq!(stats.min_opt().unwrap(), &-2.0);
2895            assert_eq!(stats.max_opt().unwrap(), &0.0);
2896            assert!(stats.max_opt().unwrap().is_sign_positive());
2897        } else {
2898            panic!("expecting Statistics::Float");
2899        }
2900    }
2901
2902    #[test]
2903    fn test_double_statistics_nan_middle() {
2904        let stats = statistics_roundtrip::<DoubleType>(&[1.0, f64::NAN, 2.0]);
2905        assert!(stats.is_min_max_backwards_compatible());
2906        if let Statistics::Double(stats) = stats {
2907            assert_eq!(stats.min_opt().unwrap(), &1.0);
2908            assert_eq!(stats.max_opt().unwrap(), &2.0);
2909        } else {
2910            panic!("expecting Statistics::Double");
2911        }
2912    }
2913
2914    #[test]
2915    fn test_double_statistics_nan_start() {
2916        let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, 1.0, 2.0]);
2917        assert!(stats.is_min_max_backwards_compatible());
2918        if let Statistics::Double(stats) = stats {
2919            assert_eq!(stats.min_opt().unwrap(), &1.0);
2920            assert_eq!(stats.max_opt().unwrap(), &2.0);
2921        } else {
2922            panic!("expecting Statistics::Double");
2923        }
2924    }
2925
2926    #[test]
2927    fn test_double_statistics_nan_only() {
2928        let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, f64::NAN]);
2929        assert!(stats.min_bytes_opt().is_none());
2930        assert!(stats.max_bytes_opt().is_none());
2931        assert!(matches!(stats, Statistics::Double(_)));
2932        assert!(stats.is_min_max_backwards_compatible());
2933    }
2934
2935    #[test]
2936    fn test_double_statistics_zero_only() {
2937        let stats = statistics_roundtrip::<DoubleType>(&[0.0]);
2938        assert!(stats.is_min_max_backwards_compatible());
2939        if let Statistics::Double(stats) = stats {
2940            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2941            assert!(stats.min_opt().unwrap().is_sign_negative());
2942            assert_eq!(stats.max_opt().unwrap(), &0.0);
2943            assert!(stats.max_opt().unwrap().is_sign_positive());
2944        } else {
2945            panic!("expecting Statistics::Double");
2946        }
2947    }
2948
2949    #[test]
2950    fn test_double_statistics_neg_zero_only() {
2951        let stats = statistics_roundtrip::<DoubleType>(&[-0.0]);
2952        assert!(stats.is_min_max_backwards_compatible());
2953        if let Statistics::Double(stats) = stats {
2954            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2955            assert!(stats.min_opt().unwrap().is_sign_negative());
2956            assert_eq!(stats.max_opt().unwrap(), &0.0);
2957            assert!(stats.max_opt().unwrap().is_sign_positive());
2958        } else {
2959            panic!("expecting Statistics::Double");
2960        }
2961    }
2962
2963    #[test]
2964    fn test_double_statistics_zero_min() {
2965        let stats = statistics_roundtrip::<DoubleType>(&[0.0, 1.0, f64::NAN, 2.0]);
2966        assert!(stats.is_min_max_backwards_compatible());
2967        if let Statistics::Double(stats) = stats {
2968            assert_eq!(stats.min_opt().unwrap(), &-0.0);
2969            assert!(stats.min_opt().unwrap().is_sign_negative());
2970            assert_eq!(stats.max_opt().unwrap(), &2.0);
2971        } else {
2972            panic!("expecting Statistics::Double");
2973        }
2974    }
2975
2976    #[test]
2977    fn test_double_statistics_neg_zero_max() {
2978        let stats = statistics_roundtrip::<DoubleType>(&[-0.0, -1.0, f64::NAN, -2.0]);
2979        assert!(stats.is_min_max_backwards_compatible());
2980        if let Statistics::Double(stats) = stats {
2981            assert_eq!(stats.min_opt().unwrap(), &-2.0);
2982            assert_eq!(stats.max_opt().unwrap(), &0.0);
2983            assert!(stats.max_opt().unwrap().is_sign_positive());
2984        } else {
2985            panic!("expecting Statistics::Double");
2986        }
2987    }
2988
2989    #[test]
2990    fn test_compare_greater_byte_array_decimals() {
2991        assert!(!compare_greater_byte_array_decimals(&[], &[],),);
2992        assert!(compare_greater_byte_array_decimals(&[1u8,], &[],),);
2993        assert!(!compare_greater_byte_array_decimals(&[], &[1u8,],),);
2994        assert!(compare_greater_byte_array_decimals(&[1u8,], &[0u8,],),);
2995        assert!(!compare_greater_byte_array_decimals(&[1u8,], &[1u8,],),);
2996        assert!(compare_greater_byte_array_decimals(&[1u8, 0u8,], &[0u8,],),);
2997        assert!(!compare_greater_byte_array_decimals(
2998            &[0u8, 1u8,],
2999            &[1u8, 0u8,],
3000        ),);
3001        assert!(!compare_greater_byte_array_decimals(
3002            &[255u8, 35u8, 0u8, 0u8,],
3003            &[0u8,],
3004        ),);
3005        assert!(compare_greater_byte_array_decimals(
3006            &[0u8,],
3007            &[255u8, 35u8, 0u8, 0u8,],
3008        ),);
3009    }
3010
3011    #[test]
3012    fn test_column_index_with_null_pages() {
3013        // write a single page of all nulls
3014        let page_writer = get_test_page_writer();
3015        let props = Default::default();
3016        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
3017        writer.write_batch(&[], Some(&[0, 0, 0, 0]), None).unwrap();
3018
3019        let r = writer.close().unwrap();
3020        assert!(r.column_index.is_some());
3021        let col_idx = r.column_index.unwrap();
3022        let col_idx = match col_idx {
3023            ColumnIndexMetaData::INT32(col_idx) => col_idx,
3024            _ => panic!("wrong stats type"),
3025        };
3026        // null_pages should be true for page 0
3027        assert!(col_idx.is_null_page(0));
3028        // min and max should be empty byte arrays
3029        assert!(col_idx.min_value(0).is_none());
3030        assert!(col_idx.max_value(0).is_none());
3031        // null_counts should be defined and be 4 for page 0
3032        assert!(col_idx.null_count(0).is_some());
3033        assert_eq!(col_idx.null_count(0), Some(4));
3034        // there is no repetition so rep histogram should be absent
3035        assert!(col_idx.repetition_level_histogram(0).is_none());
3036        // definition_level_histogram should be present and should be 0:4, 1:0
3037        assert!(col_idx.definition_level_histogram(0).is_some());
3038        assert_eq!(col_idx.definition_level_histogram(0).unwrap(), &[4, 0]);
3039    }
3040
3041    #[test]
3042    fn test_column_offset_index_metadata() {
3043        // write data
3044        // and check the offset index and column index
3045        let page_writer = get_test_page_writer();
3046        let props = Default::default();
3047        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3048        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
3049        // first page
3050        writer.flush_data_pages().unwrap();
3051        // second page
3052        writer.write_batch(&[4, 8, 2, -5], None, None).unwrap();
3053
3054        let r = writer.close().unwrap();
3055        let column_index = r.column_index.unwrap();
3056        let offset_index = r.offset_index.unwrap();
3057
3058        assert_eq!(8, r.rows_written);
3059
3060        // column index
3061        let column_index = match column_index {
3062            ColumnIndexMetaData::INT32(column_index) => column_index,
3063            _ => panic!("wrong stats type"),
3064        };
3065        assert_eq!(2, column_index.num_pages());
3066        assert_eq!(2, offset_index.page_locations.len());
3067        assert_eq!(BoundaryOrder::UNORDERED, column_index.boundary_order);
3068        for idx in 0..2 {
3069            assert!(!column_index.is_null_page(idx));
3070            assert_eq!(0, column_index.null_count(0).unwrap());
3071        }
3072
3073        if let Some(stats) = r.metadata.statistics() {
3074            assert_eq!(stats.null_count_opt(), Some(0));
3075            assert_eq!(stats.distinct_count_opt(), None);
3076            if let Statistics::Int32(stats) = stats {
3077                // first page is [1,2,3,4]
3078                // second page is [-5,2,4,8]
3079                // note that we don't increment here, as this is a non BinaryArray type.
3080                assert_eq!(stats.min_opt(), column_index.min_value(1));
3081                assert_eq!(stats.max_opt(), column_index.max_value(1));
3082            } else {
3083                panic!("expecting Statistics::Int32");
3084            }
3085        } else {
3086            panic!("metadata missing statistics");
3087        }
3088
3089        // page location
3090        assert_eq!(0, offset_index.page_locations[0].first_row_index);
3091        assert_eq!(4, offset_index.page_locations[1].first_row_index);
3092    }
3093
3094    /// Verify min/max value truncation in the column index works as expected
3095    #[test]
3096    fn test_column_offset_index_metadata_truncating() {
3097        // write data
3098        // and check the offset index and column index
3099        let page_writer = get_test_page_writer();
3100        let props = WriterProperties::builder()
3101            .set_statistics_truncate_length(None) // disable column index truncation
3102            .build()
3103            .into();
3104        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
3105
3106        let mut data = vec![FixedLenByteArray::default(); 3];
3107        // This is the expected min value - "aaa..."
3108        data[0].set_data(Bytes::from(vec![97_u8; 200]));
3109        // This is the expected max value - "ZZZ..."
3110        data[1].set_data(Bytes::from(vec![112_u8; 200]));
3111        data[2].set_data(Bytes::from(vec![98_u8; 200]));
3112
3113        writer.write_batch(&data, None, None).unwrap();
3114
3115        writer.flush_data_pages().unwrap();
3116
3117        let r = writer.close().unwrap();
3118        let column_index = r.column_index.unwrap();
3119        let offset_index = r.offset_index.unwrap();
3120
3121        let column_index = match column_index {
3122            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index,
3123            _ => panic!("wrong stats type"),
3124        };
3125
3126        assert_eq!(3, r.rows_written);
3127
3128        // column index
3129        assert_eq!(1, column_index.num_pages());
3130        assert_eq!(1, offset_index.page_locations.len());
3131        assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
3132        assert!(!column_index.is_null_page(0));
3133        assert_eq!(Some(0), column_index.null_count(0));
3134
3135        if let Some(stats) = r.metadata.statistics() {
3136            assert_eq!(stats.null_count_opt(), Some(0));
3137            assert_eq!(stats.distinct_count_opt(), None);
3138            if let Statistics::FixedLenByteArray(stats) = stats {
3139                let column_index_min_value = column_index.min_value(0).unwrap();
3140                let column_index_max_value = column_index.max_value(0).unwrap();
3141
3142                // Column index stats are truncated, while the column chunk's aren't.
3143                assert_ne!(stats.min_bytes_opt().unwrap(), column_index_min_value);
3144                assert_ne!(stats.max_bytes_opt().unwrap(), column_index_max_value);
3145
3146                assert_eq!(
3147                    column_index_min_value.len(),
3148                    DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
3149                );
3150                assert_eq!(column_index_min_value, &[97_u8; 64]);
3151                assert_eq!(
3152                    column_index_max_value.len(),
3153                    DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
3154                );
3155
3156                // We expect the last byte to be incremented
3157                assert_eq!(
3158                    *column_index_max_value.last().unwrap(),
3159                    *column_index_max_value.first().unwrap() + 1
3160                );
3161            } else {
3162                panic!("expecting Statistics::FixedLenByteArray");
3163            }
3164        } else {
3165            panic!("metadata missing statistics");
3166        }
3167    }
3168
3169    #[test]
3170    fn test_column_offset_index_truncating_spec_example() {
3171        // write data
3172        // and check the offset index and column index
3173        let page_writer = get_test_page_writer();
3174
3175        // Truncate values at 1 byte
3176        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
3177        let props = Arc::new(builder.build());
3178        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
3179
3180        let mut data = vec![FixedLenByteArray::default(); 1];
3181        // This is the expected min value
3182        data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
3183
3184        writer.write_batch(&data, None, None).unwrap();
3185
3186        writer.flush_data_pages().unwrap();
3187
3188        let r = writer.close().unwrap();
3189        let column_index = r.column_index.unwrap();
3190        let offset_index = r.offset_index.unwrap();
3191
3192        let column_index = match column_index {
3193            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index,
3194            _ => panic!("wrong stats type"),
3195        };
3196
3197        assert_eq!(1, r.rows_written);
3198
3199        // column index
3200        assert_eq!(1, column_index.num_pages());
3201        assert_eq!(1, offset_index.page_locations.len());
3202        assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
3203        assert!(!column_index.is_null_page(0));
3204        assert_eq!(Some(0), column_index.null_count(0));
3205
3206        if let Some(stats) = r.metadata.statistics() {
3207            assert_eq!(stats.null_count_opt(), Some(0));
3208            assert_eq!(stats.distinct_count_opt(), None);
3209            if let Statistics::FixedLenByteArray(_stats) = stats {
3210                let column_index_min_value = column_index.min_value(0).unwrap();
3211                let column_index_max_value = column_index.max_value(0).unwrap();
3212
3213                assert_eq!(column_index_min_value.len(), 1);
3214                assert_eq!(column_index_max_value.len(), 1);
3215
3216                assert_eq!("B".as_bytes(), column_index_min_value);
3217                assert_eq!("C".as_bytes(), column_index_max_value);
3218
3219                assert_ne!(column_index_min_value, stats.min_bytes_opt().unwrap());
3220                assert_ne!(column_index_max_value, stats.max_bytes_opt().unwrap());
3221            } else {
3222                panic!("expecting Statistics::FixedLenByteArray");
3223            }
3224        } else {
3225            panic!("metadata missing statistics");
3226        }
3227    }
3228
3229    #[test]
3230    fn test_float16_min_max_no_truncation() {
3231        // Even if we set truncation to occur at 1 byte, we should not truncate for Float16
3232        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
3233        let props = Arc::new(builder.build());
3234        let page_writer = get_test_page_writer();
3235        let mut writer = get_test_float16_column_writer(page_writer, props);
3236
3237        let expected_value = f16::PI.to_le_bytes().to_vec();
3238        let data = vec![ByteArray::from(expected_value.clone()).into()];
3239        writer.write_batch(&data, None, None).unwrap();
3240        writer.flush_data_pages().unwrap();
3241
3242        let r = writer.close().unwrap();
3243
3244        // stats should still be written
3245        // ensure bytes weren't truncated for column index
3246        let column_index = r.column_index.unwrap();
3247        let column_index = match column_index {
3248            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index,
3249            _ => panic!("wrong stats type"),
3250        };
3251        let column_index_min_bytes = column_index.min_value(0).unwrap();
3252        let column_index_max_bytes = column_index.max_value(0).unwrap();
3253        assert_eq!(expected_value, column_index_min_bytes);
3254        assert_eq!(expected_value, column_index_max_bytes);
3255
3256        // ensure bytes weren't truncated for statistics
3257        let stats = r.metadata.statistics().unwrap();
3258        if let Statistics::FixedLenByteArray(stats) = stats {
3259            let stats_min_bytes = stats.min_bytes_opt().unwrap();
3260            let stats_max_bytes = stats.max_bytes_opt().unwrap();
3261            assert_eq!(expected_value, stats_min_bytes);
3262            assert_eq!(expected_value, stats_max_bytes);
3263        } else {
3264            panic!("expecting Statistics::FixedLenByteArray");
3265        }
3266    }
3267
3268    #[test]
3269    fn test_decimal_min_max_no_truncation() {
3270        // Even if we set truncation to occur at 1 byte, we should not truncate for Decimal
3271        let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
3272        let props = Arc::new(builder.build());
3273        let page_writer = get_test_page_writer();
3274        let mut writer =
3275            get_test_decimals_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
3276
3277        let expected_value = vec![
3278            255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, 35u8,
3279            231u8, 90u8, 0u8, 0u8,
3280        ];
3281        let data = vec![ByteArray::from(expected_value.clone()).into()];
3282        writer.write_batch(&data, None, None).unwrap();
3283        writer.flush_data_pages().unwrap();
3284
3285        let r = writer.close().unwrap();
3286
3287        // stats should still be written
3288        // ensure bytes weren't truncated for column index
3289        let column_index = r.column_index.unwrap();
3290        let column_index = match column_index {
3291            ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) => column_index,
3292            _ => panic!("wrong stats type"),
3293        };
3294        let column_index_min_bytes = column_index.min_value(0).unwrap();
3295        let column_index_max_bytes = column_index.max_value(0).unwrap();
3296        assert_eq!(expected_value, column_index_min_bytes);
3297        assert_eq!(expected_value, column_index_max_bytes);
3298
3299        // ensure bytes weren't truncated for statistics
3300        let stats = r.metadata.statistics().unwrap();
3301        if let Statistics::FixedLenByteArray(stats) = stats {
3302            let stats_min_bytes = stats.min_bytes_opt().unwrap();
3303            let stats_max_bytes = stats.max_bytes_opt().unwrap();
3304            assert_eq!(expected_value, stats_min_bytes);
3305            assert_eq!(expected_value, stats_max_bytes);
3306        } else {
3307            panic!("expecting Statistics::FixedLenByteArray");
3308        }
3309    }
3310
3311    #[test]
3312    fn test_statistics_truncating_byte_array_default() {
3313        let page_writer = get_test_page_writer();
3314
3315        // The default truncate length is 64 bytes
3316        let props = WriterProperties::builder().build().into();
3317        let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
3318
3319        let mut data = vec![ByteArray::default(); 1];
3320        data[0].set_data(Bytes::from(String::from(
3321            "This string is longer than 64 bytes, so it will almost certainly be truncated.",
3322        )));
3323        writer.write_batch(&data, None, None).unwrap();
3324        writer.flush_data_pages().unwrap();
3325
3326        let r = writer.close().unwrap();
3327
3328        assert_eq!(1, r.rows_written);
3329
3330        let stats = r.metadata.statistics().expect("statistics");
3331        if let Statistics::ByteArray(_stats) = stats {
3332            let min_value = _stats.min_opt().unwrap();
3333            let max_value = _stats.max_opt().unwrap();
3334
3335            assert!(!_stats.min_is_exact());
3336            assert!(!_stats.max_is_exact());
3337
3338            let expected_len = 64;
3339            assert_eq!(min_value.len(), expected_len);
3340            assert_eq!(max_value.len(), expected_len);
3341
3342            let expected_min =
3343                "This string is longer than 64 bytes, so it will almost certainly".as_bytes();
3344            assert_eq!(expected_min, min_value.as_bytes());
3345            // note the max value is different from the min value: the last byte is incremented
3346            let expected_max =
3347                "This string is longer than 64 bytes, so it will almost certainlz".as_bytes();
3348            assert_eq!(expected_max, max_value.as_bytes());
3349        } else {
3350            panic!("expecting Statistics::ByteArray");
3351        }
3352    }
3353
3354    #[test]
3355    fn test_statistics_truncating_byte_array() {
3356        let page_writer = get_test_page_writer();
3357
3358        const TEST_TRUNCATE_LENGTH: usize = 1;
3359
3360        // Truncate values at 1 byte
3361        let builder =
3362            WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
3363        let props = Arc::new(builder.build());
3364        let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
3365
3366        let mut data = vec![ByteArray::default(); 1];
3367        // This is the expected min value
3368        data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
3369
3370        writer.write_batch(&data, None, None).unwrap();
3371
3372        writer.flush_data_pages().unwrap();
3373
3374        let r = writer.close().unwrap();
3375
3376        assert_eq!(1, r.rows_written);
3377
3378        let stats = r.metadata.statistics().expect("statistics");
3379        assert_eq!(stats.null_count_opt(), Some(0));
3380        assert_eq!(stats.distinct_count_opt(), None);
3381        if let Statistics::ByteArray(_stats) = stats {
3382            let min_value = _stats.min_opt().unwrap();
3383            let max_value = _stats.max_opt().unwrap();
3384
3385            assert!(!_stats.min_is_exact());
3386            assert!(!_stats.max_is_exact());
3387
3388            assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
3389            assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
3390
3391            assert_eq!("B".as_bytes(), min_value.as_bytes());
3392            assert_eq!("C".as_bytes(), max_value.as_bytes());
3393        } else {
3394            panic!("expecting Statistics::ByteArray");
3395        }
3396    }
3397
3398    #[test]
3399    fn test_statistics_truncating_fixed_len_byte_array() {
3400        let page_writer = get_test_page_writer();
3401
3402        const TEST_TRUNCATE_LENGTH: usize = 1;
3403
3404        // Truncate values at 1 byte
3405        let builder =
3406            WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
3407        let props = Arc::new(builder.build());
3408        let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
3409
3410        let mut data = vec![FixedLenByteArray::default(); 1];
3411
3412        const PSEUDO_DECIMAL_VALUE: i128 = 6541894651216648486512564456564654;
3413        const PSEUDO_DECIMAL_BYTES: [u8; 16] = PSEUDO_DECIMAL_VALUE.to_be_bytes();
3414
3415        const EXPECTED_MIN: [u8; TEST_TRUNCATE_LENGTH] = [PSEUDO_DECIMAL_BYTES[0]]; // parquet specifies big-endian order for decimals
3416        const EXPECTED_MAX: [u8; TEST_TRUNCATE_LENGTH] =
3417            [PSEUDO_DECIMAL_BYTES[0].overflowing_add(1).0];
3418
3419        // This is the expected min value
3420        data[0].set_data(Bytes::from(PSEUDO_DECIMAL_BYTES.as_slice()));
3421
3422        writer.write_batch(&data, None, None).unwrap();
3423
3424        writer.flush_data_pages().unwrap();
3425
3426        let r = writer.close().unwrap();
3427
3428        assert_eq!(1, r.rows_written);
3429
3430        let stats = r.metadata.statistics().expect("statistics");
3431        assert_eq!(stats.null_count_opt(), Some(0));
3432        assert_eq!(stats.distinct_count_opt(), None);
3433        if let Statistics::FixedLenByteArray(_stats) = stats {
3434            let min_value = _stats.min_opt().unwrap();
3435            let max_value = _stats.max_opt().unwrap();
3436
3437            assert!(!_stats.min_is_exact());
3438            assert!(!_stats.max_is_exact());
3439
3440            assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
3441            assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
3442
3443            assert_eq!(EXPECTED_MIN.as_slice(), min_value.as_bytes());
3444            assert_eq!(EXPECTED_MAX.as_slice(), max_value.as_bytes());
3445
3446            let reconstructed_min = i128::from_be_bytes([
3447                min_value.as_bytes()[0],
3448                0,
3449                0,
3450                0,
3451                0,
3452                0,
3453                0,
3454                0,
3455                0,
3456                0,
3457                0,
3458                0,
3459                0,
3460                0,
3461                0,
3462                0,
3463            ]);
3464
3465            let reconstructed_max = i128::from_be_bytes([
3466                max_value.as_bytes()[0],
3467                0,
3468                0,
3469                0,
3470                0,
3471                0,
3472                0,
3473                0,
3474                0,
3475                0,
3476                0,
3477                0,
3478                0,
3479                0,
3480                0,
3481                0,
3482            ]);
3483
3484            // check that the inner value is correctly bounded by the min/max
3485            println!("min: {reconstructed_min} {PSEUDO_DECIMAL_VALUE}");
3486            assert!(reconstructed_min <= PSEUDO_DECIMAL_VALUE);
3487            println!("max {reconstructed_max} {PSEUDO_DECIMAL_VALUE}");
3488            assert!(reconstructed_max >= PSEUDO_DECIMAL_VALUE);
3489        } else {
3490            panic!("expecting Statistics::FixedLenByteArray");
3491        }
3492    }
3493
3494    #[test]
3495    fn test_send() {
3496        fn test<T: Send>() {}
3497        test::<ColumnWriterImpl<Int32Type>>();
3498    }
3499
3500    #[test]
3501    fn test_increment() {
3502        let v = increment(vec![0, 0, 0]).unwrap();
3503        assert_eq!(&v, &[0, 0, 1]);
3504
3505        // Handle overflow
3506        let v = increment(vec![0, 255, 255]).unwrap();
3507        assert_eq!(&v, &[1, 0, 0]);
3508
3509        // Return `None` if all bytes are u8::MAX
3510        let v = increment(vec![255, 255, 255]);
3511        assert!(v.is_none());
3512    }
3513
3514    #[test]
3515    fn test_increment_utf8() {
3516        let test_inc = |o: &str, expected: &str| {
3517            if let Ok(v) = String::from_utf8(increment_utf8(o).unwrap()) {
3518                // Got the expected result...
3519                assert_eq!(v, expected);
3520                // and it's greater than the original string
3521                assert!(*v > *o);
3522                // Also show that BinaryArray level comparison works here
3523                let mut greater = ByteArray::new();
3524                greater.set_data(Bytes::from(v));
3525                let mut original = ByteArray::new();
3526                original.set_data(Bytes::from(o.as_bytes().to_vec()));
3527                assert!(greater > original);
3528            } else {
3529                panic!("Expected incremented UTF8 string to also be valid.");
3530            }
3531        };
3532
3533        // Basic ASCII case
3534        test_inc("hello", "hellp");
3535
3536        // 1-byte ending in max 1-byte
3537        test_inc("a\u{7f}", "b");
3538
3539        // 1-byte max should not truncate as it would need 2-byte code points
3540        assert!(increment_utf8("\u{7f}\u{7f}").is_none());
3541
3542        // UTF8 string
3543        test_inc("❤️🧡💛💚💙💜", "❤️🧡💛💚💙💝");
3544
3545        // 2-byte without overflow
3546        test_inc("éééé", "éééê");
3547
3548        // 2-byte that overflows lowest byte
3549        test_inc("\u{ff}\u{ff}", "\u{ff}\u{100}");
3550
3551        // 2-byte ending in max 2-byte
3552        test_inc("a\u{7ff}", "b");
3553
3554        // Max 2-byte should not truncate as it would need 3-byte code points
3555        assert!(increment_utf8("\u{7ff}\u{7ff}").is_none());
3556
3557        // 3-byte without overflow [U+800, U+800] -> [U+800, U+801] (note that these
3558        // characters should render right to left).
3559        test_inc("ࠀࠀ", "ࠀࠁ");
3560
3561        // 3-byte ending in max 3-byte
3562        test_inc("a\u{ffff}", "b");
3563
3564        // Max 3-byte should not truncate as it would need 4-byte code points
3565        assert!(increment_utf8("\u{ffff}\u{ffff}").is_none());
3566
3567        // 4-byte without overflow
3568        test_inc("𐀀𐀀", "𐀀𐀁");
3569
3570        // 4-byte ending in max unicode
3571        test_inc("a\u{10ffff}", "b");
3572
3573        // Max 4-byte should not truncate
3574        assert!(increment_utf8("\u{10ffff}\u{10ffff}").is_none());
3575
3576        // Skip over surrogate pair range (0xD800..=0xDFFF)
3577        //test_inc("a\u{D7FF}", "a\u{e000}");
3578        test_inc("a\u{D7FF}", "b");
3579    }
3580
3581    #[test]
3582    fn test_truncate_utf8() {
3583        // No-op
3584        let data = "❤️🧡💛💚💙💜";
3585        let r = truncate_utf8(data, data.len()).unwrap();
3586        assert_eq!(r.len(), data.len());
3587        assert_eq!(&r, data.as_bytes());
3588
3589        // We slice it away from the UTF8 boundary
3590        let r = truncate_utf8(data, 13).unwrap();
3591        assert_eq!(r.len(), 10);
3592        assert_eq!(&r, "❤️🧡".as_bytes());
3593
3594        // One multi-byte code point, and a length shorter than it, so we can't slice it
3595        let r = truncate_utf8("\u{0836}", 1);
3596        assert!(r.is_none());
3597
3598        // Test truncate and increment for max bounds on UTF-8 statistics
3599        // 7-bit (i.e. ASCII)
3600        let r = truncate_and_increment_utf8("yyyyyyyyy", 8).unwrap();
3601        assert_eq!(&r, "yyyyyyyz".as_bytes());
3602
3603        // 2-byte without overflow
3604        let r = truncate_and_increment_utf8("ééééé", 7).unwrap();
3605        assert_eq!(&r, "ééê".as_bytes());
3606
3607        // 2-byte that overflows lowest byte
3608        let r = truncate_and_increment_utf8("\u{ff}\u{ff}\u{ff}\u{ff}\u{ff}", 8).unwrap();
3609        assert_eq!(&r, "\u{ff}\u{ff}\u{ff}\u{100}".as_bytes());
3610
3611        // max 2-byte should not truncate as it would need 3-byte code points
3612        let r = truncate_and_increment_utf8("߿߿߿߿߿", 8);
3613        assert!(r.is_none());
3614
3615        // 3-byte without overflow [U+800, U+800, U+800] -> [U+800, U+801] (note that these
3616        // characters should render right to left).
3617        let r = truncate_and_increment_utf8("ࠀࠀࠀࠀ", 8).unwrap();
3618        assert_eq!(&r, "ࠀࠁ".as_bytes());
3619
3620        // max 3-byte should not truncate as it would need 4-byte code points
3621        let r = truncate_and_increment_utf8("\u{ffff}\u{ffff}\u{ffff}", 8);
3622        assert!(r.is_none());
3623
3624        // 4-byte without overflow
3625        let r = truncate_and_increment_utf8("𐀀𐀀𐀀𐀀", 9).unwrap();
3626        assert_eq!(&r, "𐀀𐀁".as_bytes());
3627
3628        // max 4-byte should not truncate
3629        let r = truncate_and_increment_utf8("\u{10ffff}\u{10ffff}", 8);
3630        assert!(r.is_none());
3631    }
3632
3633    #[test]
3634    // Check fallback truncation of statistics that should be UTF-8, but aren't
3635    // (see https://github.com/apache/arrow-rs/pull/6870).
3636    fn test_byte_array_truncate_invalid_utf8_statistics() {
3637        let message_type = "
3638            message test_schema {
3639                OPTIONAL BYTE_ARRAY a (UTF8);
3640            }
3641        ";
3642        let schema = Arc::new(parse_message_type(message_type).unwrap());
3643
3644        // Create Vec<ByteArray> containing non-UTF8 bytes
3645        let data = vec![ByteArray::from(vec![128u8; 32]); 7];
3646        let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1];
3647        let file: File = tempfile::tempfile().unwrap();
3648        let props = Arc::new(
3649            WriterProperties::builder()
3650                .set_statistics_enabled(EnabledStatistics::Chunk)
3651                .set_statistics_truncate_length(Some(8))
3652                .build(),
3653        );
3654
3655        let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
3656        let mut row_group_writer = writer.next_row_group().unwrap();
3657
3658        let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
3659        col_writer
3660            .typed::<ByteArrayType>()
3661            .write_batch(&data, Some(&def_levels), None)
3662            .unwrap();
3663        col_writer.close().unwrap();
3664        row_group_writer.close().unwrap();
3665        let file_metadata = writer.close().unwrap();
3666        let stats = file_metadata.row_group(0).column(0).statistics().unwrap();
3667        assert!(!stats.max_is_exact());
3668        // Truncation of invalid UTF-8 should fall back to binary truncation, so last byte should
3669        // be incremented by 1.
3670        assert_eq!(
3671            stats.max_bytes_opt().map(|v| v.to_vec()),
3672            Some([128, 128, 128, 128, 128, 128, 128, 129].to_vec())
3673        );
3674    }
3675
3676    #[test]
3677    fn test_increment_max_binary_chars() {
3678        let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]);
3679        assert_eq!(&r.unwrap(), &[0xFF, 0xFE, 0xFE, 0x00, 0x00]);
3680
3681        let incremented = increment(vec![0xFF, 0xFF, 0xFF]);
3682        assert!(incremented.is_none())
3683    }
3684
3685    #[test]
3686    fn test_no_column_index_when_stats_disabled() {
3687        // https://github.com/apache/arrow-rs/issues/6010
3688        // Test that column index is not created/written for all-nulls column when page
3689        // statistics are disabled.
3690        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
3691        let props = Arc::new(
3692            WriterProperties::builder()
3693                .set_statistics_enabled(EnabledStatistics::None)
3694                .build(),
3695        );
3696        let column_writer = get_column_writer(descr, props, get_test_page_writer());
3697        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
3698
3699        let data = Vec::new();
3700        let def_levels = vec![0; 10];
3701        writer.write_batch(&data, Some(&def_levels), None).unwrap();
3702        writer.flush_data_pages().unwrap();
3703
3704        let column_close_result = writer.close().unwrap();
3705        assert!(column_close_result.offset_index.is_some());
3706        assert!(column_close_result.column_index.is_none());
3707    }
3708
3709    #[test]
3710    fn test_no_offset_index_when_disabled() {
3711        // Test that offset indexes can be disabled
3712        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
3713        let props = Arc::new(
3714            WriterProperties::builder()
3715                .set_statistics_enabled(EnabledStatistics::None)
3716                .set_offset_index_disabled(true)
3717                .build(),
3718        );
3719        let column_writer = get_column_writer(descr, props, get_test_page_writer());
3720        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
3721
3722        let data = Vec::new();
3723        let def_levels = vec![0; 10];
3724        writer.write_batch(&data, Some(&def_levels), None).unwrap();
3725        writer.flush_data_pages().unwrap();
3726
3727        let column_close_result = writer.close().unwrap();
3728        assert!(column_close_result.offset_index.is_none());
3729        assert!(column_close_result.column_index.is_none());
3730    }
3731
3732    #[test]
3733    fn test_offset_index_overridden() {
3734        // Test that offset indexes are not disabled when gathering page statistics
3735        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
3736        let props = Arc::new(
3737            WriterProperties::builder()
3738                .set_statistics_enabled(EnabledStatistics::Page)
3739                .set_offset_index_disabled(true)
3740                .build(),
3741        );
3742        let column_writer = get_column_writer(descr, props, get_test_page_writer());
3743        let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
3744
3745        let data = Vec::new();
3746        let def_levels = vec![0; 10];
3747        writer.write_batch(&data, Some(&def_levels), None).unwrap();
3748        writer.flush_data_pages().unwrap();
3749
3750        let column_close_result = writer.close().unwrap();
3751        assert!(column_close_result.offset_index.is_some());
3752        assert!(column_close_result.column_index.is_some());
3753    }
3754
3755    #[test]
3756    fn test_boundary_order() -> Result<()> {
3757        let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
3758        // min max both ascending
3759        let column_close_result = write_multiple_pages::<Int32Type>(
3760            &descr,
3761            &[
3762                &[Some(-10), Some(10)],
3763                &[Some(-5), Some(11)],
3764                &[None],
3765                &[Some(-5), Some(11)],
3766            ],
3767        )?;
3768        let boundary_order = column_close_result
3769            .column_index
3770            .unwrap()
3771            .get_boundary_order();
3772        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
3773
3774        // min max both descending
3775        let column_close_result = write_multiple_pages::<Int32Type>(
3776            &descr,
3777            &[
3778                &[Some(10), Some(11)],
3779                &[Some(5), Some(11)],
3780                &[None],
3781                &[Some(-5), Some(0)],
3782            ],
3783        )?;
3784        let boundary_order = column_close_result
3785            .column_index
3786            .unwrap()
3787            .get_boundary_order();
3788        assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
3789
3790        // min max both equal
3791        let column_close_result = write_multiple_pages::<Int32Type>(
3792            &descr,
3793            &[&[Some(10), Some(11)], &[None], &[Some(10), Some(11)]],
3794        )?;
3795        let boundary_order = column_close_result
3796            .column_index
3797            .unwrap()
3798            .get_boundary_order();
3799        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
3800
3801        // only nulls
3802        let column_close_result =
3803            write_multiple_pages::<Int32Type>(&descr, &[&[None], &[None], &[None]])?;
3804        let boundary_order = column_close_result
3805            .column_index
3806            .unwrap()
3807            .get_boundary_order();
3808        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
3809
3810        // one page
3811        let column_close_result =
3812            write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)]])?;
3813        let boundary_order = column_close_result
3814            .column_index
3815            .unwrap()
3816            .get_boundary_order();
3817        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
3818
3819        // one non-null page
3820        let column_close_result =
3821            write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)], &[None]])?;
3822        let boundary_order = column_close_result
3823            .column_index
3824            .unwrap()
3825            .get_boundary_order();
3826        assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
3827
3828        // min max both unordered
3829        let column_close_result = write_multiple_pages::<Int32Type>(
3830            &descr,
3831            &[
3832                &[Some(10), Some(11)],
3833                &[Some(11), Some(16)],
3834                &[None],
3835                &[Some(-5), Some(0)],
3836            ],
3837        )?;
3838        let boundary_order = column_close_result
3839            .column_index
3840            .unwrap()
3841            .get_boundary_order();
3842        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
3843
3844        // min max both ordered in different orders
3845        let column_close_result = write_multiple_pages::<Int32Type>(
3846            &descr,
3847            &[
3848                &[Some(1), Some(9)],
3849                &[Some(2), Some(8)],
3850                &[None],
3851                &[Some(3), Some(7)],
3852            ],
3853        )?;
3854        let boundary_order = column_close_result
3855            .column_index
3856            .unwrap()
3857            .get_boundary_order();
3858        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
3859
3860        Ok(())
3861    }
3862
3863    #[test]
3864    fn test_boundary_order_logical_type() -> Result<()> {
3865        // ensure that logical types account for different sort order than underlying
3866        // physical type representation
3867        let f16_descr = Arc::new(get_test_float16_column_descr(1, 0));
3868        let fba_descr = {
3869            let tpe = SchemaType::primitive_type_builder(
3870                "col",
3871                FixedLenByteArrayType::get_physical_type(),
3872            )
3873            .with_length(2)
3874            .build()?;
3875            Arc::new(ColumnDescriptor::new(
3876                Arc::new(tpe),
3877                1,
3878                0,
3879                ColumnPath::from("col"),
3880            ))
3881        };
3882
3883        let values: &[&[Option<FixedLenByteArray>]] = &[
3884            &[Some(FixedLenByteArray::from(ByteArray::from(f16::ONE)))],
3885            &[Some(FixedLenByteArray::from(ByteArray::from(f16::ZERO)))],
3886            &[Some(FixedLenByteArray::from(ByteArray::from(
3887                f16::NEG_ZERO,
3888            )))],
3889            &[Some(FixedLenByteArray::from(ByteArray::from(f16::NEG_ONE)))],
3890        ];
3891
3892        // f16 descending
3893        let column_close_result =
3894            write_multiple_pages::<FixedLenByteArrayType>(&f16_descr, values)?;
3895        let boundary_order = column_close_result
3896            .column_index
3897            .unwrap()
3898            .get_boundary_order();
3899        assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
3900
3901        // same bytes, but fba unordered
3902        let column_close_result =
3903            write_multiple_pages::<FixedLenByteArrayType>(&fba_descr, values)?;
3904        let boundary_order = column_close_result
3905            .column_index
3906            .unwrap()
3907            .get_boundary_order();
3908        assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
3909
3910        Ok(())
3911    }
3912
3913    #[test]
3914    fn test_interval_stats_should_not_have_min_max() {
3915        let input = [
3916            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
3917            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
3918            vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2],
3919        ]
3920        .into_iter()
3921        .map(|s| ByteArray::from(s).into())
3922        .collect::<Vec<_>>();
3923
3924        let page_writer = get_test_page_writer();
3925        let mut writer = get_test_interval_column_writer(page_writer);
3926        writer.write_batch(&input, None, None).unwrap();
3927
3928        let metadata = writer.close().unwrap().metadata;
3929        let stats = if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
3930            stats.clone()
3931        } else {
3932            panic!("metadata missing statistics");
3933        };
3934        assert!(stats.min_bytes_opt().is_none());
3935        assert!(stats.max_bytes_opt().is_none());
3936    }
3937
3938    #[test]
3939    #[cfg(feature = "arrow")]
3940    fn test_column_writer_get_estimated_total_bytes() {
3941        let page_writer = get_test_page_writer();
3942        let props = Default::default();
3943        let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3944        assert_eq!(writer.get_estimated_total_bytes(), 0);
3945
3946        writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
3947        writer.add_data_page().unwrap();
3948        let size_with_one_page = writer.get_estimated_total_bytes();
3949        assert_eq!(size_with_one_page, 20);
3950
3951        writer.write_batch(&[5, 6, 7, 8], None, None).unwrap();
3952        writer.add_data_page().unwrap();
3953        let size_with_two_pages = writer.get_estimated_total_bytes();
3954        // different pages have different compressed lengths
3955        assert_eq!(size_with_two_pages, 20 + 21);
3956    }
3957
3958    fn write_multiple_pages<T: DataType>(
3959        column_descr: &Arc<ColumnDescriptor>,
3960        pages: &[&[Option<T::T>]],
3961    ) -> Result<ColumnCloseResult> {
3962        let column_writer = get_column_writer(
3963            column_descr.clone(),
3964            Default::default(),
3965            get_test_page_writer(),
3966        );
3967        let mut writer = get_typed_column_writer::<T>(column_writer);
3968
3969        for &page in pages {
3970            let values = page.iter().filter_map(Clone::clone).collect::<Vec<_>>();
3971            let def_levels = page
3972                .iter()
3973                .map(|maybe_value| if maybe_value.is_some() { 1 } else { 0 })
3974                .collect::<Vec<_>>();
3975            writer.write_batch(&values, Some(&def_levels), None)?;
3976            writer.flush_data_pages()?;
3977        }
3978
3979        writer.close()
3980    }
3981
3982    /// Performs write-read roundtrip with randomly generated values and levels.
3983    /// `max_size` is maximum number of values or levels (if `max_def_level` > 0) to write
3984    /// for a column.
3985    fn column_roundtrip_random<T: DataType>(
3986        props: WriterProperties,
3987        max_size: usize,
3988        min_value: T::T,
3989        max_value: T::T,
3990        max_def_level: i16,
3991        max_rep_level: i16,
3992    ) where
3993        T::T: PartialOrd + SampleUniform + Copy,
3994    {
3995        let mut num_values: usize = 0;
3996
3997        let mut buf: Vec<i16> = Vec::new();
3998        let def_levels = if max_def_level > 0 {
3999            random_numbers_range(max_size, 0, max_def_level + 1, &mut buf);
4000            for &dl in &buf[..] {
4001                if dl == max_def_level {
4002                    num_values += 1;
4003                }
4004            }
4005            Some(&buf[..])
4006        } else {
4007            num_values = max_size;
4008            None
4009        };
4010
4011        let mut buf: Vec<i16> = Vec::new();
4012        let rep_levels = if max_rep_level > 0 {
4013            random_numbers_range(max_size, 0, max_rep_level + 1, &mut buf);
4014            buf[0] = 0; // Must start on record boundary
4015            Some(&buf[..])
4016        } else {
4017            None
4018        };
4019
4020        let mut values: Vec<T::T> = Vec::new();
4021        random_numbers_range(num_values, min_value, max_value, &mut values);
4022
4023        column_roundtrip::<T>(props, &values[..], def_levels, rep_levels);
4024    }
4025
4026    /// Performs write-read roundtrip and asserts written values and levels.
4027    fn column_roundtrip<T: DataType>(
4028        props: WriterProperties,
4029        values: &[T::T],
4030        def_levels: Option<&[i16]>,
4031        rep_levels: Option<&[i16]>,
4032    ) {
4033        let mut file = tempfile::tempfile().unwrap();
4034        let mut write = TrackedWrite::new(&mut file);
4035        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
4036
4037        let max_def_level = match def_levels {
4038            Some(buf) => *buf.iter().max().unwrap_or(&0i16),
4039            None => 0i16,
4040        };
4041
4042        let max_rep_level = match rep_levels {
4043            Some(buf) => *buf.iter().max().unwrap_or(&0i16),
4044            None => 0i16,
4045        };
4046
4047        let mut max_batch_size = values.len();
4048        if let Some(levels) = def_levels {
4049            max_batch_size = max_batch_size.max(levels.len());
4050        }
4051        if let Some(levels) = rep_levels {
4052            max_batch_size = max_batch_size.max(levels.len());
4053        }
4054
4055        let mut writer =
4056            get_test_column_writer::<T>(page_writer, max_def_level, max_rep_level, Arc::new(props));
4057
4058        let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap();
4059        assert_eq!(values_written, values.len());
4060        let result = writer.close().unwrap();
4061
4062        drop(write);
4063
4064        let props = ReaderProperties::builder()
4065            .set_backward_compatible_lz4(false)
4066            .build();
4067        let page_reader = Box::new(
4068            SerializedPageReader::new_with_properties(
4069                Arc::new(file),
4070                &result.metadata,
4071                result.rows_written as usize,
4072                None,
4073                Arc::new(props),
4074            )
4075            .unwrap(),
4076        );
4077        let mut reader = get_test_column_reader::<T>(page_reader, max_def_level, max_rep_level);
4078
4079        let mut actual_values = Vec::with_capacity(max_batch_size);
4080        let mut actual_def_levels = def_levels.map(|_| Vec::with_capacity(max_batch_size));
4081        let mut actual_rep_levels = rep_levels.map(|_| Vec::with_capacity(max_batch_size));
4082
4083        let (_, values_read, levels_read) = reader
4084            .read_records(
4085                max_batch_size,
4086                actual_def_levels.as_mut(),
4087                actual_rep_levels.as_mut(),
4088                &mut actual_values,
4089            )
4090            .unwrap();
4091
4092        // Assert values, definition and repetition levels.
4093
4094        assert_eq!(&actual_values[..values_read], values);
4095        match actual_def_levels {
4096            Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), def_levels),
4097            None => assert_eq!(None, def_levels),
4098        }
4099        match actual_rep_levels {
4100            Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), rep_levels),
4101            None => assert_eq!(None, rep_levels),
4102        }
4103
4104        // Assert written rows.
4105
4106        if let Some(levels) = actual_rep_levels {
4107            let mut actual_rows_written = 0;
4108            for l in levels {
4109                if l == 0 {
4110                    actual_rows_written += 1;
4111                }
4112            }
4113            assert_eq!(actual_rows_written, result.rows_written);
4114        } else if actual_def_levels.is_some() {
4115            assert_eq!(levels_read as u64, result.rows_written);
4116        } else {
4117            assert_eq!(values_read as u64, result.rows_written);
4118        }
4119    }
4120
4121    /// Performs write of provided values and returns column metadata of those values.
4122    /// Used to test encoding support for column writer.
4123    fn column_write_and_get_metadata<T: DataType>(
4124        props: WriterProperties,
4125        values: &[T::T],
4126    ) -> ColumnChunkMetaData {
4127        let page_writer = get_test_page_writer();
4128        let props = Arc::new(props);
4129        let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
4130        writer.write_batch(values, None, None).unwrap();
4131        writer.close().unwrap().metadata
4132    }
4133
4134    // Helper function to more compactly create a PageEncodingStats struct.
4135    fn encoding_stats(page_type: PageType, encoding: Encoding, count: i32) -> PageEncodingStats {
4136        PageEncodingStats {
4137            page_type,
4138            encoding,
4139            count,
4140        }
4141    }
4142
4143    // Function to use in tests for EncodingWriteSupport. This checks that dictionary
4144    // offset and encodings to make sure that column writer uses provided by trait
4145    // encodings.
4146    fn check_encoding_write_support<T: DataType>(
4147        version: WriterVersion,
4148        dict_enabled: bool,
4149        data: &[T::T],
4150        dictionary_page_offset: Option<i64>,
4151        encodings: &[Encoding],
4152        page_encoding_stats: &[PageEncodingStats],
4153    ) {
4154        let props = WriterProperties::builder()
4155            .set_writer_version(version)
4156            .set_dictionary_enabled(dict_enabled)
4157            .build();
4158        let meta = column_write_and_get_metadata::<T>(props, data);
4159        assert_eq!(meta.dictionary_page_offset(), dictionary_page_offset);
4160        assert_eq!(meta.encodings().collect::<Vec<_>>(), encodings);
4161        assert_eq!(meta.page_encoding_stats().unwrap(), page_encoding_stats);
4162    }
4163
4164    /// Returns column writer.
4165    fn get_test_column_writer<'a, T: DataType>(
4166        page_writer: Box<dyn PageWriter + 'a>,
4167        max_def_level: i16,
4168        max_rep_level: i16,
4169        props: WriterPropertiesPtr,
4170    ) -> ColumnWriterImpl<'a, T> {
4171        let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
4172        let column_writer = get_column_writer(descr, props, page_writer);
4173        get_typed_column_writer::<T>(column_writer)
4174    }
4175
4176    fn get_test_column_writer_with_path<'a, T: DataType>(
4177        page_writer: Box<dyn PageWriter + 'a>,
4178        max_def_level: i16,
4179        max_rep_level: i16,
4180        props: WriterPropertiesPtr,
4181        path: ColumnPath,
4182    ) -> ColumnWriterImpl<'a, T> {
4183        let descr = Arc::new(get_test_column_descr_with_path::<T>(
4184            max_def_level,
4185            max_rep_level,
4186            path,
4187        ));
4188        let column_writer = get_column_writer(descr, props, page_writer);
4189        get_typed_column_writer::<T>(column_writer)
4190    }
4191
4192    /// Returns column reader.
4193    fn get_test_column_reader<T: DataType>(
4194        page_reader: Box<dyn PageReader>,
4195        max_def_level: i16,
4196        max_rep_level: i16,
4197    ) -> ColumnReaderImpl<T> {
4198        let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
4199        let column_reader = get_column_reader(descr, page_reader);
4200        get_typed_column_reader::<T>(column_reader)
4201    }
4202
4203    /// Returns descriptor for primitive column.
4204    fn get_test_column_descr<T: DataType>(
4205        max_def_level: i16,
4206        max_rep_level: i16,
4207    ) -> ColumnDescriptor {
4208        let path = ColumnPath::from("col");
4209        let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type())
4210            // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type,
4211            // it should be no-op for other types
4212            .with_length(1)
4213            .build()
4214            .unwrap();
4215        ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
4216    }
4217
4218    fn get_test_column_descr_with_path<T: DataType>(
4219        max_def_level: i16,
4220        max_rep_level: i16,
4221        path: ColumnPath,
4222    ) -> ColumnDescriptor {
4223        let name = path.string();
4224        let tpe = SchemaType::primitive_type_builder(&name, T::get_physical_type())
4225            // length is set for "encoding support" tests for FIXED_LEN_BYTE_ARRAY type,
4226            // it should be no-op for other types
4227            .with_length(1)
4228            .build()
4229            .unwrap();
4230        ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
4231    }
4232
4233    fn write_and_collect_page_values(
4234        path: ColumnPath,
4235        props: WriterPropertiesPtr,
4236        data: &[i32],
4237    ) -> Vec<u32> {
4238        let mut file = tempfile::tempfile().unwrap();
4239        let mut write = TrackedWrite::new(&mut file);
4240        let page_writer = Box::new(SerializedPageWriter::new(&mut write));
4241        let mut writer =
4242            get_test_column_writer_with_path::<Int32Type>(page_writer, 0, 0, props, path);
4243        writer.write_batch(data, None, None).unwrap();
4244        let r = writer.close().unwrap();
4245
4246        drop(write);
4247
4248        let props = ReaderProperties::builder()
4249            .set_backward_compatible_lz4(false)
4250            .build();
4251        let mut page_reader = Box::new(
4252            SerializedPageReader::new_with_properties(
4253                Arc::new(file),
4254                &r.metadata,
4255                r.rows_written as usize,
4256                None,
4257                Arc::new(props),
4258            )
4259            .unwrap(),
4260        );
4261
4262        let mut values_per_page = Vec::new();
4263        while let Some(page) = page_reader.get_next_page().unwrap() {
4264            assert_eq!(page.page_type(), PageType::DATA_PAGE);
4265            values_per_page.push(page.num_values());
4266        }
4267
4268        values_per_page
4269    }
4270
4271    /// Returns page writer that collects pages without serializing them.
4272    fn get_test_page_writer() -> Box<dyn PageWriter> {
4273        Box::new(TestPageWriter {})
4274    }
4275
4276    struct TestPageWriter {}
4277
4278    impl PageWriter for TestPageWriter {
4279        fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
4280            let mut res = PageWriteSpec::new();
4281            res.page_type = page.page_type();
4282            res.uncompressed_size = page.uncompressed_size();
4283            res.compressed_size = page.compressed_size();
4284            res.num_values = page.num_values();
4285            res.offset = 0;
4286            res.bytes_written = page.data().len() as u64;
4287            Ok(res)
4288        }
4289
4290        fn close(&mut self) -> Result<()> {
4291            Ok(())
4292        }
4293    }
4294
4295    /// Write data into parquet using [`get_test_page_writer`] and [`get_test_column_writer`] and returns generated statistics.
4296    fn statistics_roundtrip<T: DataType>(values: &[<T as DataType>::T]) -> Statistics {
4297        let page_writer = get_test_page_writer();
4298        let props = Default::default();
4299        let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
4300        writer.write_batch(values, None, None).unwrap();
4301
4302        let metadata = writer.close().unwrap().metadata;
4303        if let Some(stats) = metadata.statistics() {
4304            stats.clone()
4305        } else {
4306            panic!("metadata missing statistics");
4307        }
4308    }
4309
4310    /// Returns Decimals column writer.
4311    fn get_test_decimals_column_writer<T: DataType>(
4312        page_writer: Box<dyn PageWriter>,
4313        max_def_level: i16,
4314        max_rep_level: i16,
4315        props: WriterPropertiesPtr,
4316    ) -> ColumnWriterImpl<'static, T> {
4317        let descr = Arc::new(get_test_decimals_column_descr::<T>(
4318            max_def_level,
4319            max_rep_level,
4320        ));
4321        let column_writer = get_column_writer(descr, props, page_writer);
4322        get_typed_column_writer::<T>(column_writer)
4323    }
4324
4325    /// Returns descriptor for Decimal type with primitive column.
4326    fn get_test_decimals_column_descr<T: DataType>(
4327        max_def_level: i16,
4328        max_rep_level: i16,
4329    ) -> ColumnDescriptor {
4330        let path = ColumnPath::from("col");
4331        let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type())
4332            .with_length(16)
4333            .with_logical_type(Some(LogicalType::Decimal {
4334                scale: 2,
4335                precision: 3,
4336            }))
4337            .with_scale(2)
4338            .with_precision(3)
4339            .build()
4340            .unwrap();
4341        ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
4342    }
4343
4344    fn float16_statistics_roundtrip(
4345        values: &[FixedLenByteArray],
4346    ) -> ValueStatistics<FixedLenByteArray> {
4347        let page_writer = get_test_page_writer();
4348        let mut writer = get_test_float16_column_writer(page_writer, Default::default());
4349        writer.write_batch(values, None, None).unwrap();
4350
4351        let metadata = writer.close().unwrap().metadata;
4352        if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
4353            stats.clone()
4354        } else {
4355            panic!("metadata missing statistics");
4356        }
4357    }
4358
4359    fn get_test_float16_column_writer(
4360        page_writer: Box<dyn PageWriter>,
4361        props: WriterPropertiesPtr,
4362    ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
4363        let descr = Arc::new(get_test_float16_column_descr(0, 0));
4364        let column_writer = get_column_writer(descr, props, page_writer);
4365        get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
4366    }
4367
4368    fn get_test_float16_column_descr(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor {
4369        let path = ColumnPath::from("col");
4370        let tpe =
4371            SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
4372                .with_length(2)
4373                .with_logical_type(Some(LogicalType::Float16))
4374                .build()
4375                .unwrap();
4376        ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
4377    }
4378
4379    fn get_test_interval_column_writer(
4380        page_writer: Box<dyn PageWriter>,
4381    ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
4382        let descr = Arc::new(get_test_interval_column_descr());
4383        let column_writer = get_column_writer(descr, Default::default(), page_writer);
4384        get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
4385    }
4386
4387    fn get_test_interval_column_descr() -> ColumnDescriptor {
4388        let path = ColumnPath::from("col");
4389        let tpe =
4390            SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
4391                .with_length(12)
4392                .with_converted_type(ConvertedType::INTERVAL)
4393                .build()
4394                .unwrap();
4395        ColumnDescriptor::new(Arc::new(tpe), 0, 0, path)
4396    }
4397
4398    /// Returns column writer for UINT32 Column provided as ConvertedType only
4399    fn get_test_unsigned_int_given_as_converted_column_writer<'a, T: DataType>(
4400        page_writer: Box<dyn PageWriter + 'a>,
4401        max_def_level: i16,
4402        max_rep_level: i16,
4403        props: WriterPropertiesPtr,
4404    ) -> ColumnWriterImpl<'a, T> {
4405        let descr = Arc::new(get_test_converted_type_unsigned_integer_column_descr::<T>(
4406            max_def_level,
4407            max_rep_level,
4408        ));
4409        let column_writer = get_column_writer(descr, props, page_writer);
4410        get_typed_column_writer::<T>(column_writer)
4411    }
4412
4413    /// Returns column descriptor for UINT32 Column provided as ConvertedType only
4414    fn get_test_converted_type_unsigned_integer_column_descr<T: DataType>(
4415        max_def_level: i16,
4416        max_rep_level: i16,
4417    ) -> ColumnDescriptor {
4418        let path = ColumnPath::from("col");
4419        let tpe = SchemaType::primitive_type_builder("col", T::get_physical_type())
4420            .with_converted_type(ConvertedType::UINT_32)
4421            .build()
4422            .unwrap();
4423        ColumnDescriptor::new(Arc::new(tpe), max_def_level, max_rep_level, path)
4424    }
4425
4426    #[test]
4427    fn test_page_v2_snappy_compression_fallback() {
4428        // Test that PageV2 sets is_compressed to false when Snappy compression increases data size
4429        let page_writer = TestPageWriter {};
4430
4431        // Create WriterProperties with PageV2 and Snappy compression
4432        let props = WriterProperties::builder()
4433            .set_writer_version(WriterVersion::PARQUET_2_0)
4434            // Disable dictionary to ensure data is written directly
4435            .set_dictionary_enabled(false)
4436            .set_compression(Compression::SNAPPY)
4437            .build();
4438
4439        let mut column_writer =
4440            get_test_column_writer::<ByteArrayType>(Box::new(page_writer), 0, 0, Arc::new(props));
4441
4442        // Create small, simple data that Snappy compression will likely increase in size
4443        // due to compression overhead for very small data
4444        let values = vec![ByteArray::from("a")];
4445
4446        column_writer.write_batch(&values, None, None).unwrap();
4447
4448        let result = column_writer.close().unwrap();
4449        assert_eq!(
4450            result.metadata.uncompressed_size(),
4451            result.metadata.compressed_size()
4452        );
4453    }
4454}