1use bytes::Bytes;
21use half::f16;
22
23use crate::bloom_filter::Sbbf;
24use crate::file::page_index::column_index::ColumnIndexMetaData;
25use crate::file::page_index::offset_index::OffsetIndexMetaData;
26use std::cmp::Ordering;
27use std::collections::{BTreeSet, VecDeque};
28use std::str;
29
30use crate::basic::{
31 BoundaryOrder, Compression, ConvertedType, Encoding, EncodingMask, LogicalType, PageType,
32 SortOrder, Type,
33};
34use crate::column::page::{CompressedPage, Page, PageWriteSpec, PageWriter};
35use crate::column::writer::encoder::{ColumnValueEncoder, ColumnValueEncoderImpl, ColumnValues};
36use crate::compression::{Codec, CodecOptionsBuilder, create_codec};
37use crate::data_type::private::ParquetValueType;
38use crate::data_type::*;
39use crate::encodings::levels::LevelEncoder;
40#[cfg(feature = "encryption")]
41use crate::encryption::encrypt::get_column_crypto_metadata;
42use crate::errors::{ParquetError, Result};
43use crate::file::metadata::{
44 ColumnChunkMetaData, ColumnChunkMetaDataBuilder, ColumnIndexBuilder, LevelHistogram,
45 OffsetIndexBuilder, PageEncodingStats,
46};
47use crate::file::properties::{
48 EnabledStatistics, ResolvedColumnProperties, WriterProperties, WriterPropertiesPtr,
49 WriterVersion,
50};
51use crate::file::statistics::{Statistics, ValueStatistics};
52use crate::schema::types::{BasicTypeInfo, ColumnDescPtr, ColumnDescriptor};
53
54mod byte_budget_chunker;
55pub(crate) mod encoder;
56
57use byte_budget_chunker::{ByteBudgetChunker, SubBatchStrategy};
58
59macro_rules! downcast_writer {
60 ($e:expr, $i:ident, $b:expr) => {
61 match $e {
62 Self::BoolColumnWriter($i) => $b,
63 Self::Int32ColumnWriter($i) => $b,
64 Self::Int64ColumnWriter($i) => $b,
65 Self::Int96ColumnWriter($i) => $b,
66 Self::FloatColumnWriter($i) => $b,
67 Self::DoubleColumnWriter($i) => $b,
68 Self::ByteArrayColumnWriter($i) => $b,
69 Self::FixedLenByteArrayColumnWriter($i) => $b,
70 }
71 };
72}
73
74pub enum ColumnWriter<'a> {
78 BoolColumnWriter(ColumnWriterImpl<'a, BoolType>),
80 Int32ColumnWriter(ColumnWriterImpl<'a, Int32Type>),
82 Int64ColumnWriter(ColumnWriterImpl<'a, Int64Type>),
84 Int96ColumnWriter(ColumnWriterImpl<'a, Int96Type>),
86 FloatColumnWriter(ColumnWriterImpl<'a, FloatType>),
88 DoubleColumnWriter(ColumnWriterImpl<'a, DoubleType>),
90 ByteArrayColumnWriter(ColumnWriterImpl<'a, ByteArrayType>),
92 FixedLenByteArrayColumnWriter(ColumnWriterImpl<'a, FixedLenByteArrayType>),
94}
95
96impl ColumnWriter<'_> {
97 #[cfg(feature = "arrow")]
99 pub(crate) fn memory_size(&self) -> usize {
100 downcast_writer!(self, typed, typed.memory_size())
101 }
102
103 #[cfg(feature = "arrow")]
105 pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
106 downcast_writer!(self, typed, typed.get_estimated_total_bytes())
107 }
108
109 #[cfg(feature = "arrow")]
114 pub(crate) fn add_data_page(&mut self) -> Result<()> {
115 downcast_writer!(self, typed, typed.add_data_page())
116 }
117
118 #[cfg(feature = "arrow")]
122 pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
123 downcast_writer!(self, typed, typed.set_distinct_count_override(count))
124 }
125
126 pub fn close(self) -> Result<ColumnCloseResult> {
128 downcast_writer!(self, typed, typed.close())
129 }
130}
131
132pub fn get_column_writer<'a>(
134 descr: ColumnDescPtr,
135 props: WriterPropertiesPtr,
136 page_writer: Box<dyn PageWriter + 'a>,
137) -> ColumnWriter<'a> {
138 match descr.physical_type() {
139 Type::BOOLEAN => {
140 ColumnWriter::BoolColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
141 }
142 Type::INT32 => {
143 ColumnWriter::Int32ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
144 }
145 Type::INT64 => {
146 ColumnWriter::Int64ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
147 }
148 Type::INT96 => {
149 ColumnWriter::Int96ColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
150 }
151 Type::FLOAT => {
152 ColumnWriter::FloatColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
153 }
154 Type::DOUBLE => {
155 ColumnWriter::DoubleColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
156 }
157 Type::BYTE_ARRAY => {
158 ColumnWriter::ByteArrayColumnWriter(ColumnWriterImpl::new(descr, props, page_writer))
159 }
160 Type::FIXED_LEN_BYTE_ARRAY => ColumnWriter::FixedLenByteArrayColumnWriter(
161 ColumnWriterImpl::new(descr, props, page_writer),
162 ),
163 }
164}
165
166pub fn get_typed_column_writer<T: DataType>(col_writer: ColumnWriter) -> ColumnWriterImpl<T> {
173 T::get_column_writer(col_writer).unwrap_or_else(|| {
174 panic!(
175 "Failed to convert column writer into a typed column writer for `{}` type",
176 T::get_physical_type()
177 )
178 })
179}
180
181pub fn get_typed_column_writer_ref<'a, 'b: 'a, T: DataType>(
183 col_writer: &'b ColumnWriter<'a>,
184) -> &'b ColumnWriterImpl<'a, T> {
185 T::get_column_writer_ref(col_writer).unwrap_or_else(|| {
186 panic!(
187 "Failed to convert column writer into a typed column writer for `{}` type",
188 T::get_physical_type()
189 )
190 })
191}
192
193pub fn get_typed_column_writer_mut<'a, 'b: 'a, T: DataType>(
195 col_writer: &'a mut ColumnWriter<'b>,
196) -> &'a mut ColumnWriterImpl<'b, T> {
197 T::get_column_writer_mut(col_writer).unwrap_or_else(|| {
198 panic!(
199 "Failed to convert column writer into a typed column writer for `{}` type",
200 T::get_physical_type()
201 )
202 })
203}
204
205#[derive(Debug, Clone)]
209pub struct ColumnCloseResult {
210 pub bytes_written: u64,
212 pub rows_written: u64,
214 pub metadata: ColumnChunkMetaData,
216 pub bloom_filter: Option<Sbbf>,
218 pub column_index: Option<ColumnIndexMetaData>,
220 pub offset_index: Option<OffsetIndexMetaData>,
222}
223
224impl ColumnCloseResult {
225 pub fn update_dictionary_location(mut self, dictionary_len: usize) -> Result<Self> {
236 if dictionary_len > 0 {
237 self.metadata = self
238 .metadata
239 .into_builder()
240 .set_dictionary_page_offset(Some(0))
241 .set_data_page_offset(dictionary_len as i64)
242 .build()?;
243 if let Some(offset_index) = self.offset_index.as_mut() {
244 let mut offset = dictionary_len as i64;
245 for location in &mut offset_index.page_locations {
246 location.offset = offset;
247 offset += location.compressed_page_size as i64;
248 }
249 }
250 }
251 Ok(self)
252 }
253}
254
255#[derive(Default)]
257struct PageMetrics {
258 num_buffered_values: u32,
259 num_buffered_rows: u32,
260 page_size_exemption: usize,
266 num_page_nulls: u64,
267 num_page_nans: Option<u64>,
268 repetition_level_histogram: Option<LevelHistogram>,
269 definition_level_histogram: Option<LevelHistogram>,
270}
271
272impl PageMetrics {
273 fn new() -> Self {
274 Default::default()
275 }
276
277 fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
279 self.repetition_level_histogram = LevelHistogram::try_new(max_level);
280 self
281 }
282
283 fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
285 self.definition_level_histogram = LevelHistogram::try_new(max_level);
286 self
287 }
288
289 fn new_page(&mut self) {
292 self.num_buffered_values = 0;
293 self.num_buffered_rows = 0;
294 self.page_size_exemption = 0;
295 self.num_page_nulls = 0;
296 self.num_page_nans = None;
297 self.repetition_level_histogram
298 .as_mut()
299 .map(LevelHistogram::reset);
300 self.definition_level_histogram
301 .as_mut()
302 .map(LevelHistogram::reset);
303 }
304}
305
306#[derive(Default)]
308struct ColumnMetrics<T: Default> {
309 total_bytes_written: u64,
310 total_rows_written: u64,
311 total_uncompressed_size: u64,
312 total_compressed_size: u64,
313 total_num_values: u64,
314 dictionary_page_offset: Option<u64>,
315 data_page_offset: Option<u64>,
316 min_column_value: Option<T>,
317 max_column_value: Option<T>,
318 num_column_nulls: u64,
319 num_column_nans: Option<u64>,
320 column_distinct_count: Option<u64>,
321 variable_length_bytes: Option<i64>,
322 repetition_level_histogram: Option<LevelHistogram>,
323 definition_level_histogram: Option<LevelHistogram>,
324}
325
326impl<T: Default> ColumnMetrics<T> {
327 fn new() -> Self {
328 Default::default()
329 }
330
331 fn with_repetition_level_histogram(mut self, max_level: i16) -> Self {
333 self.repetition_level_histogram = LevelHistogram::try_new(max_level);
334 self
335 }
336
337 fn with_definition_level_histogram(mut self, max_level: i16) -> Self {
339 self.definition_level_histogram = LevelHistogram::try_new(max_level);
340 self
341 }
342
343 fn update_histogram(
345 chunk_histogram: &mut Option<LevelHistogram>,
346 page_histogram: Option<&LevelHistogram>,
347 ) {
348 if let (Some(page_hist), Some(chunk_hist)) = (page_histogram, chunk_histogram) {
349 chunk_hist.add(page_hist);
350 }
351 }
352
353 fn update_from_page_metrics(&mut self, page_metrics: &PageMetrics) {
356 ColumnMetrics::<T>::update_histogram(
357 &mut self.definition_level_histogram,
358 page_metrics.definition_level_histogram.as_ref(),
359 );
360 ColumnMetrics::<T>::update_histogram(
361 &mut self.repetition_level_histogram,
362 page_metrics.repetition_level_histogram.as_ref(),
363 );
364 }
365
366 fn update_variable_length_bytes(&mut self, variable_length_bytes: Option<i64>) {
368 if let Some(var_bytes) = variable_length_bytes {
369 *self.variable_length_bytes.get_or_insert(0) += var_bytes;
370 }
371 }
372}
373
374#[derive(Debug, Clone, Copy)]
382pub(crate) enum LevelDataRef<'a> {
383 Absent,
384 Materialized(&'a [i16]),
385 Uniform { value: i16, count: usize },
386}
387
388impl<'a> From<&'a [i16]> for LevelDataRef<'a> {
389 fn from(levels: &'a [i16]) -> Self {
390 Self::Materialized(levels)
391 }
392}
393
394impl<'a> From<Option<&'a [i16]>> for LevelDataRef<'a> {
395 fn from(levels: Option<&'a [i16]>) -> Self {
396 levels.map_or(Self::Absent, Self::from)
397 }
398}
399
400impl LevelDataRef<'_> {
401 pub(crate) fn len(self) -> usize {
402 match self {
403 Self::Absent => 0,
404 Self::Materialized(values) => values.len(),
405 Self::Uniform { count, .. } => count,
406 }
407 }
408
409 pub(crate) fn first(self) -> Option<i16> {
410 match self {
411 Self::Absent => None,
412 Self::Materialized(values) => values.first().copied(),
413 Self::Uniform { value, count } => (count > 0).then_some(value),
414 }
415 }
416
417 #[cfg(feature = "arrow")]
418 pub(crate) fn value_at(self, idx: usize) -> Option<i16> {
419 match self {
420 Self::Absent => None,
421 Self::Materialized(values) => values.get(idx).copied(),
422 Self::Uniform { value, count } => (idx < count).then_some(value),
423 }
424 }
425
426 pub(crate) fn slice(self, offset: usize, len: usize) -> Self {
427 match self {
428 Self::Absent => Self::Absent,
429 Self::Materialized(values) => Self::Materialized(&values[offset..offset + len]),
430 Self::Uniform { value, .. } => Self::Uniform { value, count: len },
431 }
432 }
433
434 pub(crate) fn value_count(self, total: usize, max_def: i16) -> usize {
439 match self {
440 Self::Absent => total,
441 Self::Materialized(values) => values.iter().filter(|&&d| d == max_def).count(),
442 Self::Uniform { value, count } => {
443 if value == max_def {
444 count
445 } else {
446 0
447 }
448 }
449 }
450 }
451}
452
453pub type ColumnWriterImpl<'a, T> = GenericColumnWriter<'a, ColumnValueEncoderImpl<T>>;
455
456pub struct GenericColumnWriter<'a, E: ColumnValueEncoder> {
458 descr: ColumnDescPtr,
460 props: WriterPropertiesPtr,
461 column_props: ResolvedColumnProperties,
465
466 page_writer: Box<dyn PageWriter + 'a>,
467 codec: Compression,
468 compressor: Option<Box<dyn Codec>>,
469 encoder: E,
470
471 page_metrics: PageMetrics,
472 column_metrics: ColumnMetrics<E::T>,
474
475 distinct_count_override: Option<u64>,
478
479 encodings: BTreeSet<Encoding>,
482 encoding_stats: Vec<PageEncodingStats>,
483 def_levels_encoder: LevelEncoder,
485 rep_levels_encoder: LevelEncoder,
486 data_pages: VecDeque<CompressedPage>,
487 column_index_builder: ColumnIndexBuilder,
489 offset_index_builder: Option<OffsetIndexBuilder>,
490
491 data_page_boundary_ascending: bool,
494 data_page_boundary_descending: bool,
495 last_non_null_data_page_min_max: Option<(E::T, E::T)>,
497}
498
499impl<'a, E: ColumnValueEncoder> GenericColumnWriter<'a, E> {
500 pub fn new(
502 descr: ColumnDescPtr,
503 props: WriterPropertiesPtr,
504 page_writer: Box<dyn PageWriter + 'a>,
505 ) -> Self {
506 let column_props = props.resolve_column_properties(descr.path());
507 let codec = column_props.compression;
508 let codec_options = CodecOptionsBuilder::default().build();
509 let compressor = create_codec(codec, &codec_options).unwrap();
510 let encoder = E::try_new(&descr, props.as_ref(), &column_props).unwrap();
511
512 let statistics_enabled = column_props.statistics_enabled;
513
514 let mut encodings = BTreeSet::new();
515 encodings.insert(Encoding::RLE);
517
518 let mut page_metrics = PageMetrics::new();
519 let mut column_metrics = ColumnMetrics::<E::T>::new();
520
521 if statistics_enabled != EnabledStatistics::None {
523 page_metrics = page_metrics
524 .with_repetition_level_histogram(descr.max_rep_level())
525 .with_definition_level_histogram(descr.max_def_level());
526 column_metrics = column_metrics
527 .with_repetition_level_histogram(descr.max_rep_level())
528 .with_definition_level_histogram(descr.max_def_level())
529 }
530
531 let mut column_index_builder = ColumnIndexBuilder::new(descr.physical_type());
533 if statistics_enabled != EnabledStatistics::Page {
534 column_index_builder.to_invalid()
535 }
536
537 let offset_index_builder = match props.offset_index_disabled() {
539 false => Some(OffsetIndexBuilder::new()),
540 _ => None,
541 };
542
543 Self {
544 def_levels_encoder: Self::create_level_encoder(descr.max_def_level(), &props),
545 rep_levels_encoder: Self::create_level_encoder(descr.max_rep_level(), &props),
546 descr,
547 props,
548 column_props,
549 page_writer,
550 codec,
551 compressor,
552 encoder,
553 data_pages: VecDeque::new(),
554 page_metrics,
555 column_metrics,
556 distinct_count_override: None,
557 column_index_builder,
558 offset_index_builder,
559 encodings,
560 encoding_stats: vec![],
561 data_page_boundary_ascending: true,
562 data_page_boundary_descending: true,
563 last_non_null_data_page_min_max: None,
564 }
565 }
566
567 #[cfg(feature = "arrow")]
573 pub(crate) fn set_distinct_count_override(&mut self, count: u64) {
574 self.distinct_count_override = Some(count);
575 }
576
577 #[expect(clippy::too_many_arguments)]
578 pub(crate) fn write_batch_internal(
579 &mut self,
580 values: &E::Values,
581 value_indices: Option<&[usize]>,
582 def_levels: LevelDataRef<'_>,
583 rep_levels: LevelDataRef<'_>,
584 min: Option<&E::T>,
585 max: Option<&E::T>,
586 distinct_count: Option<u64>,
587 ) -> Result<usize> {
588 if def_levels.len() != 0 && rep_levels.len() != 0 && def_levels.len() != rep_levels.len() {
590 return Err(general_err!(
591 "Inconsistent length of definition and repetition levels: {} != {}",
592 def_levels.len(),
593 rep_levels.len()
594 ));
595 }
596
597 let num_levels = def_levels.len().max(rep_levels.len());
607 let num_levels = if num_levels > 0 {
608 num_levels
609 } else {
610 value_indices.map_or_else(|| values.len(), |i| i.len())
611 };
612
613 if let Some(min) = min {
614 update_min(&self.descr, min, &mut self.column_metrics.min_column_value);
615 }
616 if let Some(max) = max {
617 update_max(&self.descr, max, &mut self.column_metrics.max_column_value);
618 }
619
620 if self.encoder.num_values() == 0 {
622 self.column_metrics.column_distinct_count = distinct_count;
623 } else {
624 self.column_metrics.column_distinct_count = None;
625 }
626
627 let mut values_offset = 0;
628 let mut levels_offset = 0;
629 let both_levels_compact = !matches!(def_levels, LevelDataRef::Materialized(_))
630 && !matches!(rep_levels, LevelDataRef::Materialized(_));
631 let has_levels = !matches!(def_levels, LevelDataRef::Absent)
632 || !matches!(rep_levels, LevelDataRef::Absent);
633
634 let base_batch_size = if both_levels_compact && has_levels {
638 self.props.data_page_row_count_limit()
639 } else {
640 self.props.write_batch_size()
641 };
642 debug_assert!(base_batch_size > 0);
643
644 let chunker = ByteBudgetChunker::new(&self.descr, &self.column_props, base_batch_size);
645 while levels_offset < num_levels {
646 let mut end_offset = num_levels.min(levels_offset + base_batch_size);
647
648 if let LevelDataRef::Materialized(levels) = rep_levels {
650 while end_offset < levels.len() && levels[end_offset] != 0 {
651 end_offset += 1;
652 }
653 }
654
655 let chunk_size = end_offset - levels_offset;
656 let chunk_def = def_levels.slice(levels_offset, chunk_size);
657 let chunk_rep = rep_levels.slice(levels_offset, chunk_size);
658
659 let sub_batch = chunker.pick_sub_batch(
666 &self.encoder,
667 values,
668 value_indices,
669 chunk_def,
670 values_offset,
671 chunk_size,
672 );
673
674 match sub_batch {
675 None => {
676 values_offset += self.write_mini_batch(
677 values,
678 values_offset,
679 value_indices,
680 chunk_size,
681 chunk_def,
682 chunk_rep,
683 )?;
684 }
685 Some(sub_batch) => {
686 values_offset += self.write_granular_chunk(
687 values,
688 values_offset,
689 value_indices,
690 chunk_size,
691 chunk_def,
692 chunk_rep,
693 sub_batch,
694 )?;
695 }
696 }
697 levels_offset = end_offset;
698 }
699
700 Ok(values_offset)
702 }
703
704 pub fn write_batch(
717 &mut self,
718 values: &E::Values,
719 def_levels: Option<&[i16]>,
720 rep_levels: Option<&[i16]>,
721 ) -> Result<usize> {
722 self.write_batch_internal(
723 values,
724 None,
725 LevelDataRef::from(def_levels),
726 LevelDataRef::from(rep_levels),
727 None,
728 None,
729 None,
730 )
731 }
732
733 pub fn write_batch_with_statistics(
741 &mut self,
742 values: &E::Values,
743 def_levels: Option<&[i16]>,
744 rep_levels: Option<&[i16]>,
745 min: Option<&E::T>,
746 max: Option<&E::T>,
747 distinct_count: Option<u64>,
748 ) -> Result<usize> {
749 self.write_batch_internal(
750 values,
751 None,
752 LevelDataRef::from(def_levels),
753 LevelDataRef::from(rep_levels),
754 min,
755 max,
756 distinct_count,
757 )
758 }
759
760 #[cfg(feature = "arrow")]
765 pub(crate) fn memory_size(&self) -> usize {
766 self.encoder.estimated_memory_size()
772 + self
773 .data_pages
774 .iter()
775 .map(|page| page.memory_usage())
776 .sum::<usize>()
777 + self.page_writer.buffered_memory_size()
778 }
779
780 pub fn get_total_bytes_written(&self) -> u64 {
786 self.column_metrics.total_bytes_written
787 }
788
789 #[cfg(feature = "arrow")]
795 pub(crate) fn get_estimated_total_bytes(&self) -> u64 {
796 self.data_pages
797 .iter()
798 .map(|page| page.data().len() as u64)
799 .sum::<u64>()
800 + self.column_metrics.total_bytes_written
801 + self.encoder.estimated_data_page_size() as u64
802 + self.encoder.estimated_dict_page_size().unwrap_or_default() as u64
803 }
804
805 pub fn get_total_rows_written(&self) -> u64 {
808 self.column_metrics.total_rows_written
809 }
810
811 pub fn get_descriptor(&self) -> &ColumnDescPtr {
813 &self.descr
814 }
815
816 pub fn close(mut self) -> Result<ColumnCloseResult> {
819 if self.page_metrics.num_buffered_values > 0 {
820 self.add_data_page()?;
821 }
822 if self.encoder.has_dictionary() {
823 self.write_dictionary_page()?;
824 }
825 self.flush_data_pages()?;
826 let metadata = self.build_column_metadata()?;
827 self.page_writer.close()?;
828
829 let write_bloom_filter = self.props.bloom_filter_for_dictionary_encoded_chunks()
830 || self.has_non_dictionary_data_page();
831 let bloom_filter = self
832 .encoder
833 .flush_bloom_filter()
834 .filter(|_| write_bloom_filter);
835
836 let boundary_order = match (
837 self.data_page_boundary_ascending,
838 self.data_page_boundary_descending,
839 ) {
840 (true, _) => BoundaryOrder::ASCENDING,
843 (false, true) => BoundaryOrder::DESCENDING,
844 (false, false) => BoundaryOrder::UNORDERED,
845 };
846 self.column_index_builder.set_boundary_order(boundary_order);
847
848 let column_index = match self.column_index_builder.valid() {
849 true => Some(self.column_index_builder.build()?),
850 false => None,
851 };
852
853 let offset_index = self.offset_index_builder.map(|b| b.build());
854
855 Ok(ColumnCloseResult {
856 bytes_written: self.column_metrics.total_bytes_written,
857 rows_written: self.column_metrics.total_rows_written,
858 bloom_filter,
859 metadata,
860 column_index,
861 offset_index,
862 })
863 }
864
865 #[expect(clippy::too_many_arguments)]
885 #[inline(never)]
886 fn write_granular_chunk(
887 &mut self,
888 values: &E::Values,
889 values_offset: usize,
890 value_indices: Option<&[usize]>,
891 chunk_size: usize,
892 chunk_def: LevelDataRef<'_>,
893 chunk_rep: LevelDataRef<'_>,
894 sub_batch: SubBatchStrategy,
895 ) -> Result<usize> {
896 debug_assert!(
900 matches!(sub_batch, SubBatchStrategy::Values(n) | SubBatchStrategy::Levels(n) if n >= 1),
901 "chunker must size at least one value or level"
902 );
903 let max_def_level = self.descr.max_def_level();
904 let mut values_consumed = 0;
905 let mut sub_start = 0;
906 while sub_start < chunk_size {
907 let window_end = match sub_batch {
908 SubBatchStrategy::Values(n) => {
909 Self::window_end_for_values(chunk_def, chunk_size, max_def_level, sub_start, n)
910 }
911 SubBatchStrategy::Levels(n) => (sub_start + n).min(chunk_size),
912 };
913 let sub_end = match chunk_rep {
914 LevelDataRef::Materialized(levels) => {
915 let mut e = window_end;
922 while e < chunk_size && levels[e] != 0 {
923 e += 1;
924 }
925 e
926 }
927 _ => window_end,
928 };
929 let sub_len = sub_end - sub_start;
930 let written = self.write_mini_batch(
931 values,
932 values_offset + values_consumed,
933 value_indices,
934 sub_len,
935 chunk_def.slice(sub_start, sub_len),
936 chunk_rep.slice(sub_start, sub_len),
937 )?;
938 values_consumed += written;
939 sub_start = sub_end;
940 }
941 Ok(values_consumed)
942 }
943
944 fn has_non_dictionary_data_page(&self) -> bool {
945 self.encoding_stats.iter().any(|stats| {
946 matches!(
947 stats.page_type,
948 PageType::DATA_PAGE | PageType::DATA_PAGE_V2
949 ) && !matches!(
950 stats.encoding,
951 Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY
952 )
953 })
954 }
955
956 fn window_end_for_values(
971 chunk_def: LevelDataRef<'_>,
972 chunk_size: usize,
973 max_def_level: i16,
974 start: usize,
975 max_values: usize,
976 ) -> usize {
977 match chunk_def {
978 LevelDataRef::Absent => (start + max_values).min(chunk_size),
980 LevelDataRef::Uniform { value, .. } => {
981 if value == max_def_level {
982 (start + max_values).min(chunk_size)
983 } else {
984 chunk_size
988 }
989 }
990 LevelDataRef::Materialized(levels) => {
991 let mut seen = 0;
992 let mut end = start;
993 while end < chunk_size && seen < max_values {
994 if levels[end] == max_def_level {
995 seen += 1;
996 }
997 end += 1;
998 }
999 end
1000 }
1001 }
1002 }
1003
1004 fn create_level_encoder(max_level: i16, props: &WriterProperties) -> LevelEncoder {
1006 match props.writer_version() {
1007 WriterVersion::PARQUET_1_0 => LevelEncoder::v1_streaming(max_level),
1008 WriterVersion::PARQUET_2_0 => LevelEncoder::v2_streaming(max_level),
1009 }
1010 }
1011
1012 fn write_mini_batch(
1016 &mut self,
1017 values: &E::Values,
1018 values_offset: usize,
1019 value_indices: Option<&[usize]>,
1020 num_levels: usize,
1021 def_levels: LevelDataRef<'_>,
1022 rep_levels: LevelDataRef<'_>,
1023 ) -> Result<usize> {
1024 let values_to_write = if self.descr.max_def_level() > 0 {
1026 let max_def = self.descr.max_def_level();
1027 match def_levels {
1028 LevelDataRef::Absent => {
1029 return Err(general_err!(
1030 "Definition levels are required, because max definition level = {}",
1031 self.descr.max_def_level()
1032 ));
1033 }
1034 LevelDataRef::Materialized(levels) => {
1035 let mut values_to_write = 0usize;
1038 let encoder = &mut self.def_levels_encoder;
1039 match self.page_metrics.definition_level_histogram.as_mut() {
1040 Some(histogram) => encoder.put_with_observer(levels, |level, count| {
1041 values_to_write += count * (level == max_def) as usize;
1042 histogram.increment_by(level, count as i64);
1043 }),
1044 None => encoder.put_with_observer(levels, |level, count| {
1045 values_to_write += count * (level == max_def) as usize;
1046 }),
1047 };
1048 self.page_metrics.num_page_nulls += (levels.len() - values_to_write) as u64;
1049 values_to_write
1050 }
1051 LevelDataRef::Uniform { value, count } => {
1052 let encoder = &mut self.def_levels_encoder;
1055 match self.page_metrics.definition_level_histogram.as_mut() {
1056 Some(histogram) => {
1057 encoder.put_n_with_observer(value, count, |level, run_len| {
1058 histogram.increment_by(level, run_len as i64);
1059 })
1060 }
1061 None => encoder.put_n_with_observer(value, count, |_, _| {}),
1062 }
1063 let values_to_write = count * (value == max_def) as usize;
1064 self.page_metrics.num_page_nulls += (count - values_to_write) as u64;
1065 values_to_write
1066 }
1067 }
1068 } else {
1069 num_levels
1070 };
1071
1072 if self.descr.max_rep_level() > 0 {
1074 let first_level = rep_levels.first().ok_or_else(|| {
1076 general_err!(
1077 "Repetition levels are required, because max repetition level = {}",
1078 self.descr.max_rep_level()
1079 )
1080 })?;
1081
1082 if first_level != 0 {
1083 return Err(general_err!(
1084 "Write must start at a record boundary, got non-zero repetition level of {}",
1085 first_level
1086 ));
1087 }
1088
1089 let mut new_rows = 0u32;
1090 match rep_levels {
1091 LevelDataRef::Absent => unreachable!(),
1092 LevelDataRef::Materialized(levels) => {
1093 let encoder = &mut self.rep_levels_encoder;
1094 match self.page_metrics.repetition_level_histogram.as_mut() {
1095 Some(histogram) => encoder.put_with_observer(levels, |level, count| {
1096 new_rows += (count as u32) * (level == 0) as u32;
1097 histogram.increment_by(level, count as i64);
1098 }),
1099 None => encoder.put_with_observer(levels, |level, count| {
1100 new_rows += (count as u32) * (level == 0) as u32;
1101 }),
1102 };
1103 }
1104 LevelDataRef::Uniform { value, count } => {
1105 let encoder = &mut self.rep_levels_encoder;
1106 match self.page_metrics.repetition_level_histogram.as_mut() {
1107 Some(histogram) => {
1108 encoder.put_n_with_observer(value, count, |level, run_len| {
1109 new_rows += (run_len as u32) * (level == 0) as u32;
1110 histogram.increment_by(level, run_len as i64);
1111 })
1112 }
1113 None => encoder.put_n_with_observer(value, count, |level, run_len| {
1114 new_rows += (run_len as u32) * (level == 0) as u32;
1115 }),
1116 }
1117 }
1118 }
1119 self.page_metrics.num_buffered_rows += new_rows;
1120 } else {
1121 self.page_metrics.num_buffered_rows += num_levels as u32;
1124 }
1125
1126 match value_indices {
1127 Some(indices) => {
1128 let indices = &indices[values_offset..values_offset + values_to_write];
1129 self.encoder.write_gather(values, indices)?;
1130 }
1131 None => self.encoder.write(values, values_offset, values_to_write)?,
1132 }
1133
1134 let page_was_empty = self.page_metrics.num_buffered_values == 0;
1135 self.page_metrics.num_buffered_values += num_levels as u32;
1136
1137 if page_was_empty && values_to_write == 1 {
1138 self.set_page_size_exemption();
1139 }
1140
1141 if self.should_add_data_page() {
1142 self.add_data_page()?;
1143 }
1144
1145 if self.should_dict_fallback() {
1146 self.dict_fallback()?;
1147 }
1148
1149 Ok(values_to_write)
1150 }
1151
1152 #[inline]
1157 fn should_dict_fallback(&self) -> bool {
1158 match self.encoder.estimated_dict_page_size() {
1159 Some(size) => size >= self.column_props.dictionary_page_size_limit,
1160 None => false,
1161 }
1162 }
1163
1164 #[cold]
1192 fn set_page_size_exemption(&mut self) {
1193 if !self.encoder.compresses_against_previous_value() {
1194 return;
1195 }
1196 let size = self.encoder.estimated_data_page_size();
1197 if size >= self.column_props.data_page_size_limit {
1198 self.page_metrics.page_size_exemption = size;
1199 }
1200 }
1201
1202 #[inline]
1204 fn should_add_data_page(&self) -> bool {
1205 if self.page_metrics.num_buffered_values == 0 {
1210 return false;
1211 }
1212
1213 self.page_metrics.num_buffered_rows as usize >= self.props.data_page_row_count_limit()
1214 || self
1215 .encoder
1216 .estimated_data_page_size()
1217 .saturating_sub(self.page_metrics.page_size_exemption)
1218 >= self.column_props.data_page_size_limit
1219 }
1220
1221 fn dict_fallback(&mut self) -> Result<()> {
1224 if self.page_metrics.num_buffered_values > 0 {
1226 self.add_data_page()?;
1227 }
1228 self.write_dictionary_page()?;
1229 self.flush_data_pages()?;
1230 Ok(())
1231 }
1232
1233 fn get_nan_count<T: ParquetValueType>(&self) -> Option<i64> {
1236 let nan_count = || {
1237 let nan_count = self.page_metrics.num_page_nans.unwrap_or(0);
1238 match i64::try_from(nan_count) {
1239 Ok(count) => Some(count),
1240 _ => Some(i64::MAX),
1241 }
1242 };
1243 match T::PHYSICAL_TYPE {
1244 Type::FLOAT | Type::DOUBLE => nan_count(),
1245 Type::FIXED_LEN_BYTE_ARRAY
1246 if matches!(self.descr.logical_type_ref(), Some(LogicalType::Float16)) =>
1247 {
1248 nan_count()
1249 }
1250 _ => None,
1251 }
1252 }
1253
1254 fn update_column_offset_index(
1256 &mut self,
1257 page_statistics: Option<&ValueStatistics<E::T>>,
1258 page_variable_length_bytes: Option<i64>,
1259 ) {
1260 let null_page =
1262 (self.page_metrics.num_buffered_rows as u64) == self.page_metrics.num_page_nulls;
1263 if null_page && self.column_index_builder.valid() {
1266 self.column_index_builder.append(
1267 null_page,
1268 vec![],
1269 vec![],
1270 self.page_metrics.num_page_nulls as i64,
1271 self.get_nan_count::<E::T>(),
1272 );
1273 } else if self.column_index_builder.valid() {
1274 match &page_statistics {
1277 None => {
1278 self.column_index_builder.to_invalid();
1279 }
1280 Some(stat) => {
1281 let new_min = stat.min_opt().unwrap();
1283 let new_max = stat.max_opt().unwrap();
1284 if let Some((last_min, last_max)) = &self.last_non_null_data_page_min_max {
1285 let basic_info = self.descr.get_basic_info();
1286 if self.data_page_boundary_ascending {
1287 let not_ascending = compare_greater(basic_info, last_min, new_min)
1289 || compare_greater(basic_info, last_max, new_max);
1290 if not_ascending {
1291 self.data_page_boundary_ascending = false;
1292 }
1293 }
1294
1295 if self.data_page_boundary_descending {
1296 let not_descending = compare_greater(basic_info, new_min, last_min)
1298 || compare_greater(basic_info, new_max, last_max);
1299 if not_descending {
1300 self.data_page_boundary_descending = false;
1301 }
1302 }
1303 }
1304 self.last_non_null_data_page_min_max = Some((new_min.clone(), new_max.clone()));
1305
1306 if self.can_truncate_value() {
1307 self.column_index_builder.append(
1308 null_page,
1309 self.truncate_min_value(
1310 self.props.column_index_truncate_length(),
1311 stat.min_bytes_opt().unwrap(),
1312 )
1313 .0,
1314 self.truncate_max_value(
1315 self.props.column_index_truncate_length(),
1316 stat.max_bytes_opt().unwrap(),
1317 )
1318 .0,
1319 self.page_metrics.num_page_nulls as i64,
1320 self.get_nan_count::<E::T>(),
1321 );
1322 } else {
1323 self.column_index_builder.append(
1324 null_page,
1325 stat.min_bytes_opt().unwrap().to_vec(),
1326 stat.max_bytes_opt().unwrap().to_vec(),
1327 self.page_metrics.num_page_nulls as i64,
1328 self.get_nan_count::<E::T>(),
1329 );
1330 }
1331 }
1332 }
1333 }
1334
1335 self.column_index_builder.append_histograms(
1337 &self.page_metrics.repetition_level_histogram,
1338 &self.page_metrics.definition_level_histogram,
1339 );
1340
1341 if let Some(builder) = self.offset_index_builder.as_mut() {
1343 builder.append_row_count(self.page_metrics.num_buffered_rows as i64);
1344 builder.append_unencoded_byte_array_data_bytes(page_variable_length_bytes);
1345 }
1346 }
1347
1348 fn can_truncate_value(&self) -> bool {
1350 match self.descr.physical_type() {
1351 Type::FIXED_LEN_BYTE_ARRAY
1355 if !matches!(
1356 self.descr.logical_type_ref(),
1357 Some(&LogicalType::Decimal { .. } | &LogicalType::Float16)
1358 ) =>
1359 {
1360 true
1361 }
1362 Type::BYTE_ARRAY => true,
1363 _ => false,
1365 }
1366 }
1367
1368 fn is_utf8(&self) -> bool {
1370 self.get_descriptor().logical_type_ref() == Some(&LogicalType::String)
1371 || self.get_descriptor().converted_type() == ConvertedType::UTF8
1372 }
1373
1374 fn truncate_min_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
1385 truncation_length
1386 .filter(|l| data.len() > *l)
1387 .and_then(|l|
1388 if self.is_utf8() {
1390 match str::from_utf8(data) {
1391 Ok(str_data) => truncate_utf8(str_data, l),
1392 Err(_) => Some(data[..l].to_vec()),
1393 }
1394 } else {
1395 Some(data[..l].to_vec())
1396 }
1397 )
1398 .map(|truncated| (truncated, true))
1399 .unwrap_or_else(|| (data.to_vec(), false))
1400 }
1401
1402 fn truncate_max_value(&self, truncation_length: Option<usize>, data: &[u8]) -> (Vec<u8>, bool) {
1416 truncation_length
1417 .filter(|l| data.len() > *l)
1418 .and_then(|l|
1419 if self.is_utf8() {
1421 match str::from_utf8(data) {
1422 Ok(str_data) => truncate_and_increment_utf8(str_data, l),
1423 Err(_) => increment(data[..l].to_vec()),
1424 }
1425 } else {
1426 increment(data[..l].to_vec())
1427 }
1428 )
1429 .map(|truncated| (truncated, true))
1430 .unwrap_or_else(|| (data.to_vec(), false))
1431 }
1432
1433 fn truncate_statistics(&self, statistics: Statistics) -> Statistics {
1436 let backwards_compatible_min_max = self.descr.sort_order().is_signed();
1437 match statistics {
1438 Statistics::ByteArray(stats) if stats._internal_has_min_max_set() => {
1439 let (min, did_truncate_min) = self.truncate_min_value(
1440 self.props.statistics_truncate_length(),
1441 stats.min_bytes_opt().unwrap(),
1442 );
1443 let (max, did_truncate_max) = self.truncate_max_value(
1444 self.props.statistics_truncate_length(),
1445 stats.max_bytes_opt().unwrap(),
1446 );
1447 Statistics::ByteArray(
1448 ValueStatistics::new(
1449 Some(min.into()),
1450 Some(max.into()),
1451 stats.distinct_count(),
1452 stats.null_count_opt(),
1453 backwards_compatible_min_max,
1454 )
1455 .with_max_is_exact(!did_truncate_max)
1456 .with_min_is_exact(!did_truncate_min),
1457 )
1458 }
1459 Statistics::FixedLenByteArray(stats)
1460 if (stats._internal_has_min_max_set() && self.can_truncate_value()) =>
1461 {
1462 let (min, did_truncate_min) = self.truncate_min_value(
1463 self.props.statistics_truncate_length(),
1464 stats.min_bytes_opt().unwrap(),
1465 );
1466 let (max, did_truncate_max) = self.truncate_max_value(
1467 self.props.statistics_truncate_length(),
1468 stats.max_bytes_opt().unwrap(),
1469 );
1470 Statistics::FixedLenByteArray(
1471 ValueStatistics::new(
1472 Some(min.into()),
1473 Some(max.into()),
1474 stats.distinct_count(),
1475 stats.null_count_opt(),
1476 backwards_compatible_min_max,
1477 )
1478 .with_max_is_exact(!did_truncate_max)
1479 .with_min_is_exact(!did_truncate_min),
1480 )
1481 }
1482 stats => stats,
1483 }
1484 }
1485
1486 pub(crate) fn add_data_page(&mut self) -> Result<()> {
1489 let values_data = self.encoder.flush_data_page()?;
1491
1492 let max_def_level = self.descr.max_def_level();
1493 let max_rep_level = self.descr.max_rep_level();
1494
1495 self.column_metrics.num_column_nulls += self.page_metrics.num_page_nulls;
1496
1497 if let Some(nan_count) = values_data.nan_count {
1498 *self.column_metrics.num_column_nans.get_or_insert(0) += nan_count;
1499 self.page_metrics.num_page_nans = Some(nan_count);
1500 }
1501
1502 let page_statistics = match (values_data.min_value, values_data.max_value) {
1503 (Some(min), Some(max)) => {
1504 update_min(&self.descr, &min, &mut self.column_metrics.min_column_value);
1506 update_max(&self.descr, &max, &mut self.column_metrics.max_column_value);
1507
1508 (self.column_props.statistics_enabled == EnabledStatistics::Page).then_some(
1509 ValueStatistics::new(
1510 Some(min),
1511 Some(max),
1512 None,
1513 Some(self.page_metrics.num_page_nulls),
1514 false,
1515 )
1516 .with_nan_count(values_data.nan_count),
1517 )
1518 }
1519 _ => None,
1520 };
1521
1522 self.update_column_offset_index(
1524 page_statistics.as_ref(),
1525 values_data.variable_length_bytes,
1526 );
1527
1528 self.column_metrics
1530 .update_from_page_metrics(&self.page_metrics);
1531 self.column_metrics
1532 .update_variable_length_bytes(values_data.variable_length_bytes);
1533
1534 let page_statistics = page_statistics
1536 .filter(|_| self.column_props.write_page_header_statistics)
1537 .map(|stats| self.truncate_statistics(Statistics::from(stats)));
1538
1539 let compressed_page = match self.props.writer_version() {
1540 WriterVersion::PARQUET_1_0 => {
1541 let mut buffer = vec![];
1542
1543 if max_rep_level > 0 {
1544 self.rep_levels_encoder
1545 .flush_to(|data| buffer.extend_from_slice(data));
1546 }
1547
1548 if max_def_level > 0 {
1549 self.def_levels_encoder
1550 .flush_to(|data| buffer.extend_from_slice(data));
1551 }
1552
1553 buffer.extend_from_slice(&values_data.buf);
1554 let uncompressed_size = buffer.len();
1555
1556 if let Some(ref mut cmpr) = self.compressor {
1557 let mut compressed_buf = Vec::with_capacity(uncompressed_size);
1558 cmpr.compress(&buffer[..], &mut compressed_buf)?;
1559 compressed_buf.shrink_to_fit();
1560 buffer = compressed_buf;
1561 }
1562
1563 let data_page = Page::DataPage {
1564 buf: buffer.into(),
1565 num_values: self.page_metrics.num_buffered_values,
1566 encoding: values_data.encoding,
1567 def_level_encoding: Encoding::RLE,
1568 rep_level_encoding: Encoding::RLE,
1569 statistics: page_statistics,
1570 };
1571
1572 CompressedPage::new(data_page, uncompressed_size)
1573 }
1574 WriterVersion::PARQUET_2_0 => {
1575 let mut rep_levels_byte_len = 0;
1576 let mut def_levels_byte_len = 0;
1577 let mut buffer = vec![];
1578
1579 if max_rep_level > 0 {
1580 self.rep_levels_encoder
1581 .flush_to(|data| buffer.extend_from_slice(data));
1582 rep_levels_byte_len = buffer.len();
1583 }
1584
1585 if max_def_level > 0 {
1586 self.def_levels_encoder
1587 .flush_to(|data| buffer.extend_from_slice(data));
1588 def_levels_byte_len = buffer.len() - rep_levels_byte_len;
1589 }
1590
1591 let uncompressed_size =
1592 rep_levels_byte_len + def_levels_byte_len + values_data.buf.len();
1593
1594 let is_compressed = match self.compressor {
1596 Some(ref mut cmpr) => {
1597 let buffer_len = buffer.len();
1598 cmpr.compress(&values_data.buf, &mut buffer)?;
1599 let compressed_values_size = buffer.len() - buffer_len;
1600 let threshold = self.column_props.data_page_v2_compression_ratio_threshold;
1601 if (compressed_values_size as f64) >= (uncompressed_size as f64) * threshold
1602 {
1603 buffer.truncate(buffer_len);
1604 buffer.extend_from_slice(&values_data.buf);
1605 false
1606 } else {
1607 true
1608 }
1609 }
1610 None => {
1611 buffer.extend_from_slice(&values_data.buf);
1612 false
1613 }
1614 };
1615
1616 let data_page = Page::DataPageV2 {
1617 buf: buffer.into(),
1618 num_values: self.page_metrics.num_buffered_values,
1619 encoding: values_data.encoding,
1620 num_nulls: self.page_metrics.num_page_nulls as u32,
1621 num_rows: self.page_metrics.num_buffered_rows,
1622 def_levels_byte_len: def_levels_byte_len as u32,
1623 rep_levels_byte_len: rep_levels_byte_len as u32,
1624 is_compressed,
1625 statistics: page_statistics,
1626 };
1627
1628 CompressedPage::new(data_page, uncompressed_size)
1629 }
1630 };
1631
1632 if self.encoder.has_dictionary() && !self.page_writer.defers_dictionary_ordering() {
1641 self.data_pages.push_back(compressed_page);
1642 } else {
1643 self.write_data_page(compressed_page)?;
1644 }
1645
1646 self.column_metrics.total_rows_written += self.page_metrics.num_buffered_rows as u64;
1648 self.page_metrics.new_page();
1649
1650 Ok(())
1651 }
1652
1653 #[inline]
1656 fn flush_data_pages(&mut self) -> Result<()> {
1657 if self.page_metrics.num_buffered_values > 0 {
1659 self.add_data_page()?;
1660 }
1661
1662 while let Some(page) = self.data_pages.pop_front() {
1663 self.write_data_page(page)?;
1664 }
1665
1666 Ok(())
1667 }
1668
1669 fn build_column_metadata(&mut self) -> Result<ColumnChunkMetaData> {
1671 let total_compressed_size = self.column_metrics.total_compressed_size as i64;
1672 let total_uncompressed_size = self.column_metrics.total_uncompressed_size as i64;
1673 let num_values = self.column_metrics.total_num_values as i64;
1674 let dict_page_offset = self.column_metrics.dictionary_page_offset.map(|v| v as i64);
1675 let data_page_offset = self.column_metrics.data_page_offset.unwrap_or(0) as i64;
1677
1678 let mut builder = ColumnChunkMetaData::builder(self.descr.clone())
1679 .set_compression(self.codec)
1680 .set_encodings_mask(EncodingMask::new_from_encodings(self.encodings.iter()))
1681 .set_page_encoding_stats(self.encoding_stats.clone())
1682 .set_total_compressed_size(total_compressed_size)
1683 .set_total_uncompressed_size(total_uncompressed_size)
1684 .set_num_values(num_values)
1685 .set_data_page_offset(data_page_offset)
1686 .set_dictionary_page_offset(dict_page_offset);
1687
1688 if self.column_props.statistics_enabled != EnabledStatistics::None {
1689 let backwards_compatible_min_max = self.descr.sort_order().is_signed();
1690
1691 let distinct_count = self
1692 .distinct_count_override
1693 .or(self.column_metrics.column_distinct_count);
1694 let statistics = ValueStatistics::<E::T>::new(
1695 self.column_metrics.min_column_value.clone(),
1696 self.column_metrics.max_column_value.clone(),
1697 distinct_count,
1698 Some(self.column_metrics.num_column_nulls),
1699 false,
1700 )
1701 .with_nan_count(self.column_metrics.num_column_nans)
1702 .with_backwards_compatible_min_max(backwards_compatible_min_max)
1703 .into();
1704
1705 let statistics = self.truncate_statistics(statistics);
1706
1707 builder = builder
1708 .set_statistics(statistics)
1709 .set_unencoded_byte_array_data_bytes(self.column_metrics.variable_length_bytes)
1710 .set_repetition_level_histogram(
1711 self.column_metrics.repetition_level_histogram.take(),
1712 )
1713 .set_definition_level_histogram(
1714 self.column_metrics.definition_level_histogram.take(),
1715 );
1716
1717 if let Some(geo_stats) = self.encoder.flush_geospatial_statistics() {
1718 builder = builder.set_geo_statistics(geo_stats);
1719 }
1720 }
1721
1722 builder = self.set_column_chunk_encryption_properties(builder);
1723
1724 let metadata = builder.build()?;
1725 Ok(metadata)
1726 }
1727
1728 #[inline]
1730 fn write_data_page(&mut self, page: CompressedPage) -> Result<()> {
1731 self.encodings.insert(page.encoding());
1732 match self.encoding_stats.last_mut() {
1733 Some(encoding_stats)
1734 if encoding_stats.page_type == page.page_type()
1735 && encoding_stats.encoding == page.encoding() =>
1736 {
1737 encoding_stats.count += 1;
1738 }
1739 _ => {
1740 self.encoding_stats.push(PageEncodingStats {
1743 page_type: page.page_type(),
1744 encoding: page.encoding(),
1745 count: 1,
1746 });
1747 }
1748 }
1749 let page_spec = self.page_writer.write_page(page)?;
1750 if let Some(builder) = self.offset_index_builder.as_mut() {
1753 builder
1754 .append_offset_and_size(page_spec.offset as i64, page_spec.compressed_size as i32)
1755 }
1756 self.update_metrics_for_page(page_spec);
1757 Ok(())
1758 }
1759
1760 #[inline]
1762 fn write_dictionary_page(&mut self) -> Result<()> {
1763 let compressed_page = {
1764 let mut page = self
1765 .encoder
1766 .flush_dict_page()?
1767 .ok_or_else(|| general_err!("Dictionary encoder is not set"))?;
1768
1769 let uncompressed_size = page.buf.len();
1770
1771 if let Some(ref mut cmpr) = self.compressor {
1772 let mut output_buf = Vec::with_capacity(uncompressed_size);
1773 cmpr.compress(&page.buf, &mut output_buf)?;
1774 page.buf = Bytes::from(output_buf);
1775 }
1776
1777 let dict_page = Page::DictionaryPage {
1778 buf: page.buf,
1779 num_values: page.num_values as u32,
1780 encoding: self.props.dictionary_page_encoding(),
1781 is_sorted: page.is_sorted,
1782 };
1783 CompressedPage::new(dict_page, uncompressed_size)
1784 };
1785
1786 self.encodings.insert(compressed_page.encoding());
1787 self.encoding_stats.push(PageEncodingStats {
1788 page_type: PageType::DICTIONARY_PAGE,
1789 encoding: compressed_page.encoding(),
1790 count: 1,
1791 });
1792 let page_spec = self.page_writer.write_page(compressed_page)?;
1793 self.update_metrics_for_page(page_spec);
1794 Ok(())
1796 }
1797
1798 #[inline]
1800 fn update_metrics_for_page(&mut self, page_spec: PageWriteSpec) {
1801 self.column_metrics.total_uncompressed_size += page_spec.uncompressed_size as u64;
1802 self.column_metrics.total_compressed_size += page_spec.compressed_size as u64;
1803 self.column_metrics.total_bytes_written += page_spec.bytes_written;
1804
1805 match page_spec.page_type {
1806 PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
1807 self.column_metrics.total_num_values += page_spec.num_values as u64;
1808 if self.column_metrics.data_page_offset.is_none() {
1809 self.column_metrics.data_page_offset = Some(page_spec.offset);
1810 }
1811 }
1812 PageType::DICTIONARY_PAGE => {
1813 assert!(
1814 self.column_metrics.dictionary_page_offset.is_none(),
1815 "Dictionary offset is already set"
1816 );
1817 self.column_metrics.dictionary_page_offset = Some(page_spec.offset);
1818 }
1819 PageType::INDEX_PAGE => {}
1820 }
1821 }
1822
1823 #[inline]
1824 #[cfg(feature = "encryption")]
1825 fn set_column_chunk_encryption_properties(
1826 &self,
1827 builder: ColumnChunkMetaDataBuilder,
1828 ) -> ColumnChunkMetaDataBuilder {
1829 if let Some(encryption_properties) = self.props.file_encryption_properties.as_ref() {
1830 builder.set_column_crypto_metadata(get_column_crypto_metadata(
1831 encryption_properties,
1832 &self.descr,
1833 ))
1834 } else {
1835 builder
1836 }
1837 }
1838
1839 #[inline]
1840 #[cfg(not(feature = "encryption"))]
1841 fn set_column_chunk_encryption_properties(
1842 &self,
1843 builder: ColumnChunkMetaDataBuilder,
1844 ) -> ColumnChunkMetaDataBuilder {
1845 builder
1846 }
1847}
1848
1849fn update_min<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, min: &mut Option<T>) {
1850 match min {
1851 None => *min = Some(val.clone()),
1852 Some(min) => {
1853 let basic_type_info = descr.get_basic_info();
1854 let is_min_nan = is_nan(basic_type_info, min);
1855 let is_val_nan = is_nan(basic_type_info, val);
1856 match (is_min_nan, is_val_nan) {
1857 (false, true) => {}
1859 (true, false) => *min = val.clone(),
1861 _ => {
1863 update_stat::<T, _>(val, min, |cur| compare_greater(basic_type_info, cur, val))
1864 }
1865 }
1866 }
1867 }
1868}
1869
1870fn update_max<T: ParquetValueType>(descr: &ColumnDescriptor, val: &T, max: &mut Option<T>) {
1871 match max {
1872 None => *max = Some(val.clone()),
1873 Some(max) => {
1874 let basic_type_info = descr.get_basic_info();
1875 let is_max_nan = is_nan(basic_type_info, max);
1876 let is_val_nan = is_nan(basic_type_info, val);
1877 match (is_max_nan, is_val_nan) {
1878 (false, true) => {}
1880 (true, false) => *max = val.clone(),
1882 _ => {
1884 update_stat::<T, _>(val, max, |cur| compare_greater(basic_type_info, val, cur))
1885 }
1886 }
1887 }
1888 }
1889}
1890
1891#[inline]
1892#[expect(clippy::eq_op)]
1893fn is_nan<T: ParquetValueType>(basic_type_info: &BasicTypeInfo, val: &T) -> bool {
1894 match T::PHYSICAL_TYPE {
1895 Type::FLOAT | Type::DOUBLE => val != val,
1896 Type::FIXED_LEN_BYTE_ARRAY
1897 if matches!(basic_type_info.sort_order(), SortOrder::TOTAL_ORDER) =>
1898 {
1899 let val = val.as_bytes();
1901 let uval = ((val[1] as u16) << 8) | val[0] as u16;
1903 uval & 0x7FFFu16 > 0x7C00u16
1904 }
1905 _ => false,
1906 }
1907}
1908
1909fn update_stat<T: ParquetValueType, F>(val: &T, cur: &mut T, should_update: F)
1914where
1915 F: Fn(&T) -> bool,
1916{
1917 if should_update(cur) {
1918 *cur = val.clone();
1919 }
1920}
1921
1922fn compare_greater<T: ParquetValueType>(basic_type_info: &BasicTypeInfo, a: &T, b: &T) -> bool {
1924 match T::PHYSICAL_TYPE {
1925 Type::FLOAT => {
1926 let a = f32::from_le_bytes(a.as_bytes().try_into().unwrap());
1927 let b = f32::from_le_bytes(b.as_bytes().try_into().unwrap());
1928 return a.total_cmp(&b) == Ordering::Greater;
1929 }
1930 Type::DOUBLE => {
1931 let a = f64::from_le_bytes(a.as_bytes().try_into().unwrap());
1932 let b = f64::from_le_bytes(b.as_bytes().try_into().unwrap());
1933 return a.total_cmp(&b) == Ordering::Greater;
1934 }
1935 Type::INT32 | Type::INT64
1936 if matches!(basic_type_info.sort_order(), SortOrder::UNSIGNED) =>
1937 {
1938 return compare_greater_unsigned_int(a, b);
1939 }
1940 Type::FIXED_LEN_BYTE_ARRAY
1941 if matches!(basic_type_info.sort_order(), SortOrder::TOTAL_ORDER) =>
1942 {
1943 return compare_greater_f16(a.as_bytes(), b.as_bytes());
1944 }
1945 Type::FIXED_LEN_BYTE_ARRAY | Type::BYTE_ARRAY
1946 if matches!(basic_type_info.converted_type(), ConvertedType::DECIMAL)
1947 || matches!(
1948 basic_type_info.logical_type_ref(),
1949 Some(LogicalType::Decimal(_))
1950 ) =>
1951 {
1952 return compare_greater_byte_array_decimals(a.as_bytes(), b.as_bytes());
1953 }
1954
1955 _ => {}
1956 }
1957
1958 a > b
1960}
1961
1962fn fallback_encoding(kind: Type, props: &WriterProperties) -> Encoding {
1970 match (kind, props.writer_version()) {
1971 (Type::BOOLEAN, WriterVersion::PARQUET_2_0) => Encoding::RLE,
1972 (Type::INT32, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1973 (Type::INT64, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BINARY_PACKED,
1974 (Type::BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1975 (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => Encoding::DELTA_BYTE_ARRAY,
1976 _ => Encoding::PLAIN,
1977 }
1978}
1979
1980fn has_dictionary_support(kind: Type, props: &WriterProperties) -> bool {
1982 match (kind, props.writer_version()) {
1983 (Type::BOOLEAN, _) => false,
1985 (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_1_0) => false,
1987 (Type::FIXED_LEN_BYTE_ARRAY, WriterVersion::PARQUET_2_0) => true,
1988 _ => true,
1989 }
1990}
1991
1992#[inline]
1993fn compare_greater_unsigned_int<T: ParquetValueType>(a: &T, b: &T) -> bool {
1994 a.as_u64().unwrap() > b.as_u64().unwrap()
1995}
1996
1997#[inline]
1998fn compare_greater_f16(a: &[u8], b: &[u8]) -> bool {
1999 let a = f16::from_le_bytes(a.try_into().unwrap());
2000 let b = f16::from_le_bytes(b.try_into().unwrap());
2001 a.total_cmp(&b) == Ordering::Greater
2002}
2003
2004fn compare_greater_byte_array_decimals(a: &[u8], b: &[u8]) -> bool {
2006 let a_length = a.len();
2007 let b_length = b.len();
2008
2009 if a_length == 0 || b_length == 0 {
2010 return a_length > 0;
2011 }
2012
2013 let first_a: u8 = a[0];
2014 let first_b: u8 = b[0];
2015
2016 if (0x80 & first_a) != (0x80 & first_b) || (a_length == b_length && first_a != first_b) {
2021 return (first_a as i8) > (first_b as i8);
2022 }
2023
2024 let extension: u8 = if (first_a as i8) < 0 { 0xFF } else { 0 };
2031
2032 if a_length != b_length {
2033 let not_equal = if a_length > b_length {
2034 let lead_length = a_length - b_length;
2035 a[0..lead_length].iter().any(|&x| x != extension)
2036 } else {
2037 let lead_length = b_length - a_length;
2038 b[0..lead_length].iter().any(|&x| x != extension)
2039 };
2040
2041 if not_equal {
2042 let negative_values: bool = (first_a as i8) < 0;
2043 let a_longer: bool = a_length > b_length;
2044 return if negative_values { !a_longer } else { a_longer };
2045 }
2046 }
2047
2048 let tail_length = a_length.min(b_length);
2049 (a[a_length - tail_length..]) > (b[b_length - tail_length..])
2050}
2051
2052fn truncate_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
2058 let split = (1..=length).rfind(|x| data.is_char_boundary(*x))?;
2059 Some(data.as_bytes()[..split].to_vec())
2060}
2061
2062fn truncate_and_increment_utf8(data: &str, length: usize) -> Option<Vec<u8>> {
2068 let lower_bound = length.saturating_sub(3);
2070 let split = (lower_bound..=length).rfind(|x| data.is_char_boundary(*x))?;
2071 increment_utf8(data.get(..split)?)
2072}
2073
2074fn increment_utf8(data: &str) -> Option<Vec<u8>> {
2081 for (idx, original_char) in data.char_indices().rev() {
2082 let original_len = original_char.len_utf8();
2083 if let Some(next_char) = char::from_u32(original_char as u32 + 1) {
2084 if next_char.len_utf8() == original_len {
2086 let mut result = data.as_bytes()[..idx + original_len].to_vec();
2087 next_char.encode_utf8(&mut result[idx..]);
2088 return Some(result);
2089 }
2090 }
2091 }
2092
2093 None
2094}
2095
2096fn increment(mut data: Vec<u8>) -> Option<Vec<u8>> {
2100 for byte in data.iter_mut().rev() {
2101 let (incremented, overflow) = byte.overflowing_add(1);
2102 *byte = incremented;
2103
2104 if !overflow {
2105 return Some(data);
2106 }
2107 }
2108
2109 None
2110}
2111
2112#[cfg(test)]
2113mod tests {
2114 use crate::{
2115 file::{properties::DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, writer::SerializedFileWriter},
2116 schema::parser::parse_message_type,
2117 };
2118 use core::str;
2119 use rand::distr::uniform::SampleUniform;
2120 use std::{fs::File, sync::Arc};
2121
2122 use crate::column::{
2123 page::PageReader,
2124 reader::{ColumnReaderImpl, get_column_reader, get_typed_column_reader},
2125 };
2126 use crate::file::writer::TrackedWrite;
2127 use crate::file::{
2128 properties::ReaderProperties, reader::SerializedPageReader, writer::SerializedPageWriter,
2129 };
2130 use crate::schema::types::{ColumnPath, Type as SchemaType};
2131 use crate::util::test_common::rand_gen::random_numbers_range;
2132
2133 use super::*;
2134
2135 #[test]
2136 fn test_column_writer_inconsistent_def_rep_length() {
2137 let page_writer = get_test_page_writer();
2138 let props = Default::default();
2139 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 1, props);
2140 let res = writer.write_batch(&[1, 2, 3, 4], Some(&[1, 1, 1]), Some(&[0, 0]));
2141 assert!(res.is_err());
2142 if let Err(err) = res {
2143 assert_eq!(
2144 format!("{err}"),
2145 "Parquet error: Inconsistent length of definition and repetition levels: 3 != 2"
2146 );
2147 }
2148 }
2149
2150 #[test]
2151 fn test_column_writer_invalid_def_levels() {
2152 let page_writer = get_test_page_writer();
2153 let props = Default::default();
2154 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2155 let res = writer.write_batch(&[1, 2, 3, 4], None, None);
2156 assert!(res.is_err());
2157 if let Err(err) = res {
2158 assert_eq!(
2159 format!("{err}"),
2160 "Parquet error: Definition levels are required, because max definition level = 1"
2161 );
2162 }
2163 }
2164
2165 #[test]
2166 fn test_column_writer_invalid_rep_levels() {
2167 let page_writer = get_test_page_writer();
2168 let props = Default::default();
2169 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 1, props);
2170 let res = writer.write_batch(&[1, 2, 3, 4], None, None);
2171 assert!(res.is_err());
2172 if let Err(err) = res {
2173 assert_eq!(
2174 format!("{err}"),
2175 "Parquet error: Repetition levels are required, because max repetition level = 1"
2176 );
2177 }
2178 }
2179
2180 #[test]
2181 fn test_column_writer_not_enough_values_to_write() {
2182 let page_writer = get_test_page_writer();
2183 let props = Default::default();
2184 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2185 let res = writer.write_batch(&[1, 2], Some(&[1, 1, 1, 1]), None);
2186 assert!(res.is_err());
2187 if let Err(err) = res {
2188 assert_eq!(
2189 format!("{err}"),
2190 "Parquet error: Expected to write 4 values, but have only 2"
2191 );
2192 }
2193 }
2194
2195 #[test]
2196 fn test_column_writer_write_only_one_dictionary_page() {
2197 let page_writer = get_test_page_writer();
2198 let props = Default::default();
2199 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2200 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2201 writer.add_data_page().unwrap();
2203 writer.write_dictionary_page().unwrap();
2204 let err = writer.write_dictionary_page().unwrap_err().to_string();
2205 assert_eq!(err, "Parquet error: Dictionary encoder is not set");
2206 }
2207
2208 #[test]
2209 fn test_column_writer_error_when_writing_disabled_dictionary() {
2210 let page_writer = get_test_page_writer();
2211 let props = Arc::new(
2212 WriterProperties::builder()
2213 .set_dictionary_enabled(false)
2214 .build(),
2215 );
2216 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2217 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2218 let err = writer.write_dictionary_page().unwrap_err().to_string();
2219 assert_eq!(err, "Parquet error: Dictionary encoder is not set");
2220 }
2221
2222 #[test]
2223 fn test_column_writer_boolean_type_does_not_support_dictionary() {
2224 let page_writer = get_test_page_writer();
2225 let props = Arc::new(
2226 WriterProperties::builder()
2227 .set_dictionary_enabled(true)
2228 .build(),
2229 );
2230 let mut writer = get_test_column_writer::<BoolType>(page_writer, 0, 0, props);
2231 writer
2232 .write_batch(&[true, false, true, false], None, None)
2233 .unwrap();
2234
2235 let r = writer.close().unwrap();
2236 assert_eq!(r.bytes_written, 1);
2239 assert_eq!(r.rows_written, 4);
2240
2241 let metadata = r.metadata;
2242 assert_eq!(
2243 metadata.encodings().collect::<Vec<_>>(),
2244 vec![Encoding::PLAIN, Encoding::RLE]
2245 );
2246 assert_eq!(metadata.num_values(), 4); assert_eq!(metadata.dictionary_page_offset(), None);
2248 }
2249
2250 #[test]
2251 fn test_bloom_filter_for_dictionary_encoded_chunks() {
2252 fn bloom_filter_written(dictionary_enabled: bool, for_dictionary_chunks: bool) -> bool {
2253 let props = Arc::new(
2254 WriterProperties::builder()
2255 .set_dictionary_enabled(dictionary_enabled)
2256 .set_bloom_filter_enabled(true)
2257 .set_bloom_filter_for_dictionary_encoded_chunks(for_dictionary_chunks)
2258 .build(),
2259 );
2260 let mut writer =
2261 get_test_column_writer::<Int32Type>(get_test_page_writer(), 0, 0, props);
2262 writer.write_batch(&[1, 2, 1, 2], None, None).unwrap();
2263 writer.close().unwrap().bloom_filter.is_some()
2264 }
2265
2266 assert!(bloom_filter_written(true, true));
2267 assert!(bloom_filter_written(false, true));
2268 assert!(bloom_filter_written(false, false));
2269 assert!(!bloom_filter_written(true, false));
2270 }
2271
2272 #[test]
2273 fn test_column_writer_default_encoding_support_bool() {
2274 check_encoding_write_support::<BoolType>(
2275 WriterVersion::PARQUET_1_0,
2276 true,
2277 &[true, false],
2278 None,
2279 &[Encoding::PLAIN, Encoding::RLE],
2280 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2281 );
2282 check_encoding_write_support::<BoolType>(
2283 WriterVersion::PARQUET_1_0,
2284 false,
2285 &[true, false],
2286 None,
2287 &[Encoding::PLAIN, Encoding::RLE],
2288 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2289 );
2290 check_encoding_write_support::<BoolType>(
2291 WriterVersion::PARQUET_2_0,
2292 true,
2293 &[true, false],
2294 None,
2295 &[Encoding::RLE],
2296 &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
2297 );
2298 check_encoding_write_support::<BoolType>(
2299 WriterVersion::PARQUET_2_0,
2300 false,
2301 &[true, false],
2302 None,
2303 &[Encoding::RLE],
2304 &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE, 1)],
2305 );
2306 }
2307
2308 #[test]
2309 fn test_column_writer_default_encoding_support_int32() {
2310 check_encoding_write_support::<Int32Type>(
2311 WriterVersion::PARQUET_1_0,
2312 true,
2313 &[1, 2],
2314 Some(0),
2315 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2316 &[
2317 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2318 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2319 ],
2320 );
2321 check_encoding_write_support::<Int32Type>(
2322 WriterVersion::PARQUET_1_0,
2323 false,
2324 &[1, 2],
2325 None,
2326 &[Encoding::PLAIN, Encoding::RLE],
2327 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2328 );
2329 check_encoding_write_support::<Int32Type>(
2330 WriterVersion::PARQUET_2_0,
2331 true,
2332 &[1, 2],
2333 Some(0),
2334 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2335 &[
2336 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2337 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2338 ],
2339 );
2340 check_encoding_write_support::<Int32Type>(
2341 WriterVersion::PARQUET_2_0,
2342 false,
2343 &[1, 2],
2344 None,
2345 &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
2346 &[encoding_stats(
2347 PageType::DATA_PAGE_V2,
2348 Encoding::DELTA_BINARY_PACKED,
2349 1,
2350 )],
2351 );
2352 }
2353
2354 #[test]
2355 fn test_column_writer_default_encoding_support_int64() {
2356 check_encoding_write_support::<Int64Type>(
2357 WriterVersion::PARQUET_1_0,
2358 true,
2359 &[1, 2],
2360 Some(0),
2361 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2362 &[
2363 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2364 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2365 ],
2366 );
2367 check_encoding_write_support::<Int64Type>(
2368 WriterVersion::PARQUET_1_0,
2369 false,
2370 &[1, 2],
2371 None,
2372 &[Encoding::PLAIN, Encoding::RLE],
2373 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2374 );
2375 check_encoding_write_support::<Int64Type>(
2376 WriterVersion::PARQUET_2_0,
2377 true,
2378 &[1, 2],
2379 Some(0),
2380 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2381 &[
2382 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2383 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2384 ],
2385 );
2386 check_encoding_write_support::<Int64Type>(
2387 WriterVersion::PARQUET_2_0,
2388 false,
2389 &[1, 2],
2390 None,
2391 &[Encoding::RLE, Encoding::DELTA_BINARY_PACKED],
2392 &[encoding_stats(
2393 PageType::DATA_PAGE_V2,
2394 Encoding::DELTA_BINARY_PACKED,
2395 1,
2396 )],
2397 );
2398 }
2399
2400 #[test]
2401 fn test_column_writer_default_encoding_support_int96() {
2402 check_encoding_write_support::<Int96Type>(
2403 WriterVersion::PARQUET_1_0,
2404 true,
2405 &[Int96::from(vec![1, 2, 3])],
2406 Some(0),
2407 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2408 &[
2409 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2410 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2411 ],
2412 );
2413 check_encoding_write_support::<Int96Type>(
2414 WriterVersion::PARQUET_1_0,
2415 false,
2416 &[Int96::from(vec![1, 2, 3])],
2417 None,
2418 &[Encoding::PLAIN, Encoding::RLE],
2419 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2420 );
2421 check_encoding_write_support::<Int96Type>(
2422 WriterVersion::PARQUET_2_0,
2423 true,
2424 &[Int96::from(vec![1, 2, 3])],
2425 Some(0),
2426 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2427 &[
2428 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2429 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2430 ],
2431 );
2432 check_encoding_write_support::<Int96Type>(
2433 WriterVersion::PARQUET_2_0,
2434 false,
2435 &[Int96::from(vec![1, 2, 3])],
2436 None,
2437 &[Encoding::PLAIN, Encoding::RLE],
2438 &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2439 );
2440 }
2441
2442 #[test]
2443 fn test_column_writer_default_encoding_support_float() {
2444 check_encoding_write_support::<FloatType>(
2445 WriterVersion::PARQUET_1_0,
2446 true,
2447 &[1.0, 2.0],
2448 Some(0),
2449 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2450 &[
2451 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2452 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2453 ],
2454 );
2455 check_encoding_write_support::<FloatType>(
2456 WriterVersion::PARQUET_1_0,
2457 false,
2458 &[1.0, 2.0],
2459 None,
2460 &[Encoding::PLAIN, Encoding::RLE],
2461 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2462 );
2463 check_encoding_write_support::<FloatType>(
2464 WriterVersion::PARQUET_2_0,
2465 true,
2466 &[1.0, 2.0],
2467 Some(0),
2468 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2469 &[
2470 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2471 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2472 ],
2473 );
2474 check_encoding_write_support::<FloatType>(
2475 WriterVersion::PARQUET_2_0,
2476 false,
2477 &[1.0, 2.0],
2478 None,
2479 &[Encoding::PLAIN, Encoding::RLE],
2480 &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2481 );
2482 }
2483
2484 #[test]
2485 fn test_column_writer_default_encoding_support_double() {
2486 check_encoding_write_support::<DoubleType>(
2487 WriterVersion::PARQUET_1_0,
2488 true,
2489 &[1.0, 2.0],
2490 Some(0),
2491 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2492 &[
2493 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2494 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2495 ],
2496 );
2497 check_encoding_write_support::<DoubleType>(
2498 WriterVersion::PARQUET_1_0,
2499 false,
2500 &[1.0, 2.0],
2501 None,
2502 &[Encoding::PLAIN, Encoding::RLE],
2503 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2504 );
2505 check_encoding_write_support::<DoubleType>(
2506 WriterVersion::PARQUET_2_0,
2507 true,
2508 &[1.0, 2.0],
2509 Some(0),
2510 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2511 &[
2512 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2513 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2514 ],
2515 );
2516 check_encoding_write_support::<DoubleType>(
2517 WriterVersion::PARQUET_2_0,
2518 false,
2519 &[1.0, 2.0],
2520 None,
2521 &[Encoding::PLAIN, Encoding::RLE],
2522 &[encoding_stats(PageType::DATA_PAGE_V2, Encoding::PLAIN, 1)],
2523 );
2524 }
2525
2526 #[test]
2527 fn test_column_writer_default_encoding_support_byte_array() {
2528 check_encoding_write_support::<ByteArrayType>(
2529 WriterVersion::PARQUET_1_0,
2530 true,
2531 &[ByteArray::from(vec![1u8])],
2532 Some(0),
2533 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2534 &[
2535 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2536 encoding_stats(PageType::DATA_PAGE, Encoding::RLE_DICTIONARY, 1),
2537 ],
2538 );
2539 check_encoding_write_support::<ByteArrayType>(
2540 WriterVersion::PARQUET_1_0,
2541 false,
2542 &[ByteArray::from(vec![1u8])],
2543 None,
2544 &[Encoding::PLAIN, Encoding::RLE],
2545 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2546 );
2547 check_encoding_write_support::<ByteArrayType>(
2548 WriterVersion::PARQUET_2_0,
2549 true,
2550 &[ByteArray::from(vec![1u8])],
2551 Some(0),
2552 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2553 &[
2554 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2555 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2556 ],
2557 );
2558 check_encoding_write_support::<ByteArrayType>(
2559 WriterVersion::PARQUET_2_0,
2560 false,
2561 &[ByteArray::from(vec![1u8])],
2562 None,
2563 &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2564 &[encoding_stats(
2565 PageType::DATA_PAGE_V2,
2566 Encoding::DELTA_BYTE_ARRAY,
2567 1,
2568 )],
2569 );
2570 }
2571
2572 #[test]
2573 fn test_column_writer_default_encoding_support_fixed_len_byte_array() {
2574 check_encoding_write_support::<FixedLenByteArrayType>(
2575 WriterVersion::PARQUET_1_0,
2576 true,
2577 &[ByteArray::from(vec![1u8]).into()],
2578 None,
2579 &[Encoding::PLAIN, Encoding::RLE],
2580 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2581 );
2582 check_encoding_write_support::<FixedLenByteArrayType>(
2583 WriterVersion::PARQUET_1_0,
2584 false,
2585 &[ByteArray::from(vec![1u8]).into()],
2586 None,
2587 &[Encoding::PLAIN, Encoding::RLE],
2588 &[encoding_stats(PageType::DATA_PAGE, Encoding::PLAIN, 1)],
2589 );
2590 check_encoding_write_support::<FixedLenByteArrayType>(
2591 WriterVersion::PARQUET_2_0,
2592 true,
2593 &[ByteArray::from(vec![1u8]).into()],
2594 Some(0),
2595 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
2596 &[
2597 encoding_stats(PageType::DICTIONARY_PAGE, Encoding::PLAIN, 1),
2598 encoding_stats(PageType::DATA_PAGE_V2, Encoding::RLE_DICTIONARY, 1),
2599 ],
2600 );
2601 check_encoding_write_support::<FixedLenByteArrayType>(
2602 WriterVersion::PARQUET_2_0,
2603 false,
2604 &[ByteArray::from(vec![1u8]).into()],
2605 None,
2606 &[Encoding::RLE, Encoding::DELTA_BYTE_ARRAY],
2607 &[encoding_stats(
2608 PageType::DATA_PAGE_V2,
2609 Encoding::DELTA_BYTE_ARRAY,
2610 1,
2611 )],
2612 );
2613 }
2614
2615 #[test]
2616 fn test_column_writer_check_metadata() {
2617 let page_writer = get_test_page_writer();
2618 let props = Default::default();
2619 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2620 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2621
2622 let r = writer.close().unwrap();
2623 assert_eq!(r.bytes_written, 20);
2624 assert_eq!(r.rows_written, 4);
2625
2626 let metadata = r.metadata;
2627 assert_eq!(
2628 metadata.encodings().collect::<Vec<_>>(),
2629 vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2630 );
2631 assert_eq!(metadata.num_values(), 4);
2632 assert_eq!(metadata.compressed_size(), 20);
2633 assert_eq!(metadata.uncompressed_size(), 20);
2634 assert_eq!(metadata.data_page_offset(), 0);
2635 assert_eq!(metadata.dictionary_page_offset(), Some(0));
2636 if let Some(stats) = metadata.statistics() {
2637 assert_eq!(stats.null_count_opt(), Some(0));
2638 assert_eq!(stats.distinct_count_opt(), None);
2639 if let Statistics::Int32(stats) = stats {
2640 assert_eq!(stats.min_opt().unwrap(), &1);
2641 assert_eq!(stats.max_opt().unwrap(), &4);
2642 } else {
2643 panic!("expecting Statistics::Int32");
2644 }
2645 } else {
2646 panic!("metadata missing statistics");
2647 }
2648 }
2649
2650 #[test]
2651 fn test_column_writer_check_byte_array_min_max() {
2652 let page_writer = get_test_page_writer();
2653 let props = Default::default();
2654 let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
2655 writer
2656 .write_batch(
2657 &[
2658 ByteArray::from(vec![
2659 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2660 35u8, 231u8, 90u8, 0u8, 0u8,
2661 ]),
2662 ByteArray::from(vec![
2663 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 228u8, 62u8, 146u8,
2664 152u8, 177u8, 56u8, 0u8, 0u8,
2665 ]),
2666 ByteArray::from(vec![
2667 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8,
2668 0u8,
2669 ]),
2670 ByteArray::from(vec![
2671 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2672 44u8, 0u8, 0u8,
2673 ]),
2674 ],
2675 None,
2676 None,
2677 )
2678 .unwrap();
2679 let metadata = writer.close().unwrap().metadata;
2680 if let Some(stats) = metadata.statistics() {
2681 if let Statistics::ByteArray(stats) = stats {
2682 assert_eq!(
2683 stats.min_opt().unwrap(),
2684 &ByteArray::from(vec![
2685 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8,
2686 35u8, 231u8, 90u8, 0u8, 0u8,
2687 ])
2688 );
2689 assert_eq!(
2690 stats.max_opt().unwrap(),
2691 &ByteArray::from(vec![
2692 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 41u8, 162u8, 36u8, 26u8, 246u8,
2693 44u8, 0u8, 0u8,
2694 ])
2695 );
2696 } else {
2697 panic!("expecting Statistics::ByteArray");
2698 }
2699 } else {
2700 panic!("metadata missing statistics");
2701 }
2702 }
2703
2704 #[test]
2705 fn test_column_writer_byte_array_min_max_unequal_lengths() {
2706 let page_writer = get_test_page_writer();
2709 let props = Default::default();
2710 let mut writer = get_test_decimals_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
2711 writer
2712 .write_batch(
2713 &[
2714 ByteArray::from(vec![0u8, 255u8]), ByteArray::from(vec![0u8, 128u8, 0u8]), ByteArray::from(vec![255u8, 127u8]), ByteArray::from(vec![128u8]), ],
2719 None,
2720 None,
2721 )
2722 .unwrap();
2723 let metadata = writer.close().unwrap().metadata;
2724 let stats = metadata.statistics().expect("metadata missing statistics");
2725 let Statistics::ByteArray(stats) = stats else {
2726 panic!("expecting Statistics::ByteArray");
2727 };
2728 assert_eq!(
2730 stats.min_opt().unwrap(),
2731 &ByteArray::from(vec![255u8, 127u8])
2732 );
2733 assert_eq!(
2735 stats.max_opt().unwrap(),
2736 &ByteArray::from(vec![0u8, 128u8, 0u8])
2737 );
2738 }
2739
2740 #[test]
2741 fn test_column_writer_uint32_converted_type_min_max() {
2742 let page_writer = get_test_page_writer();
2743 let props = Default::default();
2744 let mut writer = get_test_unsigned_int_given_as_converted_column_writer::<Int32Type>(
2745 page_writer,
2746 0,
2747 0,
2748 props,
2749 );
2750 writer.write_batch(&[0, 1, 2, 3, 4, 5], None, None).unwrap();
2751 let metadata = writer.close().unwrap().metadata;
2752 if let Some(stats) = metadata.statistics() {
2753 if let Statistics::Int32(stats) = stats {
2754 assert_eq!(stats.min_opt().unwrap(), &0,);
2755 assert_eq!(stats.max_opt().unwrap(), &5,);
2756 } else {
2757 panic!("expecting Statistics::Int32");
2758 }
2759 } else {
2760 panic!("metadata missing statistics");
2761 }
2762 }
2763
2764 #[test]
2765 fn test_column_writer_precalculated_statistics() {
2766 let page_writer = get_test_page_writer();
2767 let props = Arc::new(
2768 WriterProperties::builder()
2769 .set_statistics_enabled(EnabledStatistics::Chunk)
2770 .build(),
2771 );
2772 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2773 writer
2774 .write_batch_with_statistics(
2775 &[1, 2, 3, 4],
2776 None,
2777 None,
2778 Some(&-17),
2779 Some(&9000),
2780 Some(55),
2781 )
2782 .unwrap();
2783
2784 let r = writer.close().unwrap();
2785 assert_eq!(r.bytes_written, 20);
2786 assert_eq!(r.rows_written, 4);
2787
2788 let metadata = r.metadata;
2789 assert_eq!(
2790 metadata.encodings().collect::<Vec<_>>(),
2791 vec![Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY]
2792 );
2793 assert_eq!(metadata.num_values(), 4);
2794 assert_eq!(metadata.compressed_size(), 20);
2795 assert_eq!(metadata.uncompressed_size(), 20);
2796 assert_eq!(metadata.data_page_offset(), 0);
2797 assert_eq!(metadata.dictionary_page_offset(), Some(0));
2798 if let Some(stats) = metadata.statistics() {
2799 assert_eq!(stats.null_count_opt(), Some(0));
2800 assert_eq!(stats.distinct_count_opt().unwrap_or(0), 55);
2801 if let Statistics::Int32(stats) = stats {
2802 assert_eq!(stats.min_opt().unwrap(), &-17);
2803 assert_eq!(stats.max_opt().unwrap(), &9000);
2804 } else {
2805 panic!("expecting Statistics::Int32");
2806 }
2807 } else {
2808 panic!("metadata missing statistics");
2809 }
2810 }
2811
2812 #[test]
2813 fn test_mixed_precomputed_statistics() {
2814 let mut buf = Vec::with_capacity(100);
2815 let mut write = TrackedWrite::new(&mut buf);
2816 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2817 let props = Arc::new(
2818 WriterProperties::builder()
2819 .set_write_page_header_statistics(true)
2820 .build(),
2821 );
2822 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
2823
2824 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
2825 writer
2826 .write_batch_with_statistics(&[5, 6, 7], None, None, Some(&5), Some(&7), Some(3))
2827 .unwrap();
2828
2829 let r = writer.close().unwrap();
2830
2831 let stats = r.metadata.statistics().unwrap();
2832 assert_eq!(stats.min_bytes_opt().unwrap(), 1_i32.to_le_bytes());
2833 assert_eq!(stats.max_bytes_opt().unwrap(), 7_i32.to_le_bytes());
2834 assert_eq!(stats.null_count_opt(), Some(0));
2835 assert!(stats.distinct_count_opt().is_none());
2836
2837 drop(write);
2838
2839 let props = ReaderProperties::builder()
2840 .set_backward_compatible_lz4(false)
2841 .set_read_page_statistics(true)
2842 .build();
2843 let reader = SerializedPageReader::new_with_properties(
2844 Arc::new(Bytes::from(buf)),
2845 &r.metadata,
2846 r.rows_written as usize,
2847 None,
2848 Arc::new(props),
2849 )
2850 .unwrap();
2851
2852 let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2853 assert_eq!(pages.len(), 2);
2854
2855 assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2856 assert_eq!(pages[1].page_type(), PageType::DATA_PAGE);
2857
2858 let page_statistics = pages[1].statistics().unwrap();
2859 assert_eq!(
2860 page_statistics.min_bytes_opt().unwrap(),
2861 1_i32.to_le_bytes()
2862 );
2863 assert_eq!(
2864 page_statistics.max_bytes_opt().unwrap(),
2865 7_i32.to_le_bytes()
2866 );
2867 assert_eq!(page_statistics.null_count_opt(), Some(0));
2868 assert!(page_statistics.distinct_count_opt().is_none());
2869 }
2870
2871 #[test]
2872 fn test_disabled_statistics() {
2873 let mut buf = Vec::with_capacity(100);
2874 let mut write = TrackedWrite::new(&mut buf);
2875 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
2876 let props = WriterProperties::builder()
2877 .set_statistics_enabled(EnabledStatistics::None)
2878 .set_writer_version(WriterVersion::PARQUET_2_0)
2879 .build();
2880 let props = Arc::new(props);
2881
2882 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
2883 writer
2884 .write_batch(&[1, 2, 3, 4], Some(&[1, 0, 0, 1, 1, 1]), None)
2885 .unwrap();
2886
2887 let r = writer.close().unwrap();
2888 assert!(r.metadata.statistics().is_none());
2889
2890 drop(write);
2891
2892 let props = ReaderProperties::builder()
2893 .set_backward_compatible_lz4(false)
2894 .build();
2895 let reader = SerializedPageReader::new_with_properties(
2896 Arc::new(Bytes::from(buf)),
2897 &r.metadata,
2898 r.rows_written as usize,
2899 None,
2900 Arc::new(props),
2901 )
2902 .unwrap();
2903
2904 let pages = reader.collect::<Result<Vec<_>>>().unwrap();
2905 assert_eq!(pages.len(), 2);
2906
2907 assert_eq!(pages[0].page_type(), PageType::DICTIONARY_PAGE);
2908 assert_eq!(pages[1].page_type(), PageType::DATA_PAGE_V2);
2909
2910 match &pages[1] {
2911 Page::DataPageV2 {
2912 num_values,
2913 num_nulls,
2914 num_rows,
2915 statistics,
2916 ..
2917 } => {
2918 assert_eq!(*num_values, 6);
2919 assert_eq!(*num_nulls, 2);
2920 assert_eq!(*num_rows, 6);
2921 assert!(statistics.is_none());
2922 }
2923 _ => unreachable!(),
2924 }
2925 }
2926
2927 #[test]
2928 fn test_column_writer_empty_column_roundtrip() {
2929 let props = Default::default();
2930 column_roundtrip::<Int32Type>(props, &[], None, None);
2931 }
2932
2933 #[test]
2934 fn test_column_writer_non_nullable_values_roundtrip() {
2935 let props = Default::default();
2936 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 0, 0);
2937 }
2938
2939 #[test]
2940 fn test_column_writer_nullable_non_repeated_values_roundtrip() {
2941 let props = Default::default();
2942 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 0);
2943 }
2944
2945 #[test]
2946 fn test_column_writer_nullable_repeated_values_roundtrip() {
2947 let props = Default::default();
2948 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2949 }
2950
2951 #[test]
2952 fn test_column_writer_dictionary_fallback_small_data_page() {
2953 let props = WriterProperties::builder()
2954 .set_dictionary_page_size_limit(32)
2955 .set_data_page_size_limit(32)
2956 .build();
2957 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2958 }
2959
2960 #[test]
2961 #[cfg_attr(miri, ignore)] fn test_column_writer_small_write_batch_size() {
2963 for i in &[1usize, 2, 5, 10, 11, 1023] {
2964 let props = WriterProperties::builder().set_write_batch_size(*i).build();
2965
2966 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2967 }
2968 }
2969
2970 #[test]
2971 fn test_column_writer_dictionary_disabled_v1() {
2972 let props = WriterProperties::builder()
2973 .set_writer_version(WriterVersion::PARQUET_1_0)
2974 .set_dictionary_enabled(false)
2975 .build();
2976 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2977 }
2978
2979 #[test]
2980 fn test_column_writer_dictionary_disabled_v2() {
2981 let props = WriterProperties::builder()
2982 .set_writer_version(WriterVersion::PARQUET_2_0)
2983 .set_dictionary_enabled(false)
2984 .build();
2985 column_roundtrip_random::<Int32Type>(props, 1024, i32::MIN, i32::MAX, 10, 10);
2986 }
2987
2988 #[test]
2989 fn test_column_writer_compression_v1() {
2990 let props = WriterProperties::builder()
2991 .set_writer_version(WriterVersion::PARQUET_1_0)
2992 .set_compression(Compression::SNAPPY)
2993 .build();
2994 column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
2995 }
2996
2997 #[test]
2998 fn test_column_writer_compression_v2() {
2999 let props = WriterProperties::builder()
3000 .set_writer_version(WriterVersion::PARQUET_2_0)
3001 .set_compression(Compression::SNAPPY)
3002 .build();
3003 column_roundtrip_random::<Int32Type>(props, 2048, i32::MIN, i32::MAX, 10, 10);
3004 }
3005
3006 #[test]
3007 fn test_column_writer_v2_compression_ratio_threshold() {
3008 fn write_v2_page(threshold: f64) -> bool {
3009 let mut buf = Vec::with_capacity(4096);
3010 let mut write = TrackedWrite::new(&mut buf);
3011 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
3012 let props = Arc::new(
3013 WriterProperties::builder()
3014 .set_writer_version(WriterVersion::PARQUET_2_0)
3015 .set_compression(Compression::SNAPPY)
3016 .set_dictionary_enabled(false)
3017 .set_data_page_v2_compression_ratio_threshold(threshold)
3018 .build(),
3019 );
3020
3021 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3022 let values: Vec<i32> = vec![42; 4096];
3023 writer.write_batch(&values, None, None).unwrap();
3024 let r = writer.close().unwrap();
3025 drop(write);
3026
3027 let reader_props = ReaderProperties::builder()
3028 .set_backward_compatible_lz4(false)
3029 .build();
3030 let reader = SerializedPageReader::new_with_properties(
3031 Arc::new(Bytes::from(buf)),
3032 &r.metadata,
3033 r.rows_written as usize,
3034 None,
3035 Arc::new(reader_props),
3036 )
3037 .unwrap();
3038 let pages = reader.collect::<Result<Vec<_>>>().unwrap();
3039 let data_page = pages
3040 .iter()
3041 .find(|p| p.page_type() == PageType::DATA_PAGE_V2)
3042 .expect("expected a v2 data page");
3043 match data_page {
3044 Page::DataPageV2 { is_compressed, .. } => *is_compressed,
3045 _ => unreachable!(),
3046 }
3047 }
3048
3049 assert!(write_v2_page(1.0));
3051 assert!(!write_v2_page(0.001));
3053 }
3054
3055 #[test]
3056 fn test_column_writer_add_data_pages_with_dict() {
3057 let mut file = tempfile::tempfile().unwrap();
3060 let mut write = TrackedWrite::new(&mut file);
3061 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
3062 let props = Arc::new(
3063 WriterProperties::builder()
3064 .set_data_page_size_limit(10)
3065 .set_write_batch_size(3) .build(),
3067 );
3068 let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
3069 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
3070 writer.write_batch(data, None, None).unwrap();
3071 let r = writer.close().unwrap();
3072
3073 drop(write);
3074
3075 let props = ReaderProperties::builder()
3077 .set_backward_compatible_lz4(false)
3078 .build();
3079 let mut page_reader = Box::new(
3080 SerializedPageReader::new_with_properties(
3081 Arc::new(file),
3082 &r.metadata,
3083 r.rows_written as usize,
3084 None,
3085 Arc::new(props),
3086 )
3087 .unwrap(),
3088 );
3089 let mut res = Vec::new();
3090 while let Some(page) = page_reader.get_next_page().unwrap() {
3091 res.push((page.page_type(), page.num_values(), page.buffer().len()));
3092 }
3093 assert_eq!(
3094 res,
3095 vec![
3096 (PageType::DICTIONARY_PAGE, 10, 40),
3097 (PageType::DATA_PAGE, 9, 10),
3098 (PageType::DATA_PAGE, 1, 3),
3099 ]
3100 );
3101 assert_eq!(
3102 r.metadata.page_encoding_stats(),
3103 Some(&vec![
3104 PageEncodingStats {
3105 page_type: PageType::DICTIONARY_PAGE,
3106 encoding: Encoding::PLAIN,
3107 count: 1
3108 },
3109 PageEncodingStats {
3110 page_type: PageType::DATA_PAGE,
3111 encoding: Encoding::RLE_DICTIONARY,
3112 count: 2,
3113 }
3114 ])
3115 );
3116 }
3117
3118 #[test]
3119 fn test_column_writer_column_data_page_size_limit() {
3120 let props = Arc::new(
3121 WriterProperties::builder()
3122 .set_writer_version(WriterVersion::PARQUET_1_0)
3123 .set_dictionary_enabled(false)
3124 .set_data_page_size_limit(1000)
3125 .set_column_data_page_size_limit(ColumnPath::from("col"), 10)
3126 .set_write_batch_size(3)
3127 .build(),
3128 );
3129 let data = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
3130
3131 let col_values =
3132 write_and_collect_page_values(ColumnPath::from("col"), Arc::clone(&props), data);
3133 let other_values = write_and_collect_page_values(ColumnPath::from("other"), props, data);
3134
3135 assert_eq!(col_values, vec![3, 3, 3, 1]);
3136 assert_eq!(other_values, vec![10]);
3137 }
3138
3139 #[test]
3140 fn test_column_writer_caps_page_size_for_large_byte_array_values() {
3141 let value_size = 64 * 1024; let page_byte_limit = 16 * 1024; let num_rows = 64;
3149
3150 let props = WriterProperties::builder()
3151 .set_writer_version(WriterVersion::PARQUET_1_0)
3152 .set_dictionary_enabled(false)
3153 .set_encoding(Encoding::PLAIN)
3154 .set_data_page_size_limit(page_byte_limit)
3155 .build();
3158
3159 let data: Vec<_> = (0..num_rows)
3160 .map(|i| ByteArray::from(vec![i as u8; value_size]))
3161 .collect();
3162 let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3163
3164 let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3166 assert_eq!(total_values as usize, num_rows);
3167 assert!(
3171 pages.data_pages.len() >= num_rows / 2,
3172 "expected pages to be cut close to one per value, got {:?}",
3173 pages.data_pages,
3174 );
3175 for (size, _) in &pages.data_pages {
3179 assert!(
3180 *size <= value_size + 64,
3181 "page size {size} exceeds one-value bound ({}B) — pages {:?}",
3182 value_size + 64,
3183 pages.data_pages,
3184 );
3185 }
3186 }
3187
3188 #[test]
3189 fn test_column_writer_delta_byte_array_dedups_large_shared_prefix_values() {
3190 let value_size = 64 * 1024; let page_byte_limit = 16 * 1024;
3196 let num_rows = 16;
3197
3198 let props = WriterProperties::builder()
3199 .set_writer_version(WriterVersion::PARQUET_1_0)
3200 .set_dictionary_enabled(false)
3201 .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3202 .set_data_page_size_limit(page_byte_limit)
3203 .set_statistics_enabled(EnabledStatistics::None)
3204 .build();
3205
3206 let data: Vec<_> = (0..num_rows)
3209 .map(|_| ByteArray::from(vec![b'a'; value_size]))
3210 .collect();
3211 let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3212
3213 let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3215 assert_eq!(total_values as usize, num_rows);
3216
3217 let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum();
3220 assert!(
3221 total_bytes < 2 * value_size,
3222 "expected under 2x a single value ({}B) for {num_rows} identical \
3223 values, got {total_bytes}B across pages {:?}",
3224 2 * value_size,
3225 pages.data_pages,
3226 );
3227 }
3228
3229 #[test]
3230 fn test_column_writer_delta_byte_array_bounds_pages_without_shared_prefix() {
3231 let value_size = 64 * 1024;
3236 let page_byte_limit = 16 * 1024;
3237 let num_rows = 16;
3238
3239 let props = WriterProperties::builder()
3240 .set_writer_version(WriterVersion::PARQUET_1_0)
3241 .set_dictionary_enabled(false)
3242 .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3243 .set_data_page_size_limit(page_byte_limit)
3244 .set_statistics_enabled(EnabledStatistics::None)
3245 .build();
3246
3247 let data: Vec<_> = (0..num_rows)
3249 .map(|i| ByteArray::from(vec![i as u8; value_size]))
3250 .collect();
3251 let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3252
3253 let total_values: u32 = pages.data_pages.iter().map(|(_, n)| n).sum();
3254 assert_eq!(total_values as usize, num_rows);
3255
3256 let upper_bound = 2 * value_size + 64;
3259 for (size, n_values) in &pages.data_pages {
3260 assert!(
3261 *size <= upper_bound,
3262 "page size {size} exceeds two-value bound ({upper_bound}B); pages {:?}",
3263 pages.data_pages,
3264 );
3265 assert!(
3266 *n_values <= 2,
3267 "page holds {n_values} values, expected at most 2; pages {:?}",
3268 pages.data_pages,
3269 );
3270 }
3271 }
3272
3273 #[test]
3274 fn test_column_writer_caps_page_size_with_sparse_nulls() {
3275 let value_size = 64 * 1024;
3290 let page_byte_limit = 16 * 1024;
3291 let num_values = 16;
3292
3293 let props = WriterProperties::builder()
3294 .set_dictionary_enabled(false)
3295 .set_encoding(Encoding::PLAIN)
3296 .set_data_page_size_limit(page_byte_limit)
3297 .set_statistics_enabled(EnabledStatistics::None)
3298 .build();
3299
3300 let data: Vec<_> = (0..num_values)
3301 .map(|_| ByteArray::from(vec![b'a'; value_size]))
3302 .collect();
3303 let def_levels: Vec<i16> = (0..num_values as i16 + 1)
3305 .map(|i| i16::from(i != 8))
3306 .collect();
3307 let pages =
3308 write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3309
3310 let upper_bound = 2 * value_size + 64;
3313 for (size, n_levels) in &pages.data_pages {
3314 assert!(
3315 *size <= upper_bound,
3316 "page size {size} exceeds two-value bound ({upper_bound}B); pages {:?}",
3317 pages.data_pages,
3318 );
3319 assert!(
3320 *n_levels <= 3,
3321 "page holds {n_levels} levels, expected at most 3 (two values + a null); \
3322 pages {:?}",
3323 pages.data_pages,
3324 );
3325 }
3326 let num_levels: usize = num_values + 1;
3329 assert_eq!(pages.data_pages.len(), num_levels.div_ceil(2));
3330 }
3331
3332 #[test]
3333 fn test_column_writer_delta_byte_array_nullable_shared_prefix_dedup() {
3334 let value_size = 64 * 1024;
3346 let page_byte_limit = 16 * 1024;
3347 let num_values = 16;
3348
3349 let props = WriterProperties::builder()
3350 .set_writer_version(WriterVersion::PARQUET_1_0)
3351 .set_dictionary_enabled(false)
3352 .set_encoding(Encoding::DELTA_BYTE_ARRAY)
3353 .set_data_page_size_limit(page_byte_limit)
3354 .set_statistics_enabled(EnabledStatistics::None)
3355 .build();
3356
3357 let data: Vec<_> = (0..num_values)
3358 .map(|_| ByteArray::from(vec![b'a'; value_size]))
3359 .collect();
3360 let def_levels: Vec<i16> = (0..num_values as i16 + 1)
3362 .map(|i| i16::from(i != 8))
3363 .collect();
3364 let pages =
3365 write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3366
3367 let per_page_values: Vec<u32> = pages.data_pages.iter().map(|(_, n)| *n).collect();
3370 assert_eq!(per_page_values, vec![num_values as u32 + 1]);
3371
3372 let total_bytes: usize = pages.data_pages.iter().map(|(size, _)| size).sum();
3373 assert!(
3374 total_bytes < 2 * value_size,
3375 "expected ~one full value's worth of bytes (full dedup), \
3376 got {total_bytes}B across pages {:?}",
3377 pages.data_pages,
3378 );
3379 }
3380
3381 #[test]
3382 fn test_column_writer_caps_page_size_for_large_values_in_list() {
3383 let value_size = 32 * 1024;
3394 let page_byte_limit = 16 * 1024;
3395 let values_per_record = 3;
3396 let num_records = 3;
3397 let num_values = values_per_record * num_records;
3398
3399 let mut rep_levels = Vec::with_capacity(num_values);
3401 for _ in 0..num_records {
3402 rep_levels.push(0i16);
3403 rep_levels.extend(std::iter::repeat_n(1i16, values_per_record - 1));
3404 }
3405 let def_levels = vec![1i16; num_values];
3406
3407 let props = WriterProperties::builder()
3408 .set_writer_version(WriterVersion::PARQUET_1_0)
3409 .set_dictionary_enabled(false)
3410 .set_encoding(Encoding::PLAIN)
3411 .set_data_page_size_limit(page_byte_limit)
3412 .build();
3413
3414 let data: Vec<_> = (0..num_values)
3415 .map(|i| ByteArray::from(vec![i as u8; value_size]))
3416 .collect();
3417 let pages = write_and_collect_pages::<ByteArrayType>(
3418 props,
3419 1,
3420 1,
3421 &data,
3422 Some(&def_levels),
3423 Some(&rep_levels),
3424 );
3425 let data_pages = pages.data_pages;
3426
3427 assert_eq!(
3431 data_pages.len(),
3432 num_records,
3433 "expected one data page per record, got {data_pages:?}"
3434 );
3435 for (bytes, n_values) in &data_pages {
3436 assert_eq!(
3437 *n_values as usize, values_per_record,
3438 "each page must hold a whole record's leaves, got {data_pages:?}"
3439 );
3440 let upper_bound = values_per_record * (value_size + 16);
3444 assert!(
3445 *bytes <= upper_bound,
3446 "page size {bytes} exceeds whole-record bound ({upper_bound}); pages {data_pages:?}"
3447 );
3448 }
3449 }
3450
3451 #[test]
3452 fn test_column_writer_caps_page_size_with_nullable_large_values() {
3453 let value_size = 32 * 1024;
3459 let page_byte_limit = 16 * 1024;
3460 let num_levels = 32;
3461
3462 let def_levels: Vec<i16> = (0..num_levels as i16).map(|i| i % 2).collect();
3464 let num_values = def_levels.iter().filter(|&&d| d == 1).count();
3465
3466 let props = WriterProperties::builder()
3467 .set_writer_version(WriterVersion::PARQUET_1_0)
3468 .set_dictionary_enabled(false)
3469 .set_encoding(Encoding::PLAIN)
3470 .set_data_page_size_limit(page_byte_limit)
3471 .build();
3472
3473 let data: Vec<_> = (0..num_values)
3474 .map(|i| ByteArray::from(vec![i as u8; value_size]))
3475 .collect();
3476 let pages =
3477 write_and_collect_pages::<ByteArrayType>(props, 1, 0, &data, Some(&def_levels), None);
3478 let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3479
3480 assert!(
3485 data_pages.len() >= num_values / 2,
3486 "expected at least {} pages for {num_values} large values, got {} pages: {data_pages:?}",
3487 num_values / 2,
3488 data_pages.len(),
3489 );
3490 for size in &data_pages {
3492 assert!(
3493 *size <= value_size + 64,
3494 "page size {size} exceeds one-value bound; pages {data_pages:?}"
3495 );
3496 }
3497 }
3498
3499 #[test]
3500 fn test_column_writer_dict_enabled_large_values_post_spill() {
3501 let value_size = 64 * 1024;
3511 let page_byte_limit = 16 * 1024;
3512 let num_rows = 32;
3513
3514 let props = WriterProperties::builder()
3515 .set_writer_version(WriterVersion::PARQUET_1_0)
3516 .set_dictionary_enabled(true)
3517 .set_dictionary_page_size_limit(1024)
3520 .set_data_page_size_limit(page_byte_limit)
3521 .set_write_batch_size(4)
3526 .build();
3527
3528 let data: Vec<_> = (0..num_rows)
3529 .map(|i| ByteArray::from(vec![i as u8; value_size]))
3530 .collect();
3531 let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3532 let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3533
3534 assert!(
3538 data_pages.len() >= num_rows / 2,
3539 "expected >= {} data pages after dict spill, got {} ({data_pages:?})",
3540 num_rows / 2,
3541 data_pages.len(),
3542 );
3543 for size in &data_pages {
3544 assert!(
3545 *size <= value_size + 64,
3546 "page size {size} exceeds one-value bound; pages {data_pages:?}"
3547 );
3548 }
3549 }
3550
3551 #[test]
3552 fn test_column_writer_caps_dictionary_page_size() {
3553 let value_size = 8 * 1024;
3561 let dict_page_limit = 64 * 1024;
3562 let num_rows = 2048;
3563
3564 let props = WriterProperties::builder()
3565 .set_writer_version(WriterVersion::PARQUET_1_0)
3566 .set_dictionary_enabled(true)
3567 .set_dictionary_page_size_limit(dict_page_limit)
3568 .build();
3569
3570 let data: Vec<_> = (0..num_rows)
3571 .map(|i| {
3572 let mut v = vec![0u8; value_size];
3574 v[..8].copy_from_slice(&(i as u64).to_le_bytes());
3575 ByteArray::from(v)
3576 })
3577 .collect();
3578 let pages = write_and_collect_pages::<ByteArrayType>(props, 0, 0, &data, None, None);
3579 let dict_page_size = pages.dict_page_size;
3580
3581 assert!(
3582 dict_page_size > 0,
3583 "expected the column to dictionary-encode"
3584 );
3585 assert!(
3589 dict_page_size <= 3 * dict_page_limit,
3590 "dictionary page {dict_page_size} exceeds 3x the {dict_page_limit} limit",
3591 );
3592 }
3593
3594 #[test]
3595 fn test_column_writer_caps_page_size_for_fixed_len_byte_array() {
3596 let page_byte_limit = 4;
3602 let num_values = 128;
3603
3604 let props = WriterProperties::builder()
3605 .set_writer_version(WriterVersion::PARQUET_1_0)
3606 .set_dictionary_enabled(false)
3607 .set_encoding(Encoding::PLAIN)
3608 .set_data_page_size_limit(page_byte_limit)
3609 .build();
3610
3611 let data: Vec<_> = (0..num_values)
3612 .map(|i| {
3613 let mut fla = FixedLenByteArray::default();
3614 fla.set_data(Bytes::from(vec![i as u8]));
3615 fla
3616 })
3617 .collect();
3618 let pages =
3619 write_and_collect_pages::<FixedLenByteArrayType>(props, 0, 0, &data, None, None);
3620 let data_pages: Vec<_> = pages.data_pages.iter().map(|(size, _)| *size).collect();
3621
3622 assert!(
3625 data_pages.len() >= num_values / 8,
3626 "expected pages capped by byte budget, got {data_pages:?}"
3627 );
3628 for size in &data_pages {
3629 assert!(
3630 *size <= page_byte_limit * 4,
3631 "page size {size} larger than expected; pages {data_pages:?}"
3632 );
3633 }
3634 }
3635
3636 #[test]
3637 fn test_bool_statistics() {
3638 let stats = statistics_roundtrip::<BoolType>(&[true, false, false, true]);
3639 assert!(!stats.is_min_max_backwards_compatible());
3642 if let Statistics::Boolean(stats) = stats {
3643 assert_eq!(stats.min_opt().unwrap(), &false);
3644 assert_eq!(stats.max_opt().unwrap(), &true);
3645 } else {
3646 panic!("expecting Statistics::Boolean, got {stats:?}");
3647 }
3648 }
3649
3650 #[test]
3651 fn test_int32_statistics() {
3652 let stats = statistics_roundtrip::<Int32Type>(&[-1, 3, -2, 2]);
3653 assert!(stats.is_min_max_backwards_compatible());
3654 if let Statistics::Int32(stats) = stats {
3655 assert_eq!(stats.min_opt().unwrap(), &-2);
3656 assert_eq!(stats.max_opt().unwrap(), &3);
3657 } else {
3658 panic!("expecting Statistics::Int32, got {stats:?}");
3659 }
3660 }
3661
3662 #[test]
3663 fn test_int64_statistics() {
3664 let stats = statistics_roundtrip::<Int64Type>(&[-1, 3, -2, 2]);
3665 assert!(stats.is_min_max_backwards_compatible());
3666 if let Statistics::Int64(stats) = stats {
3667 assert_eq!(stats.min_opt().unwrap(), &-2);
3668 assert_eq!(stats.max_opt().unwrap(), &3);
3669 } else {
3670 panic!("expecting Statistics::Int64, got {stats:?}");
3671 }
3672 }
3673
3674 #[test]
3675 fn test_int96_statistics() {
3676 let input = vec![
3677 Int96::from(vec![1, 20, 30]),
3678 Int96::from(vec![3, 20, 10]),
3679 Int96::from(vec![0, 20, 30]),
3680 Int96::from(vec![2, 20, 30]),
3681 ]
3682 .into_iter()
3683 .collect::<Vec<Int96>>();
3684
3685 let stats = statistics_roundtrip::<Int96Type>(&input);
3686 assert!(!stats.is_min_max_backwards_compatible());
3687 if let Statistics::Int96(stats) = stats {
3688 assert_eq!(stats.min_opt().unwrap(), &Int96::from(vec![3, 20, 10]));
3689 assert_eq!(stats.max_opt().unwrap(), &Int96::from(vec![2, 20, 30]));
3690 } else {
3691 panic!("expecting Statistics::Int96, got {stats:?}");
3692 }
3693 }
3694
3695 #[test]
3696 fn test_float_statistics() {
3697 let stats = statistics_roundtrip::<FloatType>(&[-1.0, 3.0, -2.0, 2.0]);
3698 assert!(!stats.is_min_max_backwards_compatible());
3699 if let Statistics::Float(stats) = stats {
3700 assert_eq!(stats.min_opt().unwrap(), &-2.0);
3701 assert_eq!(stats.max_opt().unwrap(), &3.0);
3702 } else {
3703 panic!("expecting Statistics::Float, got {stats:?}");
3704 }
3705 }
3706
3707 #[test]
3708 fn test_double_statistics() {
3709 let stats = statistics_roundtrip::<DoubleType>(&[-1.0, 3.0, -2.0, 2.0]);
3710 assert!(!stats.is_min_max_backwards_compatible());
3711 if let Statistics::Double(stats) = stats {
3712 assert_eq!(stats.min_opt().unwrap(), &-2.0);
3713 assert_eq!(stats.max_opt().unwrap(), &3.0);
3714 } else {
3715 panic!("expecting Statistics::Double, got {stats:?}");
3716 }
3717 }
3718
3719 #[test]
3720 fn test_byte_array_statistics() {
3721 let input = ["aawaa", "zz", "aaw", "m", "qrs"]
3722 .iter()
3723 .map(|&s| s.into())
3724 .collect::<Vec<_>>();
3725
3726 let stats = statistics_roundtrip::<ByteArrayType>(&input);
3727 assert!(!stats.is_min_max_backwards_compatible());
3728 if let Statistics::ByteArray(stats) = stats {
3729 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from("aaw"));
3730 assert_eq!(stats.max_opt().unwrap(), &ByteArray::from("zz"));
3731 } else {
3732 panic!("expecting Statistics::ByteArray, got {stats:?}");
3733 }
3734 }
3735
3736 #[test]
3737 fn test_fixed_len_byte_array_statistics() {
3738 let input = ["aawaa", "zz ", "aaw ", "m ", "qrs "]
3739 .iter()
3740 .map(|&s| ByteArray::from(s).into())
3741 .collect::<Vec<_>>();
3742
3743 let stats = statistics_roundtrip::<FixedLenByteArrayType>(&input);
3744 assert!(!stats.is_min_max_backwards_compatible());
3745 if let Statistics::FixedLenByteArray(stats) = stats {
3746 let expected_min: FixedLenByteArray = ByteArray::from("aaw ").into();
3747 assert_eq!(stats.min_opt().unwrap(), &expected_min);
3748 let expected_max: FixedLenByteArray = ByteArray::from("zz ").into();
3749 assert_eq!(stats.max_opt().unwrap(), &expected_max);
3750 } else {
3751 panic!("expecting Statistics::FixedLenByteArray, got {stats:?}");
3752 }
3753 }
3754
3755 #[test]
3756 fn test_ieee754_total_order_float() {
3757 let neg_nan = f32::from_bits(0xffc00000); let neg_inf = f32::NEG_INFINITY;
3761 let neg_one = -1.0_f32;
3762 let neg_zero = -0.0_f32;
3763 let pos_zero = 0.0_f32;
3764 let pos_one = 1.0_f32;
3765 let pos_inf = f32::INFINITY;
3766 let pos_nan = f32::from_bits(0x7fc00000); let values = vec![
3769 pos_nan, neg_zero, pos_inf, neg_one, neg_nan, pos_one, neg_inf, pos_zero,
3770 ];
3771
3772 let stats = statistics_roundtrip::<FloatType>(&values);
3773 if let Statistics::Float(stats) = stats {
3774 assert_eq!(stats.min_opt().unwrap(), &neg_inf);
3777 assert_eq!(stats.max_opt().unwrap(), &pos_inf);
3778 assert_eq!(stats.nan_count_opt(), Some(2)); } else {
3780 panic!("Expected float statistics");
3781 }
3782 }
3783
3784 #[test]
3785 fn test_ieee754_total_order_float_only_nan() {
3786 let neg_nan1 = f32::from_bits(0xffc00000); let neg_nan2 = f32::from_bits(0xffc00001); let neg_nan3 = f32::from_bits(0xffc00002); let pos_nan1 = f32::from_bits(0x7fc00000); let pos_nan2 = f32::from_bits(0x7fc00001); let pos_nan3 = f32::from_bits(0x7fc00002); let values = vec![neg_nan1, neg_nan2, neg_nan3, pos_nan1, pos_nan2, pos_nan3];
3796
3797 let stats = statistics_roundtrip::<FloatType>(&values);
3798 if let Statistics::Float(stats) = stats {
3799 assert_eq!(
3801 stats.min_opt().unwrap().total_cmp(&neg_nan3),
3802 Ordering::Equal
3803 );
3804 assert_eq!(
3805 stats.max_opt().unwrap().total_cmp(&pos_nan3),
3806 Ordering::Equal
3807 );
3808 assert_eq!(stats.nan_count_opt(), Some(6));
3809 } else {
3810 panic!("Expected float statistics");
3811 }
3812 }
3813
3814 #[test]
3815 fn test_ieee754_total_order_double() {
3816 let neg_nan = f64::from_bits(0xfff8000000000000);
3818 let neg_inf = f64::NEG_INFINITY;
3819 let neg_one = -1.0_f64;
3820 let neg_zero = -0.0_f64;
3821 let pos_zero = 0.0_f64;
3822 let pos_one = 1.0_f64;
3823 let pos_inf = f64::INFINITY;
3824 let pos_nan = f64::from_bits(0x7ff8000000000000);
3825
3826 let values = vec![
3827 pos_nan, neg_zero, pos_inf, neg_one, neg_nan, pos_one, neg_inf, pos_zero,
3828 ];
3829
3830 let stats = statistics_roundtrip::<DoubleType>(&values);
3831 if let Statistics::Double(stats) = stats {
3832 assert_eq!(stats.min_opt().unwrap(), &neg_inf);
3834 assert_eq!(stats.max_opt().unwrap(), &pos_inf);
3835 assert_eq!(stats.nan_count_opt(), Some(2));
3836 } else {
3837 panic!("Expected double statistics");
3838 }
3839 }
3840
3841 #[test]
3842 fn test_ieee754_total_order_double_only_nan() {
3843 let neg_nan1 = f64::from_bits(0xfff8000000000000);
3846 let neg_nan2 = f64::from_bits(0xfff8000000000001);
3847 let neg_nan3 = f64::from_bits(0xfff8000000000002);
3848 let pos_nan1 = f64::from_bits(0x7ff8000000000000);
3849 let pos_nan2 = f64::from_bits(0x7ff8000000000001);
3850 let pos_nan3 = f64::from_bits(0x7ff8000000000002);
3851
3852 let values = vec![neg_nan1, neg_nan2, neg_nan3, pos_nan1, pos_nan2, pos_nan3];
3853
3854 let stats = statistics_roundtrip::<DoubleType>(&values);
3855 if let Statistics::Double(stats) = stats {
3856 assert_eq!(
3858 stats.min_opt().unwrap().total_cmp(&neg_nan3),
3859 Ordering::Equal
3860 );
3861 assert_eq!(
3862 stats.max_opt().unwrap().total_cmp(&pos_nan3),
3863 Ordering::Equal
3864 );
3865 assert_eq!(stats.nan_count_opt(), Some(6));
3866 } else {
3867 panic!("Expected float statistics");
3868 }
3869 }
3870
3871 #[test]
3872 fn test_ieee754_total_order_zeros() {
3873 let values = vec![-0.0_f32, 0.0_f32, -0.0_f32, 0.0_f32];
3875
3876 let stats = statistics_roundtrip::<FloatType>(&values);
3877 if let Statistics::Float(stats) = stats {
3878 assert_eq!(stats.min_opt().unwrap().to_bits(), (-0.0_f32).to_bits());
3880 assert_eq!(stats.max_opt().unwrap().to_bits(), 0.0_f32.to_bits());
3881 } else {
3882 panic!("Expected float statistics");
3883 }
3884 }
3885
3886 #[test]
3887 #[cfg_attr(miri, ignore)] fn test_column_writer_check_float16_min_max() {
3889 let input = [
3890 -f16::ONE,
3891 f16::from_f32(3.0),
3892 -f16::from_f32(2.0),
3893 f16::from_f32(2.0),
3894 ]
3895 .into_iter()
3896 .map(|s| ByteArray::from(s).into())
3897 .collect::<Vec<_>>();
3898
3899 let stats = float16_statistics_roundtrip(&input);
3900 assert!(!stats.is_min_max_backwards_compatible());
3901 assert_eq!(
3902 stats.min_opt().unwrap(),
3903 &ByteArray::from(-f16::from_f32(2.0))
3904 );
3905 assert_eq!(
3906 stats.max_opt().unwrap(),
3907 &ByteArray::from(f16::from_f32(3.0))
3908 );
3909 }
3910
3911 #[test]
3912 #[cfg_attr(miri, ignore)] fn test_column_writer_check_float16_nan_middle() {
3914 let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
3915 .into_iter()
3916 .map(|s| ByteArray::from(s).into())
3917 .collect::<Vec<_>>();
3918
3919 let stats = float16_statistics_roundtrip(&input);
3920 assert!(!stats.is_min_max_backwards_compatible());
3921 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
3922 assert_eq!(
3923 stats.max_opt().unwrap(),
3924 &ByteArray::from(f16::ONE + f16::ONE)
3925 );
3926 assert_eq!(stats.nan_count_opt(), Some(1));
3927 }
3928
3929 #[test]
3930 #[cfg_attr(miri, ignore)] fn test_float16_statistics_nan_middle() {
3932 let input = [f16::ONE, f16::NAN, f16::ONE + f16::ONE]
3933 .into_iter()
3934 .map(|s| ByteArray::from(s).into())
3935 .collect::<Vec<_>>();
3936
3937 let stats = float16_statistics_roundtrip(&input);
3938 assert!(!stats.is_min_max_backwards_compatible());
3939 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
3940 assert_eq!(
3941 stats.max_opt().unwrap(),
3942 &ByteArray::from(f16::ONE + f16::ONE)
3943 );
3944 assert_eq!(stats.nan_count_opt(), Some(1));
3945 }
3946
3947 #[test]
3948 #[cfg_attr(miri, ignore)] fn test_float16_statistics_nan_start() {
3950 let input = [f16::NAN, f16::ONE, f16::ONE + f16::ONE]
3951 .into_iter()
3952 .map(|s| ByteArray::from(s).into())
3953 .collect::<Vec<_>>();
3954
3955 let stats = float16_statistics_roundtrip(&input);
3956 assert!(!stats.is_min_max_backwards_compatible());
3957 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ONE));
3958 assert_eq!(
3959 stats.max_opt().unwrap(),
3960 &ByteArray::from(f16::ONE + f16::ONE)
3961 );
3962 assert_eq!(stats.nan_count_opt(), Some(1));
3963 }
3964
3965 #[test]
3966 fn test_float16_statistics_nan_only() {
3967 let input = [f16::NAN, f16::NAN]
3968 .into_iter()
3969 .map(|s| ByteArray::from(s).into())
3970 .collect::<Vec<_>>();
3971
3972 let stats = float16_statistics_roundtrip(&input);
3973 assert_eq!(
3974 stats.min_bytes_opt(),
3975 Some(ByteArray::from(f16::NAN).as_bytes())
3976 );
3977 assert_eq!(
3978 stats.max_bytes_opt(),
3979 Some(ByteArray::from(f16::NAN).as_bytes())
3980 );
3981 assert!(!stats.is_min_max_backwards_compatible());
3982 assert_eq!(stats.nan_count_opt(), Some(2));
3983 }
3984
3985 #[test]
3986 fn test_float16_statistics_zero_only() {
3987 let input = std::iter::once(f16::ZERO)
3988 .map(|s| ByteArray::from(s).into())
3989 .collect::<Vec<_>>();
3990
3991 let stats = float16_statistics_roundtrip(&input);
3992 assert!(!stats.is_min_max_backwards_compatible());
3993 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ZERO));
3994 assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::ZERO));
3995 }
3996
3997 #[test]
3998 fn test_float16_statistics_neg_zero_only() {
3999 let input = std::iter::once(f16::NEG_ZERO)
4000 .map(|s| ByteArray::from(s).into())
4001 .collect::<Vec<_>>();
4002
4003 let stats = float16_statistics_roundtrip(&input);
4004 assert!(!stats.is_min_max_backwards_compatible());
4005 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4006 assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4007 }
4008
4009 #[test]
4010 fn test_float16_statistics_zero_min() {
4011 let input = [f16::ZERO, f16::ONE, f16::NAN, f16::PI]
4012 .into_iter()
4013 .map(|s| ByteArray::from(s).into())
4014 .collect::<Vec<_>>();
4015
4016 let stats = float16_statistics_roundtrip(&input);
4017 assert!(!stats.is_min_max_backwards_compatible());
4018 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(f16::ZERO));
4019 assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::PI));
4020 }
4021
4022 #[test]
4023 fn test_float16_statistics_neg_zero_max() {
4024 let input = [f16::NEG_ZERO, f16::NEG_ONE, f16::NAN, -f16::PI]
4025 .into_iter()
4026 .map(|s| ByteArray::from(s).into())
4027 .collect::<Vec<_>>();
4028
4029 let stats = float16_statistics_roundtrip(&input);
4030 assert!(!stats.is_min_max_backwards_compatible());
4031 assert_eq!(stats.min_opt().unwrap(), &ByteArray::from(-f16::PI));
4032 assert_eq!(stats.max_opt().unwrap(), &ByteArray::from(f16::NEG_ZERO));
4033 }
4034
4035 #[test]
4036 fn test_float_statistics_nan_middle() {
4037 let stats = statistics_roundtrip::<FloatType>(&[1.0, f32::NAN, 2.0]);
4038 assert!(!stats.is_min_max_backwards_compatible());
4039 if let Statistics::Float(stats) = stats {
4040 assert_eq!(stats.min_opt().unwrap(), &1.0);
4041 assert_eq!(stats.max_opt().unwrap(), &2.0);
4042 assert_eq!(stats.nan_count_opt(), Some(1))
4043 } else {
4044 panic!("expecting Statistics::Float");
4045 }
4046 }
4047
4048 #[test]
4049 fn test_float_statistics_nan_start() {
4050 let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, 1.0, 2.0]);
4051 assert!(!stats.is_min_max_backwards_compatible());
4052 if let Statistics::Float(stats) = stats {
4053 assert_eq!(stats.min_opt().unwrap(), &1.0);
4054 assert_eq!(stats.max_opt().unwrap(), &2.0);
4055 assert_eq!(stats.nan_count_opt(), Some(1))
4056 } else {
4057 panic!("expecting Statistics::Float");
4058 }
4059 }
4060
4061 #[test]
4062 fn test_float_statistics_nan_only() {
4063 let stats = statistics_roundtrip::<FloatType>(&[f32::NAN, f32::NAN]);
4064 assert_eq!(stats.min_bytes_opt(), Some(f32::NAN.as_bytes()));
4065 assert_eq!(stats.max_bytes_opt(), Some(f32::NAN.as_bytes()));
4066 assert_eq!(stats.nan_count_opt(), Some(2));
4067 assert!(!stats.is_min_max_backwards_compatible());
4068 assert!(matches!(stats, Statistics::Float(_)));
4069 }
4070
4071 #[test]
4072 fn test_float_statistics_zero_only() {
4073 let stats = statistics_roundtrip::<FloatType>(&[0.0]);
4074 assert!(!stats.is_min_max_backwards_compatible());
4075 if let Statistics::Float(stats) = stats {
4076 assert_eq!(stats.min_opt().unwrap(), &0.0);
4077 assert!(stats.min_opt().unwrap().is_sign_positive());
4078 assert_eq!(stats.max_opt().unwrap(), &0.0);
4079 assert!(stats.max_opt().unwrap().is_sign_positive());
4080 } else {
4081 panic!("expecting Statistics::Float");
4082 }
4083 }
4084
4085 #[test]
4086 fn test_float_statistics_neg_zero_only() {
4087 let stats = statistics_roundtrip::<FloatType>(&[-0.0]);
4088 assert!(!stats.is_min_max_backwards_compatible());
4089 if let Statistics::Float(stats) = stats {
4090 assert_eq!(stats.min_opt().unwrap(), &-0.0);
4091 assert!(stats.min_opt().unwrap().is_sign_negative());
4092 assert_eq!(stats.max_opt().unwrap(), &-0.0);
4093 assert!(stats.max_opt().unwrap().is_sign_negative());
4094 } else {
4095 panic!("expecting Statistics::Float");
4096 }
4097 }
4098
4099 #[test]
4100 fn test_float_statistics_zero_min() {
4101 let stats = statistics_roundtrip::<FloatType>(&[0.0, 1.0, f32::NAN, 2.0]);
4102 assert!(!stats.is_min_max_backwards_compatible());
4103 if let Statistics::Float(stats) = stats {
4104 assert_eq!(stats.min_opt().unwrap(), &0.0);
4105 assert!(stats.min_opt().unwrap().is_sign_positive());
4106 assert_eq!(stats.max_opt().unwrap(), &2.0);
4107 } else {
4108 panic!("expecting Statistics::Float");
4109 }
4110 }
4111
4112 #[test]
4113 fn test_float_statistics_neg_zero_max() {
4114 let stats = statistics_roundtrip::<FloatType>(&[-0.0, -1.0, f32::NAN, -2.0]);
4115 assert!(!stats.is_min_max_backwards_compatible());
4116 if let Statistics::Float(stats) = stats {
4117 assert_eq!(stats.min_opt().unwrap(), &-2.0);
4118 assert_eq!(stats.max_opt().unwrap(), &-0.0);
4119 assert!(stats.max_opt().unwrap().is_sign_negative());
4120 } else {
4121 panic!("expecting Statistics::Float");
4122 }
4123 }
4124
4125 #[test]
4126 fn test_double_statistics_nan_middle() {
4127 let stats = statistics_roundtrip::<DoubleType>(&[1.0, f64::NAN, 2.0]);
4128 assert!(!stats.is_min_max_backwards_compatible());
4129 if let Statistics::Double(stats) = stats {
4130 assert_eq!(stats.min_opt().unwrap(), &1.0);
4131 assert_eq!(stats.max_opt().unwrap(), &2.0);
4132 assert_eq!(stats.nan_count_opt(), Some(1))
4133 } else {
4134 panic!("expecting Statistics::Double");
4135 }
4136 }
4137
4138 #[test]
4139 fn test_double_statistics_nan_start() {
4140 let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, 1.0, 2.0]);
4141 assert!(!stats.is_min_max_backwards_compatible());
4142 if let Statistics::Double(stats) = stats {
4143 assert_eq!(stats.min_opt().unwrap(), &1.0);
4144 assert_eq!(stats.max_opt().unwrap(), &2.0);
4145 assert_eq!(stats.nan_count_opt(), Some(1))
4146 } else {
4147 panic!("expecting Statistics::Double");
4148 }
4149 }
4150
4151 #[test]
4152 fn test_double_statistics_nan_only() {
4153 let stats = statistics_roundtrip::<DoubleType>(&[f64::NAN, f64::NAN]);
4154 assert_eq!(stats.min_bytes_opt(), Some(f64::NAN.as_bytes()));
4155 assert_eq!(stats.max_bytes_opt(), Some(f64::NAN.as_bytes()));
4156 assert_eq!(stats.nan_count_opt(), Some(2));
4157 assert!(matches!(stats, Statistics::Double(_)));
4158 assert!(!stats.is_min_max_backwards_compatible());
4159 }
4160
4161 #[test]
4162 fn test_double_statistics_zero_only() {
4163 let stats = statistics_roundtrip::<DoubleType>(&[0.0]);
4164 assert!(!stats.is_min_max_backwards_compatible());
4165 if let Statistics::Double(stats) = stats {
4166 assert_eq!(stats.min_opt().unwrap(), &0.0);
4167 assert!(stats.min_opt().unwrap().is_sign_positive());
4168 assert_eq!(stats.max_opt().unwrap(), &0.0);
4169 assert!(stats.max_opt().unwrap().is_sign_positive());
4170 } else {
4171 panic!("expecting Statistics::Double");
4172 }
4173 }
4174
4175 #[test]
4176 fn test_double_statistics_neg_zero_only() {
4177 let stats = statistics_roundtrip::<DoubleType>(&[-0.0]);
4178 assert!(!stats.is_min_max_backwards_compatible());
4179 if let Statistics::Double(stats) = stats {
4180 assert_eq!(stats.min_opt().unwrap(), &-0.0);
4181 assert!(stats.min_opt().unwrap().is_sign_negative());
4182 assert_eq!(stats.max_opt().unwrap(), &-0.0);
4183 assert!(stats.max_opt().unwrap().is_sign_negative());
4184 } else {
4185 panic!("expecting Statistics::Double");
4186 }
4187 }
4188
4189 #[test]
4190 fn test_double_statistics_zero_min() {
4191 let stats = statistics_roundtrip::<DoubleType>(&[0.0, 1.0, f64::NAN, 2.0]);
4192 assert!(!stats.is_min_max_backwards_compatible());
4193 if let Statistics::Double(stats) = stats {
4194 assert_eq!(stats.min_opt().unwrap(), &0.0);
4195 assert!(stats.min_opt().unwrap().is_sign_positive());
4196 assert_eq!(stats.max_opt().unwrap(), &2.0);
4197 } else {
4198 panic!("expecting Statistics::Double");
4199 }
4200 }
4201
4202 #[test]
4203 fn test_double_statistics_neg_zero_max() {
4204 let stats = statistics_roundtrip::<DoubleType>(&[-0.0, -1.0, f64::NAN, -2.0]);
4205 assert!(!stats.is_min_max_backwards_compatible());
4206 if let Statistics::Double(stats) = stats {
4207 assert_eq!(stats.min_opt().unwrap(), &-2.0);
4208 assert_eq!(stats.max_opt().unwrap(), &-0.0);
4209 assert!(stats.max_opt().unwrap().is_sign_negative());
4210 } else {
4211 panic!("expecting Statistics::Double");
4212 }
4213 }
4214
4215 #[test]
4216 fn test_compare_greater_byte_array_decimals() {
4217 assert!(!compare_greater_byte_array_decimals(&[], &[],),);
4218 assert!(compare_greater_byte_array_decimals(&[1u8,], &[],),);
4219 assert!(!compare_greater_byte_array_decimals(&[], &[1u8,],),);
4220 assert!(compare_greater_byte_array_decimals(&[1u8,], &[0u8,],),);
4221 assert!(!compare_greater_byte_array_decimals(&[1u8,], &[1u8,],),);
4222 assert!(compare_greater_byte_array_decimals(&[1u8, 0u8,], &[0u8,],),);
4223 assert!(!compare_greater_byte_array_decimals(
4224 &[0u8, 1u8,],
4225 &[1u8, 0u8,],
4226 ),);
4227 assert!(!compare_greater_byte_array_decimals(
4228 &[255u8, 35u8, 0u8, 0u8,],
4229 &[0u8,],
4230 ),);
4231 assert!(compare_greater_byte_array_decimals(
4232 &[0u8,],
4233 &[255u8, 35u8, 0u8, 0u8,],
4234 ),);
4235
4236 assert!(compare_greater_byte_array_decimals(
4242 &[0u8, 128u8, 0u8,],
4243 &[0u8, 255u8,],
4244 ),);
4245 assert!(!compare_greater_byte_array_decimals(
4246 &[0u8, 255u8,],
4247 &[0u8, 128u8, 0u8,],
4248 ),);
4249 assert!(compare_greater_byte_array_decimals(
4251 &[128u8,],
4252 &[255u8, 127u8,],
4253 ),);
4254 assert!(!compare_greater_byte_array_decimals(
4255 &[255u8, 127u8,],
4256 &[128u8,],
4257 ),);
4258 assert!(compare_greater_byte_array_decimals(
4260 &[128u8,],
4261 &[255u8, 0u8,],
4262 ),);
4263 assert!(compare_greater_byte_array_decimals(&[0u8, 10u8,], &[5u8,],),);
4265 assert!(compare_greater_byte_array_decimals(&[10u8,], &[0u8, 5u8,],),);
4266 assert!(!compare_greater_byte_array_decimals(
4268 &[255u8, 128u8,],
4269 &[128u8,],
4270 ),);
4271 assert!(!compare_greater_byte_array_decimals(
4272 &[128u8,],
4273 &[255u8, 128u8,],
4274 ),);
4275 }
4276
4277 #[test]
4278 fn test_column_index_with_null_pages() {
4279 let page_writer = get_test_page_writer();
4281 let props = Default::default();
4282 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 1, 0, props);
4283 writer.write_batch(&[], Some(&[0, 0, 0, 0]), None).unwrap();
4284
4285 let r = writer.close().unwrap();
4286 assert!(r.column_index.is_some());
4287 let col_idx = r.column_index.unwrap();
4288 let ColumnIndexMetaData::INT32(col_idx) = col_idx else {
4289 panic!("wrong stats type")
4290 };
4291 assert!(col_idx.is_null_page(0));
4293 assert!(col_idx.min_value(0).is_none());
4295 assert!(col_idx.max_value(0).is_none());
4296 assert!(col_idx.null_count(0).is_some());
4298 assert_eq!(col_idx.null_count(0), Some(4));
4299 assert!(col_idx.repetition_level_histogram(0).is_none());
4301 assert!(col_idx.definition_level_histogram(0).is_some());
4303 assert_eq!(col_idx.definition_level_histogram(0).unwrap(), &[4, 0]);
4304 }
4305
4306 #[test]
4307 fn test_column_offset_index_metadata() {
4308 let page_writer = get_test_page_writer();
4311 let props = Default::default();
4312 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
4313 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
4314 writer.flush_data_pages().unwrap();
4316 writer.write_batch(&[4, 8, 2, -5], None, None).unwrap();
4318
4319 let r = writer.close().unwrap();
4320 let column_index = r.column_index.unwrap();
4321 let offset_index = r.offset_index.unwrap();
4322
4323 assert_eq!(8, r.rows_written);
4324
4325 let ColumnIndexMetaData::INT32(column_index) = column_index else {
4327 panic!("wrong stats type")
4328 };
4329 assert_eq!(2, column_index.num_pages());
4330 assert_eq!(2, offset_index.page_locations.len());
4331 assert_eq!(BoundaryOrder::UNORDERED, column_index.boundary_order);
4332 for idx in 0..2 {
4333 assert!(!column_index.is_null_page(idx));
4334 assert_eq!(0, column_index.null_count(0).unwrap());
4335 }
4336
4337 if let Some(stats) = r.metadata.statistics() {
4338 assert_eq!(stats.null_count_opt(), Some(0));
4339 assert_eq!(stats.distinct_count_opt(), None);
4340 if let Statistics::Int32(stats) = stats {
4341 assert_eq!(stats.min_opt(), column_index.min_value(1));
4345 assert_eq!(stats.max_opt(), column_index.max_value(1));
4346 } else {
4347 panic!("expecting Statistics::Int32");
4348 }
4349 } else {
4350 panic!("metadata missing statistics");
4351 }
4352
4353 assert_eq!(0, offset_index.page_locations[0].first_row_index);
4355 assert_eq!(4, offset_index.page_locations[1].first_row_index);
4356 }
4357
4358 #[test]
4360 fn test_column_offset_index_metadata_truncating() {
4361 let page_writer = get_test_page_writer();
4364 let props = WriterProperties::builder()
4365 .set_statistics_truncate_length(None) .build()
4367 .into();
4368 let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4369
4370 let mut data = vec![FixedLenByteArray::default(); 3];
4371 data[0].set_data(Bytes::from(vec![97_u8; 200]));
4373 data[1].set_data(Bytes::from(vec![112_u8; 200]));
4375 data[2].set_data(Bytes::from(vec![98_u8; 200]));
4376
4377 writer.write_batch(&data, None, None).unwrap();
4378
4379 writer.flush_data_pages().unwrap();
4380
4381 let r = writer.close().unwrap();
4382 let column_index = r.column_index.unwrap();
4383 let offset_index = r.offset_index.unwrap();
4384
4385 let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4386 panic!("wrong stats type")
4387 };
4388
4389 assert_eq!(3, r.rows_written);
4390
4391 assert_eq!(1, column_index.num_pages());
4393 assert_eq!(1, offset_index.page_locations.len());
4394 assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
4395 assert!(!column_index.is_null_page(0));
4396 assert_eq!(Some(0), column_index.null_count(0));
4397
4398 if let Some(stats) = r.metadata.statistics() {
4399 assert_eq!(stats.null_count_opt(), Some(0));
4400 assert_eq!(stats.distinct_count_opt(), None);
4401 if let Statistics::FixedLenByteArray(stats) = stats {
4402 let column_index_min_value = column_index.min_value(0).unwrap();
4403 let column_index_max_value = column_index.max_value(0).unwrap();
4404
4405 assert_ne!(stats.min_bytes_opt().unwrap(), column_index_min_value);
4407 assert_ne!(stats.max_bytes_opt().unwrap(), column_index_max_value);
4408
4409 assert_eq!(
4410 column_index_min_value.len(),
4411 DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
4412 );
4413 assert_eq!(column_index_min_value, &[97_u8; 64]);
4414 assert_eq!(
4415 column_index_max_value.len(),
4416 DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH.unwrap()
4417 );
4418
4419 assert_eq!(
4421 *column_index_max_value.last().unwrap(),
4422 *column_index_max_value.first().unwrap() + 1
4423 );
4424 } else {
4425 panic!("expecting Statistics::FixedLenByteArray");
4426 }
4427 } else {
4428 panic!("metadata missing statistics");
4429 }
4430 }
4431
4432 #[test]
4433 fn test_column_offset_index_truncating_spec_example() {
4434 let page_writer = get_test_page_writer();
4437
4438 let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4440 let props = Arc::new(builder.build());
4441 let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4442
4443 let mut data = vec![FixedLenByteArray::default(); 1];
4444 data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
4446
4447 writer.write_batch(&data, None, None).unwrap();
4448
4449 writer.flush_data_pages().unwrap();
4450
4451 let r = writer.close().unwrap();
4452 let column_index = r.column_index.unwrap();
4453 let offset_index = r.offset_index.unwrap();
4454
4455 let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4456 panic!("wrong stats type")
4457 };
4458
4459 assert_eq!(1, r.rows_written);
4460
4461 assert_eq!(1, column_index.num_pages());
4463 assert_eq!(1, offset_index.page_locations.len());
4464 assert_eq!(BoundaryOrder::ASCENDING, column_index.boundary_order);
4465 assert!(!column_index.is_null_page(0));
4466 assert_eq!(Some(0), column_index.null_count(0));
4467
4468 if let Some(stats) = r.metadata.statistics() {
4469 assert_eq!(stats.null_count_opt(), Some(0));
4470 assert_eq!(stats.distinct_count_opt(), None);
4471 if let Statistics::FixedLenByteArray(_stats) = stats {
4472 let column_index_min_value = column_index.min_value(0).unwrap();
4473 let column_index_max_value = column_index.max_value(0).unwrap();
4474
4475 assert_eq!(column_index_min_value.len(), 1);
4476 assert_eq!(column_index_max_value.len(), 1);
4477
4478 assert_eq!(b"B", column_index_min_value);
4479 assert_eq!(b"C", column_index_max_value);
4480
4481 assert_ne!(column_index_min_value, stats.min_bytes_opt().unwrap());
4482 assert_ne!(column_index_max_value, stats.max_bytes_opt().unwrap());
4483 } else {
4484 panic!("expecting Statistics::FixedLenByteArray");
4485 }
4486 } else {
4487 panic!("metadata missing statistics");
4488 }
4489 }
4490
4491 #[test]
4492 fn test_float16_min_max_no_truncation() {
4493 let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4495 let props = Arc::new(builder.build());
4496 let page_writer = get_test_page_writer();
4497 let mut writer = get_test_float16_column_writer(page_writer, props);
4498
4499 let expected_value = f16::PI.to_le_bytes().to_vec();
4500 let data = vec![ByteArray::from(expected_value.clone()).into()];
4501 writer.write_batch(&data, None, None).unwrap();
4502 writer.flush_data_pages().unwrap();
4503
4504 let r = writer.close().unwrap();
4505
4506 let column_index = r.column_index.unwrap();
4509 let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4510 panic!("wrong stats type")
4511 };
4512 let column_index_min_bytes = column_index.min_value(0).unwrap();
4513 let column_index_max_bytes = column_index.max_value(0).unwrap();
4514 assert_eq!(expected_value, column_index_min_bytes);
4515 assert_eq!(expected_value, column_index_max_bytes);
4516
4517 let stats = r.metadata.statistics().unwrap();
4519 if let Statistics::FixedLenByteArray(stats) = stats {
4520 let stats_min_bytes = stats.min_bytes_opt().unwrap();
4521 let stats_max_bytes = stats.max_bytes_opt().unwrap();
4522 assert_eq!(expected_value, stats_min_bytes);
4523 assert_eq!(expected_value, stats_max_bytes);
4524 } else {
4525 panic!("expecting Statistics::FixedLenByteArray");
4526 }
4527 }
4528
4529 #[test]
4530 fn test_decimal_min_max_no_truncation() {
4531 let builder = WriterProperties::builder().set_column_index_truncate_length(Some(1));
4533 let props = Arc::new(builder.build());
4534 let page_writer = get_test_page_writer();
4535 let mut writer =
4536 get_test_decimals_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4537
4538 let expected_value = vec![
4539 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 179u8, 172u8, 19u8, 35u8,
4540 231u8, 90u8, 0u8, 0u8,
4541 ];
4542 let data = vec![ByteArray::from(expected_value.clone()).into()];
4543 writer.write_batch(&data, None, None).unwrap();
4544 writer.flush_data_pages().unwrap();
4545
4546 let r = writer.close().unwrap();
4547
4548 let column_index = r.column_index.unwrap();
4551 let ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(column_index) = column_index else {
4552 panic!("wrong stats type")
4553 };
4554 let column_index_min_bytes = column_index.min_value(0).unwrap();
4555 let column_index_max_bytes = column_index.max_value(0).unwrap();
4556 assert_eq!(expected_value, column_index_min_bytes);
4557 assert_eq!(expected_value, column_index_max_bytes);
4558
4559 let stats = r.metadata.statistics().unwrap();
4561 if let Statistics::FixedLenByteArray(stats) = stats {
4562 let stats_min_bytes = stats.min_bytes_opt().unwrap();
4563 let stats_max_bytes = stats.max_bytes_opt().unwrap();
4564 assert_eq!(expected_value, stats_min_bytes);
4565 assert_eq!(expected_value, stats_max_bytes);
4566 } else {
4567 panic!("expecting Statistics::FixedLenByteArray");
4568 }
4569 }
4570
4571 #[test]
4572 fn test_statistics_truncating_byte_array_default() {
4573 let page_writer = get_test_page_writer();
4574
4575 let props = WriterProperties::builder().build().into();
4577 let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
4578
4579 let mut data = vec![ByteArray::default(); 1];
4580 data[0].set_data(Bytes::from(String::from(
4581 "This string is longer than 64 bytes, so it will almost certainly be truncated.",
4582 )));
4583 writer.write_batch(&data, None, None).unwrap();
4584 writer.flush_data_pages().unwrap();
4585
4586 let r = writer.close().unwrap();
4587
4588 assert_eq!(1, r.rows_written);
4589
4590 let stats = r.metadata.statistics().expect("statistics");
4591 if let Statistics::ByteArray(stats) = stats {
4592 let min_value = stats.min_opt().unwrap();
4593 let max_value = stats.max_opt().unwrap();
4594
4595 assert!(!stats.min_is_exact());
4596 assert!(!stats.max_is_exact());
4597
4598 let expected_len = 64;
4599 assert_eq!(min_value.len(), expected_len);
4600 assert_eq!(max_value.len(), expected_len);
4601
4602 let expected_min =
4603 "This string is longer than 64 bytes, so it will almost certainly".as_bytes();
4604 assert_eq!(expected_min, min_value.as_bytes());
4605 let expected_max =
4607 "This string is longer than 64 bytes, so it will almost certainlz".as_bytes();
4608 assert_eq!(expected_max, max_value.as_bytes());
4609 } else {
4610 panic!("expecting Statistics::ByteArray");
4611 }
4612 }
4613
4614 #[test]
4615 fn test_statistics_truncating_byte_array() {
4616 let page_writer = get_test_page_writer();
4617
4618 const TEST_TRUNCATE_LENGTH: usize = 1;
4619
4620 let builder =
4622 WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
4623 let props = Arc::new(builder.build());
4624 let mut writer = get_test_column_writer::<ByteArrayType>(page_writer, 0, 0, props);
4625
4626 let mut data = vec![ByteArray::default(); 1];
4627 data[0].set_data(Bytes::from(String::from("Blart Versenwald III")));
4629
4630 writer.write_batch(&data, None, None).unwrap();
4631
4632 writer.flush_data_pages().unwrap();
4633
4634 let r = writer.close().unwrap();
4635
4636 assert_eq!(1, r.rows_written);
4637
4638 let stats = r.metadata.statistics().expect("statistics");
4639 assert_eq!(stats.null_count_opt(), Some(0));
4640 assert_eq!(stats.distinct_count_opt(), None);
4641 if let Statistics::ByteArray(stats) = stats {
4642 let min_value = stats.min_opt().unwrap();
4643 let max_value = stats.max_opt().unwrap();
4644
4645 assert!(!stats.min_is_exact());
4646 assert!(!stats.max_is_exact());
4647
4648 assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
4649 assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
4650
4651 assert_eq!(b"B", min_value.as_bytes());
4652 assert_eq!(b"C", max_value.as_bytes());
4653 } else {
4654 panic!("expecting Statistics::ByteArray");
4655 }
4656 }
4657
4658 #[test]
4659 fn test_statistics_truncating_fixed_len_byte_array() {
4660 let page_writer = get_test_page_writer();
4661
4662 const TEST_TRUNCATE_LENGTH: usize = 1;
4663
4664 let builder =
4666 WriterProperties::builder().set_statistics_truncate_length(Some(TEST_TRUNCATE_LENGTH));
4667 let props = Arc::new(builder.build());
4668 let mut writer = get_test_column_writer::<FixedLenByteArrayType>(page_writer, 0, 0, props);
4669
4670 let mut data = vec![FixedLenByteArray::default(); 1];
4671
4672 const PSEUDO_DECIMAL_VALUE: i128 = 6541894651216648486512564456564654;
4673 const PSEUDO_DECIMAL_BYTES: [u8; 16] = PSEUDO_DECIMAL_VALUE.to_be_bytes();
4674
4675 const EXPECTED_MIN: [u8; TEST_TRUNCATE_LENGTH] = [PSEUDO_DECIMAL_BYTES[0]]; const EXPECTED_MAX: [u8; TEST_TRUNCATE_LENGTH] =
4677 [PSEUDO_DECIMAL_BYTES[0].overflowing_add(1).0];
4678
4679 data[0].set_data(Bytes::from(PSEUDO_DECIMAL_BYTES.as_slice()));
4681
4682 writer.write_batch(&data, None, None).unwrap();
4683
4684 writer.flush_data_pages().unwrap();
4685
4686 let r = writer.close().unwrap();
4687
4688 assert_eq!(1, r.rows_written);
4689
4690 let stats = r.metadata.statistics().expect("statistics");
4691 assert_eq!(stats.null_count_opt(), Some(0));
4692 assert_eq!(stats.distinct_count_opt(), None);
4693 if let Statistics::FixedLenByteArray(stats) = stats {
4694 let min_value = stats.min_opt().unwrap();
4695 let max_value = stats.max_opt().unwrap();
4696
4697 assert!(!stats.min_is_exact());
4698 assert!(!stats.max_is_exact());
4699
4700 assert_eq!(min_value.len(), TEST_TRUNCATE_LENGTH);
4701 assert_eq!(max_value.len(), TEST_TRUNCATE_LENGTH);
4702
4703 assert_eq!(EXPECTED_MIN.as_slice(), min_value.as_bytes());
4704 assert_eq!(EXPECTED_MAX.as_slice(), max_value.as_bytes());
4705
4706 let reconstructed_min = i128::from_be_bytes([
4707 min_value.as_bytes()[0],
4708 0,
4709 0,
4710 0,
4711 0,
4712 0,
4713 0,
4714 0,
4715 0,
4716 0,
4717 0,
4718 0,
4719 0,
4720 0,
4721 0,
4722 0,
4723 ]);
4724
4725 let reconstructed_max = i128::from_be_bytes([
4726 max_value.as_bytes()[0],
4727 0,
4728 0,
4729 0,
4730 0,
4731 0,
4732 0,
4733 0,
4734 0,
4735 0,
4736 0,
4737 0,
4738 0,
4739 0,
4740 0,
4741 0,
4742 ]);
4743
4744 println!("min: {reconstructed_min} {PSEUDO_DECIMAL_VALUE}");
4746 assert!(reconstructed_min <= PSEUDO_DECIMAL_VALUE);
4747 println!("max {reconstructed_max} {PSEUDO_DECIMAL_VALUE}");
4748 assert!(reconstructed_max >= PSEUDO_DECIMAL_VALUE);
4749 } else {
4750 panic!("expecting Statistics::FixedLenByteArray");
4751 }
4752 }
4753
4754 #[test]
4755 fn test_send() {
4756 fn test<T: Send>() {}
4757 test::<ColumnWriterImpl<Int32Type>>();
4758 }
4759
4760 #[test]
4761 fn test_increment() {
4762 let v = increment(vec![0, 0, 0]).unwrap();
4763 assert_eq!(&v, &[0, 0, 1]);
4764
4765 let v = increment(vec![0, 255, 255]).unwrap();
4767 assert_eq!(&v, &[1, 0, 0]);
4768
4769 let v = increment(vec![255, 255, 255]);
4771 assert!(v.is_none());
4772 }
4773
4774 #[test]
4775 fn test_increment_utf8() {
4776 let test_inc = |o: &str, expected: &str| {
4777 if let Ok(v) = String::from_utf8(increment_utf8(o).unwrap()) {
4778 assert_eq!(v, expected);
4780 assert!(*v > *o);
4782 let mut greater = ByteArray::new();
4784 greater.set_data(Bytes::from(v));
4785 let mut original = ByteArray::new();
4786 original.set_data(Bytes::from(o.as_bytes().to_vec()));
4787 assert!(greater > original);
4788 } else {
4789 panic!("Expected incremented UTF8 string to also be valid.");
4790 }
4791 };
4792
4793 test_inc("hello", "hellp");
4795
4796 test_inc("a\u{7f}", "b");
4798
4799 assert!(increment_utf8("\u{7f}\u{7f}").is_none());
4801
4802 test_inc("❤️🧡💛💚💙💜", "❤️🧡💛💚💙💝");
4804
4805 test_inc("éééé", "éééê");
4807
4808 test_inc("\u{ff}\u{ff}", "\u{ff}\u{100}");
4810
4811 test_inc("a\u{7ff}", "b");
4813
4814 assert!(increment_utf8("\u{7ff}\u{7ff}").is_none());
4816
4817 test_inc("ࠀࠀ", "ࠀࠁ");
4820
4821 test_inc("a\u{ffff}", "b");
4823
4824 assert!(increment_utf8("\u{ffff}\u{ffff}").is_none());
4826
4827 test_inc("𐀀𐀀", "𐀀𐀁");
4829
4830 test_inc("a\u{10ffff}", "b");
4832
4833 assert!(increment_utf8("\u{10ffff}\u{10ffff}").is_none());
4835
4836 test_inc("a\u{D7FF}", "b");
4839 }
4840
4841 #[test]
4842 fn test_truncate_utf8() {
4843 let data = "❤️🧡💛💚💙💜";
4845 let r = truncate_utf8(data, data.len()).unwrap();
4846 assert_eq!(r.len(), data.len());
4847 assert_eq!(&r, data.as_bytes());
4848
4849 let r = truncate_utf8(data, 13).unwrap();
4851 assert_eq!(r.len(), 10);
4852 assert_eq!(&r, "❤️🧡".as_bytes());
4853
4854 let r = truncate_utf8("\u{0836}", 1);
4856 assert!(r.is_none());
4857
4858 let r = truncate_and_increment_utf8("yyyyyyyyy", 8).unwrap();
4861 assert_eq!(&r, b"yyyyyyyz");
4862
4863 let r = truncate_and_increment_utf8("ééééé", 7).unwrap();
4865 assert_eq!(&r, "ééê".as_bytes());
4866
4867 let r = truncate_and_increment_utf8("\u{ff}\u{ff}\u{ff}\u{ff}\u{ff}", 8).unwrap();
4869 assert_eq!(&r, "\u{ff}\u{ff}\u{ff}\u{100}".as_bytes());
4870
4871 let r = truncate_and_increment_utf8("߿߿߿߿߿", 8);
4873 assert!(r.is_none());
4874
4875 let r = truncate_and_increment_utf8("ࠀࠀࠀࠀ", 8).unwrap();
4878 assert_eq!(&r, "ࠀࠁ".as_bytes());
4879
4880 let r = truncate_and_increment_utf8("\u{ffff}\u{ffff}\u{ffff}", 8);
4882 assert!(r.is_none());
4883
4884 let r = truncate_and_increment_utf8("𐀀𐀀𐀀𐀀", 9).unwrap();
4886 assert_eq!(&r, "𐀀𐀁".as_bytes());
4887
4888 let r = truncate_and_increment_utf8("\u{10ffff}\u{10ffff}", 8);
4890 assert!(r.is_none());
4891 }
4892
4893 #[test]
4894 fn test_byte_array_truncate_invalid_utf8_statistics() {
4897 let message_type = "
4898 message test_schema {
4899 OPTIONAL BYTE_ARRAY a (UTF8);
4900 }
4901 ";
4902 let schema = Arc::new(parse_message_type(message_type).unwrap());
4903
4904 let data = vec![ByteArray::from(vec![128u8; 32]); 7];
4906 let def_levels = [1, 1, 1, 1, 0, 1, 0, 1, 0, 1];
4907 let file: File = tempfile::tempfile().unwrap();
4908 let props = Arc::new(
4909 WriterProperties::builder()
4910 .set_statistics_enabled(EnabledStatistics::Chunk)
4911 .set_statistics_truncate_length(Some(8))
4912 .build(),
4913 );
4914
4915 let mut writer = SerializedFileWriter::new(&file, schema, props).unwrap();
4916 let mut row_group_writer = writer.next_row_group().unwrap();
4917
4918 let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
4919 col_writer
4920 .typed::<ByteArrayType>()
4921 .write_batch(&data, Some(&def_levels), None)
4922 .unwrap();
4923 col_writer.close().unwrap();
4924 row_group_writer.close().unwrap();
4925 let file_metadata = writer.close().unwrap();
4926 let stats = file_metadata.row_group(0).column(0).statistics().unwrap();
4927 assert!(!stats.max_is_exact());
4928 assert_eq!(
4931 stats.max_bytes_opt().map(|v| v.to_vec()),
4932 Some([128, 128, 128, 128, 128, 128, 128, 129].to_vec())
4933 );
4934 }
4935
4936 #[test]
4937 fn test_increment_max_binary_chars() {
4938 let r = increment(vec![0xFF, 0xFE, 0xFD, 0xFF, 0xFF]);
4939 assert_eq!(&r.unwrap(), &[0xFF, 0xFE, 0xFE, 0x00, 0x00]);
4940
4941 let incremented = increment(vec![0xFF, 0xFF, 0xFF]);
4942 assert!(incremented.is_none())
4943 }
4944
4945 #[test]
4946 fn test_no_column_index_when_stats_disabled() {
4947 let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
4951 let props = Arc::new(
4952 WriterProperties::builder()
4953 .set_statistics_enabled(EnabledStatistics::None)
4954 .build(),
4955 );
4956 let column_writer = get_column_writer(descr, props, get_test_page_writer());
4957 let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
4958
4959 let data = Vec::new();
4960 let def_levels = vec![0; 10];
4961 writer.write_batch(&data, Some(&def_levels), None).unwrap();
4962 writer.flush_data_pages().unwrap();
4963
4964 let column_close_result = writer.close().unwrap();
4965 assert!(column_close_result.offset_index.is_some());
4966 assert!(column_close_result.column_index.is_none());
4967 }
4968
4969 #[test]
4970 fn test_no_offset_index_when_disabled() {
4971 let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
4973 let props = Arc::new(
4974 WriterProperties::builder()
4975 .set_statistics_enabled(EnabledStatistics::None)
4976 .set_offset_index_disabled(true)
4977 .build(),
4978 );
4979 let column_writer = get_column_writer(descr, props, get_test_page_writer());
4980 let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
4981
4982 let data = Vec::new();
4983 let def_levels = vec![0; 10];
4984 writer.write_batch(&data, Some(&def_levels), None).unwrap();
4985 writer.flush_data_pages().unwrap();
4986
4987 let column_close_result = writer.close().unwrap();
4988 assert!(column_close_result.offset_index.is_none());
4989 assert!(column_close_result.column_index.is_none());
4990 }
4991
4992 #[test]
4993 fn test_offset_index_overridden() {
4994 let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
4996 let props = Arc::new(
4997 WriterProperties::builder()
4998 .set_statistics_enabled(EnabledStatistics::Page)
4999 .set_offset_index_disabled(true)
5000 .build(),
5001 );
5002 let column_writer = get_column_writer(descr, props, get_test_page_writer());
5003 let mut writer = get_typed_column_writer::<Int32Type>(column_writer);
5004
5005 let data = Vec::new();
5006 let def_levels = vec![0; 10];
5007 writer.write_batch(&data, Some(&def_levels), None).unwrap();
5008 writer.flush_data_pages().unwrap();
5009
5010 let column_close_result = writer.close().unwrap();
5011 assert!(column_close_result.offset_index.is_some());
5012 assert!(column_close_result.column_index.is_some());
5013 }
5014
5015 #[test]
5016 fn test_boundary_order() -> Result<()> {
5017 let descr = Arc::new(get_test_column_descr::<Int32Type>(1, 0));
5018 let column_close_result = write_multiple_pages::<Int32Type>(
5020 &descr,
5021 &[
5022 &[Some(-10), Some(10)],
5023 &[Some(-5), Some(11)],
5024 &[None],
5025 &[Some(-5), Some(11)],
5026 ],
5027 )?;
5028 let boundary_order = column_close_result
5029 .column_index
5030 .unwrap()
5031 .get_boundary_order();
5032 assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5033
5034 let column_close_result = write_multiple_pages::<Int32Type>(
5036 &descr,
5037 &[
5038 &[Some(10), Some(11)],
5039 &[Some(5), Some(11)],
5040 &[None],
5041 &[Some(-5), Some(0)],
5042 ],
5043 )?;
5044 let boundary_order = column_close_result
5045 .column_index
5046 .unwrap()
5047 .get_boundary_order();
5048 assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
5049
5050 let column_close_result = write_multiple_pages::<Int32Type>(
5052 &descr,
5053 &[&[Some(10), Some(11)], &[None], &[Some(10), Some(11)]],
5054 )?;
5055 let boundary_order = column_close_result
5056 .column_index
5057 .unwrap()
5058 .get_boundary_order();
5059 assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5060
5061 let column_close_result =
5063 write_multiple_pages::<Int32Type>(&descr, &[&[None], &[None], &[None]])?;
5064 let boundary_order = column_close_result
5065 .column_index
5066 .unwrap()
5067 .get_boundary_order();
5068 assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5069
5070 let column_close_result =
5072 write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)]])?;
5073 let boundary_order = column_close_result
5074 .column_index
5075 .unwrap()
5076 .get_boundary_order();
5077 assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5078
5079 let column_close_result =
5081 write_multiple_pages::<Int32Type>(&descr, &[&[Some(-10), Some(10)], &[None]])?;
5082 let boundary_order = column_close_result
5083 .column_index
5084 .unwrap()
5085 .get_boundary_order();
5086 assert_eq!(boundary_order, Some(BoundaryOrder::ASCENDING));
5087
5088 let column_close_result = write_multiple_pages::<Int32Type>(
5090 &descr,
5091 &[
5092 &[Some(10), Some(11)],
5093 &[Some(11), Some(16)],
5094 &[None],
5095 &[Some(-5), Some(0)],
5096 ],
5097 )?;
5098 let boundary_order = column_close_result
5099 .column_index
5100 .unwrap()
5101 .get_boundary_order();
5102 assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5103
5104 let column_close_result = write_multiple_pages::<Int32Type>(
5106 &descr,
5107 &[
5108 &[Some(1), Some(9)],
5109 &[Some(2), Some(8)],
5110 &[None],
5111 &[Some(3), Some(7)],
5112 ],
5113 )?;
5114 let boundary_order = column_close_result
5115 .column_index
5116 .unwrap()
5117 .get_boundary_order();
5118 assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5119
5120 Ok(())
5121 }
5122
5123 #[test]
5124 fn test_boundary_order_logical_type() -> Result<()> {
5125 let f16_descr = Arc::new(get_test_float16_column_descr(1, 0));
5128 let fba_descr = {
5129 let type_ = SchemaType::primitive_type_builder(
5130 "col",
5131 FixedLenByteArrayType::get_physical_type(),
5132 )
5133 .with_length(2)
5134 .build()?;
5135 Arc::new(ColumnDescriptor::new(
5136 Arc::new(type_),
5137 1,
5138 0,
5139 ColumnPath::from("col"),
5140 ))
5141 };
5142
5143 let values: &[&[Option<FixedLenByteArray>]] = &[
5144 &[Some(FixedLenByteArray::from(ByteArray::from(f16::ONE)))],
5145 &[Some(FixedLenByteArray::from(ByteArray::from(f16::ZERO)))],
5146 &[Some(FixedLenByteArray::from(ByteArray::from(
5147 f16::NEG_ZERO,
5148 )))],
5149 &[Some(FixedLenByteArray::from(ByteArray::from(f16::NEG_ONE)))],
5150 ];
5151
5152 let column_close_result =
5154 write_multiple_pages::<FixedLenByteArrayType>(&f16_descr, values)?;
5155 let boundary_order = column_close_result
5156 .column_index
5157 .unwrap()
5158 .get_boundary_order();
5159 assert_eq!(boundary_order, Some(BoundaryOrder::DESCENDING));
5160
5161 let column_close_result =
5163 write_multiple_pages::<FixedLenByteArrayType>(&fba_descr, values)?;
5164 let boundary_order = column_close_result
5165 .column_index
5166 .unwrap()
5167 .get_boundary_order();
5168 assert_eq!(boundary_order, Some(BoundaryOrder::UNORDERED));
5169
5170 Ok(())
5171 }
5172
5173 #[test]
5174 fn test_interval_stats_should_not_have_min_max() {
5175 let input = [
5176 vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
5177 vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],
5178 vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2],
5179 ]
5180 .into_iter()
5181 .map(|s| ByteArray::from(s).into())
5182 .collect::<Vec<_>>();
5183
5184 let page_writer = get_test_page_writer();
5185 let mut writer = get_test_interval_column_writer(page_writer);
5186 writer.write_batch(&input, None, None).unwrap();
5187
5188 let metadata = writer.close().unwrap().metadata;
5189 let stats = if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
5190 stats.clone()
5191 } else {
5192 panic!("metadata missing statistics");
5193 };
5194 assert!(stats.min_bytes_opt().is_none());
5195 assert!(stats.max_bytes_opt().is_none());
5196 }
5197
5198 #[test]
5199 #[cfg(feature = "arrow")]
5200 fn test_column_writer_get_estimated_total_bytes() {
5201 let page_writer = get_test_page_writer();
5202 let props = Default::default();
5203 let mut writer = get_test_column_writer::<Int32Type>(page_writer, 0, 0, props);
5204 assert_eq!(writer.get_estimated_total_bytes(), 0);
5205
5206 writer.write_batch(&[1, 2, 3, 4], None, None).unwrap();
5207 writer.add_data_page().unwrap();
5208 let size_with_one_page = writer.get_estimated_total_bytes();
5209 assert_eq!(size_with_one_page, 20);
5210
5211 writer.write_batch(&[5, 6, 7, 8], None, None).unwrap();
5212 writer.add_data_page().unwrap();
5213 let size_with_two_pages = writer.get_estimated_total_bytes();
5214 assert_eq!(size_with_two_pages, 20 + 21);
5216 }
5217
5218 fn write_multiple_pages<T: DataType>(
5219 column_descr: &Arc<ColumnDescriptor>,
5220 pages: &[&[Option<T::T>]],
5221 ) -> Result<ColumnCloseResult> {
5222 let column_writer = get_column_writer(
5223 column_descr.clone(),
5224 Default::default(),
5225 get_test_page_writer(),
5226 );
5227 let mut writer = get_typed_column_writer::<T>(column_writer);
5228
5229 for &page in pages {
5230 let values = page.iter().filter_map(Clone::clone).collect::<Vec<_>>();
5231 let def_levels = page
5232 .iter()
5233 .map(|maybe_value| i16::from(maybe_value.is_some()))
5234 .collect::<Vec<_>>();
5235 writer.write_batch(&values, Some(&def_levels), None)?;
5236 writer.flush_data_pages()?;
5237 }
5238
5239 writer.close()
5240 }
5241
5242 fn column_roundtrip_random<T: DataType>(
5246 props: WriterProperties,
5247 max_size: usize,
5248 min_value: T::T,
5249 max_value: T::T,
5250 max_def_level: i16,
5251 max_rep_level: i16,
5252 ) where
5253 T::T: PartialOrd + SampleUniform + Copy,
5254 {
5255 let mut num_values: usize = 0;
5256
5257 let mut buf: Vec<i16> = Vec::new();
5258 let def_levels = if max_def_level > 0 {
5259 random_numbers_range(max_size, 0, max_def_level + 1, &mut buf);
5260 for &dl in &buf[..] {
5261 if dl == max_def_level {
5262 num_values += 1;
5263 }
5264 }
5265 Some(&buf[..])
5266 } else {
5267 num_values = max_size;
5268 None
5269 };
5270
5271 let mut buf: Vec<i16> = Vec::new();
5272 let rep_levels = if max_rep_level > 0 {
5273 random_numbers_range(max_size, 0, max_rep_level + 1, &mut buf);
5274 buf[0] = 0; Some(&buf[..])
5276 } else {
5277 None
5278 };
5279
5280 let mut values: Vec<T::T> = Vec::new();
5281 random_numbers_range(num_values, min_value, max_value, &mut values);
5282
5283 column_roundtrip::<T>(props, &values[..], def_levels, rep_levels);
5284 }
5285
5286 fn column_roundtrip<T: DataType>(
5288 props: WriterProperties,
5289 values: &[T::T],
5290 def_levels: Option<&[i16]>,
5291 rep_levels: Option<&[i16]>,
5292 ) {
5293 let mut file = tempfile::tempfile().unwrap();
5294 let mut write = TrackedWrite::new(&mut file);
5295 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5296
5297 let max_def_level = match def_levels {
5298 Some(buf) => *buf.iter().max().unwrap_or(&0i16),
5299 None => 0i16,
5300 };
5301
5302 let max_rep_level = match rep_levels {
5303 Some(buf) => *buf.iter().max().unwrap_or(&0i16),
5304 None => 0i16,
5305 };
5306
5307 let mut max_batch_size = values.len();
5308 if let Some(levels) = def_levels {
5309 max_batch_size = max_batch_size.max(levels.len());
5310 }
5311 if let Some(levels) = rep_levels {
5312 max_batch_size = max_batch_size.max(levels.len());
5313 }
5314
5315 let mut writer =
5316 get_test_column_writer::<T>(page_writer, max_def_level, max_rep_level, Arc::new(props));
5317
5318 let values_written = writer.write_batch(values, def_levels, rep_levels).unwrap();
5319 assert_eq!(values_written, values.len());
5320 let result = writer.close().unwrap();
5321
5322 drop(write);
5323
5324 let props = ReaderProperties::builder()
5325 .set_backward_compatible_lz4(false)
5326 .build();
5327 let page_reader = Box::new(
5328 SerializedPageReader::new_with_properties(
5329 Arc::new(file),
5330 &result.metadata,
5331 result.rows_written as usize,
5332 None,
5333 Arc::new(props),
5334 )
5335 .unwrap(),
5336 );
5337 let mut reader = get_test_column_reader::<T>(page_reader, max_def_level, max_rep_level);
5338
5339 let mut actual_values = Vec::with_capacity(max_batch_size);
5340 let mut actual_def_levels = def_levels.map(|_| Vec::with_capacity(max_batch_size));
5341 let mut actual_rep_levels = rep_levels.map(|_| Vec::with_capacity(max_batch_size));
5342
5343 let (_, values_read, levels_read) = reader
5344 .read_records(
5345 max_batch_size,
5346 actual_def_levels.as_mut(),
5347 actual_rep_levels.as_mut(),
5348 &mut actual_values,
5349 )
5350 .unwrap();
5351
5352 assert_eq!(&actual_values[..values_read], values);
5355 match actual_def_levels {
5356 Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), def_levels),
5357 None => assert_eq!(None, def_levels),
5358 }
5359 match actual_rep_levels {
5360 Some(ref vec) => assert_eq!(Some(&vec[..levels_read]), rep_levels),
5361 None => assert_eq!(None, rep_levels),
5362 }
5363
5364 if let Some(levels) = actual_rep_levels {
5367 let mut actual_rows_written = 0;
5368 for l in levels {
5369 if l == 0 {
5370 actual_rows_written += 1;
5371 }
5372 }
5373 assert_eq!(actual_rows_written, result.rows_written);
5374 } else if actual_def_levels.is_some() {
5375 assert_eq!(levels_read as u64, result.rows_written);
5376 } else {
5377 assert_eq!(values_read as u64, result.rows_written);
5378 }
5379 }
5380
5381 fn column_write_and_get_metadata<T: DataType>(
5384 props: WriterProperties,
5385 values: &[T::T],
5386 ) -> ColumnChunkMetaData {
5387 let page_writer = get_test_page_writer();
5388 let props = Arc::new(props);
5389 let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
5390 writer.write_batch(values, None, None).unwrap();
5391 writer.close().unwrap().metadata
5392 }
5393
5394 fn encoding_stats(page_type: PageType, encoding: Encoding, count: i32) -> PageEncodingStats {
5396 PageEncodingStats {
5397 page_type,
5398 encoding,
5399 count,
5400 }
5401 }
5402
5403 fn check_encoding_write_support<T: DataType>(
5407 version: WriterVersion,
5408 dict_enabled: bool,
5409 data: &[T::T],
5410 dictionary_page_offset: Option<i64>,
5411 encodings: &[Encoding],
5412 page_encoding_stats: &[PageEncodingStats],
5413 ) {
5414 let props = WriterProperties::builder()
5415 .set_writer_version(version)
5416 .set_dictionary_enabled(dict_enabled)
5417 .build();
5418 let meta = column_write_and_get_metadata::<T>(props, data);
5419 assert_eq!(meta.dictionary_page_offset(), dictionary_page_offset);
5420 assert_eq!(meta.encodings().collect::<Vec<_>>(), encodings);
5421 assert_eq!(meta.page_encoding_stats().unwrap(), page_encoding_stats);
5422 }
5423
5424 fn get_test_column_writer<'a, T: DataType>(
5426 page_writer: Box<dyn PageWriter + 'a>,
5427 max_def_level: i16,
5428 max_rep_level: i16,
5429 props: WriterPropertiesPtr,
5430 ) -> ColumnWriterImpl<'a, T> {
5431 let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
5432 let column_writer = get_column_writer(descr, props, page_writer);
5433 get_typed_column_writer::<T>(column_writer)
5434 }
5435
5436 fn get_test_column_writer_with_path<'a, T: DataType>(
5437 page_writer: Box<dyn PageWriter + 'a>,
5438 max_def_level: i16,
5439 max_rep_level: i16,
5440 props: WriterPropertiesPtr,
5441 path: ColumnPath,
5442 ) -> ColumnWriterImpl<'a, T> {
5443 let descr = Arc::new(get_test_column_descr_with_path::<T>(
5444 max_def_level,
5445 max_rep_level,
5446 path,
5447 ));
5448 let column_writer = get_column_writer(descr, props, page_writer);
5449 get_typed_column_writer::<T>(column_writer)
5450 }
5451
5452 struct CollectedPages {
5454 data_pages: Vec<(usize, u32)>,
5456 dict_page_size: usize,
5458 }
5459
5460 fn write_and_collect_pages<T: DataType>(
5465 props: WriterProperties,
5466 max_def_level: i16,
5467 max_rep_level: i16,
5468 data: &[T::T],
5469 def_levels: Option<&[i16]>,
5470 rep_levels: Option<&[i16]>,
5471 ) -> CollectedPages {
5472 let mut file = tempfile::tempfile().unwrap();
5473 let mut write = TrackedWrite::new(&mut file);
5474 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5475 let mut writer =
5476 get_test_column_writer::<T>(page_writer, max_def_level, max_rep_level, Arc::new(props));
5477 writer.write_batch(data, def_levels, rep_levels).unwrap();
5478 let r = writer.close().unwrap();
5479 drop(write);
5480
5481 let read_props = ReaderProperties::builder()
5482 .set_backward_compatible_lz4(false)
5483 .build();
5484 let mut page_reader = Box::new(
5485 SerializedPageReader::new_with_properties(
5486 Arc::new(file),
5487 &r.metadata,
5488 r.rows_written as usize,
5489 None,
5490 Arc::new(read_props),
5491 )
5492 .unwrap(),
5493 );
5494
5495 let mut collected = CollectedPages {
5496 data_pages: Vec::new(),
5497 dict_page_size: 0,
5498 };
5499 while let Some(page) = page_reader.get_next_page().unwrap() {
5500 match page.page_type() {
5501 PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => {
5502 collected
5503 .data_pages
5504 .push((page.buffer().len(), page.num_values()));
5505 }
5506 PageType::DICTIONARY_PAGE => {
5507 collected.dict_page_size = collected.dict_page_size.max(page.buffer().len());
5508 }
5509 PageType::INDEX_PAGE => {}
5510 }
5511 }
5512 collected
5513 }
5514
5515 fn get_test_column_reader<T: DataType>(
5517 page_reader: Box<dyn PageReader>,
5518 max_def_level: i16,
5519 max_rep_level: i16,
5520 ) -> ColumnReaderImpl<T> {
5521 let descr = Arc::new(get_test_column_descr::<T>(max_def_level, max_rep_level));
5522 let column_reader = get_column_reader(descr, page_reader);
5523 get_typed_column_reader::<T>(column_reader)
5524 }
5525
5526 fn get_test_column_descr<T: DataType>(
5528 max_def_level: i16,
5529 max_rep_level: i16,
5530 ) -> ColumnDescriptor {
5531 let path = ColumnPath::from("col");
5532 let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5533 .with_length(1)
5536 .build()
5537 .unwrap();
5538 ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5539 }
5540
5541 fn get_test_column_descr_with_path<T: DataType>(
5542 max_def_level: i16,
5543 max_rep_level: i16,
5544 path: ColumnPath,
5545 ) -> ColumnDescriptor {
5546 let name = path.string();
5547 let type_ = SchemaType::primitive_type_builder(&name, T::get_physical_type())
5548 .with_length(1)
5551 .build()
5552 .unwrap();
5553 ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5554 }
5555
5556 fn write_and_collect_page_values(
5557 path: ColumnPath,
5558 props: WriterPropertiesPtr,
5559 data: &[i32],
5560 ) -> Vec<u32> {
5561 let mut file = tempfile::tempfile().unwrap();
5562 let mut write = TrackedWrite::new(&mut file);
5563 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5564 let mut writer =
5565 get_test_column_writer_with_path::<Int32Type>(page_writer, 0, 0, props, path);
5566 writer.write_batch(data, None, None).unwrap();
5567 let r = writer.close().unwrap();
5568
5569 drop(write);
5570
5571 let props = ReaderProperties::builder()
5572 .set_backward_compatible_lz4(false)
5573 .build();
5574 let mut page_reader = Box::new(
5575 SerializedPageReader::new_with_properties(
5576 Arc::new(file),
5577 &r.metadata,
5578 r.rows_written as usize,
5579 None,
5580 Arc::new(props),
5581 )
5582 .unwrap(),
5583 );
5584
5585 let mut values_per_page = Vec::new();
5586 while let Some(page) = page_reader.get_next_page().unwrap() {
5587 assert_eq!(page.page_type(), PageType::DATA_PAGE);
5588 values_per_page.push(page.num_values());
5589 }
5590
5591 values_per_page
5592 }
5593
5594 fn get_test_page_writer() -> Box<dyn PageWriter> {
5596 Box::new(TestPageWriter {})
5597 }
5598
5599 struct TestPageWriter {}
5600
5601 impl PageWriter for TestPageWriter {
5602 fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
5603 let mut res = PageWriteSpec::new();
5604 res.page_type = page.page_type();
5605 res.uncompressed_size = page.uncompressed_size();
5606 res.compressed_size = page.compressed_size();
5607 res.num_values = page.num_values();
5608 res.offset = 0;
5609 res.bytes_written = page.data().len() as u64;
5610 Ok(res)
5611 }
5612
5613 fn close(&mut self) -> Result<()> {
5614 Ok(())
5615 }
5616 }
5617
5618 fn statistics_roundtrip<T: DataType>(values: &[<T as DataType>::T]) -> Statistics {
5620 let page_writer = get_test_page_writer();
5621 let props = Default::default();
5622 let mut writer = get_test_column_writer::<T>(page_writer, 0, 0, props);
5623 writer.write_batch(values, None, None).unwrap();
5624
5625 let metadata = writer.close().unwrap().metadata;
5626 if let Some(stats) = metadata.statistics() {
5627 stats.clone()
5628 } else {
5629 panic!("metadata missing statistics");
5630 }
5631 }
5632
5633 fn get_test_decimals_column_writer<T: DataType>(
5635 page_writer: Box<dyn PageWriter>,
5636 max_def_level: i16,
5637 max_rep_level: i16,
5638 props: WriterPropertiesPtr,
5639 ) -> ColumnWriterImpl<'static, T> {
5640 let descr = Arc::new(get_test_decimals_column_descr::<T>(
5641 max_def_level,
5642 max_rep_level,
5643 ));
5644 let column_writer = get_column_writer(descr, props, page_writer);
5645 get_typed_column_writer::<T>(column_writer)
5646 }
5647
5648 fn get_test_decimals_column_descr<T: DataType>(
5650 max_def_level: i16,
5651 max_rep_level: i16,
5652 ) -> ColumnDescriptor {
5653 let path = ColumnPath::from("col");
5654 let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5655 .with_length(16)
5656 .with_logical_type(Some(LogicalType::decimal(2, 3)))
5657 .with_scale(2)
5658 .with_precision(3)
5659 .build()
5660 .unwrap();
5661 ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5662 }
5663
5664 fn float16_statistics_roundtrip(
5665 values: &[FixedLenByteArray],
5666 ) -> ValueStatistics<FixedLenByteArray> {
5667 let page_writer = get_test_page_writer();
5668 let mut writer = get_test_float16_column_writer(page_writer, Default::default());
5669 writer.write_batch(values, None, None).unwrap();
5670
5671 let metadata = writer.close().unwrap().metadata;
5672 if let Some(Statistics::FixedLenByteArray(stats)) = metadata.statistics() {
5673 stats.clone()
5674 } else {
5675 panic!("metadata missing statistics");
5676 }
5677 }
5678
5679 fn get_test_float16_column_writer(
5680 page_writer: Box<dyn PageWriter>,
5681 props: WriterPropertiesPtr,
5682 ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
5683 let descr = Arc::new(get_test_float16_column_descr(0, 0));
5684 let column_writer = get_column_writer(descr, props, page_writer);
5685 get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
5686 }
5687
5688 fn get_test_float16_column_descr(max_def_level: i16, max_rep_level: i16) -> ColumnDescriptor {
5689 let path = ColumnPath::from("col");
5690 let type_ =
5691 SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
5692 .with_length(2)
5693 .with_logical_type(Some(LogicalType::Float16))
5694 .build()
5695 .unwrap();
5696 ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5697 }
5698
5699 fn get_test_interval_column_writer(
5700 page_writer: Box<dyn PageWriter>,
5701 ) -> ColumnWriterImpl<'static, FixedLenByteArrayType> {
5702 let descr = Arc::new(get_test_interval_column_descr());
5703 let column_writer = get_column_writer(descr, Default::default(), page_writer);
5704 get_typed_column_writer::<FixedLenByteArrayType>(column_writer)
5705 }
5706
5707 fn get_test_interval_column_descr() -> ColumnDescriptor {
5708 let path = ColumnPath::from("col");
5709 let type_ =
5710 SchemaType::primitive_type_builder("col", FixedLenByteArrayType::get_physical_type())
5711 .with_length(12)
5712 .with_converted_type(ConvertedType::INTERVAL)
5713 .build()
5714 .unwrap();
5715 ColumnDescriptor::new(Arc::new(type_), 0, 0, path)
5716 }
5717
5718 fn get_test_unsigned_int_given_as_converted_column_writer<'a, T: DataType>(
5720 page_writer: Box<dyn PageWriter + 'a>,
5721 max_def_level: i16,
5722 max_rep_level: i16,
5723 props: WriterPropertiesPtr,
5724 ) -> ColumnWriterImpl<'a, T> {
5725 let descr = Arc::new(get_test_converted_type_unsigned_integer_column_descr::<T>(
5726 max_def_level,
5727 max_rep_level,
5728 ));
5729 let column_writer = get_column_writer(descr, props, page_writer);
5730 get_typed_column_writer::<T>(column_writer)
5731 }
5732
5733 fn get_test_converted_type_unsigned_integer_column_descr<T: DataType>(
5735 max_def_level: i16,
5736 max_rep_level: i16,
5737 ) -> ColumnDescriptor {
5738 let path = ColumnPath::from("col");
5739 let type_ = SchemaType::primitive_type_builder("col", T::get_physical_type())
5740 .with_converted_type(ConvertedType::UINT_32)
5741 .build()
5742 .unwrap();
5743 ColumnDescriptor::new(Arc::new(type_), max_def_level, max_rep_level, path)
5744 }
5745
5746 #[test]
5747 fn test_page_v2_snappy_compression_fallback() {
5748 let page_writer = TestPageWriter {};
5750
5751 let props = WriterProperties::builder()
5753 .set_writer_version(WriterVersion::PARQUET_2_0)
5754 .set_dictionary_enabled(false)
5756 .set_compression(Compression::SNAPPY)
5757 .build();
5758
5759 let mut column_writer =
5760 get_test_column_writer::<ByteArrayType>(Box::new(page_writer), 0, 0, Arc::new(props));
5761
5762 let values = vec![ByteArray::from("a")];
5765
5766 column_writer.write_batch(&values, None, None).unwrap();
5767
5768 let result = column_writer.close().unwrap();
5769 assert_eq!(
5770 result.metadata.uncompressed_size(),
5771 result.metadata.compressed_size()
5772 );
5773 }
5774
5775 struct ColumnRoundTripUniform<'a, T: DataType> {
5776 props: WriterProperties,
5777 values: &'a [T::T],
5778 def_levels: LevelDataRef<'a>,
5779 rep_levels: LevelDataRef<'a>,
5780 max_def_level: i16,
5781 max_rep_level: i16,
5782 expected_values: &'a [T::T],
5783 expected_def_levels: Option<&'a [i16]>,
5784 expected_rep_levels: Option<&'a [i16]>,
5785 }
5786
5787 impl<'a, T: DataType> ColumnRoundTripUniform<'a, T>
5788 where
5789 T::T: PartialEq + std::fmt::Debug,
5790 {
5791 fn new() -> Self {
5792 Self {
5793 props: Default::default(),
5794 values: &[],
5795 def_levels: LevelDataRef::Absent,
5796 rep_levels: LevelDataRef::Absent,
5797 max_def_level: 0,
5798 max_rep_level: 0,
5799 expected_values: &[],
5800 expected_def_levels: None,
5801 expected_rep_levels: None,
5802 }
5803 }
5804
5805 fn with_props(mut self, props: WriterProperties) -> Self {
5806 self.props = props;
5807 self
5808 }
5809
5810 fn with_values(mut self, values: &'a [T::T]) -> Self {
5811 self.values = values;
5812 self
5813 }
5814
5815 fn with_def_levels(mut self, def_levels: LevelDataRef<'a>) -> Self {
5816 self.def_levels = def_levels;
5817 self
5818 }
5819
5820 fn with_rep_levels(mut self, rep_levels: LevelDataRef<'a>) -> Self {
5821 self.rep_levels = rep_levels;
5822 self
5823 }
5824
5825 fn with_max_def_level(mut self, max_def_level: i16) -> Self {
5826 self.max_def_level = max_def_level;
5827 self
5828 }
5829
5830 fn with_max_rep_level(mut self, max_rep_level: i16) -> Self {
5831 self.max_rep_level = max_rep_level;
5832 self
5833 }
5834
5835 fn with_expected_values(mut self, expected_values: &'a [T::T]) -> Self {
5836 self.expected_values = expected_values;
5837 self
5838 }
5839
5840 fn with_expected_def_levels(mut self, expected_def_levels: &'a [i16]) -> Self {
5841 self.expected_def_levels = Some(expected_def_levels);
5842 self
5843 }
5844
5845 fn with_expected_rep_levels(mut self, expected_rep_levels: &'a [i16]) -> Self {
5846 self.expected_rep_levels = Some(expected_rep_levels);
5847 self
5848 }
5849
5850 fn run(self) {
5853 let mut file = tempfile::tempfile().unwrap();
5854 let mut write = TrackedWrite::new(&mut file);
5855 let page_writer = Box::new(SerializedPageWriter::new(&mut write));
5856 let mut writer = get_test_column_writer::<T>(
5857 page_writer,
5858 self.max_def_level,
5859 self.max_rep_level,
5860 Arc::new(self.props),
5861 );
5862
5863 writer
5864 .write_batch_internal(
5865 self.values,
5866 None,
5867 self.def_levels,
5868 self.rep_levels,
5869 None,
5870 None,
5871 None,
5872 )
5873 .unwrap();
5874 let result = writer.close().unwrap();
5875 drop(write);
5876
5877 let props = ReaderProperties::builder()
5878 .set_backward_compatible_lz4(false)
5879 .build();
5880 let page_reader = Box::new(
5881 SerializedPageReader::new_with_properties(
5882 Arc::new(file),
5883 &result.metadata,
5884 result.rows_written as usize,
5885 None,
5886 Arc::new(props),
5887 )
5888 .unwrap(),
5889 );
5890 let mut reader =
5891 get_test_column_reader::<T>(page_reader, self.max_def_level, self.max_rep_level);
5892
5893 let batch_size = self
5894 .expected_def_levels
5895 .map_or(self.expected_values.len(), |l| l.len());
5896 let mut actual_values = Vec::with_capacity(batch_size);
5897 let mut actual_def = self
5898 .expected_def_levels
5899 .map(|_| Vec::with_capacity(batch_size));
5900 let mut actual_rep = self
5901 .expected_rep_levels
5902 .map(|_| Vec::with_capacity(batch_size));
5903
5904 let (_, values_read, levels_read) = reader
5905 .read_records(
5906 batch_size,
5907 actual_def.as_mut(),
5908 actual_rep.as_mut(),
5909 &mut actual_values,
5910 )
5911 .unwrap();
5912
5913 assert_eq!(&actual_values[..values_read], self.expected_values);
5914 if let Some(ref v) = actual_def {
5915 assert_eq!(&v[..levels_read], self.expected_def_levels.unwrap());
5916 }
5917 if let Some(ref v) = actual_rep {
5918 assert_eq!(&v[..levels_read], self.expected_rep_levels.unwrap());
5919 }
5920 }
5921 }
5922
5923 #[test]
5924 fn test_level_data_ref_value_count() {
5925 let max_def = 2;
5930 assert_eq!(LevelDataRef::Absent.value_count(64, max_def), 64);
5933 assert_eq!(
5935 LevelDataRef::Uniform {
5936 value: max_def,
5937 count: 40
5938 }
5939 .value_count(40, max_def),
5940 40
5941 );
5942 assert_eq!(
5943 LevelDataRef::Uniform {
5944 value: max_def - 1,
5945 count: 40
5946 }
5947 .value_count(40, max_def),
5948 0
5949 );
5950 let levels = [2i16, 0, 2, 1, 2, 2, 0];
5953 assert_eq!(
5954 LevelDataRef::Materialized(&levels).value_count(levels.len(), max_def),
5955 4
5956 );
5957 }
5958
5959 #[test]
5960 fn test_uniform_def_levels_all_null() {
5961 let max_def_level = 1;
5963 let count = 100;
5964 let expected_def_levels = vec![0i16; count];
5965 ColumnRoundTripUniform::<Int32Type>::new()
5966 .with_def_levels(LevelDataRef::Uniform { value: 0, count })
5967 .with_max_def_level(max_def_level)
5968 .with_expected_def_levels(&expected_def_levels)
5969 .run();
5970 }
5971
5972 #[test]
5973 fn test_uniform_def_levels_all_valid() {
5974 let max_def_level = 1;
5976 let values: Vec<i32> = (0..50).collect();
5977 let expected_def_levels = vec![max_def_level; values.len()];
5978 ColumnRoundTripUniform::<Int32Type>::new()
5979 .with_values(&values)
5980 .with_def_levels(LevelDataRef::Uniform {
5981 value: max_def_level,
5982 count: values.len(),
5983 })
5984 .with_max_def_level(max_def_level)
5985 .with_expected_values(&values)
5986 .with_expected_def_levels(&expected_def_levels)
5987 .run();
5988 }
5989
5990 #[test]
5991 fn test_uniform_def_and_rep_levels() {
5992 let max_def_level = 2;
5995 let max_rep_level = 1;
5996 let count = 200;
5997 let expected_def_levels = vec![0i16; count];
5998 let expected_rep_levels = vec![0i16; count];
5999 ColumnRoundTripUniform::<Int32Type>::new()
6000 .with_def_levels(LevelDataRef::Uniform { value: 0, count })
6001 .with_rep_levels(LevelDataRef::Uniform { value: 0, count })
6002 .with_max_def_level(max_def_level)
6003 .with_max_rep_level(max_rep_level)
6004 .with_expected_def_levels(&expected_def_levels)
6005 .with_expected_rep_levels(&expected_rep_levels)
6006 .run();
6007 }
6008
6009 #[test]
6010 fn test_uniform_levels_v1_and_v2() {
6011 for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
6013 let props = WriterProperties::builder()
6014 .set_writer_version(version)
6015 .build();
6016 let max_def = 1;
6017 let count = 100;
6018 let expected_def_levels = vec![0i16; count];
6019 ColumnRoundTripUniform::<Int32Type>::new()
6020 .with_props(props)
6021 .with_def_levels(LevelDataRef::Uniform { value: 0, count })
6022 .with_max_def_level(max_def)
6023 .with_expected_def_levels(&expected_def_levels)
6024 .run();
6025 }
6026 }
6027}