1use std::io::Write;
14use std::sync::Arc;
15
16#[cfg(feature = "encryption")]
17pub(crate) mod encryption;
18
19#[cfg(feature = "encryption")]
20use crate::file::{
21 column_crypto_metadata::ColumnCryptoMetaData, metadata::thrift::encryption::EncryptionAlgorithm,
22};
23use crate::{
24 basic::{
25 ColumnOrder, CompressionCodec, ConvertedType, Encoding, EncodingMask, LogicalType,
26 PageType, Repetition, Type,
27 },
28 data_type::{ByteArray, FixedLenByteArray, Int96},
29 errors::{ParquetError, Result},
30 file::{
31 metadata::{
32 ColumnChunkMetaData, ColumnChunkMetaDataBuilder, FileMetaData, KeyValue,
33 LevelHistogram, PageEncodingStats, ParquetMetaData, ParquetMetaDataOptions,
34 ParquetPageEncodingStats, RowGroupMetaData, RowGroupMetaDataBuilder, SortingColumn,
35 },
36 statistics::ValueStatistics,
37 },
38 parquet_thrift::{
39 ElementType, FieldType, ReadThrift, ThriftCompactInputProtocol,
40 ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift, WriteThriftField,
41 check_list_size, read_thrift_vec, validate_list_type,
42 },
43 schema::types::{
44 ColumnDescriptor, SchemaDescriptor, TypePtr, num_nodes, parquet_schema_from_array,
45 },
46 thrift_struct,
47 util::bit_util::FromBytes,
48 write_thrift_field,
49};
50
51thrift_struct!(
53pub(crate) struct SchemaElement<'a> {
54 1: optional Type r#type;
56 2: optional i32 type_length;
61 3: optional Repetition repetition_type;
64 4: required string<'a> name;
66 5: optional i32 num_children;
71 6: optional ConvertedType converted_type;
76 7: optional i32 scale
81 8: optional i32 precision
82 9: optional i32 field_id;
85 10: optional LogicalType logical_type
90}
91);
92
93thrift_struct!(
94struct Statistics<'a> {
95 1: optional binary<'a> max;
96 2: optional binary<'a> min;
97 3: optional i64 null_count;
98 4: optional i64 distinct_count;
99 5: optional binary<'a> max_value;
100 6: optional binary<'a> min_value;
101 7: optional bool is_max_value_exact;
102 8: optional bool is_min_value_exact;
103 9: optional i64 nan_count;
104}
105);
106
107thrift_struct!(
108struct BoundingBox {
109 1: required double xmin;
110 2: required double xmax;
111 3: required double ymin;
112 4: required double ymax;
113 5: optional double zmin;
114 6: optional double zmax;
115 7: optional double mmin;
116 8: optional double mmax;
117}
118);
119
120thrift_struct!(
121struct GeospatialStatistics {
122 1: optional BoundingBox bbox;
123 2: optional list<i32> geospatial_types;
124}
125);
126
127thrift_struct!(
128struct SizeStatistics {
129 1: optional i64 unencoded_byte_array_data_bytes;
130 2: optional list<i64> repetition_level_histogram;
131 3: optional list<i64> definition_level_histogram;
132}
133);
134
135fn convert_geo_stats(
136 st: GeospatialStatistics,
137) -> crate::geospatial::statistics::GeospatialStatistics {
138 let bbox = st.bbox.map(convert_bounding_box);
139 let geospatial_types: Option<Vec<i32>> = st.geospatial_types.filter(|v| !v.is_empty());
140 crate::geospatial::statistics::GeospatialStatistics::new(bbox, geospatial_types)
141}
142
143fn convert_bounding_box(bb: BoundingBox) -> crate::geospatial::bounding_box::BoundingBox {
144 let mut newbb = crate::geospatial::bounding_box::BoundingBox::new(
145 bb.xmin.into(),
146 bb.xmax.into(),
147 bb.ymin.into(),
148 bb.ymax.into(),
149 );
150
151 newbb = match (bb.zmin, bb.zmax) {
152 (Some(zmin), Some(zmax)) => newbb.with_zrange(zmin.into(), zmax.into()),
153 _ => newbb,
155 };
156
157 newbb = match (bb.mmin, bb.mmax) {
158 (Some(mmin), Some(mmax)) => newbb.with_mrange(mmin.into(), mmax.into()),
159 _ => newbb,
161 };
162
163 newbb
164}
165
166fn convert_stats(
168 column_descr: &Arc<ColumnDescriptor>,
169 thrift_stats: Option<Statistics>,
170) -> Result<Option<crate::file::statistics::Statistics>> {
171 use crate::file::statistics::Statistics as FStatistics;
172 Ok(match thrift_stats {
173 Some(stats) => {
174 let null_count = stats
176 .null_count
177 .map(|null_count| {
178 if null_count < 0 {
179 return Err(general_err!(
180 "Statistics null count is negative {}",
181 null_count
182 ));
183 }
184 Ok(null_count as u64)
185 })
186 .transpose()?;
187 let distinct_count = stats.distinct_count.map(|value| value as u64);
189 let nan_count = stats
191 .nan_count
192 .map(|nan_count| {
193 if nan_count < 0 {
194 return Err(general_err!(
195 "Statistics NaN count is negative {}",
196 nan_count
197 ));
198 }
199 Ok(nan_count as u64)
200 })
201 .transpose()?;
202 let old_format = stats.min_value.is_none() && stats.max_value.is_none();
204 let min = if old_format {
206 stats.min
207 } else {
208 stats.min_value
209 };
210 let max = if old_format {
212 stats.max
213 } else {
214 stats.max_value
215 };
216
217 fn check_len(min: Option<&[u8]>, max: Option<&[u8]>, len: usize) -> Result<()> {
218 if let Some(min) = min
219 && min.len() < len
220 {
221 return Err(general_err!("Insufficient bytes to parse min statistic",));
222 }
223 if let Some(max) = max
224 && max.len() < len
225 {
226 return Err(general_err!("Insufficient bytes to parse max statistic",));
227 }
228 Ok(())
229 }
230
231 let physical_type = column_descr.physical_type();
232 match physical_type {
233 Type::BOOLEAN => check_len(min, max, 1),
234 Type::INT32 | Type::FLOAT => check_len(min, max, 4),
235 Type::INT64 | Type::DOUBLE => check_len(min, max, 8),
236 Type::INT96 => check_len(min, max, 12),
237 _ => Ok(()),
238 }?;
239
240 let res = match physical_type {
245 Type::BOOLEAN => FStatistics::boolean(
246 min.map(|data| data[0] != 0),
247 max.map(|data| data[0] != 0),
248 distinct_count,
249 null_count,
250 old_format,
251 ),
252 Type::INT32 => FStatistics::int32(
253 min.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())),
254 max.map(|data| i32::from_le_bytes(data[..4].try_into().unwrap())),
255 distinct_count,
256 null_count,
257 old_format,
258 ),
259 Type::INT64 => FStatistics::int64(
260 min.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())),
261 max.map(|data| i64::from_le_bytes(data[..8].try_into().unwrap())),
262 distinct_count,
263 null_count,
264 old_format,
265 ),
266 Type::INT96 => {
267 let min = if let Some(data) = min {
269 if data.len() != 12 {
270 return Err(general_err!("Incorrect Int96 min statistics"));
271 }
272 Some(Int96::try_from_le_slice(data)?)
273 } else {
274 None
275 };
276 let max = if let Some(data) = max {
277 if data.len() != 12 {
278 return Err(general_err!("Incorrect Int96 max statistics"));
279 }
280 Some(Int96::try_from_le_slice(data)?)
281 } else {
282 None
283 };
284 FStatistics::int96(min, max, distinct_count, null_count, old_format)
285 }
286 Type::FLOAT => FStatistics::Float(
287 ValueStatistics::new(
288 min.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())),
289 max.map(|data| f32::from_le_bytes(data[..4].try_into().unwrap())),
290 distinct_count,
291 null_count,
292 old_format,
293 )
294 .with_nan_count(nan_count),
295 ),
296 Type::DOUBLE => FStatistics::Double(
297 ValueStatistics::new(
298 min.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())),
299 max.map(|data| f64::from_le_bytes(data[..8].try_into().unwrap())),
300 distinct_count,
301 null_count,
302 old_format,
303 )
304 .with_nan_count(nan_count),
305 ),
306 Type::BYTE_ARRAY => FStatistics::ByteArray(
307 ValueStatistics::new(
308 min.map(ByteArray::from),
309 max.map(ByteArray::from),
310 distinct_count,
311 null_count,
312 old_format,
313 )
314 .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false))
315 .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)),
316 ),
317 Type::FIXED_LEN_BYTE_ARRAY => FStatistics::FixedLenByteArray(
318 ValueStatistics::new(
319 min.map(ByteArray::from).map(FixedLenByteArray::from),
320 max.map(ByteArray::from).map(FixedLenByteArray::from),
321 distinct_count,
322 null_count,
323 old_format,
324 )
325 .with_nan_count(nan_count)
326 .with_max_is_exact(stats.is_max_value_exact.unwrap_or(false))
327 .with_min_is_exact(stats.is_min_value_exact.unwrap_or(false)),
328 ),
329 };
330
331 Some(res)
332 }
333 None => None,
334 })
335}
336
337const COL_META_TYPE: u16 = 1 << 1;
339const COL_META_ENCODINGS: u16 = 1 << 2;
340const COL_META_CODEC: u16 = 1 << 4;
341const COL_META_NUM_VALUES: u16 = 1 << 5;
342const COL_META_TOTAL_UNCOMP_SZ: u16 = 1 << 6;
343const COL_META_TOTAL_COMP_SZ: u16 = 1 << 7;
344const COL_META_DATA_PAGE_OFFSET: u16 = 1 << 9;
345
346const COL_META_ALL_REQUIRED: u16 = COL_META_TYPE
348 | COL_META_ENCODINGS
349 | COL_META_CODEC
350 | COL_META_NUM_VALUES
351 | COL_META_TOTAL_UNCOMP_SZ
352 | COL_META_TOTAL_COMP_SZ
353 | COL_META_DATA_PAGE_OFFSET;
354
355fn validate_column_metadata(mask: u16) -> Result<()> {
358 if mask != COL_META_ALL_REQUIRED {
359 if mask & COL_META_ENCODINGS == 0 {
360 return Err(general_err!("Required field encodings is missing"));
361 }
362
363 if mask & COL_META_CODEC == 0 {
364 return Err(general_err!("Required field codec is missing"));
365 }
366 if mask & COL_META_NUM_VALUES == 0 {
367 return Err(general_err!("Required field num_values is missing"));
368 }
369 if mask & COL_META_TOTAL_UNCOMP_SZ == 0 {
370 return Err(general_err!(
371 "Required field total_uncompressed_size is missing"
372 ));
373 }
374 if mask & COL_META_TOTAL_COMP_SZ == 0 {
375 return Err(general_err!(
376 "Required field total_compressed_size is missing"
377 ));
378 }
379 if mask & COL_META_DATA_PAGE_OFFSET == 0 {
380 return Err(general_err!("Required field data_page_offset is missing"));
381 }
382 }
383
384 Ok(())
385}
386
387fn read_encoding_stats_as_mask(prot: &mut ThriftSliceInputProtocol<'_>) -> Result<EncodingMask> {
388 let mut mask = 0i32;
390 let list_ident = prot.read_list_begin()?;
391 validate_list_type(ElementType::Struct, &list_ident)?;
393
394 for _ in 0..list_ident.size {
395 let pes = PageEncodingStats::read_thrift(prot)?;
396 match pes.page_type {
397 PageType::DATA_PAGE | PageType::DATA_PAGE_V2 => mask |= 1 << pes.encoding as i32,
398 _ => {}
399 }
400 }
401 EncodingMask::try_new(mask)
402}
403
404#[expect(clippy::useless_let_if_seq)] fn read_column_metadata(
408 prot: &mut ThriftSliceInputProtocol<'_>,
409 column: &mut ColumnChunkMetaData,
410 col_index: usize,
411 options: Option<&ParquetMetaDataOptions>,
412) -> Result<u16> {
413 let mut seen_mask = 0u16;
415
416 let mut skip_pes = false;
417 let mut pes_mask = true;
418 let mut skip_col_stats = false;
419 let mut skip_size_stats = false;
420
421 if let Some(opts) = options {
422 skip_pes = opts.skip_encoding_stats(col_index);
423 pes_mask = opts.encoding_stats_as_mask();
424 skip_col_stats = opts.skip_column_stats(col_index);
425 skip_size_stats = opts.skip_size_stats(col_index);
426 }
427
428 let column_descr = &column.column_descr;
448
449 let mut last_field_id = 0i16;
450 loop {
451 let field_ident = prot.read_field_begin(last_field_id)?;
452 if field_ident.field_type == FieldType::Stop {
453 break;
454 }
455 match field_ident.id {
456 1 => {
458 Type::read_thrift(&mut *prot)?;
460 seen_mask |= COL_META_TYPE;
461 }
462 2 => {
463 column.encodings = EncodingMask::read_thrift(&mut *prot)?;
464 seen_mask |= COL_META_ENCODINGS;
465 }
466 4 => {
468 column.compression = CompressionCodec::read_thrift(&mut *prot)?;
469 seen_mask |= COL_META_CODEC;
470 }
471 5 => {
472 column.num_values = i64::read_thrift(&mut *prot)?;
473 seen_mask |= COL_META_NUM_VALUES;
474 }
475 6 => {
476 column.total_uncompressed_size = i64::read_thrift(&mut *prot)?;
477 seen_mask |= COL_META_TOTAL_UNCOMP_SZ;
478 }
479 7 => {
480 column.total_compressed_size = i64::read_thrift(&mut *prot)?;
481 seen_mask |= COL_META_TOTAL_COMP_SZ;
482 }
483 9 => {
485 column.data_page_offset = i64::read_thrift(&mut *prot)?;
486 seen_mask |= COL_META_DATA_PAGE_OFFSET;
487 }
488 10 => {
489 column.index_page_offset = Some(i64::read_thrift(&mut *prot)?);
490 }
491 11 => {
492 column.dictionary_page_offset = Some(i64::read_thrift(&mut *prot)?);
493 }
494 12 if !skip_col_stats => {
495 column.statistics =
496 convert_stats(column_descr, Some(Statistics::read_thrift(&mut *prot)?))?;
497 }
498 13 if !skip_pes => {
499 if pes_mask {
500 let val = read_encoding_stats_as_mask(&mut *prot)?;
501 column.encoding_stats = Some(ParquetPageEncodingStats::Mask(val));
502 } else {
503 let val =
504 read_thrift_vec::<PageEncodingStats, ThriftSliceInputProtocol>(&mut *prot)?;
505 column.encoding_stats = Some(ParquetPageEncodingStats::Full(val));
506 }
507 }
508 14 => {
509 column.bloom_filter_offset = Some(i64::read_thrift(&mut *prot)?);
510 }
511 15 => {
512 column.bloom_filter_length = Some(i32::read_thrift(&mut *prot)?);
513 }
514 16 if !skip_size_stats => {
515 let val = SizeStatistics::read_thrift(&mut *prot)?;
516 column.unencoded_byte_array_data_bytes = val.unencoded_byte_array_data_bytes;
517 column.repetition_level_histogram =
518 val.repetition_level_histogram.map(LevelHistogram::from);
519 column.definition_level_histogram =
520 val.definition_level_histogram.map(LevelHistogram::from);
521 }
522 17 => {
523 let val = GeospatialStatistics::read_thrift(&mut *prot)?;
524 column.geo_statistics = Some(Box::new(convert_geo_stats(val)));
525 }
526 _ => {
527 prot.skip(field_ident.field_type)?;
528 }
529 }
530 last_field_id = field_ident.id;
531 }
532
533 Ok(seen_mask)
534}
535
536fn read_column_chunk(
539 prot: &mut ThriftSliceInputProtocol<'_>,
540 column_descr: &Arc<ColumnDescriptor>,
541 col_index: usize,
542 options: Option<&ParquetMetaDataOptions>,
543) -> Result<ColumnChunkMetaData> {
544 let mut col = ColumnChunkMetaDataBuilder::new(column_descr.clone()).build()?;
546
547 let mut has_file_offset = false;
549
550 let mut col_meta_mask = 0u16;
552
553 let mut last_field_id = 0i16;
565 loop {
566 let field_ident = prot.read_field_begin(last_field_id)?;
567 if field_ident.field_type == FieldType::Stop {
568 break;
569 }
570 match field_ident.id {
571 1 => {
572 col.file_path = Some(String::read_thrift(&mut *prot)?);
573 }
574 2 => {
575 col.file_offset = i64::read_thrift(&mut *prot)?;
576 has_file_offset = true;
577 }
578 3 => {
579 col_meta_mask = read_column_metadata(&mut *prot, &mut col, col_index, options)?;
580 }
581 4 => {
582 col.offset_index_offset = Some(i64::read_thrift(&mut *prot)?);
583 }
584 5 => {
585 col.offset_index_length = Some(i32::read_thrift(&mut *prot)?);
586 }
587 6 => {
588 col.column_index_offset = Some(i64::read_thrift(&mut *prot)?);
589 }
590 7 => {
591 col.column_index_length = Some(i32::read_thrift(&mut *prot)?);
592 }
593 #[cfg(feature = "encryption")]
594 8 => {
595 let val = ColumnCryptoMetaData::read_thrift(&mut *prot)?;
596 col.column_crypto_metadata = Some(Box::new(val));
597 }
598 #[cfg(feature = "encryption")]
599 9 => {
600 col.encrypted_column_metadata = Some(<&[u8]>::read_thrift(&mut *prot)?.to_vec());
601 }
602 _ => {
603 prot.skip(field_ident.field_type)?;
604 }
605 }
606 last_field_id = field_ident.id;
607 }
608
609 if !has_file_offset {
611 return Err(general_err!("Required field file_offset is missing"));
612 }
613
614 #[cfg(feature = "encryption")]
616 if col.encrypted_column_metadata.is_some() {
617 return Ok(col);
618 }
619
620 validate_column_metadata(col_meta_mask)?;
622
623 Ok(col)
624}
625
626fn read_row_group(
627 prot: &mut ThriftSliceInputProtocol,
628 schema_descr: &Arc<SchemaDescriptor>,
629 options: Option<&ParquetMetaDataOptions>,
630) -> Result<RowGroupMetaData> {
631 let mut row_group = RowGroupMetaDataBuilder::new(schema_descr.clone()).build_unchecked();
633
634 const RG_COLUMNS: u8 = 1 << 1;
636 const RG_TOT_BYTE_SIZE: u8 = 1 << 2;
637 const RG_NUM_ROWS: u8 = 1 << 3;
638 const RG_ALL_REQUIRED: u8 = RG_COLUMNS | RG_TOT_BYTE_SIZE | RG_NUM_ROWS;
639
640 let mut mask = 0u8;
641
642 let mut last_field_id = 0i16;
652 loop {
653 let field_ident = prot.read_field_begin(last_field_id)?;
654 if field_ident.field_type == FieldType::Stop {
655 break;
656 }
657 match field_ident.id {
658 1 => {
659 let list_ident = prot.read_list_begin()?;
660 validate_list_type(ElementType::Struct, &list_ident)?;
662 if schema_descr.num_columns() != list_ident.size as usize {
663 return Err(general_err!(
664 "Column count mismatch. Schema has {} columns while Row Group has {}",
665 schema_descr.num_columns(),
666 list_ident.size
667 ));
668 }
669 for i in 0..list_ident.size as usize {
670 let col = read_column_chunk(prot, &schema_descr.columns()[i], i, options)?;
671 row_group.columns.push(col);
672 }
673 mask |= RG_COLUMNS;
674 }
675 2 => {
676 row_group.total_byte_size = i64::read_thrift(&mut *prot)?;
677 mask |= RG_TOT_BYTE_SIZE;
678 }
679 3 => {
680 row_group.num_rows = i64::read_thrift(&mut *prot)?;
681 mask |= RG_NUM_ROWS;
682 }
683 4 => {
684 let val = read_thrift_vec::<SortingColumn, ThriftSliceInputProtocol>(&mut *prot)?;
685 row_group.sorting_columns = Some(val);
686 }
687 5 => {
688 row_group.file_offset = Some(i64::read_thrift(&mut *prot)?);
689 }
690 7 => {
692 row_group.ordinal = Some(i16::read_thrift(&mut *prot)? as i32);
693 }
694 _ => {
695 prot.skip(field_ident.field_type)?;
696 }
697 }
698 last_field_id = field_ident.id;
699 }
700
701 if mask != RG_ALL_REQUIRED {
702 if mask & RG_COLUMNS == 0 {
703 return Err(general_err!("Required field columns is missing"));
704 }
705 if mask & RG_TOT_BYTE_SIZE == 0 {
706 return Err(general_err!("Required field total_byte_size is missing"));
707 }
708 if mask & RG_NUM_ROWS == 0 {
709 return Err(general_err!("Required field num_rows is missing"));
710 }
711 }
712
713 Ok(row_group)
714}
715
716pub(crate) fn parquet_schema_from_bytes(buf: &[u8]) -> Result<SchemaDescriptor> {
719 let mut prot = ThriftSliceInputProtocol::new(buf);
720
721 let mut last_field_id = 0i16;
722 loop {
723 let field_ident = prot.read_field_begin(last_field_id)?;
724 if field_ident.field_type == FieldType::Stop {
725 break;
726 }
727 match field_ident.id {
728 2 => {
729 let val = read_thrift_vec::<SchemaElement, ThriftSliceInputProtocol>(&mut prot)?;
731 let val = parquet_schema_from_array(val)?;
732 return Ok(SchemaDescriptor::new(val));
733 }
734 _ => prot.skip(field_ident.field_type)?,
735 }
736 last_field_id = field_ident.id;
737 }
738 Err(general_err!("Input does not contain a schema"))
739}
740
741pub(crate) fn parquet_metadata_from_bytes(
744 buf: &[u8],
745 options: Option<&ParquetMetaDataOptions>,
746) -> Result<ParquetMetaData> {
747 let (file_metadata, row_groups) = file_metadata_and_row_groups_from_bytes(buf, options)?;
748 Ok(ParquetMetaData::new(file_metadata, row_groups))
749}
750
751pub(crate) fn file_metadata_and_row_groups_from_bytes(
752 buf: &[u8],
753 options: Option<&ParquetMetaDataOptions>,
754) -> Result<(FileMetaData, Vec<RowGroupMetaData>)> {
755 let mut prot = ThriftSliceInputProtocol::new(buf);
756
757 let mut version: Option<i32> = None;
759 let mut num_rows: Option<i64> = None;
760 let mut row_groups: Option<Vec<RowGroupMetaData>> = None;
761 let mut key_value_metadata: Option<Vec<KeyValue>> = None;
762 let mut created_by: Option<&str> = None;
763 let mut column_orders: Option<Vec<ColumnOrder>> = None;
764 #[cfg(feature = "encryption")]
765 let mut encryption_algorithm: Option<EncryptionAlgorithm> = None;
766 #[cfg(feature = "encryption")]
767 let mut footer_signing_key_metadata: Option<&[u8]> = None;
768
769 let mut schema_descr: Option<Arc<SchemaDescriptor>> =
772 options.and_then(|options| options.schema().cloned());
773
774 let mut last_field_id = 0i16;
786 loop {
787 let field_ident = prot.read_field_begin(last_field_id)?;
788 if field_ident.field_type == FieldType::Stop {
789 break;
790 }
791 match field_ident.id {
792 1 => {
793 version = Some(i32::read_thrift(&mut prot)?);
794 }
795 2 => {
796 if schema_descr.is_some() {
798 prot.skip(field_ident.field_type)?;
799 } else {
800 let val =
802 read_thrift_vec::<SchemaElement, ThriftSliceInputProtocol>(&mut prot)?;
803 let val = parquet_schema_from_array(val)?;
804 schema_descr = Some(Arc::new(SchemaDescriptor::new(val)));
805 }
806 }
807 3 => {
808 num_rows = Some(i64::read_thrift(&mut prot)?);
809 }
810 4 => {
811 if schema_descr.is_none() {
812 return Err(general_err!("Required field schema is missing"));
813 }
814 let schema_descr = schema_descr.as_ref().unwrap();
815 let list_ident = prot.read_list_begin()?;
816 validate_list_type(ElementType::Struct, &list_ident)?;
818 let size = list_ident.size as usize;
819 check_list_size(prot.remaining_bytes(), size)?;
821 let mut rg_vec = Vec::with_capacity(size);
822
823 for _ in 0..list_ident.size {
824 rg_vec.push(read_row_group(&mut prot, schema_descr, options)?);
825 }
826 ensure_row_group_ordinals(&mut rg_vec)?;
827 row_groups = Some(rg_vec);
828 }
829 5 => {
830 let val = read_thrift_vec::<KeyValue, ThriftSliceInputProtocol>(&mut prot)?;
831 key_value_metadata = Some(val);
832 }
833 6 => {
834 created_by = Some(<&str>::read_thrift(&mut prot)?);
835 }
836 7 => {
837 let val = read_thrift_vec::<ColumnOrder, ThriftSliceInputProtocol>(&mut prot)?;
838 column_orders = Some(val);
839 }
840 #[cfg(feature = "encryption")]
841 8 => {
842 let val = EncryptionAlgorithm::read_thrift(&mut prot)?;
843 encryption_algorithm = Some(val);
844 }
845 #[cfg(feature = "encryption")]
846 9 => {
847 footer_signing_key_metadata = Some(<&[u8]>::read_thrift(&mut prot)?);
848 }
849 _ => {
850 prot.skip(field_ident.field_type)?;
851 }
852 }
853 last_field_id = field_ident.id;
854 }
855 let Some(version) = version else {
856 return Err(general_err!("Required field version is missing"));
857 };
858 let Some(num_rows) = num_rows else {
859 return Err(general_err!("Required field num_rows is missing"));
860 };
861 let Some(row_groups) = row_groups else {
862 return Err(general_err!("Required field row_groups is missing"));
863 };
864
865 let created_by = created_by.map(|c| c.to_owned());
866
867 let schema_descr = schema_descr.unwrap();
869
870 if column_orders
872 .as_ref()
873 .is_some_and(|cos| cos.len() != schema_descr.num_columns())
874 {
875 return Err(general_err!("Column order length mismatch"));
876 }
877 let column_orders = column_orders.map(|mut cos| {
879 for (i, column) in schema_descr.columns().iter().enumerate() {
880 if let ColumnOrder::TYPE_DEFINED_ORDER(_) = cos[i] {
881 let sort_order = ColumnOrder::get_sort_order_for_type(
885 column.logical_type_ref(),
886 column.converted_type(),
887 column.physical_type(),
888 true,
889 );
890 cos[i] = ColumnOrder::TYPE_DEFINED_ORDER(sort_order);
891 }
892 }
893 cos
894 });
895
896 #[cfg(not(feature = "encryption"))]
897 let fmd = FileMetaData::new(
898 version,
899 num_rows,
900 created_by,
901 key_value_metadata,
902 schema_descr,
903 column_orders,
904 );
905 #[cfg(feature = "encryption")]
906 let fmd = FileMetaData::new(
907 version,
908 num_rows,
909 created_by,
910 key_value_metadata,
911 schema_descr,
912 column_orders,
913 )
914 .with_encryption_algorithm(encryption_algorithm)
915 .with_footer_signing_key_metadata(footer_signing_key_metadata.map(|v| v.to_vec()));
916
917 Ok((fmd, row_groups))
918}
919
920fn ensure_row_group_ordinals(row_groups: &mut [RowGroupMetaData]) -> Result<()> {
937 if row_groups.iter().any(|rg| rg.ordinal.is_some()) {
941 return Ok(());
942 }
943 for (idx, rg) in row_groups.iter_mut().enumerate() {
944 let ordinal: i32 = idx
945 .try_into()
946 .map_err(|_| general_err!("Row group ordinal {} exceeds i32 max value", idx))?;
947 rg.ordinal = Some(ordinal);
948 }
949 Ok(())
950}
951
952thrift_struct!(
953 pub(crate) struct IndexPageHeader {}
954);
955
956thrift_struct!(
957pub(crate) struct DictionaryPageHeader {
958 1: required i32 num_values;
960
961 2: required Encoding encoding
963
964 3: optional bool is_sorted;
966}
967);
968
969thrift_struct!(
970pub(crate) struct PageStatistics {
978 1: optional binary max;
979 2: optional binary min;
980 3: optional i64 null_count;
981 4: optional i64 distinct_count;
982 5: optional binary max_value;
983 6: optional binary min_value;
984 7: optional bool is_max_value_exact;
985 8: optional bool is_min_value_exact;
986 9: optional i64 nan_count;
987}
988);
989
990thrift_struct!(
991pub(crate) struct DataPageHeader {
992 1: required i32 num_values
993 2: required Encoding encoding
994 3: required Encoding definition_level_encoding;
995 4: required Encoding repetition_level_encoding;
996 5: optional PageStatistics statistics;
997}
998);
999
1000impl DataPageHeader {
1001 fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result<Self>
1003 where
1004 R: ThriftCompactInputProtocol<'a>,
1005 {
1006 let mut num_values: Option<i32> = None;
1007 let mut encoding: Option<Encoding> = None;
1008 let mut definition_level_encoding: Option<Encoding> = None;
1009 let mut repetition_level_encoding: Option<Encoding> = None;
1010 let statistics: Option<PageStatistics> = None;
1011 let mut last_field_id = 0i16;
1012 loop {
1013 let field_ident = prot.read_field_begin(last_field_id)?;
1014 if field_ident.field_type == FieldType::Stop {
1015 break;
1016 }
1017 match field_ident.id {
1018 1 => {
1019 let val = i32::read_thrift(&mut *prot)?;
1020 num_values = Some(val);
1021 }
1022 2 => {
1023 let val = Encoding::read_thrift(&mut *prot)?;
1024 encoding = Some(val);
1025 }
1026 3 => {
1027 let val = Encoding::read_thrift(&mut *prot)?;
1028 definition_level_encoding = Some(val);
1029 }
1030 4 => {
1031 let val = Encoding::read_thrift(&mut *prot)?;
1032 repetition_level_encoding = Some(val);
1033 }
1034 _ => {
1035 prot.skip(field_ident.field_type)?;
1036 }
1037 }
1038 last_field_id = field_ident.id;
1039 }
1040 let Some(num_values) = num_values else {
1041 return Err(general_err!("Required field num_values is missing"));
1042 };
1043 let Some(encoding) = encoding else {
1044 return Err(general_err!("Required field encoding is missing"));
1045 };
1046 let Some(definition_level_encoding) = definition_level_encoding else {
1047 return Err(general_err!(
1048 "Required field definition_level_encoding is missing"
1049 ));
1050 };
1051 let Some(repetition_level_encoding) = repetition_level_encoding else {
1052 return Err(general_err!(
1053 "Required field repetition_level_encoding is missing"
1054 ));
1055 };
1056 Ok(Self {
1057 num_values,
1058 encoding,
1059 definition_level_encoding,
1060 repetition_level_encoding,
1061 statistics,
1062 })
1063 }
1064}
1065
1066thrift_struct!(
1067pub(crate) struct DataPageHeaderV2 {
1068 1: required i32 num_values
1069 2: required i32 num_nulls
1070 3: required i32 num_rows
1071 4: required Encoding encoding
1072 5: required i32 definition_levels_byte_length;
1073 6: required i32 repetition_levels_byte_length;
1074 7: optional bool is_compressed = true;
1075 8: optional PageStatistics statistics;
1076}
1077);
1078
1079impl DataPageHeaderV2 {
1080 fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result<Self>
1082 where
1083 R: ThriftCompactInputProtocol<'a>,
1084 {
1085 let mut num_values: Option<i32> = None;
1086 let mut num_nulls: Option<i32> = None;
1087 let mut num_rows: Option<i32> = None;
1088 let mut encoding: Option<Encoding> = None;
1089 let mut definition_levels_byte_length: Option<i32> = None;
1090 let mut repetition_levels_byte_length: Option<i32> = None;
1091 let mut is_compressed: Option<bool> = None;
1092 let statistics: Option<PageStatistics> = None;
1093 let mut last_field_id = 0i16;
1094 loop {
1095 let field_ident = prot.read_field_begin(last_field_id)?;
1096 if field_ident.field_type == FieldType::Stop {
1097 break;
1098 }
1099 match field_ident.id {
1100 1 => {
1101 let val = i32::read_thrift(&mut *prot)?;
1102 num_values = Some(val);
1103 }
1104 2 => {
1105 let val = i32::read_thrift(&mut *prot)?;
1106 num_nulls = Some(val);
1107 }
1108 3 => {
1109 let val = i32::read_thrift(&mut *prot)?;
1110 num_rows = Some(val);
1111 }
1112 4 => {
1113 let val = Encoding::read_thrift(&mut *prot)?;
1114 encoding = Some(val);
1115 }
1116 5 => {
1117 let val = i32::read_thrift(&mut *prot)?;
1118 definition_levels_byte_length = Some(val);
1119 }
1120 6 => {
1121 let val = i32::read_thrift(&mut *prot)?;
1122 repetition_levels_byte_length = Some(val);
1123 }
1124 7 => {
1125 is_compressed = Some(field_ident.bool_val()?);
1126 }
1127 _ => {
1128 prot.skip(field_ident.field_type)?;
1129 }
1130 }
1131 last_field_id = field_ident.id;
1132 }
1133 let Some(num_values) = num_values else {
1134 return Err(general_err!("Required field num_values is missing"));
1135 };
1136 let Some(num_nulls) = num_nulls else {
1137 return Err(general_err!("Required field num_nulls is missing"));
1138 };
1139 let Some(num_rows) = num_rows else {
1140 return Err(general_err!("Required field num_rows is missing"));
1141 };
1142 let Some(encoding) = encoding else {
1143 return Err(general_err!("Required field encoding is missing"));
1144 };
1145 let Some(definition_levels_byte_length) = definition_levels_byte_length else {
1146 return Err(general_err!(
1147 "Required field definition_levels_byte_length is missing"
1148 ));
1149 };
1150 let Some(repetition_levels_byte_length) = repetition_levels_byte_length else {
1151 return Err(general_err!(
1152 "Required field repetition_levels_byte_length is missing"
1153 ));
1154 };
1155 Ok(Self {
1156 num_values,
1157 num_nulls,
1158 num_rows,
1159 encoding,
1160 definition_levels_byte_length,
1161 repetition_levels_byte_length,
1162 is_compressed,
1163 statistics,
1164 })
1165 }
1166}
1167
1168thrift_struct!(
1169pub(crate) struct PageHeader {
1170 1: required PageType r#type
1172
1173 2: required i32 uncompressed_page_size
1175
1176 3: required i32 compressed_page_size
1178
1179 4: optional i32 crc
1181
1182 5: optional DataPageHeader data_page_header;
1184 6: optional IndexPageHeader index_page_header;
1185 7: optional DictionaryPageHeader dictionary_page_header;
1186 8: optional DataPageHeaderV2 data_page_header_v2;
1187}
1188);
1189
1190impl PageHeader {
1191 pub(crate) fn read_thrift_without_stats<'a, R>(prot: &mut R) -> Result<Self>
1195 where
1196 R: ThriftCompactInputProtocol<'a>,
1197 {
1198 let mut type_: Option<PageType> = None;
1199 let mut uncompressed_page_size: Option<i32> = None;
1200 let mut compressed_page_size: Option<i32> = None;
1201 let mut crc: Option<i32> = None;
1202 let mut data_page_header: Option<DataPageHeader> = None;
1203 let mut index_page_header: Option<IndexPageHeader> = None;
1204 let mut dictionary_page_header: Option<DictionaryPageHeader> = None;
1205 let mut data_page_header_v2: Option<DataPageHeaderV2> = None;
1206 let mut last_field_id = 0i16;
1207 loop {
1208 let field_ident = prot.read_field_begin(last_field_id)?;
1209 if field_ident.field_type == FieldType::Stop {
1210 break;
1211 }
1212 match field_ident.id {
1213 1 => {
1214 let val = PageType::read_thrift(&mut *prot)?;
1215 type_ = Some(val);
1216 }
1217 2 => {
1218 let val = i32::read_thrift(&mut *prot)?;
1219 uncompressed_page_size = Some(val);
1220 }
1221 3 => {
1222 let val = i32::read_thrift(&mut *prot)?;
1223 compressed_page_size = Some(val);
1224 }
1225 4 => {
1226 let val = i32::read_thrift(&mut *prot)?;
1227 crc = Some(val);
1228 }
1229 5 => {
1230 let val = DataPageHeader::read_thrift_without_stats(&mut *prot)?;
1231 data_page_header = Some(val);
1232 }
1233 6 => {
1234 let val = IndexPageHeader::read_thrift(&mut *prot)?;
1235 index_page_header = Some(val);
1236 }
1237 7 => {
1238 let val = DictionaryPageHeader::read_thrift(&mut *prot)?;
1239 dictionary_page_header = Some(val);
1240 }
1241 8 => {
1242 let val = DataPageHeaderV2::read_thrift_without_stats(&mut *prot)?;
1243 data_page_header_v2 = Some(val);
1244 }
1245 _ => {
1246 prot.skip(field_ident.field_type)?;
1247 }
1248 }
1249 last_field_id = field_ident.id;
1250 }
1251 let Some(type_) = type_ else {
1252 return Err(general_err!("Required field type_ is missing"));
1253 };
1254 let Some(uncompressed_page_size) = uncompressed_page_size else {
1255 return Err(general_err!(
1256 "Required field uncompressed_page_size is missing"
1257 ));
1258 };
1259 let Some(compressed_page_size) = compressed_page_size else {
1260 return Err(general_err!(
1261 "Required field compressed_page_size is missing"
1262 ));
1263 };
1264 Ok(Self {
1265 r#type: type_,
1266 uncompressed_page_size,
1267 compressed_page_size,
1268 crc,
1269 data_page_header,
1270 index_page_header,
1271 dictionary_page_header,
1272 data_page_header_v2,
1273 })
1274 }
1275}
1276
1277#[cfg(feature = "encryption")]
1281fn should_write_column_stats(column_chunk: &ColumnChunkMetaData) -> bool {
1282 column_chunk.encrypted_column_metadata.is_none()
1286}
1287
1288#[cfg(not(feature = "encryption"))]
1289fn should_write_column_stats(_column_chunk: &ColumnChunkMetaData) -> bool {
1290 true
1291}
1292
1293pub(super) fn serialize_column_meta_data<W: Write>(
1314 column_chunk: &ColumnChunkMetaData,
1315 w: &mut ThriftCompactOutputProtocol<W>,
1316) -> Result<()> {
1317 use crate::file::statistics::page_stats_to_thrift;
1318
1319 column_chunk.column_type().write_thrift_field(w, 1, 0)?;
1320 column_chunk
1321 .encodings()
1322 .collect::<Vec<_>>()
1323 .write_thrift_field(w, 2, 1)?;
1324 if w.write_path_in_schema() {
1325 let path = column_chunk.column_descr.path().parts();
1326 let path: Vec<&str> = path.iter().map(|v| v.as_str()).collect();
1327 path.write_thrift_field(w, 3, 2)?;
1328 column_chunk.compression.write_thrift_field(w, 4, 3)?;
1329 } else {
1330 column_chunk.compression.write_thrift_field(w, 4, 2)?;
1331 }
1332
1333 column_chunk.num_values.write_thrift_field(w, 5, 4)?;
1334 column_chunk
1335 .total_uncompressed_size
1336 .write_thrift_field(w, 6, 5)?;
1337 column_chunk
1338 .total_compressed_size
1339 .write_thrift_field(w, 7, 6)?;
1340 let mut last_field_id = column_chunk.data_page_offset.write_thrift_field(w, 9, 7)?;
1342 if let Some(index_page_offset) = column_chunk.index_page_offset {
1343 last_field_id = index_page_offset.write_thrift_field(w, 10, last_field_id)?;
1344 }
1345 if let Some(dictionary_page_offset) = column_chunk.dictionary_page_offset {
1346 last_field_id = dictionary_page_offset.write_thrift_field(w, 11, last_field_id)?;
1347 }
1348
1349 if should_write_column_stats(column_chunk) {
1350 let stats = page_stats_to_thrift(column_chunk.statistics());
1352 if let Some(stats) = stats {
1353 last_field_id = stats.write_thrift_field(w, 12, last_field_id)?;
1354 }
1355 if let Some(page_encoding_stats) = column_chunk.page_encoding_stats() {
1356 last_field_id = page_encoding_stats.write_thrift_field(w, 13, last_field_id)?;
1357 }
1358 if let Some(bloom_filter_offset) = column_chunk.bloom_filter_offset {
1359 last_field_id = bloom_filter_offset.write_thrift_field(w, 14, last_field_id)?;
1360 }
1361 if let Some(bloom_filter_length) = column_chunk.bloom_filter_length {
1362 last_field_id = bloom_filter_length.write_thrift_field(w, 15, last_field_id)?;
1363 }
1364
1365 let size_stats = if column_chunk.unencoded_byte_array_data_bytes.is_some()
1367 || column_chunk.repetition_level_histogram.is_some()
1368 || column_chunk.definition_level_histogram.is_some()
1369 {
1370 let repetition_level_histogram = column_chunk
1371 .repetition_level_histogram()
1372 .map(|hist| hist.clone().into_inner());
1373
1374 let definition_level_histogram = column_chunk
1375 .definition_level_histogram()
1376 .map(|hist| hist.clone().into_inner());
1377
1378 Some(SizeStatistics {
1379 unencoded_byte_array_data_bytes: column_chunk.unencoded_byte_array_data_bytes,
1380 repetition_level_histogram,
1381 definition_level_histogram,
1382 })
1383 } else {
1384 None
1385 };
1386 if let Some(size_stats) = size_stats {
1387 last_field_id = size_stats.write_thrift_field(w, 16, last_field_id)?;
1388 }
1389
1390 if let Some(geo_stats) = column_chunk.geo_statistics() {
1391 geo_stats.write_thrift_field(w, 17, last_field_id)?;
1392 }
1393 }
1394
1395 w.write_struct_end()
1396}
1397
1398pub(super) struct FileMeta<'a> {
1400 pub(super) file_metadata: &'a FileMetaData,
1401 pub(super) row_groups: &'a Vec<RowGroupMetaData>,
1402 pub(super) write_path_in_schema: bool,
1404}
1405
1406impl WriteThrift for FileMeta<'_> {
1418 const ELEMENT_TYPE: ElementType = ElementType::Struct;
1419
1420 #[cfg_attr(not(feature = "encryption"), expect(unused_assignments))]
1422 fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1423 writer.set_write_path_in_schema(self.write_path_in_schema);
1424 writer.set_write_row_group_ordinal(i16::try_from(self.row_groups.len()).is_ok());
1426
1427 self.file_metadata
1428 .version
1429 .write_thrift_field(writer, 1, 0)?;
1430
1431 let root = self.file_metadata.schema_descr().root_schema_ptr();
1434 let schema_len = num_nodes(&root)?;
1435 writer.write_field_begin(FieldType::List, 2, 1)?;
1436 writer.write_list_begin(ElementType::Struct, schema_len)?;
1437 write_schema(&root, writer)?;
1439
1440 self.file_metadata
1441 .num_rows
1442 .write_thrift_field(writer, 3, 2)?;
1443
1444 let mut last_field_id = self.row_groups.write_thrift_field(writer, 4, 3)?;
1446
1447 if let Some(kv_metadata) = self.file_metadata.key_value_metadata() {
1448 last_field_id = kv_metadata.write_thrift_field(writer, 5, last_field_id)?;
1449 }
1450 if let Some(created_by) = self.file_metadata.created_by() {
1451 last_field_id = created_by.write_thrift_field(writer, 6, last_field_id)?;
1452 }
1453 if let Some(column_orders) = self.file_metadata.column_orders() {
1454 last_field_id = column_orders.write_thrift_field(writer, 7, last_field_id)?;
1455 }
1456 #[cfg(feature = "encryption")]
1457 if let Some(algo) = self.file_metadata.encryption_algorithm.as_ref() {
1458 last_field_id = algo.write_thrift_field(writer, 8, last_field_id)?;
1459 }
1460 #[cfg(feature = "encryption")]
1461 if let Some(key) = self.file_metadata.footer_signing_key_metadata.as_ref() {
1462 key.as_slice()
1463 .write_thrift_field(writer, 9, last_field_id)?;
1464 }
1465
1466 writer.write_struct_end()
1467 }
1468}
1469
1470fn write_schema<W: Write>(
1471 schema: &TypePtr,
1472 writer: &mut ThriftCompactOutputProtocol<W>,
1473) -> Result<()> {
1474 if !schema.is_group() {
1475 return Err(general_err!("Root schema must be Group type"));
1476 }
1477 write_schema_helper(schema, writer)
1478}
1479
1480fn write_schema_helper<W: Write>(
1481 node: &TypePtr,
1482 writer: &mut ThriftCompactOutputProtocol<W>,
1483) -> Result<()> {
1484 match node.as_ref() {
1485 crate::schema::types::Type::PrimitiveType {
1486 basic_info,
1487 physical_type,
1488 type_length,
1489 scale,
1490 precision,
1491 } => {
1492 let element = SchemaElement {
1493 r#type: Some(*physical_type),
1494 type_length: if *type_length >= 0 {
1495 Some(*type_length)
1496 } else {
1497 None
1498 },
1499 repetition_type: Some(basic_info.repetition()),
1500 name: basic_info.name(),
1501 num_children: None,
1502 converted_type: match basic_info.converted_type() {
1503 ConvertedType::NONE => None,
1504 other => Some(other),
1505 },
1506 scale: if *scale >= 0 { Some(*scale) } else { None },
1507 precision: if *precision >= 0 {
1508 Some(*precision)
1509 } else {
1510 None
1511 },
1512 field_id: if basic_info.has_id() {
1513 Some(basic_info.id())
1514 } else {
1515 None
1516 },
1517 logical_type: basic_info.logical_type_ref().cloned(),
1518 };
1519 element.write_thrift(writer)
1520 }
1521 crate::schema::types::Type::GroupType { basic_info, fields } => {
1522 let repetition = if basic_info.has_repetition() {
1523 Some(basic_info.repetition())
1524 } else {
1525 None
1526 };
1527
1528 let element = SchemaElement {
1529 r#type: None,
1530 type_length: None,
1531 repetition_type: repetition,
1532 name: basic_info.name(),
1533 num_children: Some(fields.len().try_into()?),
1534 converted_type: match basic_info.converted_type() {
1535 ConvertedType::NONE => None,
1536 other => Some(other),
1537 },
1538 scale: None,
1539 precision: None,
1540 field_id: if basic_info.has_id() {
1541 Some(basic_info.id())
1542 } else {
1543 None
1544 },
1545 logical_type: basic_info.logical_type_ref().cloned(),
1546 };
1547
1548 element.write_thrift(writer)?;
1549
1550 for field in fields {
1552 write_schema_helper(field, writer)?;
1553 }
1554 Ok(())
1555 }
1556 }
1557}
1558
1559impl WriteThrift for RowGroupMetaData {
1569 const ELEMENT_TYPE: ElementType = ElementType::Struct;
1570
1571 fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1572 self.columns.write_thrift_field(writer, 1, 0)?;
1574 self.total_byte_size.write_thrift_field(writer, 2, 1)?;
1575 let mut last_field_id = self.num_rows.write_thrift_field(writer, 3, 2)?;
1576 if let Some(sorting_columns) = self.sorting_columns() {
1577 last_field_id = sorting_columns.write_thrift_field(writer, 4, last_field_id)?;
1578 }
1579 if let Some(file_offset) = self.file_offset() {
1580 last_field_id = file_offset.write_thrift_field(writer, 5, last_field_id)?;
1581 }
1582 last_field_id = self
1584 .compressed_size()
1585 .write_thrift_field(writer, 6, last_field_id)?;
1586
1587 if writer.write_row_group_ordinal()
1589 && let Some(ordinal) = self.ordinal()
1590 && let Ok(ordinal) = i16::try_from(ordinal)
1591 {
1592 ordinal.write_thrift_field(writer, 7, last_field_id)?;
1593 }
1594 writer.write_struct_end()
1595 }
1596}
1597
1598impl WriteThrift for ColumnChunkMetaData {
1610 const ELEMENT_TYPE: ElementType = ElementType::Struct;
1611
1612 #[cfg_attr(not(feature = "encryption"), expect(unused_assignments))]
1613 fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1614 let mut last_field_id = 0i16;
1615 if let Some(file_path) = self.file_path() {
1616 last_field_id = file_path.write_thrift_field(writer, 1, last_field_id)?;
1617 }
1618 last_field_id = self
1619 .file_offset()
1620 .write_thrift_field(writer, 2, last_field_id)?;
1621
1622 #[cfg(feature = "encryption")]
1623 let write_meta_data =
1624 self.encrypted_column_metadata.is_none() || self.plaintext_footer_mode;
1625 #[cfg(not(feature = "encryption"))]
1626 let write_meta_data = true;
1627
1628 if write_meta_data {
1634 writer.write_field_begin(FieldType::Struct, 3, last_field_id)?;
1635 serialize_column_meta_data(self, writer)?;
1636 last_field_id = 3;
1637 }
1638
1639 if let Some(offset_idx_off) = self.offset_index_offset() {
1640 last_field_id = offset_idx_off.write_thrift_field(writer, 4, last_field_id)?;
1641 }
1642 if let Some(offset_idx_len) = self.offset_index_length() {
1643 last_field_id = offset_idx_len.write_thrift_field(writer, 5, last_field_id)?;
1644 }
1645 if let Some(column_idx_off) = self.column_index_offset() {
1646 last_field_id = column_idx_off.write_thrift_field(writer, 6, last_field_id)?;
1647 }
1648 if let Some(column_idx_len) = self.column_index_length() {
1649 last_field_id = column_idx_len.write_thrift_field(writer, 7, last_field_id)?;
1650 }
1651 #[cfg(feature = "encryption")]
1652 {
1653 if let Some(crypto_metadata) = self.crypto_metadata() {
1654 last_field_id = crypto_metadata.write_thrift_field(writer, 8, last_field_id)?;
1655 }
1656 if let Some(encrypted_meta) = self.encrypted_column_metadata.as_ref() {
1657 encrypted_meta
1658 .as_slice()
1659 .write_thrift_field(writer, 9, last_field_id)?;
1660 }
1661 }
1662
1663 writer.write_struct_end()
1664 }
1665}
1666
1667impl WriteThrift for crate::geospatial::statistics::GeospatialStatistics {
1672 const ELEMENT_TYPE: ElementType = ElementType::Struct;
1673
1674 fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1675 let mut last_field_id = 0i16;
1676 if let Some(bbox) = self.bounding_box() {
1677 last_field_id = bbox.write_thrift_field(writer, 1, last_field_id)?;
1678 }
1679 if let Some(geo_types) = self.geospatial_types() {
1680 geo_types.write_thrift_field(writer, 2, last_field_id)?;
1681 }
1682
1683 writer.write_struct_end()
1684 }
1685}
1686
1687use crate::geospatial::statistics::GeospatialStatistics as RustGeospatialStatistics;
1689write_thrift_field!(RustGeospatialStatistics, FieldType::Struct);
1690
1691impl WriteThrift for crate::geospatial::bounding_box::BoundingBox {
1702 const ELEMENT_TYPE: ElementType = ElementType::Struct;
1703
1704 fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1705 self.get_xmin().write_thrift_field(writer, 1, 0)?;
1706 self.get_xmax().write_thrift_field(writer, 2, 1)?;
1707 self.get_ymin().write_thrift_field(writer, 3, 2)?;
1708 let mut last_field_id = self.get_ymax().write_thrift_field(writer, 4, 3)?;
1709
1710 if let Some(zmin) = self.get_zmin() {
1711 last_field_id = zmin.write_thrift_field(writer, 5, last_field_id)?;
1712 }
1713 if let Some(zmax) = self.get_zmax() {
1714 last_field_id = zmax.write_thrift_field(writer, 6, last_field_id)?;
1715 }
1716 if let Some(mmin) = self.get_mmin() {
1717 last_field_id = mmin.write_thrift_field(writer, 7, last_field_id)?;
1718 }
1719 if let Some(mmax) = self.get_mmax() {
1720 mmax.write_thrift_field(writer, 8, last_field_id)?;
1721 }
1722
1723 writer.write_struct_end()
1724 }
1725}
1726
1727use crate::geospatial::bounding_box::BoundingBox as RustBoundingBox;
1729write_thrift_field!(RustBoundingBox, FieldType::Struct);
1730
1731#[cfg(test)]
1732pub(crate) mod tests {
1733 use crate::basic::{Encoding, PageType, Type as PhysicalType};
1734 use crate::errors::Result;
1735 use crate::file::metadata::thrift::{
1736 BoundingBox, DataPageHeaderV2, DictionaryPageHeader, PageHeader, SchemaElement,
1737 write_schema,
1738 };
1739 use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaDataOptions, RowGroupMetaData};
1740 use crate::parquet_thrift::tests::test_roundtrip;
1741 use crate::parquet_thrift::{
1742 ElementType, ThriftCompactOutputProtocol, ThriftSliceInputProtocol, WriteThrift,
1743 read_thrift_vec,
1744 };
1745 use crate::schema::types::{
1746 ColumnDescriptor, ColumnPath, SchemaDescriptor, TypePtr, num_nodes,
1747 parquet_schema_from_array,
1748 };
1749 use std::sync::Arc;
1750
1751 pub(crate) fn read_row_group(
1753 buf: &[u8],
1754 schema_descr: Arc<SchemaDescriptor>,
1755 ) -> Result<RowGroupMetaData> {
1756 let mut reader = ThriftSliceInputProtocol::new(buf);
1757 crate::file::metadata::thrift::read_row_group(&mut reader, &schema_descr, None)
1758 }
1759
1760 pub(crate) fn read_column_chunk(
1761 buf: &[u8],
1762 column_descr: Arc<ColumnDescriptor>,
1763 ) -> Result<ColumnChunkMetaData> {
1764 read_column_chunk_with_options(buf, column_descr, None)
1765 }
1766
1767 pub(crate) fn read_column_chunk_with_options(
1768 buf: &[u8],
1769 column_descr: Arc<ColumnDescriptor>,
1770 options: Option<&ParquetMetaDataOptions>,
1771 ) -> Result<ColumnChunkMetaData> {
1772 let mut reader = ThriftSliceInputProtocol::new(buf);
1773 crate::file::metadata::thrift::read_column_chunk(&mut reader, &column_descr, 0, options)
1774 }
1775
1776 pub(crate) fn roundtrip_schema(schema: TypePtr) -> Result<TypePtr> {
1777 let num_nodes = num_nodes(&schema)?;
1778 let mut buf = Vec::new();
1779 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1780
1781 writer.write_list_begin(ElementType::Struct, num_nodes)?;
1783
1784 write_schema(&schema, &mut writer)?;
1786
1787 let mut prot = ThriftSliceInputProtocol::new(&buf);
1788 let se: Vec<SchemaElement> = read_thrift_vec(&mut prot)?;
1789 parquet_schema_from_array(se)
1790 }
1791
1792 pub(crate) fn schema_to_buf(schema: &TypePtr) -> Result<Vec<u8>> {
1793 let num_nodes = num_nodes(schema)?;
1794 let mut buf = Vec::new();
1795 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1796
1797 writer.write_list_begin(ElementType::Struct, num_nodes)?;
1799
1800 write_schema(schema, &mut writer)?;
1802 Ok(buf)
1803 }
1804
1805 pub(crate) fn buf_to_schema_list(buf: &mut Vec<u8>) -> Result<Vec<SchemaElement<'_>>> {
1806 let mut prot = ThriftSliceInputProtocol::new(buf.as_mut_slice());
1807 read_thrift_vec(&mut prot)
1808 }
1809
1810 fn thrift_bytes<T: WriteThrift>(value: &T) -> Vec<u8> {
1811 let mut buf = Vec::new();
1812 let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1813 value.write_thrift(&mut writer).unwrap();
1814 buf
1815 }
1816
1817 fn change_false_bool_field_to_i32(buf: &mut [u8]) {
1818 let pos = buf
1819 .iter()
1820 .rposition(|byte| *byte == 0x12)
1821 .expect("expected BOOL_FALSE field header byte");
1822 buf[pos] = 0x15;
1823 }
1824
1825 fn assert_malformed_bool_error(err: crate::errors::ParquetError) {
1826 let msg = err.to_string();
1827 assert!(
1828 msg.contains("Unexpected struct field type"),
1829 "unexpected error message: {msg}"
1830 );
1831 }
1832
1833 #[test]
1834 fn test_bounding_box_roundtrip() {
1835 test_roundtrip(BoundingBox {
1836 xmin: 0.1.into(),
1837 xmax: 10.3.into(),
1838 ymin: 0.001.into(),
1839 ymax: 128.5.into(),
1840 zmin: None,
1841 zmax: None,
1842 mmin: None,
1843 mmax: None,
1844 });
1845
1846 test_roundtrip(BoundingBox {
1847 xmin: 0.1.into(),
1848 xmax: 10.3.into(),
1849 ymin: 0.001.into(),
1850 ymax: 128.5.into(),
1851 zmin: Some(11.0.into()),
1852 zmax: Some(1300.0.into()),
1853 mmin: None,
1854 mmax: None,
1855 });
1856
1857 test_roundtrip(BoundingBox {
1858 xmin: 0.1.into(),
1859 xmax: 10.3.into(),
1860 ymin: 0.001.into(),
1861 ymax: 128.5.into(),
1862 zmin: Some(11.0.into()),
1863 zmax: Some(1300.0.into()),
1864 mmin: Some(3.7.into()),
1865 mmax: Some(42.0.into()),
1866 });
1867 }
1868
1869 #[test]
1870 fn test_convert_stats_preserves_missing_null_count() {
1871 let primitive =
1872 crate::schema::types::Type::primitive_type_builder("col", PhysicalType::INT32)
1873 .build()
1874 .unwrap();
1875 let column_descr = Arc::new(ColumnDescriptor::new(
1876 Arc::new(primitive),
1877 0,
1878 0,
1879 ColumnPath::new(vec![]),
1880 ));
1881
1882 let none_null_count = super::Statistics {
1883 max: None,
1884 min: None,
1885 null_count: None,
1886 distinct_count: None,
1887 max_value: None,
1888 min_value: None,
1889 is_max_value_exact: None,
1890 is_min_value_exact: None,
1891 nan_count: None,
1892 };
1893 let decoded_none = super::convert_stats(&column_descr, Some(none_null_count))
1894 .unwrap()
1895 .unwrap();
1896 assert_eq!(decoded_none.null_count_opt(), None);
1897
1898 let zero_null_count = super::Statistics {
1899 max: None,
1900 min: None,
1901 null_count: Some(0),
1902 distinct_count: None,
1903 max_value: None,
1904 min_value: None,
1905 is_max_value_exact: None,
1906 is_min_value_exact: None,
1907 nan_count: None,
1908 };
1909 let decoded_zero = super::convert_stats(&column_descr, Some(zero_null_count))
1910 .unwrap()
1911 .unwrap();
1912 assert_eq!(decoded_zero.null_count_opt(), Some(0));
1913 }
1914
1915 #[test]
1916 fn test_convert_stats_returns_error_for_overlong_int96_statistics() {
1917 let primitive =
1918 crate::schema::types::Type::primitive_type_builder("col", PhysicalType::INT96)
1919 .build()
1920 .unwrap();
1921 let column_descr = Arc::new(ColumnDescriptor::new(
1922 Arc::new(primitive),
1923 0,
1924 0,
1925 ColumnPath::new(vec![]),
1926 ));
1927 let invalid = (0..13).collect::<Vec<_>>();
1928
1929 let make_stats = |min, max| super::Statistics {
1930 max,
1931 min,
1932 null_count: Some(0),
1933 distinct_count: None,
1934 max_value: None,
1935 min_value: None,
1936 is_max_value_exact: None,
1937 is_min_value_exact: None,
1938 nan_count: None,
1939 };
1940
1941 let err = super::convert_stats(&column_descr, Some(make_stats(Some(&invalid), None)))
1942 .unwrap_err();
1943 assert_eq!(
1944 err.to_string(),
1945 "Parquet error: Incorrect Int96 min statistics"
1946 );
1947
1948 let err = super::convert_stats(&column_descr, Some(make_stats(None, Some(&invalid))))
1949 .unwrap_err();
1950 assert_eq!(
1951 err.to_string(),
1952 "Parquet error: Incorrect Int96 max statistics"
1953 );
1954 }
1955
1956 #[test]
1957 fn malformed_bool_field_returns_error_not_panic() {
1958 let page_header = PageHeader {
1959 r#type: PageType::DICTIONARY_PAGE,
1960 uncompressed_page_size: 1,
1961 compressed_page_size: 1,
1962 crc: None,
1963 data_page_header: None,
1964 index_page_header: None,
1965 dictionary_page_header: Some(DictionaryPageHeader {
1966 num_values: 1,
1967 encoding: Encoding::PLAIN,
1968 is_sorted: Some(false),
1969 }),
1970 data_page_header_v2: None,
1971 };
1972
1973 let mut buf = thrift_bytes(&page_header);
1974 change_false_bool_field_to_i32(&mut buf);
1975
1976 let mut prot = ThriftSliceInputProtocol::new(&buf);
1977 let err = PageHeader::read_thrift_without_stats(&mut prot)
1978 .expect_err("malformed bool field should return an error");
1979 assert_malformed_bool_error(err);
1980 }
1981
1982 #[test]
1983 fn malformed_data_page_v2_bool_field_returns_error_not_panic() {
1984 let data_page_header_v2 = DataPageHeaderV2 {
1985 num_values: 1,
1986 num_nulls: 0,
1987 num_rows: 1,
1988 encoding: Encoding::PLAIN,
1989 definition_levels_byte_length: 0,
1990 repetition_levels_byte_length: 0,
1991 is_compressed: Some(false),
1992 statistics: None,
1993 };
1994
1995 let mut buf = thrift_bytes(&data_page_header_v2);
1996 change_false_bool_field_to_i32(&mut buf);
1997
1998 let mut prot = ThriftSliceInputProtocol::new(&buf);
1999 let err = DataPageHeaderV2::read_thrift_without_stats(&mut prot)
2000 .expect_err("malformed bool field should return an error");
2001 assert_malformed_bool_error(err);
2002 }
2003
2004 fn roundtrip_rg_ordinals(ordinals: &[Option<i32>]) -> Vec<Option<i32>> {
2008 use crate::file::metadata::ParquetMetaDataWriter;
2009 use crate::file::metadata::{FileMetaData, ParquetMetaData, ParquetMetaDataReader};
2010 use crate::schema::types::Type as SchemaType;
2011
2012 let field = SchemaType::primitive_type_builder("c", PhysicalType::INT32)
2013 .build()
2014 .unwrap();
2015 let schema = SchemaType::group_type_builder("schema")
2016 .with_fields(vec![Arc::new(field)])
2017 .build()
2018 .unwrap();
2019 let schema_descr = Arc::new(SchemaDescriptor::new(Arc::new(schema)));
2020
2021 let row_groups = ordinals
2022 .iter()
2023 .map(|ordinal| {
2024 let columns = schema_descr
2025 .columns()
2026 .iter()
2027 .map(|col| ColumnChunkMetaData::builder(col.clone()).build().unwrap())
2028 .collect();
2029 let mut builder =
2030 crate::file::metadata::RowGroupMetaData::builder(schema_descr.clone())
2031 .set_num_rows(10)
2032 .set_total_byte_size(100)
2033 .set_column_metadata(columns);
2034 if let Some(ordinal) = ordinal {
2035 builder = builder.set_ordinal(*ordinal);
2036 }
2037 builder.build().unwrap()
2038 })
2039 .collect();
2040
2041 let file_metadata = FileMetaData::new(
2042 1,
2043 10 * ordinals.len() as i64,
2044 None,
2045 None,
2046 schema_descr,
2047 None,
2048 );
2049 let metadata = ParquetMetaData::new(file_metadata, row_groups);
2050
2051 let mut buffer = Vec::new();
2052 ParquetMetaDataWriter::new(&mut buffer, &metadata)
2053 .finish()
2054 .unwrap();
2055 let decoded = ParquetMetaDataReader::decode_metadata(&buffer[..buffer.len() - 8]).unwrap();
2057 decoded.row_groups().iter().map(|rg| rg.ordinal()).collect()
2058 }
2059
2060 #[test]
2063 fn ordinals_all_present_are_honored() {
2064 assert_eq!(
2065 roundtrip_rg_ordinals(&[Some(5), Some(1), Some(3)]),
2066 vec![Some(5), Some(1), Some(3)],
2067 );
2068 }
2069
2070 #[test]
2073 fn ordinals_none_present_are_sequentially_filled() {
2074 assert_eq!(
2075 roundtrip_rg_ordinals(&[None, None, None]),
2076 vec![Some(0), Some(1), Some(2)],
2077 );
2078 }
2079
2080 #[test]
2085 fn ordinals_mixed_decode_succeeds_untouched() {
2086 assert_eq!(
2087 roundtrip_rg_ordinals(&[Some(0), None, Some(2)]),
2088 vec![Some(0), None, Some(2)],
2089 );
2090 assert_eq!(
2092 roundtrip_rg_ordinals(&[None, Some(1), Some(2)]),
2093 vec![None, Some(1), Some(2)],
2094 );
2095 }
2096
2097 #[test]
2099 fn row_group_count_exceeding_remaining_input_is_rejected() {
2100 use crate::file::metadata::thrift::parquet_metadata_from_bytes;
2101 use crate::parquet_thrift::{FieldType, WriteThriftField};
2102 use crate::schema::types::Type as SchemaType;
2103
2104 let field = SchemaType::primitive_type_builder("c", PhysicalType::INT32)
2106 .build()
2107 .unwrap();
2108 let schema = SchemaType::group_type_builder("schema")
2109 .with_fields(vec![Arc::new(field)])
2110 .build()
2111 .unwrap();
2112 let root: TypePtr = Arc::new(schema);
2113 let schema_len = num_nodes(&root).unwrap();
2114
2115 let mut buf = Vec::new();
2116 let mut w = ThriftCompactOutputProtocol::new(&mut buf);
2117 1i32.write_thrift_field(&mut w, 1, 0).unwrap();
2119 w.write_field_begin(FieldType::List, 2, 1).unwrap();
2121 w.write_list_begin(ElementType::Struct, schema_len).unwrap();
2122 write_schema(&root, &mut w).unwrap();
2123 0i64.write_thrift_field(&mut w, 3, 2).unwrap();
2125 w.write_field_begin(FieldType::List, 4, 3).unwrap();
2128 w.write_list_begin(ElementType::Struct, 1024).unwrap();
2129 w.write_struct_end().unwrap();
2130
2131 let err = parquet_metadata_from_bytes(&buf, None)
2132 .expect_err("over-declared row-group count must be rejected")
2133 .to_string();
2134 assert!(
2135 err.contains("exceeds remaining input length"),
2136 "expected a remaining-input bound error, got: {err}"
2137 );
2138 }
2139}