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