1use std::collections::HashMap;
21use std::sync::Arc;
22
23use arrow_schema::SchemaRef as ArrowSchemaRef;
24use bytes::Bytes;
25use futures::future::BoxFuture;
26use itertools::Itertools;
27use parquet::arrow::AsyncArrowWriter;
28use parquet::arrow::async_reader::AsyncFileReader;
29use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter;
30use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
31use parquet::encryption::encrypt::FileEncryptionProperties;
32use parquet::file::metadata::ParquetMetaData;
33use parquet::file::properties::{CdcOptions, WriterProperties};
34use parquet::file::statistics::Statistics;
35
36use super::{FileWriter, FileWriterBuilder};
37use crate::arrow::{
38 ArrowFileReader, DEFAULT_MAP_FIELD_NAME, FieldMatchMode, NanValueCountVisitor,
39 get_parquet_stat_max_as_datum, get_parquet_stat_min_as_datum,
40};
41use crate::compression::CompressionCodec;
42use crate::encryption::{EncryptionManager, StandardKeyMetadata};
43use crate::error::invalid_data;
44use crate::io::{FileIO, FileWrite, OutputFile};
45use crate::spec::{
46 DataContentType, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, MapType,
47 NestedFieldRef, PartitionSpec, PrimitiveType, Schema, SchemaRef, SchemaVisitor, Struct,
48 StructType, TableMetadata, TableProperties, Type, VariantType, visit_schema,
49};
50use crate::transform::create_transform_function;
51use crate::writer::{CurrentFileStatus, DataFile};
52use crate::{Error, ErrorKind, Result};
53
54#[derive(Clone, Debug)]
56pub struct ParquetWriterBuilder {
57 props: WriterProperties,
58 schema: SchemaRef,
59 match_mode: FieldMatchMode,
60 encryption_manager: Option<Arc<EncryptionManager>>,
61}
62
63impl ParquetWriterBuilder {
64 pub fn new(props: WriterProperties, schema: SchemaRef) -> Self {
71 Self::new_with_match_mode(props, schema, FieldMatchMode::Id)
72 }
73
74 pub fn new_with_match_mode(
76 props: WriterProperties,
77 schema: SchemaRef,
78 match_mode: FieldMatchMode,
79 ) -> Self {
80 Self {
81 props,
82 schema,
83 match_mode,
84 encryption_manager: None,
85 }
86 }
87
88 pub fn from_table_properties(
92 table_props: &TableProperties<'_>,
93 schema: SchemaRef,
94 ) -> Result<Self> {
95 let cdc = if table_props.cdc_enabled()? {
96 Some(CdcOptions {
97 min_chunk_size: table_props.cdc_min_chunk_size()?,
98 max_chunk_size: table_props.cdc_max_chunk_size()?,
99 norm_level: table_props.cdc_norm_level()?,
100 })
101 } else {
102 None
103 };
104 let compression = parquet_compression(table_props.parquet_compression_codec()?)?;
105 let props = WriterProperties::builder()
106 .set_content_defined_chunking(cdc)
107 .set_compression(compression)
108 .set_max_row_group_bytes(Some(table_props.parquet_row_group_size_bytes()?))
109 .set_data_page_size_limit(table_props.parquet_page_size_bytes()?)
110 .set_data_page_row_count_limit(table_props.parquet_page_row_limit()?)
111 .set_dictionary_page_size_limit(table_props.parquet_dict_size_bytes()?)
112 .build();
113 Ok(Self::new_with_match_mode(props, schema, FieldMatchMode::Id))
114 }
115
116 pub fn with_match_mode(mut self, match_mode: FieldMatchMode) -> Self {
121 self.match_mode = match_mode;
122 self
123 }
124
125 pub fn with_encryption_manager(mut self, encryption_manager: Arc<EncryptionManager>) -> Self {
127 self.encryption_manager = Some(encryption_manager);
128 self
129 }
130}
131
132fn parquet_compression(codec: CompressionCodec) -> Result<Compression> {
133 let compression = match codec {
134 CompressionCodec::None => Compression::UNCOMPRESSED,
135 CompressionCodec::Snappy => Compression::SNAPPY,
136 CompressionCodec::Lzo => Compression::LZO,
137 CompressionCodec::Lz4 => Compression::LZ4,
138 CompressionCodec::Lz4Raw => Compression::LZ4_RAW,
139 CompressionCodec::Zstd(level) => {
140 let level =
141 ZstdLevel::try_new(level as i32).map_err(|e| invalid_level_error("zstd", e))?;
142 Compression::ZSTD(level)
143 }
144 CompressionCodec::Gzip(level) => {
145 let level =
146 GzipLevel::try_new(level as u32).map_err(|e| invalid_level_error("gzip", e))?;
147 Compression::GZIP(level)
148 }
149 CompressionCodec::Brotli(level) => {
150 let level =
151 BrotliLevel::try_new(level as u32).map_err(|e| invalid_level_error("brotli", e))?;
152 Compression::BROTLI(level)
153 }
154 };
155 Ok(compression)
156}
157
158fn invalid_level_error(codec: &str, source: impl Into<anyhow::Error>) -> Error {
159 invalid_data!("Invalid {codec} compression level").with_source(source)
160}
161
162impl FileWriterBuilder for ParquetWriterBuilder {
163 type R = ParquetWriter;
164
165 async fn build(&self, output_file: OutputFile) -> Result<Self::R> {
166 let key_metadata = self
167 .encryption_manager
168 .as_ref()
169 .map(|em| em.generate_key_metadata());
170 let writer_properties =
171 resolve_writer_properties(self.props.clone(), key_metadata.as_ref())?;
172 Ok(ParquetWriter {
173 schema: self.schema.clone(),
174 inner_writer: None,
175 writer_properties,
176 current_row_num: 0,
177 output_file,
178 nan_value_count_visitor: NanValueCountVisitor::new_with_match_mode(self.match_mode),
179 key_metadata,
180 })
181 }
182}
183
184struct IndexByParquetPathName {
186 name_to_id: HashMap<String, i32>,
187
188 field_names: Vec<String>,
189
190 field_id: i32,
191}
192
193impl IndexByParquetPathName {
194 pub fn new() -> Self {
196 Self {
197 name_to_id: HashMap::new(),
198 field_names: Vec::new(),
199 field_id: 0,
200 }
201 }
202
203 pub fn get(&self, name: &str) -> Option<&i32> {
205 self.name_to_id.get(name)
206 }
207}
208
209impl Default for IndexByParquetPathName {
210 fn default() -> Self {
211 Self::new()
212 }
213}
214
215impl SchemaVisitor for IndexByParquetPathName {
216 type T = ();
217
218 fn before_struct_field(&mut self, field: &NestedFieldRef) -> Result<()> {
219 self.field_names.push(field.name.to_string());
220 self.field_id = field.id;
221 Ok(())
222 }
223
224 fn after_struct_field(&mut self, _field: &NestedFieldRef) -> Result<()> {
225 self.field_names.pop();
226 Ok(())
227 }
228
229 fn before_list_element(&mut self, field: &NestedFieldRef) -> Result<()> {
230 self.field_names.push(format!("list.{}", field.name));
231 self.field_id = field.id;
232 Ok(())
233 }
234
235 fn after_list_element(&mut self, _field: &NestedFieldRef) -> Result<()> {
236 self.field_names.pop();
237 Ok(())
238 }
239
240 fn before_map_key(&mut self, field: &NestedFieldRef) -> Result<()> {
241 self.field_names
242 .push(format!("{DEFAULT_MAP_FIELD_NAME}.key"));
243 self.field_id = field.id;
244 Ok(())
245 }
246
247 fn after_map_key(&mut self, _field: &NestedFieldRef) -> Result<()> {
248 self.field_names.pop();
249 Ok(())
250 }
251
252 fn before_map_value(&mut self, field: &NestedFieldRef) -> Result<()> {
253 self.field_names
254 .push(format!("{DEFAULT_MAP_FIELD_NAME}.value"));
255 self.field_id = field.id;
256 Ok(())
257 }
258
259 fn after_map_value(&mut self, _field: &NestedFieldRef) -> Result<()> {
260 self.field_names.pop();
261 Ok(())
262 }
263
264 fn schema(&mut self, _schema: &Schema, _value: Self::T) -> Result<Self::T> {
265 Ok(())
266 }
267
268 fn field(&mut self, _field: &NestedFieldRef, _value: Self::T) -> Result<Self::T> {
269 Ok(())
270 }
271
272 fn r#struct(&mut self, _struct: &StructType, _results: Vec<Self::T>) -> Result<Self::T> {
273 Ok(())
274 }
275
276 fn list(&mut self, _list: &ListType, _value: Self::T) -> Result<Self::T> {
277 Ok(())
278 }
279
280 fn map(&mut self, _map: &MapType, _key_value: Self::T, _value: Self::T) -> Result<Self::T> {
281 Ok(())
282 }
283
284 fn primitive(&mut self, _p: &PrimitiveType) -> Result<Self::T> {
285 let full_name = self.field_names.iter().map(String::as_str).join(".");
286 let field_id = self.field_id;
287 if let Some(existing_field_id) = self.name_to_id.get(full_name.as_str()) {
288 return Err(invalid_data!(
289 "Invalid schema: multiple fields for name {full_name}: {field_id} and {existing_field_id}"
290 ));
291 } else {
292 self.name_to_id.insert(full_name, field_id);
293 }
294
295 Ok(())
296 }
297
298 fn variant(&mut self, _v: &VariantType) -> Result<Self::T> {
299 Err(Error::new(
300 ErrorKind::FeatureUnsupported,
301 "Writing variant columns to Parquet is not supported yet",
302 ))
303 }
304}
305
306pub struct ParquetWriter {
308 schema: SchemaRef,
309 output_file: OutputFile,
310 inner_writer: Option<AsyncArrowWriter<AsyncFileWriter>>,
311 writer_properties: WriterProperties,
312 current_row_num: usize,
313 nan_value_count_visitor: NanValueCountVisitor,
314 key_metadata: Option<StandardKeyMetadata>,
315}
316
317struct MinMaxColAggregator {
319 lower_bounds: HashMap<i32, Datum>,
320 upper_bounds: HashMap<i32, Datum>,
321 schema: SchemaRef,
322}
323
324impl MinMaxColAggregator {
325 fn new(schema: SchemaRef) -> Self {
327 Self {
328 lower_bounds: HashMap::new(),
329 upper_bounds: HashMap::new(),
330 schema,
331 }
332 }
333
334 fn update_state_min(&mut self, field_id: i32, datum: Datum) {
335 self.lower_bounds
336 .entry(field_id)
337 .and_modify(|e| {
338 if *e > datum {
339 *e = datum.clone()
340 }
341 })
342 .or_insert(datum);
343 }
344
345 fn update_state_max(&mut self, field_id: i32, datum: Datum) {
346 self.upper_bounds
347 .entry(field_id)
348 .and_modify(|e| {
349 if *e < datum {
350 *e = datum.clone()
351 }
352 })
353 .or_insert(datum);
354 }
355
356 fn update(&mut self, field_id: i32, value: &Statistics) -> Result<()> {
358 let Some(ty) = self
359 .schema
360 .field_by_id(field_id)
361 .map(|f| f.field_type.as_ref())
362 else {
363 return Ok(());
366 };
367 let Type::Primitive(ty) = ty.clone() else {
368 return Err(Error::new(
369 ErrorKind::Unexpected,
370 format!("Composed type {ty} is not supported for min max aggregation."),
371 ));
372 };
373
374 if value.min_is_exact() {
375 let Some(min_datum) = get_parquet_stat_min_as_datum(&ty, value)? else {
376 return Err(Error::new(
377 ErrorKind::Unexpected,
378 format!("Statistics {value} is not match with field type {ty}."),
379 ));
380 };
381
382 if !min_datum.is_nan() {
384 self.update_state_min(field_id, min_datum);
385 }
386 }
387
388 if value.max_is_exact() {
389 let Some(max_datum) = get_parquet_stat_max_as_datum(&ty, value)? else {
390 return Err(Error::new(
391 ErrorKind::Unexpected,
392 format!("Statistics {value} is not match with field type {ty}."),
393 ));
394 };
395
396 if !max_datum.is_nan() {
397 self.update_state_max(field_id, max_datum);
398 }
399 }
400
401 Ok(())
402 }
403
404 fn produce(self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) {
406 (self.lower_bounds, self.upper_bounds)
407 }
408}
409
410impl ParquetWriter {
411 #[allow(dead_code)]
413 pub(crate) async fn parquet_files_to_data_files(
414 file_io: &FileIO,
415 file_paths: Vec<String>,
416 table_metadata: &TableMetadata,
417 ) -> Result<Vec<DataFile>> {
418 let mut data_files: Vec<DataFile> = Vec::new();
420
421 for file_path in file_paths {
422 let input_file = file_io.new_input(&file_path)?;
423 let file_metadata = input_file.metadata().await?;
424 let file_size_in_bytes = file_metadata.size as usize;
425 let reader = input_file.reader().await?;
426
427 let mut parquet_reader = ArrowFileReader::new(file_metadata, reader);
428 let parquet_metadata = parquet_reader
429 .get_metadata(None)
430 .await
431 .map_err(|err| invalid_data!("Error reading Parquet metadata: {err}"))?;
432 let mut builder = ParquetWriter::parquet_to_data_file_builder(
433 table_metadata.current_schema().clone(),
434 parquet_metadata,
435 file_size_in_bytes,
436 file_path,
437 HashMap::new(),
439 None,
440 )?;
441 builder.partition_spec_id(table_metadata.default_partition_spec_id());
442 let data_file = builder.build().unwrap();
443 data_files.push(data_file);
444 }
445
446 Ok(data_files)
447 }
448
449 pub(crate) fn parquet_to_data_file_builder(
451 schema: SchemaRef,
452 metadata: Arc<ParquetMetaData>,
453 written_size: usize,
454 file_path: String,
455 nan_value_counts: HashMap<i32, u64>,
456 key_metadata: Option<StandardKeyMetadata>,
457 ) -> Result<DataFileBuilder> {
458 let index_by_parquet_path = {
459 let mut visitor = IndexByParquetPathName::new();
460 visit_schema(&schema, &mut visitor)?;
461 visitor
462 };
463
464 let (column_sizes, value_counts, null_value_counts, (lower_bounds, upper_bounds)) = {
465 let mut per_col_size: HashMap<i32, u64> = HashMap::new();
466 let mut per_col_val_num: HashMap<i32, u64> = HashMap::new();
467 let mut per_col_null_val_num: HashMap<i32, u64> = HashMap::new();
468 let mut min_max_agg = MinMaxColAggregator::new(schema);
469
470 for row_group in metadata.row_groups() {
471 for column_chunk_metadata in row_group.columns() {
472 let parquet_path = column_chunk_metadata.column_descr().path().string();
473
474 let Some(&field_id) = index_by_parquet_path.get(&parquet_path) else {
475 continue;
476 };
477
478 *per_col_size.entry(field_id).or_insert(0) +=
479 column_chunk_metadata.compressed_size() as u64;
480 *per_col_val_num.entry(field_id).or_insert(0) +=
481 column_chunk_metadata.num_values() as u64;
482
483 if let Some(statistics) = column_chunk_metadata.statistics() {
484 if let Some(null_count) = statistics.null_count_opt() {
485 *per_col_null_val_num.entry(field_id).or_insert(0) += null_count;
486 }
487
488 min_max_agg.update(field_id, statistics)?;
489 }
490 }
491 }
492 (
493 per_col_size,
494 per_col_val_num,
495 per_col_null_val_num,
496 min_max_agg.produce(),
497 )
498 };
499
500 let key_metadata = match key_metadata {
501 Some(m) => Some(m.encode()?.into_vec()),
502 None => None,
503 };
504
505 let mut builder = DataFileBuilder::default();
506 builder
507 .content(DataContentType::Data)
508 .file_path(file_path)
509 .file_format(DataFileFormat::Parquet)
510 .partition(Struct::empty())
511 .record_count(metadata.file_metadata().num_rows() as u64)
512 .file_size_in_bytes(written_size as u64)
513 .column_sizes(column_sizes)
514 .value_counts(value_counts)
515 .null_value_counts(null_value_counts)
516 .nan_value_counts(nan_value_counts)
517 .lower_bounds(lower_bounds)
520 .upper_bounds(upper_bounds)
521 .key_metadata(key_metadata)
522 .split_offsets(Some(
523 metadata
524 .row_groups()
525 .iter()
526 .filter_map(|group| group.file_offset())
527 .collect(),
528 ));
529
530 Ok(builder)
531 }
532
533 #[allow(dead_code)]
534 fn partition_value_from_bounds(
535 table_spec: Arc<PartitionSpec>,
536 lower_bounds: &HashMap<i32, Datum>,
537 upper_bounds: &HashMap<i32, Datum>,
538 ) -> Result<Struct> {
539 let mut partition_literals: Vec<Option<Literal>> = Vec::new();
540
541 for field in table_spec.fields() {
542 if let (Some(lower), Some(upper)) = (
543 lower_bounds.get(&field.source_id),
544 upper_bounds.get(&field.source_id),
545 ) {
546 if !field.transform.preserves_order() {
547 return Err(invalid_data!(
548 "cannot infer partition value for non linear partition field (needs to preserve order): {} with transform {}",
549 field.name,
550 field.transform
551 ));
552 }
553
554 if lower != upper {
555 return Err(invalid_data!(
556 "multiple partition values for field {}: lower: {:?}, upper: {:?}",
557 field.name,
558 lower,
559 upper
560 ));
561 }
562
563 let transform_fn = create_transform_function(&field.transform)?;
564 let transform_literal =
565 Literal::from(transform_fn.transform_literal_result(lower)?);
566
567 partition_literals.push(Some(transform_literal));
568 } else {
569 partition_literals.push(None);
570 }
571 }
572
573 let partition_struct = Struct::from_iter(partition_literals);
574
575 Ok(partition_struct)
576 }
577}
578
579fn resolve_writer_properties(
580 writer_properties: WriterProperties,
581 key_metadata: Option<&StandardKeyMetadata>,
582) -> Result<WriterProperties> {
583 let Some(key_metadata) = key_metadata else {
584 return Ok(writer_properties);
585 };
586
587 if writer_properties.file_encryption_properties().is_some() {
588 return Err(Error::new(
589 ErrorKind::Unexpected,
590 "Parquet writer properties already have file encryption properties set",
591 ));
592 }
593
594 let mut builder =
595 FileEncryptionProperties::builder(key_metadata.encryption_key().as_bytes().to_vec());
596 if let Some(aad) = key_metadata.aad_prefix() {
597 builder = builder.with_aad_prefix(aad.to_vec());
598 }
599 let file_encryption_properties = builder.build().map_err(|e| {
600 Error::new(
601 ErrorKind::Unexpected,
602 "Failed to build parquet file encryption properties",
603 )
604 .with_source(e)
605 })?;
606
607 Ok(writer_properties
608 .into_builder()
609 .with_file_encryption_properties(file_encryption_properties)
610 .build())
611}
612
613impl FileWriter for ParquetWriter {
614 async fn write(&mut self, batch: &arrow_array::RecordBatch) -> Result<()> {
615 if batch.num_rows() == 0 {
617 return Ok(());
618 }
619
620 self.current_row_num += batch.num_rows();
621
622 let batch_c = batch.clone();
623 self.nan_value_count_visitor
624 .compute(self.schema.clone(), batch_c)?;
625
626 let writer = if let Some(writer) = &mut self.inner_writer {
628 writer
629 } else {
630 let arrow_schema: ArrowSchemaRef = Arc::new(self.schema.as_ref().try_into()?);
631 let inner_writer = self.output_file.writer().await?;
632 let async_writer = AsyncFileWriter::new(inner_writer);
633 let writer = AsyncArrowWriter::try_new(
634 async_writer,
635 arrow_schema.clone(),
636 Some(self.writer_properties.clone()),
637 )
638 .map_err(|err| {
639 Error::new(ErrorKind::Unexpected, "Failed to build parquet writer.")
640 .with_source(err)
641 })?;
642 self.inner_writer = Some(writer);
643 self.inner_writer.as_mut().unwrap()
644 };
645
646 writer.write(batch).await.map_err(|err| {
647 Error::new(
648 ErrorKind::Unexpected,
649 "Failed to write using parquet writer.",
650 )
651 .with_source(err)
652 })?;
653
654 Ok(())
655 }
656
657 async fn close(mut self) -> Result<Vec<DataFileBuilder>> {
658 let mut writer = match self.inner_writer.take() {
659 Some(writer) => writer,
660 None => return Ok(vec![]),
661 };
662
663 let metadata = writer.finish().await.map_err(|err| {
664 Error::new(ErrorKind::Unexpected, "Failed to finish parquet writer.").with_source(err)
665 })?;
666
667 let written_size = writer.bytes_written();
668
669 if self.current_row_num == 0 {
670 self.output_file.delete().await.map_err(|err| {
671 Error::new(
672 ErrorKind::Unexpected,
673 "Failed to delete empty parquet file.",
674 )
675 .with_source(err)
676 })?;
677 Ok(vec![])
678 } else {
679 let parquet_metadata = Arc::new(metadata);
680
681 Ok(vec![Self::parquet_to_data_file_builder(
682 self.schema,
683 parquet_metadata,
684 written_size,
685 self.output_file.location().to_string(),
686 self.nan_value_count_visitor.nan_value_counts,
687 self.key_metadata,
688 )?])
689 }
690 }
691}
692
693impl CurrentFileStatus for ParquetWriter {
694 fn current_file_path(&self) -> String {
695 self.output_file.location().to_string()
696 }
697
698 fn current_row_num(&self) -> usize {
699 self.current_row_num
700 }
701
702 fn current_written_size(&self) -> usize {
703 if let Some(inner) = self.inner_writer.as_ref() {
704 inner.bytes_written() + inner.in_progress_size()
707 } else {
708 0
710 }
711 }
712}
713
714struct AsyncFileWriter(Box<dyn FileWrite>);
720
721impl AsyncFileWriter {
722 pub fn new(writer: Box<dyn FileWrite>) -> Self {
724 Self(writer)
725 }
726}
727
728impl ArrowAsyncFileWriter for AsyncFileWriter {
729 fn write(&mut self, bs: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> {
730 Box::pin(async {
731 self.0
732 .write(bs)
733 .await
734 .map_err(|err| parquet::errors::ParquetError::External(Box::new(err)))
735 })
736 }
737
738 fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> {
739 Box::pin(async {
740 self.0
742 .close()
743 .await
744 .map(|_| ())
745 .map_err(|err| parquet::errors::ParquetError::External(Box::new(err)))
746 })
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use std::collections::HashMap;
753 use std::sync::Arc;
754
755 use anyhow::Result;
756 use arrow_array::builder::{Float32Builder, Int32Builder, MapBuilder};
757 use arrow_array::types::{Float32Type, Int64Type};
758 use arrow_array::{
759 Array, ArrayRef, BooleanArray, Decimal128Array, Float32Array, Float64Array, Int32Array,
760 Int64Array, ListArray, MapArray, RecordBatch, StructArray,
761 };
762 use arrow_schema::{DataType, Field, Fields, SchemaRef as ArrowSchemaRef};
763 use arrow_select::concat::concat_batches;
764 use futures::TryStreamExt;
765 use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
766 use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
767 use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
768 use parquet::file::statistics::ValueStatistics;
769 use parquet::schema::types::ColumnPath;
770 use tempfile::TempDir;
771 use uuid::Uuid;
772
773 use super::*;
774 use crate::Runtime;
775 use crate::arrow::{ArrowReaderBuilder, schema_to_arrow_schema};
776 use crate::io::FileIO;
777 use crate::scan::{FileScanTask, FileScanTaskStream};
778 use crate::spec::decimal_utils::{decimal_mantissa, decimal_new, decimal_scale};
779 use crate::spec::{PrimitiveLiteral, Struct, *};
780 use crate::test_utils::make_encryption_manager;
781 use crate::writer::file_writer::location_generator::{
782 DefaultFileNameGenerator, DefaultLocationGenerator, FileNameGenerator, LocationGenerator,
783 };
784 use crate::writer::tests::check_parquet_data_file;
785
786 fn schema_for_all_type() -> Schema {
787 Schema::builder()
788 .with_schema_id(1)
789 .with_fields(vec![
790 NestedField::optional(0, "boolean", Type::Primitive(PrimitiveType::Boolean)).into(),
791 NestedField::optional(1, "int", Type::Primitive(PrimitiveType::Int)).into(),
792 NestedField::optional(2, "long", Type::Primitive(PrimitiveType::Long)).into(),
793 NestedField::optional(3, "float", Type::Primitive(PrimitiveType::Float)).into(),
794 NestedField::optional(4, "double", Type::Primitive(PrimitiveType::Double)).into(),
795 NestedField::optional(5, "string", Type::Primitive(PrimitiveType::String)).into(),
796 NestedField::optional(6, "binary", Type::Primitive(PrimitiveType::Binary)).into(),
797 NestedField::optional(7, "date", Type::Primitive(PrimitiveType::Date)).into(),
798 NestedField::optional(8, "time", Type::Primitive(PrimitiveType::Time)).into(),
799 NestedField::optional(9, "timestamp", Type::Primitive(PrimitiveType::Timestamp))
800 .into(),
801 NestedField::optional(
802 10,
803 "timestamptz",
804 Type::Primitive(PrimitiveType::Timestamptz),
805 )
806 .into(),
807 NestedField::optional(
808 11,
809 "timestamp_ns",
810 Type::Primitive(PrimitiveType::TimestampNs),
811 )
812 .into(),
813 NestedField::optional(
814 12,
815 "timestamptz_ns",
816 Type::Primitive(PrimitiveType::TimestamptzNs),
817 )
818 .into(),
819 NestedField::optional(
820 13,
821 "decimal",
822 Type::Primitive(PrimitiveType::Decimal {
823 precision: 10,
824 scale: 5,
825 }),
826 )
827 .into(),
828 NestedField::optional(14, "uuid", Type::Primitive(PrimitiveType::Uuid)).into(),
829 NestedField::optional(15, "fixed", Type::Primitive(PrimitiveType::Fixed(10)))
830 .into(),
831 NestedField::optional(
834 16,
835 "decimal_38",
836 Type::Primitive(PrimitiveType::Decimal {
837 precision: 38,
838 scale: 5,
839 }),
840 )
841 .into(),
842 ])
843 .build()
844 .unwrap()
845 }
846
847 fn nested_schema_for_test() -> Schema {
848 Schema::builder()
850 .with_schema_id(1)
851 .with_fields(vec![
852 NestedField::required(0, "col0", Type::Primitive(PrimitiveType::Long)).into(),
853 NestedField::required(
854 1,
855 "col1",
856 Type::Struct(StructType::new(vec![
857 NestedField::required(5, "col_1_5", Type::Primitive(PrimitiveType::Long))
858 .into(),
859 NestedField::required(6, "col_1_6", Type::Primitive(PrimitiveType::Long))
860 .into(),
861 ])),
862 )
863 .into(),
864 NestedField::required(2, "col2", Type::Primitive(PrimitiveType::String)).into(),
865 NestedField::required(
866 3,
867 "col3",
868 Type::List(ListType::new(
869 NestedField::required(7, "element", Type::Primitive(PrimitiveType::Long))
870 .into(),
871 )),
872 )
873 .into(),
874 NestedField::required(
875 4,
876 "col4",
877 Type::Struct(StructType::new(vec![
878 NestedField::required(
879 8,
880 "col_4_8",
881 Type::Struct(StructType::new(vec![
882 NestedField::required(
883 9,
884 "col_4_8_9",
885 Type::Primitive(PrimitiveType::Long),
886 )
887 .into(),
888 ])),
889 )
890 .into(),
891 ])),
892 )
893 .into(),
894 NestedField::required(
895 10,
896 "col5",
897 Type::Map(MapType::new(
898 NestedField::required(11, "key", Type::Primitive(PrimitiveType::String))
899 .into(),
900 NestedField::required(
901 12,
902 "value",
903 Type::List(ListType::new(
904 NestedField::required(
905 13,
906 "item",
907 Type::Primitive(PrimitiveType::Long),
908 )
909 .into(),
910 )),
911 )
912 .into(),
913 )),
914 )
915 .into(),
916 ])
917 .build()
918 .unwrap()
919 }
920
921 #[tokio::test]
922 async fn test_index_by_parquet_path() {
923 let expect = HashMap::from([
924 ("col0".to_string(), 0),
925 ("col1.col_1_5".to_string(), 5),
926 ("col1.col_1_6".to_string(), 6),
927 ("col2".to_string(), 2),
928 ("col3.list.element".to_string(), 7),
929 ("col4.col_4_8.col_4_8_9".to_string(), 9),
930 ("col5.key_value.key".to_string(), 11),
931 ("col5.key_value.value.list.item".to_string(), 13),
932 ]);
933 let mut visitor = IndexByParquetPathName::new();
934 visit_schema(&nested_schema_for_test(), &mut visitor).unwrap();
935 assert_eq!(visitor.name_to_id, expect);
936 }
937
938 #[test]
939 fn test_index_by_parquet_path_variant_is_unsupported() {
940 let schema = Schema::builder()
943 .with_fields(vec![
944 NestedField::optional(1, "v", Type::Variant(VariantType)).into(),
945 ])
946 .build()
947 .unwrap();
948 let mut visitor = IndexByParquetPathName::new();
949 let err = visit_schema(&schema, &mut visitor).unwrap_err();
950 assert_eq!(err.kind(), ErrorKind::FeatureUnsupported, "{err}");
951 }
952
953 #[tokio::test]
954 async fn test_parquet_writer() -> Result<()> {
955 let temp_dir = TempDir::new().unwrap();
956 let file_io = FileIO::new_with_fs();
957 let location_gen = DefaultLocationGenerator::with_data_location(
958 temp_dir.path().to_str().unwrap().to_string(),
959 );
960 let file_name_gen =
961 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
962
963 let schema = {
965 let fields =
966 vec![
967 Field::new("col", DataType::Int64, true).with_metadata(HashMap::from([(
968 PARQUET_FIELD_ID_META_KEY.to_string(),
969 "0".to_string(),
970 )])),
971 ];
972 Arc::new(arrow_schema::Schema::new(fields))
973 };
974 let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
975 let null_col = Arc::new(Int64Array::new_null(1024)) as ArrayRef;
976 let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
977 let to_write_null = RecordBatch::try_new(schema.clone(), vec![null_col]).unwrap();
978
979 let output_file = file_io.new_output(
980 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
981 )?;
982
983 let mut pw = ParquetWriterBuilder::new(
985 WriterProperties::builder()
986 .set_max_row_group_row_count(Some(128))
987 .build(),
988 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
989 )
990 .build(output_file)
991 .await?;
992 pw.write(&to_write).await?;
993 pw.write(&to_write_null).await?;
994 let res = pw.close().await?;
995 assert_eq!(res.len(), 1);
996 let data_file = res
997 .into_iter()
998 .next()
999 .unwrap()
1000 .content(DataContentType::Data)
1002 .partition(Struct::empty())
1003 .partition_spec_id(0)
1004 .build()
1005 .unwrap();
1006
1007 assert_eq!(data_file.record_count(), 2048);
1009 assert_eq!(*data_file.value_counts(), HashMap::from([(0, 2048)]));
1010 assert_eq!(
1011 *data_file.lower_bounds(),
1012 HashMap::from([(0, Datum::long(0))])
1013 );
1014 assert_eq!(
1015 *data_file.upper_bounds(),
1016 HashMap::from([(0, Datum::long(1023))])
1017 );
1018 assert_eq!(*data_file.null_value_counts(), HashMap::from([(0, 1024)]));
1019
1020 let expect_batch = concat_batches(&schema, vec![&to_write, &to_write_null]).unwrap();
1022 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
1023
1024 Ok(())
1025 }
1026
1027 #[tokio::test]
1028 async fn test_parquet_writer_encrypted_roundtrip() -> Result<()> {
1029 let temp_dir = TempDir::new().unwrap();
1030 let file_io = FileIO::new_with_fs();
1031 let location_gen = DefaultLocationGenerator::with_data_location(
1032 temp_dir.path().to_str().unwrap().to_string(),
1033 );
1034 let file_name_gen =
1035 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1036
1037 let arrow_schema = Arc::new(arrow_schema::Schema::new(vec![
1038 Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
1039 PARQUET_FIELD_ID_META_KEY.to_string(),
1040 "1".to_string(),
1041 )])),
1042 ]));
1043 let iceberg_schema: SchemaRef = Arc::new(arrow_schema.as_ref().try_into().unwrap());
1044 let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(Int32Array::from(vec![
1045 10, 20, 30,
1046 ])) as ArrayRef])
1047 .unwrap();
1048
1049 let output_file = file_io.new_output(
1050 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1051 )?;
1052 let raw_properties = HashMap::from([(
1053 TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(),
1054 "test-key".to_string(),
1055 )]);
1056 let table_properties = TableProperties::new(&raw_properties);
1057 let mut parquet_writer =
1058 ParquetWriterBuilder::from_table_properties(&table_properties, iceberg_schema.clone())?
1059 .with_encryption_manager(make_encryption_manager("test-key"))
1060 .build(output_file)
1061 .await?;
1062 parquet_writer.write(&batch).await?;
1063 let data_file = parquet_writer
1064 .close()
1065 .await?
1066 .into_iter()
1067 .next()
1068 .unwrap()
1069 .content(DataContentType::Data)
1070 .partition(Struct::empty())
1071 .partition_spec_id(0)
1072 .build()
1073 .unwrap();
1074
1075 let raw = file_io
1076 .new_input(data_file.file_path.clone())?
1077 .read()
1078 .await?;
1079 assert!(
1080 ParquetRecordBatchReaderBuilder::try_new(raw).is_err(),
1081 "an encrypted parquet file must not be readable without decryption"
1082 );
1083
1084 let reader = ArrowReaderBuilder::new(file_io, Runtime::current()).build();
1085 let task = FileScanTask::builder()
1086 .with_file_size_in_bytes(data_file.file_size_in_bytes())
1087 .with_start(0)
1088 .with_length(0)
1089 .with_data_file_path(data_file.file_path.clone())
1090 .with_data_file_format(DataFileFormat::Parquet)
1091 .with_schema(iceberg_schema)
1092 .with_project_field_ids(vec![1])
1093 .with_case_sensitive(false)
1094 .with_key_metadata(data_file.key_metadata().map(Box::from))
1095 .build()
1096 .unwrap();
1097 let tasks = Box::pin(futures::stream::iter(vec![Ok(task)])) as FileScanTaskStream;
1098 let batches: Vec<RecordBatch> = reader
1099 .read(tasks)
1100 .unwrap()
1101 .stream()
1102 .try_collect()
1103 .await
1104 .unwrap();
1105
1106 assert_eq!(batches.len(), 1);
1107 let ids = batches[0]
1108 .column(0)
1109 .as_any()
1110 .downcast_ref::<Int32Array>()
1111 .unwrap();
1112 assert_eq!(ids.values(), &[10, 20, 30]);
1113
1114 Ok(())
1115 }
1116
1117 #[tokio::test]
1118 async fn test_parquet_writer_with_complex_schema() -> Result<()> {
1119 let temp_dir = TempDir::new().unwrap();
1120 let file_io = FileIO::new_with_fs();
1121 let location_gen = DefaultLocationGenerator::with_data_location(
1122 temp_dir.path().to_str().unwrap().to_string(),
1123 );
1124 let file_name_gen =
1125 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1126
1127 let schema = nested_schema_for_test();
1129 let arrow_schema: ArrowSchemaRef = Arc::new((&schema).try_into().unwrap());
1130 let col0 = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
1131 let col1 = Arc::new(StructArray::new(
1132 {
1133 if let DataType::Struct(fields) = arrow_schema.field(1).data_type() {
1134 fields.clone()
1135 } else {
1136 unreachable!()
1137 }
1138 },
1139 vec![
1140 Arc::new(Int64Array::from_iter_values(0..1024)),
1141 Arc::new(Int64Array::from_iter_values(0..1024)),
1142 ],
1143 None,
1144 ));
1145 let col2 = Arc::new(arrow_array::StringArray::from_iter_values(
1146 (0..1024).map(|n| n.to_string()),
1147 )) as ArrayRef;
1148 let col3 = Arc::new({
1149 let list_parts = ListArray::from_iter_primitive::<Int64Type, _, _>(
1150 (0..1024).map(|n| Some(vec![Some(n)])),
1151 )
1152 .into_parts();
1153 ListArray::new(
1154 {
1155 if let DataType::List(field) = arrow_schema.field(3).data_type() {
1156 field.clone()
1157 } else {
1158 unreachable!()
1159 }
1160 },
1161 list_parts.1,
1162 list_parts.2,
1163 list_parts.3,
1164 )
1165 }) as ArrayRef;
1166 let col4 = Arc::new(StructArray::new(
1167 {
1168 if let DataType::Struct(fields) = arrow_schema.field(4).data_type() {
1169 fields.clone()
1170 } else {
1171 unreachable!()
1172 }
1173 },
1174 vec![Arc::new(StructArray::new(
1175 {
1176 if let DataType::Struct(fields) = arrow_schema.field(4).data_type() {
1177 if let DataType::Struct(fields) = fields[0].data_type() {
1178 fields.clone()
1179 } else {
1180 unreachable!()
1181 }
1182 } else {
1183 unreachable!()
1184 }
1185 },
1186 vec![Arc::new(Int64Array::from_iter_values(0..1024))],
1187 None,
1188 ))],
1189 None,
1190 ));
1191 let col5 = Arc::new({
1192 let mut map_array_builder = MapBuilder::new(
1193 None,
1194 arrow_array::builder::StringBuilder::new(),
1195 arrow_array::builder::ListBuilder::new(arrow_array::builder::PrimitiveBuilder::<
1196 Int64Type,
1197 >::new()),
1198 );
1199 for i in 0..1024 {
1200 map_array_builder.keys().append_value(i.to_string());
1201 map_array_builder
1202 .values()
1203 .append_value(vec![Some(i as i64); i + 1]);
1204 map_array_builder.append(true)?;
1205 }
1206 let (_, offset_buffer, struct_array, null_buffer, ordered) =
1207 map_array_builder.finish().into_parts();
1208 let struct_array = {
1209 let (_, mut arrays, nulls) = struct_array.into_parts();
1210 let list_array = {
1211 let list_array = arrays[1]
1212 .as_any()
1213 .downcast_ref::<ListArray>()
1214 .unwrap()
1215 .clone();
1216 let (_, offsets, array, nulls) = list_array.into_parts();
1217 let list_field = {
1218 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1219 if let DataType::Struct(fields) = map_field.data_type() {
1220 if let DataType::List(list_field) = fields[1].data_type() {
1221 list_field.clone()
1222 } else {
1223 unreachable!()
1224 }
1225 } else {
1226 unreachable!()
1227 }
1228 } else {
1229 unreachable!()
1230 }
1231 };
1232 ListArray::new(list_field, offsets, array, nulls)
1233 };
1234 arrays[1] = Arc::new(list_array) as ArrayRef;
1235 StructArray::new(
1236 {
1237 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1238 if let DataType::Struct(fields) = map_field.data_type() {
1239 fields.clone()
1240 } else {
1241 unreachable!()
1242 }
1243 } else {
1244 unreachable!()
1245 }
1246 },
1247 arrays,
1248 nulls,
1249 )
1250 };
1251 MapArray::new(
1252 {
1253 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1254 map_field.clone()
1255 } else {
1256 unreachable!()
1257 }
1258 },
1259 offset_buffer,
1260 struct_array,
1261 null_buffer,
1262 ordered,
1263 )
1264 }) as ArrayRef;
1265 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
1266 col0, col1, col2, col3, col4, col5,
1267 ])
1268 .unwrap();
1269 let output_file = file_io.new_output(
1270 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1271 )?;
1272
1273 let mut pw =
1275 ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(schema))
1276 .build(output_file)
1277 .await?;
1278 pw.write(&to_write).await?;
1279 let res = pw.close().await?;
1280 assert_eq!(res.len(), 1);
1281 let data_file = res
1282 .into_iter()
1283 .next()
1284 .unwrap()
1285 .content(DataContentType::Data)
1287 .partition(Struct::empty())
1288 .partition_spec_id(0)
1289 .build()
1290 .unwrap();
1291
1292 assert_eq!(data_file.record_count(), 1024);
1294 assert_eq!(
1295 *data_file.value_counts(),
1296 HashMap::from([
1297 (0, 1024),
1298 (5, 1024),
1299 (6, 1024),
1300 (2, 1024),
1301 (7, 1024),
1302 (9, 1024),
1303 (11, 1024),
1304 (13, (1..1025).sum()),
1305 ])
1306 );
1307 assert_eq!(
1308 *data_file.lower_bounds(),
1309 HashMap::from([
1310 (0, Datum::long(0)),
1311 (5, Datum::long(0)),
1312 (6, Datum::long(0)),
1313 (2, Datum::string("0")),
1314 (7, Datum::long(0)),
1315 (9, Datum::long(0)),
1316 (11, Datum::string("0")),
1317 (13, Datum::long(0))
1318 ])
1319 );
1320 assert_eq!(
1321 *data_file.upper_bounds(),
1322 HashMap::from([
1323 (0, Datum::long(1023)),
1324 (5, Datum::long(1023)),
1325 (6, Datum::long(1023)),
1326 (2, Datum::string("999")),
1327 (7, Datum::long(1023)),
1328 (9, Datum::long(1023)),
1329 (11, Datum::string("999")),
1330 (13, Datum::long(1023))
1331 ])
1332 );
1333
1334 check_parquet_data_file(&file_io, &data_file, &to_write).await;
1336
1337 Ok(())
1338 }
1339
1340 #[tokio::test]
1341 async fn test_all_type_for_write() -> Result<()> {
1342 let temp_dir = TempDir::new().unwrap();
1343 let file_io = FileIO::new_with_fs();
1344 let location_gen = DefaultLocationGenerator::with_data_location(
1345 temp_dir.path().to_str().unwrap().to_string(),
1346 );
1347 let file_name_gen =
1348 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1349
1350 let schema = schema_for_all_type();
1353 let arrow_schema: ArrowSchemaRef = Arc::new((&schema).try_into().unwrap());
1354 let col0 = Arc::new(BooleanArray::from(vec![
1355 Some(true),
1356 Some(false),
1357 None,
1358 Some(true),
1359 ])) as ArrayRef;
1360 let col1 = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])) as ArrayRef;
1361 let col2 = Arc::new(Int64Array::from(vec![Some(1), Some(2), None, Some(4)])) as ArrayRef;
1362 let col3 = Arc::new(Float32Array::from(vec![
1363 Some(0.5),
1364 Some(2.0),
1365 None,
1366 Some(3.5),
1367 ])) as ArrayRef;
1368 let col4 = Arc::new(Float64Array::from(vec![
1369 Some(0.5),
1370 Some(2.0),
1371 None,
1372 Some(3.5),
1373 ])) as ArrayRef;
1374 let col5 = Arc::new(arrow_array::StringArray::from(vec![
1375 Some("a"),
1376 Some("b"),
1377 None,
1378 Some("d"),
1379 ])) as ArrayRef;
1380 let col6 = Arc::new(arrow_array::LargeBinaryArray::from_opt_vec(vec![
1381 Some(b"one"),
1382 None,
1383 Some(b""),
1384 Some(b"zzzz"),
1385 ])) as ArrayRef;
1386 let col7 = Arc::new(arrow_array::Date32Array::from(vec![
1387 Some(0),
1388 Some(1),
1389 None,
1390 Some(3),
1391 ])) as ArrayRef;
1392 let col8 = Arc::new(arrow_array::Time64MicrosecondArray::from(vec![
1393 Some(0),
1394 Some(1),
1395 None,
1396 Some(3),
1397 ])) as ArrayRef;
1398 let col9 = Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![
1399 Some(0),
1400 Some(1),
1401 None,
1402 Some(3),
1403 ])) as ArrayRef;
1404 let col10 = Arc::new(
1405 arrow_array::TimestampMicrosecondArray::from(vec![Some(0), Some(1), None, Some(3)])
1406 .with_timezone_utc(),
1407 ) as ArrayRef;
1408 let col11 = Arc::new(arrow_array::TimestampNanosecondArray::from(vec![
1409 Some(0),
1410 Some(1),
1411 None,
1412 Some(3),
1413 ])) as ArrayRef;
1414 let col12 = Arc::new(
1415 arrow_array::TimestampNanosecondArray::from(vec![Some(0), Some(1), None, Some(3)])
1416 .with_timezone_utc(),
1417 ) as ArrayRef;
1418 let col13 = Arc::new(
1419 Decimal128Array::from(vec![Some(1), Some(2), None, Some(100)])
1420 .with_precision_and_scale(10, 5)
1421 .unwrap(),
1422 ) as ArrayRef;
1423 let col14 = Arc::new(
1424 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1425 vec![
1426 Some(Uuid::from_u128(0).as_bytes().to_vec()),
1427 Some(Uuid::from_u128(1).as_bytes().to_vec()),
1428 None,
1429 Some(Uuid::from_u128(3).as_bytes().to_vec()),
1430 ]
1431 .into_iter(),
1432 16,
1433 )
1434 .unwrap(),
1435 ) as ArrayRef;
1436 let col15 = Arc::new(
1437 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1438 vec![
1439 Some(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
1440 Some(vec![11, 12, 13, 14, 15, 16, 17, 18, 19, 20]),
1441 None,
1442 Some(vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30]),
1443 ]
1444 .into_iter(),
1445 10,
1446 )
1447 .unwrap(),
1448 ) as ArrayRef;
1449 let col16 = Arc::new(
1450 Decimal128Array::from(vec![Some(1), Some(2), None, Some(100)])
1451 .with_precision_and_scale(38, 5)
1452 .unwrap(),
1453 ) as ArrayRef;
1454 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
1455 col0, col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12, col13,
1456 col14, col15, col16,
1457 ])
1458 .unwrap();
1459 let output_file = file_io.new_output(
1460 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1461 )?;
1462
1463 let mut pw =
1465 ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(schema))
1466 .build(output_file)
1467 .await?;
1468 pw.write(&to_write).await?;
1469 let res = pw.close().await?;
1470 assert_eq!(res.len(), 1);
1471 let data_file = res
1472 .into_iter()
1473 .next()
1474 .unwrap()
1475 .content(DataContentType::Data)
1477 .partition(Struct::empty())
1478 .partition_spec_id(0)
1479 .build()
1480 .unwrap();
1481
1482 assert_eq!(data_file.record_count(), 4);
1484 assert!(data_file.value_counts().iter().all(|(_, &v)| { v == 4 }));
1485 assert!(
1486 data_file
1487 .null_value_counts()
1488 .iter()
1489 .all(|(_, &v)| { v == 1 })
1490 );
1491 assert_eq!(
1492 *data_file.lower_bounds(),
1493 HashMap::from([
1494 (0, Datum::bool(false)),
1495 (1, Datum::int(1)),
1496 (2, Datum::long(1)),
1497 (3, Datum::float(0.5_f32)),
1498 (4, Datum::double(0.5)),
1499 (5, Datum::string("a")),
1500 (6, Datum::binary(vec![])),
1501 (7, Datum::date(0)),
1502 (8, Datum::time_micros(0).unwrap()),
1503 (9, Datum::timestamp_micros(0)),
1504 (10, Datum::timestamptz_micros(0)),
1505 (11, Datum::timestamp_nanos(0)),
1506 (12, Datum::timestamptz_nanos(0)),
1507 (
1508 13,
1509 Datum::new(
1510 PrimitiveType::Decimal {
1511 precision: 10,
1512 scale: 5
1513 },
1514 PrimitiveLiteral::Int128(1)
1515 )
1516 ),
1517 (14, Datum::uuid(Uuid::from_u128(0))),
1518 (15, Datum::fixed(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
1519 (
1520 16,
1521 Datum::new(
1522 PrimitiveType::Decimal {
1523 precision: 38,
1524 scale: 5
1525 },
1526 PrimitiveLiteral::Int128(1)
1527 )
1528 ),
1529 ])
1530 );
1531 assert_eq!(
1532 *data_file.upper_bounds(),
1533 HashMap::from([
1534 (0, Datum::bool(true)),
1535 (1, Datum::int(4)),
1536 (2, Datum::long(4)),
1537 (3, Datum::float(3.5_f32)),
1538 (4, Datum::double(3.5)),
1539 (5, Datum::string("d")),
1540 (6, Datum::binary(vec![122, 122, 122, 122])),
1541 (7, Datum::date(3)),
1542 (8, Datum::time_micros(3).unwrap()),
1543 (9, Datum::timestamp_micros(3)),
1544 (10, Datum::timestamptz_micros(3)),
1545 (11, Datum::timestamp_nanos(3)),
1546 (12, Datum::timestamptz_nanos(3)),
1547 (
1548 13,
1549 Datum::new(
1550 PrimitiveType::Decimal {
1551 precision: 10,
1552 scale: 5
1553 },
1554 PrimitiveLiteral::Int128(100)
1555 )
1556 ),
1557 (14, Datum::uuid(Uuid::from_u128(3))),
1558 (
1559 15,
1560 Datum::fixed(vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30])
1561 ),
1562 (
1563 16,
1564 Datum::new(
1565 PrimitiveType::Decimal {
1566 precision: 38,
1567 scale: 5
1568 },
1569 PrimitiveLiteral::Int128(100)
1570 )
1571 ),
1572 ])
1573 );
1574
1575 check_parquet_data_file(&file_io, &data_file, &to_write).await;
1577
1578 Ok(())
1579 }
1580
1581 #[tokio::test]
1582 async fn test_decimal_bound() -> Result<()> {
1583 let temp_dir = TempDir::new().unwrap();
1584 let file_io = FileIO::new_with_fs();
1585 let location_gen = DefaultLocationGenerator::with_data_location(
1586 temp_dir.path().to_str().unwrap().to_string(),
1587 );
1588 let file_name_gen =
1589 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1590
1591 let schema = Arc::new(
1593 Schema::builder()
1594 .with_fields(vec![
1595 NestedField::optional(
1596 0,
1597 "decimal",
1598 Type::Primitive(PrimitiveType::Decimal {
1599 precision: 28,
1600 scale: 10,
1601 }),
1602 )
1603 .into(),
1604 ])
1605 .build()
1606 .unwrap(),
1607 );
1608 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1609 let output_file = file_io.new_output(
1610 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1611 )?;
1612 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone())
1613 .build(output_file)
1614 .await?;
1615 let col0 = Arc::new(
1616 Decimal128Array::from(vec![Some(22000000000), Some(11000000000)])
1617 .with_data_type(DataType::Decimal128(28, 10)),
1618 ) as ArrayRef;
1619 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1620 pw.write(&to_write).await?;
1621 let res = pw.close().await?;
1622 assert_eq!(res.len(), 1);
1623 let data_file = res
1624 .into_iter()
1625 .next()
1626 .unwrap()
1627 .content(DataContentType::Data)
1628 .partition(Struct::empty())
1629 .partition_spec_id(0)
1630 .build()
1631 .unwrap();
1632 assert_eq!(
1633 data_file.upper_bounds().get(&0),
1634 Some(Datum::decimal_with_precision(decimal_new(22000000000_i64, 10), 28).unwrap())
1635 .as_ref()
1636 );
1637 assert_eq!(
1638 data_file.lower_bounds().get(&0),
1639 Some(Datum::decimal_with_precision(decimal_new(11000000000_i64, 10), 28).unwrap())
1640 .as_ref()
1641 );
1642
1643 let schema = Arc::new(
1645 Schema::builder()
1646 .with_fields(vec![
1647 NestedField::optional(
1648 0,
1649 "decimal",
1650 Type::Primitive(PrimitiveType::Decimal {
1651 precision: 28,
1652 scale: 10,
1653 }),
1654 )
1655 .into(),
1656 ])
1657 .build()
1658 .unwrap(),
1659 );
1660 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1661 let output_file = file_io.new_output(
1662 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1663 )?;
1664 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone())
1665 .build(output_file)
1666 .await?;
1667 let col0 = Arc::new(
1668 Decimal128Array::from(vec![Some(-22000000000), Some(-11000000000)])
1669 .with_data_type(DataType::Decimal128(28, 10)),
1670 ) as ArrayRef;
1671 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1672 pw.write(&to_write).await?;
1673 let res = pw.close().await?;
1674 assert_eq!(res.len(), 1);
1675 let data_file = res
1676 .into_iter()
1677 .next()
1678 .unwrap()
1679 .content(DataContentType::Data)
1680 .partition(Struct::empty())
1681 .partition_spec_id(0)
1682 .build()
1683 .unwrap();
1684 assert_eq!(
1685 data_file.upper_bounds().get(&0),
1686 Some(Datum::decimal_with_precision(decimal_new(-11000000000_i64, 10), 28).unwrap())
1687 .as_ref()
1688 );
1689 assert_eq!(
1690 data_file.lower_bounds().get(&0),
1691 Some(Datum::decimal_with_precision(decimal_new(-22000000000_i64, 10), 28).unwrap())
1692 .as_ref()
1693 );
1694
1695 use crate::spec::decimal_utils::decimal_from_str_exact;
1698 let decimal_max = decimal_from_str_exact("99999999999999999999999999999999999999").unwrap();
1699 let decimal_min =
1700 decimal_from_str_exact("-99999999999999999999999999999999999999").unwrap();
1701 assert_eq!(decimal_scale(&decimal_max), decimal_scale(&decimal_min));
1702 let schema = Arc::new(
1703 Schema::builder()
1704 .with_fields(vec![
1705 NestedField::optional(
1706 0,
1707 "decimal",
1708 Type::Primitive(PrimitiveType::Decimal {
1709 precision: 38,
1710 scale: decimal_scale(&decimal_max),
1711 }),
1712 )
1713 .into(),
1714 ])
1715 .build()
1716 .unwrap(),
1717 );
1718 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1719 let output_file = file_io.new_output(
1720 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1721 )?;
1722 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema)
1723 .build(output_file)
1724 .await?;
1725 let col0 = Arc::new(
1726 Decimal128Array::from(vec![
1727 Some(decimal_mantissa(&decimal_max)),
1728 Some(decimal_mantissa(&decimal_min)),
1729 ])
1730 .with_data_type(DataType::Decimal128(38, 0)),
1731 ) as ArrayRef;
1732 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1733 pw.write(&to_write).await?;
1734 let res = pw.close().await?;
1735 assert_eq!(res.len(), 1);
1736 let data_file = res
1737 .into_iter()
1738 .next()
1739 .unwrap()
1740 .content(DataContentType::Data)
1741 .partition(Struct::empty())
1742 .partition_spec_id(0)
1743 .build()
1744 .unwrap();
1745 assert_eq!(
1746 data_file.upper_bounds().get(&0),
1747 Some(Datum::decimal(decimal_max).unwrap()).as_ref()
1748 );
1749 assert_eq!(
1750 data_file.lower_bounds().get(&0),
1751 Some(Datum::decimal(decimal_min).unwrap()).as_ref()
1752 );
1753
1754 Ok(())
1824 }
1825
1826 #[tokio::test]
1827 async fn test_empty_write() -> Result<()> {
1828 let temp_dir = TempDir::new().unwrap();
1829 let file_io = FileIO::new_with_fs();
1830 let location_gen = DefaultLocationGenerator::with_data_location(
1831 temp_dir.path().to_str().unwrap().to_string(),
1832 );
1833 let file_name_gen =
1834 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1835
1836 let schema = {
1838 let fields =
1839 vec![
1840 Field::new("col", DataType::Int64, true).with_metadata(HashMap::from([(
1841 PARQUET_FIELD_ID_META_KEY.to_string(),
1842 "0".to_string(),
1843 )])),
1844 ];
1845 Arc::new(arrow_schema::Schema::new(fields))
1846 };
1847 let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
1848 let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
1849 let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
1850 let output_file = file_io.new_output(&file_path)?;
1851 let mut pw = ParquetWriterBuilder::new(
1852 WriterProperties::builder().build(),
1853 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1854 )
1855 .build(output_file)
1856 .await?;
1857 pw.write(&to_write).await?;
1858 pw.close().await.unwrap();
1859 assert!(file_io.exists(&file_path).await.unwrap());
1860
1861 let file_name_gen =
1863 DefaultFileNameGenerator::new("test_empty".to_string(), None, DataFileFormat::Parquet);
1864 let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
1865 let output_file = file_io.new_output(&file_path)?;
1866 let pw = ParquetWriterBuilder::new(
1867 WriterProperties::builder().build(),
1868 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1869 )
1870 .build(output_file)
1871 .await?;
1872 pw.close().await.unwrap();
1873 assert!(!file_io.exists(&file_path).await.unwrap());
1874
1875 Ok(())
1876 }
1877
1878 #[tokio::test]
1879 async fn test_nan_val_cnts_primitive_type() -> Result<()> {
1880 let temp_dir = TempDir::new().unwrap();
1881 let file_io = FileIO::new_with_fs();
1882 let location_gen = DefaultLocationGenerator::with_data_location(
1883 temp_dir.path().to_str().unwrap().to_string(),
1884 );
1885 let file_name_gen =
1886 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1887
1888 let arrow_schema = {
1890 let fields = vec![
1891 Field::new("col", DataType::Float32, false).with_metadata(HashMap::from([(
1892 PARQUET_FIELD_ID_META_KEY.to_string(),
1893 "0".to_string(),
1894 )])),
1895 Field::new("col2", DataType::Float64, false).with_metadata(HashMap::from([(
1896 PARQUET_FIELD_ID_META_KEY.to_string(),
1897 "1".to_string(),
1898 )])),
1899 ];
1900 Arc::new(arrow_schema::Schema::new(fields))
1901 };
1902
1903 let float_32_col = Arc::new(Float32Array::from_iter_values_with_nulls(
1904 [1.0_f32, f32::NAN, 2.0, 2.0].into_iter(),
1905 None,
1906 )) as ArrayRef;
1907
1908 let float_64_col = Arc::new(Float64Array::from_iter_values_with_nulls(
1909 [1.0_f64, f64::NAN, 2.0, 2.0].into_iter(),
1910 None,
1911 )) as ArrayRef;
1912
1913 let to_write =
1914 RecordBatch::try_new(arrow_schema.clone(), vec![float_32_col, float_64_col]).unwrap();
1915 let output_file = file_io.new_output(
1916 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1917 )?;
1918
1919 let mut pw = ParquetWriterBuilder::new(
1921 WriterProperties::builder().build(),
1922 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1923 )
1924 .build(output_file)
1925 .await?;
1926
1927 pw.write(&to_write).await?;
1928 let res = pw.close().await?;
1929 assert_eq!(res.len(), 1);
1930 let data_file = res
1931 .into_iter()
1932 .next()
1933 .unwrap()
1934 .content(DataContentType::Data)
1936 .partition(Struct::empty())
1937 .partition_spec_id(0)
1938 .build()
1939 .unwrap();
1940
1941 assert_eq!(data_file.record_count(), 4);
1943 assert_eq!(*data_file.value_counts(), HashMap::from([(0, 4), (1, 4)]));
1944 assert_eq!(
1945 *data_file.lower_bounds(),
1946 HashMap::from([(0, Datum::float(1.0_f32)), (1, Datum::double(1.0)),])
1947 );
1948 assert_eq!(
1949 *data_file.upper_bounds(),
1950 HashMap::from([(0, Datum::float(2.0_f32)), (1, Datum::double(2.0)),])
1951 );
1952 assert_eq!(
1953 *data_file.null_value_counts(),
1954 HashMap::from([(0, 0), (1, 0)])
1955 );
1956 assert_eq!(
1957 *data_file.nan_value_counts(),
1958 HashMap::from([(0, 1), (1, 1)])
1959 );
1960
1961 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
1963 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
1964
1965 Ok(())
1966 }
1967
1968 #[tokio::test]
1969 async fn test_nan_val_cnts_struct_type() -> Result<()> {
1970 let temp_dir = TempDir::new().unwrap();
1971 let file_io = FileIO::new_with_fs();
1972 let location_gen = DefaultLocationGenerator::with_data_location(
1973 temp_dir.path().to_str().unwrap().to_string(),
1974 );
1975 let file_name_gen =
1976 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1977
1978 let schema_struct_float_fields = Fields::from(vec![
1979 Field::new("col4", DataType::Float32, false).with_metadata(HashMap::from([(
1980 PARQUET_FIELD_ID_META_KEY.to_string(),
1981 "4".to_string(),
1982 )])),
1983 ]);
1984
1985 let schema_struct_nested_float_fields = Fields::from(vec![
1986 Field::new("col7", DataType::Float32, false).with_metadata(HashMap::from([(
1987 PARQUET_FIELD_ID_META_KEY.to_string(),
1988 "7".to_string(),
1989 )])),
1990 ]);
1991
1992 let schema_struct_nested_fields = Fields::from(vec![
1993 Field::new(
1994 "col6",
1995 DataType::Struct(schema_struct_nested_float_fields.clone()),
1996 false,
1997 )
1998 .with_metadata(HashMap::from([(
1999 PARQUET_FIELD_ID_META_KEY.to_string(),
2000 "6".to_string(),
2001 )])),
2002 ]);
2003
2004 let arrow_schema = {
2006 let fields = vec![
2007 Field::new(
2008 "col3",
2009 DataType::Struct(schema_struct_float_fields.clone()),
2010 false,
2011 )
2012 .with_metadata(HashMap::from([(
2013 PARQUET_FIELD_ID_META_KEY.to_string(),
2014 "3".to_string(),
2015 )])),
2016 Field::new(
2017 "col5",
2018 DataType::Struct(schema_struct_nested_fields.clone()),
2019 false,
2020 )
2021 .with_metadata(HashMap::from([(
2022 PARQUET_FIELD_ID_META_KEY.to_string(),
2023 "5".to_string(),
2024 )])),
2025 ];
2026 Arc::new(arrow_schema::Schema::new(fields))
2027 };
2028
2029 let float_32_col = Arc::new(Float32Array::from_iter_values_with_nulls(
2030 [1.0_f32, f32::NAN, 2.0, 2.0].into_iter(),
2031 None,
2032 )) as ArrayRef;
2033
2034 let struct_float_field_col = Arc::new(StructArray::new(
2035 schema_struct_float_fields,
2036 vec![float_32_col.clone()],
2037 None,
2038 )) as ArrayRef;
2039
2040 let struct_nested_float_field_col = Arc::new(StructArray::new(
2041 schema_struct_nested_fields,
2042 vec![Arc::new(StructArray::new(
2043 schema_struct_nested_float_fields,
2044 vec![float_32_col.clone()],
2045 None,
2046 )) as ArrayRef],
2047 None,
2048 )) as ArrayRef;
2049
2050 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
2051 struct_float_field_col,
2052 struct_nested_float_field_col,
2053 ])
2054 .unwrap();
2055 let output_file = file_io.new_output(
2056 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
2057 )?;
2058
2059 let mut pw = ParquetWriterBuilder::new(
2061 WriterProperties::builder().build(),
2062 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
2063 )
2064 .build(output_file)
2065 .await?;
2066
2067 pw.write(&to_write).await?;
2068 let res = pw.close().await?;
2069 assert_eq!(res.len(), 1);
2070 let data_file = res
2071 .into_iter()
2072 .next()
2073 .unwrap()
2074 .content(DataContentType::Data)
2076 .partition(Struct::empty())
2077 .partition_spec_id(0)
2078 .build()
2079 .unwrap();
2080
2081 assert_eq!(data_file.record_count(), 4);
2083 assert_eq!(*data_file.value_counts(), HashMap::from([(4, 4), (7, 4)]));
2084 assert_eq!(
2085 *data_file.lower_bounds(),
2086 HashMap::from([(4, Datum::float(1.0_f32)), (7, Datum::float(1.0_f32)),])
2087 );
2088 assert_eq!(
2089 *data_file.upper_bounds(),
2090 HashMap::from([(4, Datum::float(2.0_f32)), (7, Datum::float(2.0_f32)),])
2091 );
2092 assert_eq!(
2093 *data_file.null_value_counts(),
2094 HashMap::from([(4, 0), (7, 0)])
2095 );
2096 assert_eq!(
2097 *data_file.nan_value_counts(),
2098 HashMap::from([(4, 1), (7, 1)])
2099 );
2100
2101 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2103 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2104
2105 Ok(())
2106 }
2107
2108 #[tokio::test]
2109 async fn test_nan_val_cnts_list_type() -> Result<()> {
2110 let temp_dir = TempDir::new().unwrap();
2111 let file_io = FileIO::new_with_fs();
2112 let location_gen = DefaultLocationGenerator::with_data_location(
2113 temp_dir.path().to_str().unwrap().to_string(),
2114 );
2115 let file_name_gen =
2116 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2117
2118 let schema_list_float_field = Field::new("element", DataType::Float32, true).with_metadata(
2119 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2120 );
2121
2122 let schema_struct_list_float_field = Field::new("element", DataType::Float32, true)
2123 .with_metadata(HashMap::from([(
2124 PARQUET_FIELD_ID_META_KEY.to_string(),
2125 "4".to_string(),
2126 )]));
2127
2128 let schema_struct_list_field = Fields::from(vec![
2129 Field::new_list("col2", schema_struct_list_float_field.clone(), true).with_metadata(
2130 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())]),
2131 ),
2132 ]);
2133
2134 let arrow_schema = {
2135 let fields = vec![
2136 Field::new_list("col0", schema_list_float_field.clone(), true).with_metadata(
2137 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())]),
2138 ),
2139 Field::new_struct("col1", schema_struct_list_field.clone(), true)
2140 .with_metadata(HashMap::from([(
2141 PARQUET_FIELD_ID_META_KEY.to_string(),
2142 "2".to_string(),
2143 )]))
2144 .clone(),
2145 ];
2149 Arc::new(arrow_schema::Schema::new(fields))
2150 };
2151
2152 let list_parts = ListArray::from_iter_primitive::<Float32Type, _, _>(vec![Some(vec![
2153 Some(1.0_f32),
2154 Some(f32::NAN),
2155 Some(2.0),
2156 Some(2.0),
2157 ])])
2158 .into_parts();
2159
2160 let list_float_field_col = Arc::new({
2161 let list_parts = list_parts.clone();
2162 ListArray::new(
2163 {
2164 if let DataType::List(field) = arrow_schema.field(0).data_type() {
2165 field.clone()
2166 } else {
2167 unreachable!()
2168 }
2169 },
2170 list_parts.1,
2171 list_parts.2,
2172 list_parts.3,
2173 )
2174 }) as ArrayRef;
2175
2176 let struct_list_fields_schema =
2177 if let DataType::Struct(fields) = arrow_schema.field(1).data_type() {
2178 fields.clone()
2179 } else {
2180 unreachable!()
2181 };
2182
2183 let struct_list_float_field_col = Arc::new({
2184 ListArray::new(
2185 {
2186 if let DataType::List(field) = struct_list_fields_schema
2187 .first()
2188 .expect("could not find first list field")
2189 .data_type()
2190 {
2191 field.clone()
2192 } else {
2193 unreachable!()
2194 }
2195 },
2196 list_parts.1,
2197 list_parts.2,
2198 list_parts.3,
2199 )
2200 }) as ArrayRef;
2201
2202 let struct_list_float_field_col = Arc::new(StructArray::new(
2203 struct_list_fields_schema,
2204 vec![struct_list_float_field_col.clone()],
2205 None,
2206 )) as ArrayRef;
2207
2208 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
2209 list_float_field_col,
2210 struct_list_float_field_col,
2211 ])
2213 .expect("Could not form record batch");
2214 let output_file = file_io.new_output(
2215 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
2216 )?;
2217
2218 let mut pw = ParquetWriterBuilder::new(
2220 WriterProperties::builder().build(),
2221 Arc::new(
2222 to_write
2223 .schema()
2224 .as_ref()
2225 .try_into()
2226 .expect("Could not convert iceberg schema"),
2227 ),
2228 )
2229 .build(output_file)
2230 .await?;
2231
2232 pw.write(&to_write).await?;
2233 let res = pw.close().await?;
2234 assert_eq!(res.len(), 1);
2235 let data_file = res
2236 .into_iter()
2237 .next()
2238 .unwrap()
2239 .content(DataContentType::Data)
2240 .partition(Struct::empty())
2241 .partition_spec_id(0)
2242 .build()
2243 .unwrap();
2244
2245 assert_eq!(data_file.record_count(), 1);
2247 assert_eq!(*data_file.value_counts(), HashMap::from([(1, 4), (4, 4)]));
2248 assert_eq!(
2249 *data_file.lower_bounds(),
2250 HashMap::from([(1, Datum::float(1.0_f32)), (4, Datum::float(1.0_f32))])
2251 );
2252 assert_eq!(
2253 *data_file.upper_bounds(),
2254 HashMap::from([(1, Datum::float(2.0_f32)), (4, Datum::float(2.0_f32))])
2255 );
2256 assert_eq!(
2257 *data_file.null_value_counts(),
2258 HashMap::from([(1, 0), (4, 0)])
2259 );
2260 assert_eq!(
2261 *data_file.nan_value_counts(),
2262 HashMap::from([(1, 1), (4, 1)])
2263 );
2264
2265 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2267 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2268
2269 Ok(())
2270 }
2271
2272 macro_rules! construct_map_arr {
2273 ($map_key_field_schema:ident, $map_value_field_schema:ident) => {{
2274 let int_builder = Int32Builder::new();
2275 let float_builder = Float32Builder::with_capacity(4);
2276 let mut builder = MapBuilder::new(None, int_builder, float_builder);
2277 builder.keys().append_value(1);
2278 builder.values().append_value(1.0_f32);
2279 builder.append(true).unwrap();
2280 builder.keys().append_value(2);
2281 builder.values().append_value(f32::NAN);
2282 builder.append(true).unwrap();
2283 builder.keys().append_value(3);
2284 builder.values().append_value(2.0);
2285 builder.append(true).unwrap();
2286 builder.keys().append_value(4);
2287 builder.values().append_value(2.0);
2288 builder.append(true).unwrap();
2289 let array = builder.finish();
2290
2291 let (_field, offsets, entries, nulls, ordered) = array.into_parts();
2292 let new_struct_fields_schema =
2293 Fields::from(vec![$map_key_field_schema, $map_value_field_schema]);
2294
2295 let entries = {
2296 let (_, arrays, nulls) = entries.into_parts();
2297 StructArray::new(new_struct_fields_schema.clone(), arrays, nulls)
2298 };
2299
2300 let field = Arc::new(Field::new(
2301 DEFAULT_MAP_FIELD_NAME,
2302 DataType::Struct(new_struct_fields_schema),
2303 false,
2304 ));
2305
2306 Arc::new(MapArray::new(field, offsets, entries, nulls, ordered))
2307 }};
2308 }
2309
2310 #[tokio::test]
2311 async fn test_nan_val_cnts_map_type() -> Result<()> {
2312 let temp_dir = TempDir::new().unwrap();
2313 let file_io = FileIO::new_with_fs();
2314 let location_gen = DefaultLocationGenerator::with_data_location(
2315 temp_dir.path().to_str().unwrap().to_string(),
2316 );
2317 let file_name_gen =
2318 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2319
2320 let map_key_field_schema =
2321 Field::new(MAP_KEY_FIELD_NAME, DataType::Int32, false).with_metadata(HashMap::from([
2322 (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()),
2323 ]));
2324
2325 let map_value_field_schema =
2326 Field::new(MAP_VALUE_FIELD_NAME, DataType::Float32, true).with_metadata(HashMap::from(
2327 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
2328 ));
2329
2330 let struct_map_key_field_schema =
2331 Field::new(MAP_KEY_FIELD_NAME, DataType::Int32, false).with_metadata(HashMap::from([
2332 (PARQUET_FIELD_ID_META_KEY.to_string(), "6".to_string()),
2333 ]));
2334
2335 let struct_map_value_field_schema =
2336 Field::new(MAP_VALUE_FIELD_NAME, DataType::Float32, true).with_metadata(HashMap::from(
2337 [(PARQUET_FIELD_ID_META_KEY.to_string(), "7".to_string())],
2338 ));
2339
2340 let schema_struct_map_field = Fields::from(vec![
2341 Field::new_map(
2342 "col3",
2343 DEFAULT_MAP_FIELD_NAME,
2344 struct_map_key_field_schema.clone(),
2345 struct_map_value_field_schema.clone(),
2346 false,
2347 false,
2348 )
2349 .with_metadata(HashMap::from([(
2350 PARQUET_FIELD_ID_META_KEY.to_string(),
2351 "5".to_string(),
2352 )])),
2353 ]);
2354
2355 let arrow_schema = {
2356 let fields = vec![
2357 Field::new_map(
2358 "col0",
2359 DEFAULT_MAP_FIELD_NAME,
2360 map_key_field_schema.clone(),
2361 map_value_field_schema.clone(),
2362 false,
2363 false,
2364 )
2365 .with_metadata(HashMap::from([(
2366 PARQUET_FIELD_ID_META_KEY.to_string(),
2367 "0".to_string(),
2368 )])),
2369 Field::new_struct("col1", schema_struct_map_field.clone(), true)
2370 .with_metadata(HashMap::from([(
2371 PARQUET_FIELD_ID_META_KEY.to_string(),
2372 "3".to_string(),
2373 )]))
2374 .clone(),
2375 ];
2376 Arc::new(arrow_schema::Schema::new(fields))
2377 };
2378
2379 let map_array = construct_map_arr!(map_key_field_schema, map_value_field_schema);
2380
2381 let struct_map_arr =
2382 construct_map_arr!(struct_map_key_field_schema, struct_map_value_field_schema);
2383
2384 let struct_list_float_field_col = Arc::new(StructArray::new(
2385 schema_struct_map_field,
2386 vec![struct_map_arr],
2387 None,
2388 )) as ArrayRef;
2389
2390 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
2391 map_array,
2392 struct_list_float_field_col,
2393 ])
2394 .expect("Could not form record batch");
2395 let output_file = file_io.new_output(
2396 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
2397 )?;
2398
2399 let mut pw = ParquetWriterBuilder::new(
2401 WriterProperties::builder().build(),
2402 Arc::new(
2403 to_write
2404 .schema()
2405 .as_ref()
2406 .try_into()
2407 .expect("Could not convert iceberg schema"),
2408 ),
2409 )
2410 .build(output_file)
2411 .await?;
2412
2413 pw.write(&to_write).await?;
2414 let res = pw.close().await?;
2415 assert_eq!(res.len(), 1);
2416 let data_file = res
2417 .into_iter()
2418 .next()
2419 .unwrap()
2420 .content(DataContentType::Data)
2421 .partition(Struct::empty())
2422 .partition_spec_id(0)
2423 .build()
2424 .unwrap();
2425
2426 assert_eq!(data_file.record_count(), 4);
2428 assert_eq!(
2429 *data_file.value_counts(),
2430 HashMap::from([(1, 4), (2, 4), (6, 4), (7, 4)])
2431 );
2432 assert_eq!(
2433 *data_file.lower_bounds(),
2434 HashMap::from([
2435 (1, Datum::int(1)),
2436 (2, Datum::float(1.0_f32)),
2437 (6, Datum::int(1)),
2438 (7, Datum::float(1.0_f32))
2439 ])
2440 );
2441 assert_eq!(
2442 *data_file.upper_bounds(),
2443 HashMap::from([
2444 (1, Datum::int(4)),
2445 (2, Datum::float(2.0_f32)),
2446 (6, Datum::int(4)),
2447 (7, Datum::float(2.0_f32))
2448 ])
2449 );
2450 assert_eq!(
2451 *data_file.null_value_counts(),
2452 HashMap::from([(1, 0), (2, 0), (6, 0), (7, 0)])
2453 );
2454 assert_eq!(
2455 *data_file.nan_value_counts(),
2456 HashMap::from([(2, 1), (7, 1)])
2457 );
2458
2459 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2461 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2462
2463 Ok(())
2464 }
2465
2466 #[tokio::test]
2467 async fn test_write_empty_parquet_file() {
2468 let temp_dir = TempDir::new().unwrap();
2469 let file_io = FileIO::new_with_fs();
2470 let location_gen = DefaultLocationGenerator::with_data_location(
2471 temp_dir.path().to_str().unwrap().to_string(),
2472 );
2473 let file_name_gen =
2474 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2475 let output_file = file_io
2476 .new_output(location_gen.generate_location(None, &file_name_gen.generate_file_name()))
2477 .unwrap();
2478
2479 let pw = ParquetWriterBuilder::new(
2481 WriterProperties::builder().build(),
2482 Arc::new(
2483 Schema::builder()
2484 .with_schema_id(1)
2485 .with_fields(vec![
2486 NestedField::required(0, "col", Type::Primitive(PrimitiveType::Long))
2487 .with_id(0)
2488 .into(),
2489 ])
2490 .build()
2491 .expect("Failed to create schema"),
2492 ),
2493 )
2494 .build(output_file)
2495 .await
2496 .unwrap();
2497
2498 let res = pw.close().await.unwrap();
2499 assert_eq!(res.len(), 0);
2500
2501 assert_eq!(std::fs::read_dir(temp_dir.path()).unwrap().count(), 0);
2503 }
2504
2505 #[tokio::test]
2506 async fn test_parquet_writer_uuid_logical_type() {
2507 let temp_dir = TempDir::new().unwrap();
2508 let file_io = FileIO::new_with_fs();
2509 let location_gen = DefaultLocationGenerator::with_data_location(
2510 temp_dir.path().to_str().unwrap().to_string(),
2511 );
2512 let file_name_gen =
2513 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2514 let output_file = file_io
2515 .new_output(location_gen.generate_location(None, &file_name_gen.generate_file_name()))
2516 .unwrap();
2517
2518 let schema = Arc::new(
2519 Schema::builder()
2520 .with_schema_id(1)
2521 .with_fields(vec![
2522 NestedField::required(0, "uuid", Type::Primitive(PrimitiveType::Uuid)).into(),
2523 ])
2524 .build()
2525 .unwrap(),
2526 );
2527 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
2528 let col = Arc::new(
2529 arrow_array::FixedSizeBinaryArray::try_from_iter(
2530 [Uuid::from_u128(0), Uuid::from_u128(1)]
2531 .into_iter()
2532 .map(|uuid| uuid.into_bytes()),
2533 )
2534 .unwrap(),
2535 ) as ArrayRef;
2536 let to_write = RecordBatch::try_new(arrow_schema, vec![col]).unwrap();
2537
2538 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema)
2539 .build(output_file)
2540 .await
2541 .unwrap();
2542 pw.write(&to_write).await.unwrap();
2543 let data_file = pw
2544 .close()
2545 .await
2546 .unwrap()
2547 .into_iter()
2548 .next()
2549 .unwrap()
2550 .content(DataContentType::Data)
2551 .partition(Struct::empty())
2552 .partition_spec_id(0)
2553 .build()
2554 .unwrap();
2555
2556 let raw = file_io
2557 .new_input(data_file.file_path.clone())
2558 .unwrap()
2559 .read()
2560 .await
2561 .unwrap();
2562 let reader = ParquetRecordBatchReaderBuilder::try_new(raw).unwrap();
2563 let column = reader.metadata().file_metadata().schema_descr().column(0);
2564 assert_eq!(
2565 (
2566 column.physical_type(),
2567 column.type_length(),
2568 column.logical_type_ref()
2569 ),
2570 (
2571 parquet::basic::Type::FIXED_LEN_BYTE_ARRAY,
2572 16,
2573 Some(&parquet::basic::LogicalType::Uuid)
2574 )
2575 );
2576 }
2577
2578 #[test]
2579 fn test_min_max_aggregator() {
2580 let schema = Arc::new(
2581 Schema::builder()
2582 .with_schema_id(1)
2583 .with_fields(vec![
2584 NestedField::required(0, "col", Type::Primitive(PrimitiveType::Int))
2585 .with_id(0)
2586 .into(),
2587 ])
2588 .build()
2589 .expect("Failed to create schema"),
2590 );
2591 let mut min_max_agg = MinMaxColAggregator::new(schema);
2592 let create_statistics =
2593 |min, max| Statistics::Int32(ValueStatistics::new(min, max, None, None, false));
2594 min_max_agg
2595 .update(0, &create_statistics(None, Some(42)))
2596 .unwrap();
2597 min_max_agg
2598 .update(0, &create_statistics(Some(0), Some(i32::MAX)))
2599 .unwrap();
2600 min_max_agg
2601 .update(0, &create_statistics(Some(i32::MIN), None))
2602 .unwrap();
2603 min_max_agg
2604 .update(0, &create_statistics(None, None))
2605 .unwrap();
2606
2607 let (lower_bounds, upper_bounds) = min_max_agg.produce();
2608
2609 assert_eq!(lower_bounds, HashMap::from([(0, Datum::int(i32::MIN))]));
2610 assert_eq!(upper_bounds, HashMap::from([(0, Datum::int(i32::MAX))]));
2611 }
2612
2613 #[test]
2614 fn test_min_max_aggregator_skips_nan_bounds() {
2615 let schema = Arc::new(
2616 Schema::builder()
2617 .with_schema_id(1)
2618 .with_fields(vec![
2619 NestedField::required(0, "f", Type::Primitive(PrimitiveType::Float)).into(),
2620 NestedField::required(1, "d", Type::Primitive(PrimitiveType::Double)).into(),
2621 ])
2622 .build()
2623 .expect("Failed to create schema"),
2624 );
2625 let mut min_max_agg = MinMaxColAggregator::new(schema);
2626 let float_stats =
2627 |min, max| Statistics::Float(ValueStatistics::new(min, max, None, None, false));
2628 let double_stats =
2629 |min, max| Statistics::Double(ValueStatistics::new(min, max, None, None, false));
2630
2631 min_max_agg
2632 .update(0, &float_stats(Some(-f32::NAN), Some(f32::NAN)))
2633 .unwrap();
2634 min_max_agg
2635 .update(0, &float_stats(Some(1.5), Some(2.5)))
2636 .unwrap();
2637 min_max_agg
2638 .update(1, &double_stats(Some(f64::NAN), Some(f64::NAN)))
2639 .unwrap();
2640
2641 let (lower_bounds, upper_bounds) = min_max_agg.produce();
2642
2643 assert_eq!(lower_bounds, HashMap::from([(0, Datum::float(1.5_f32))]));
2644 assert_eq!(upper_bounds, HashMap::from([(0, Datum::float(2.5_f32))]));
2645 }
2646
2647 fn cdc_test_schema() -> SchemaRef {
2648 Arc::new(
2649 Schema::builder()
2650 .with_schema_id(1)
2651 .with_fields(vec![
2652 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(),
2653 NestedField::required(2, "payload", Type::Primitive(PrimitiveType::String))
2654 .into(),
2655 ])
2656 .build()
2657 .unwrap(),
2658 )
2659 }
2660
2661 #[test]
2662 fn test_from_table_properties_no_cdc_by_default() {
2663 let raw_properties = HashMap::new();
2664 let tp = TableProperties::new(&raw_properties);
2665 let builder = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema()).unwrap();
2666 assert!(builder.props.content_defined_chunking().is_none());
2667 }
2668
2669 #[tokio::test]
2670 async fn test_from_table_properties_without_encryption_writes_plaintext() {
2671 let raw_properties = HashMap::new();
2672 let tp = TableProperties::new(&raw_properties);
2673 let tmp = TempDir::new().unwrap();
2674 let output = FileIO::new_with_fs()
2675 .new_output(format!("{}/plain.parquet", tmp.path().to_str().unwrap()))
2676 .unwrap();
2677 let writer = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema())
2678 .unwrap()
2679 .build(output)
2680 .await
2681 .unwrap();
2682
2683 assert!(writer.key_metadata.is_none());
2684 assert!(
2685 writer
2686 .writer_properties
2687 .file_encryption_properties()
2688 .is_none()
2689 );
2690 }
2691
2692 #[tokio::test]
2693 async fn test_from_table_properties_propagate_to_writer() {
2694 let raw_properties = HashMap::from([
2703 (
2704 TableProperties::PROPERTY_PARQUET_CDC_ENABLED.to_string(),
2705 "true".to_string(),
2706 ),
2707 (
2708 TableProperties::PROPERTY_PARQUET_CDC_MIN_CHUNK_SIZE.to_string(),
2709 "4096".to_string(),
2710 ),
2711 (
2712 TableProperties::PROPERTY_PARQUET_CDC_MAX_CHUNK_SIZE.to_string(),
2713 "8192".to_string(),
2714 ),
2715 (
2716 TableProperties::PROPERTY_PARQUET_CDC_NORM_LEVEL.to_string(),
2717 "2".to_string(),
2718 ),
2719 ]);
2720 let tp = TableProperties::new(&raw_properties);
2721
2722 let tmp = TempDir::new().unwrap();
2723 let output = FileIO::new_with_fs()
2724 .new_output(format!("{}/cdc.parquet", tmp.path().to_str().unwrap()))
2725 .unwrap();
2726 let writer = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema())
2727 .unwrap()
2728 .build(output)
2729 .await
2730 .unwrap();
2731
2732 let cdc = writer
2733 .writer_properties
2734 .content_defined_chunking()
2735 .copied()
2736 .expect("CDC should be enabled on the built writer");
2737 assert_eq!(cdc.min_chunk_size, 4096);
2738 assert_eq!(cdc.max_chunk_size, 8192);
2739 assert_eq!(cdc.norm_level, 2);
2740 }
2741
2742 #[test]
2743 fn test_from_table_properties_sizing_defaults() {
2744 let raw_properties = HashMap::new();
2747 let tp = TableProperties::new(&raw_properties);
2748 let props = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema())
2749 .unwrap()
2750 .props;
2751
2752 assert_eq!(
2753 props.max_row_group_bytes(),
2754 Some(TableProperties::PROPERTY_PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT)
2755 );
2756 assert_eq!(
2757 props.data_page_size_limit(),
2758 TableProperties::PROPERTY_PARQUET_PAGE_SIZE_BYTES_DEFAULT
2759 );
2760 assert_eq!(
2761 props.data_page_row_count_limit(),
2762 TableProperties::PROPERTY_PARQUET_PAGE_ROW_LIMIT_DEFAULT
2763 );
2764 assert_eq!(
2765 props.dictionary_page_size_limit(),
2766 TableProperties::PROPERTY_PARQUET_DICT_SIZE_BYTES_DEFAULT
2767 );
2768 assert_eq!(
2770 props.compression(&ColumnPath::from("id")),
2771 Compression::ZSTD(ZstdLevel::try_new(3).unwrap())
2772 );
2773 }
2774
2775 #[test]
2776 fn test_from_table_properties_sizing_and_compression_overrides() {
2777 let raw_properties = HashMap::from([
2778 (
2779 TableProperties::PROPERTY_PARQUET_ROW_GROUP_SIZE_BYTES.to_string(),
2780 "1048576".to_string(),
2781 ),
2782 (
2783 TableProperties::PROPERTY_PARQUET_PAGE_SIZE_BYTES.to_string(),
2784 "65536".to_string(),
2785 ),
2786 (
2787 TableProperties::PROPERTY_PARQUET_PAGE_ROW_LIMIT.to_string(),
2788 "5000".to_string(),
2789 ),
2790 (
2791 TableProperties::PROPERTY_PARQUET_DICT_SIZE_BYTES.to_string(),
2792 "131072".to_string(),
2793 ),
2794 (
2795 TableProperties::PROPERTY_PARQUET_COMPRESSION_CODEC.to_string(),
2796 "gzip".to_string(),
2797 ),
2798 (
2799 TableProperties::PROPERTY_PARQUET_COMPRESSION_LEVEL.to_string(),
2800 "9".to_string(),
2801 ),
2802 ]);
2803 let tp = TableProperties::new(&raw_properties);
2804 let props = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema())
2805 .unwrap()
2806 .props;
2807
2808 assert_eq!(props.max_row_group_bytes(), Some(1048576));
2809 assert_eq!(props.data_page_size_limit(), 65536);
2810 assert_eq!(props.data_page_row_count_limit(), 5000);
2811 assert_eq!(props.dictionary_page_size_limit(), 131072);
2812 assert_eq!(
2813 props.compression(&ColumnPath::from("id")),
2814 Compression::GZIP(GzipLevel::try_new(9).unwrap())
2815 );
2816 }
2817
2818 #[test]
2819 fn test_from_table_properties_invalid_codec_errors() {
2820 let entries = HashMap::from([(
2821 TableProperties::PROPERTY_PARQUET_COMPRESSION_CODEC.to_string(),
2822 "bogus".to_string(),
2823 )]);
2824 let err = ParquetWriterBuilder::from_table_properties(
2825 &TableProperties::new(&entries),
2826 cdc_test_schema(),
2827 )
2828 .unwrap_err();
2829 assert_eq!(err.kind(), ErrorKind::DataInvalid);
2830 assert!(err.to_string().contains("bogus"));
2831 }
2832
2833 #[test]
2834 fn test_parquet_compression_mapping() {
2835 assert_eq!(
2837 parquet_compression(CompressionCodec::None).unwrap(),
2838 Compression::UNCOMPRESSED
2839 );
2840 assert_eq!(
2841 parquet_compression(CompressionCodec::Snappy).unwrap(),
2842 Compression::SNAPPY
2843 );
2844 assert_eq!(
2845 parquet_compression(CompressionCodec::Lz4).unwrap(),
2846 Compression::LZ4
2847 );
2848 assert_eq!(
2849 parquet_compression(CompressionCodec::Lz4Raw).unwrap(),
2850 Compression::LZ4_RAW
2851 );
2852 assert_eq!(
2853 parquet_compression(CompressionCodec::Lzo).unwrap(),
2854 Compression::LZO
2855 );
2856
2857 assert_eq!(
2859 parquet_compression(CompressionCodec::zstd_default()).unwrap(),
2860 Compression::ZSTD(ZstdLevel::try_new(3).unwrap())
2861 );
2862 assert_eq!(
2863 parquet_compression(CompressionCodec::gzip_default()).unwrap(),
2864 Compression::GZIP(GzipLevel::try_new(6).unwrap())
2865 );
2866 assert_eq!(
2867 parquet_compression(CompressionCodec::brotli_default()).unwrap(),
2868 Compression::BROTLI(BrotliLevel::try_new(1).unwrap())
2869 );
2870
2871 assert_eq!(
2873 parquet_compression(CompressionCodec::Zstd(10)).unwrap(),
2874 Compression::ZSTD(ZstdLevel::try_new(10).unwrap())
2875 );
2876 }
2877
2878 #[test]
2879 fn test_parquet_compression_invalid_level() {
2880 let err = parquet_compression(CompressionCodec::Zstd(99)).unwrap_err();
2882 assert_eq!(err.kind(), ErrorKind::DataInvalid);
2883 assert!(err.to_string().contains("zstd"));
2884 }
2885}