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