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