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