1use std::collections::HashMap;
21use std::sync::Arc;
22
23use arrow_array::RecordBatch;
24use arrow_schema::{DataType, Field, Schema as ArrowSchema};
25use bytes::Bytes;
26use futures::future::BoxFuture;
27use itertools::Itertools;
28use parquet::arrow::AsyncArrowWriter;
29use parquet::arrow::async_reader::AsyncFileReader;
30use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter;
31use parquet::file::metadata::ParquetMetaData;
32use parquet::file::properties::{CdcOptions, WriterProperties};
33use parquet::file::statistics::Statistics;
34
35use super::{FileWriter, FileWriterBuilder};
36use crate::arrow::{
37 ArrowFileReader, DEFAULT_MAP_FIELD_NAME, FieldMatchMode, NanValueCountVisitor,
38 get_parquet_stat_max_as_datum, get_parquet_stat_min_as_datum,
39};
40use crate::io::{FileIO, FileWrite, OutputFile};
41use crate::spec::{
42 DataContentType, DataFileBuilder, DataFileFormat, Datum, ListType, Literal, MapType,
43 NestedFieldRef, PartitionSpec, PrimitiveType, Schema, SchemaRef, SchemaVisitor, Struct,
44 StructType, TableMetadata, TableProperties, Type, visit_schema,
45};
46use crate::transform::create_transform_function;
47use crate::writer::{CurrentFileStatus, DataFile};
48use crate::{Error, ErrorKind, Result};
49
50#[derive(Clone, Debug)]
52pub struct ParquetWriterBuilder {
53 props: WriterProperties,
54 schema: SchemaRef,
55 match_mode: FieldMatchMode,
56 arrow_schema: Option<Arc<ArrowSchema>>,
60}
61
62impl ParquetWriterBuilder {
63 pub fn new(props: WriterProperties, schema: SchemaRef) -> Self {
70 Self::new_with_match_mode(props, schema, FieldMatchMode::Id)
71 }
72
73 pub fn new_with_match_mode(
75 props: WriterProperties,
76 schema: SchemaRef,
77 match_mode: FieldMatchMode,
78 ) -> Self {
79 Self {
80 props,
81 schema,
82 match_mode,
83 arrow_schema: None,
84 }
85 }
86
87 pub fn from_table_properties(table_props: &TableProperties, schema: SchemaRef) -> Self {
95 let cdc = table_props.cdc_enabled.then_some(CdcOptions {
96 min_chunk_size: table_props.cdc_min_chunk_size,
97 max_chunk_size: table_props.cdc_max_chunk_size,
98 norm_level: table_props.cdc_norm_level,
99 });
100 let props = WriterProperties::builder()
104 .set_content_defined_chunking(cdc)
105 .build();
106 Self::new_with_match_mode(props, schema, FieldMatchMode::Id)
107 }
108
109 pub fn with_match_mode(mut self, match_mode: FieldMatchMode) -> Self {
114 self.match_mode = match_mode;
115 self
116 }
117
118 pub fn with_arrow_schema(mut self, arrow_schema: Arc<ArrowSchema>) -> Result<Self> {
123 validate_arrow_schema_matches_iceberg(&arrow_schema, &self.schema)?;
124 self.arrow_schema = Some(arrow_schema);
125 Ok(self)
126 }
127}
128
129impl FileWriterBuilder for ParquetWriterBuilder {
130 type R = ParquetWriter;
131
132 async fn build(&self, output_file: OutputFile) -> Result<Self::R> {
133 Ok(ParquetWriter {
134 schema: self.schema.clone(),
135 arrow_schema: self.arrow_schema.clone(),
136 inner_writer: None,
137 writer_properties: self.props.clone(),
138 current_row_num: 0,
139 output_file,
140 nan_value_count_visitor: NanValueCountVisitor::new_with_match_mode(self.match_mode),
141 })
142 }
143}
144
145struct IndexByParquetPathName {
147 name_to_id: HashMap<String, i32>,
148
149 field_names: Vec<String>,
150
151 field_id: i32,
152}
153
154impl IndexByParquetPathName {
155 pub fn new() -> Self {
157 Self {
158 name_to_id: HashMap::new(),
159 field_names: Vec::new(),
160 field_id: 0,
161 }
162 }
163
164 pub fn get(&self, name: &str) -> Option<&i32> {
166 self.name_to_id.get(name)
167 }
168}
169
170impl Default for IndexByParquetPathName {
171 fn default() -> Self {
172 Self::new()
173 }
174}
175
176impl SchemaVisitor for IndexByParquetPathName {
177 type T = ();
178
179 fn before_struct_field(&mut self, field: &NestedFieldRef) -> Result<()> {
180 self.field_names.push(field.name.to_string());
181 self.field_id = field.id;
182 Ok(())
183 }
184
185 fn after_struct_field(&mut self, _field: &NestedFieldRef) -> Result<()> {
186 self.field_names.pop();
187 Ok(())
188 }
189
190 fn before_list_element(&mut self, field: &NestedFieldRef) -> Result<()> {
191 self.field_names.push(format!("list.{}", field.name));
192 self.field_id = field.id;
193 Ok(())
194 }
195
196 fn after_list_element(&mut self, _field: &NestedFieldRef) -> Result<()> {
197 self.field_names.pop();
198 Ok(())
199 }
200
201 fn before_map_key(&mut self, field: &NestedFieldRef) -> Result<()> {
202 self.field_names
203 .push(format!("{DEFAULT_MAP_FIELD_NAME}.key"));
204 self.field_id = field.id;
205 Ok(())
206 }
207
208 fn after_map_key(&mut self, _field: &NestedFieldRef) -> Result<()> {
209 self.field_names.pop();
210 Ok(())
211 }
212
213 fn before_map_value(&mut self, field: &NestedFieldRef) -> Result<()> {
214 self.field_names
215 .push(format!("{DEFAULT_MAP_FIELD_NAME}.value"));
216 self.field_id = field.id;
217 Ok(())
218 }
219
220 fn after_map_value(&mut self, _field: &NestedFieldRef) -> Result<()> {
221 self.field_names.pop();
222 Ok(())
223 }
224
225 fn schema(&mut self, _schema: &Schema, _value: Self::T) -> Result<Self::T> {
226 Ok(())
227 }
228
229 fn field(&mut self, _field: &NestedFieldRef, _value: Self::T) -> Result<Self::T> {
230 Ok(())
231 }
232
233 fn r#struct(&mut self, _struct: &StructType, _results: Vec<Self::T>) -> Result<Self::T> {
234 Ok(())
235 }
236
237 fn list(&mut self, _list: &ListType, _value: Self::T) -> Result<Self::T> {
238 Ok(())
239 }
240
241 fn map(&mut self, _map: &MapType, _key_value: Self::T, _value: Self::T) -> Result<Self::T> {
242 Ok(())
243 }
244
245 fn primitive(&mut self, _p: &PrimitiveType) -> Result<Self::T> {
246 let full_name = self.field_names.iter().map(String::as_str).join(".");
247 let field_id = self.field_id;
248 if let Some(existing_field_id) = self.name_to_id.get(full_name.as_str()) {
249 return Err(Error::new(
250 ErrorKind::DataInvalid,
251 format!(
252 "Invalid schema: multiple fields for name {full_name}: {field_id} and {existing_field_id}"
253 ),
254 ));
255 } else {
256 self.name_to_id.insert(full_name, field_id);
257 }
258
259 Ok(())
260 }
261}
262
263pub struct ParquetWriter {
265 schema: SchemaRef,
266 arrow_schema: Option<Arc<ArrowSchema>>,
269 output_file: OutputFile,
270 inner_writer: Option<AsyncArrowWriter<AsyncFileWriter>>,
271 writer_properties: WriterProperties,
272 current_row_num: usize,
273 nan_value_count_visitor: NanValueCountVisitor,
274}
275
276struct MinMaxColAggregator {
278 lower_bounds: HashMap<i32, Datum>,
279 upper_bounds: HashMap<i32, Datum>,
280 schema: SchemaRef,
281}
282
283impl MinMaxColAggregator {
284 fn new(schema: SchemaRef) -> Self {
286 Self {
287 lower_bounds: HashMap::new(),
288 upper_bounds: HashMap::new(),
289 schema,
290 }
291 }
292
293 fn update_state_min(&mut self, field_id: i32, datum: Datum) {
294 self.lower_bounds
295 .entry(field_id)
296 .and_modify(|e| {
297 if *e > datum {
298 *e = datum.clone()
299 }
300 })
301 .or_insert(datum);
302 }
303
304 fn update_state_max(&mut self, field_id: i32, datum: Datum) {
305 self.upper_bounds
306 .entry(field_id)
307 .and_modify(|e| {
308 if *e < datum {
309 *e = datum.clone()
310 }
311 })
312 .or_insert(datum);
313 }
314
315 fn update(&mut self, field_id: i32, value: Statistics) -> Result<()> {
317 let Some(ty) = self
318 .schema
319 .field_by_id(field_id)
320 .map(|f| f.field_type.as_ref())
321 else {
322 return Ok(());
325 };
326 let Type::Primitive(ty) = ty.clone() else {
327 return Err(Error::new(
328 ErrorKind::Unexpected,
329 format!("Composed type {ty} is not supported for min max aggregation."),
330 ));
331 };
332
333 if value.min_is_exact() {
334 let Some(min_datum) = get_parquet_stat_min_as_datum(&ty, &value)? else {
335 return Err(Error::new(
336 ErrorKind::Unexpected,
337 format!("Statistics {value} is not match with field type {ty}."),
338 ));
339 };
340
341 self.update_state_min(field_id, min_datum);
342 }
343
344 if value.max_is_exact() {
345 let Some(max_datum) = get_parquet_stat_max_as_datum(&ty, &value)? else {
346 return Err(Error::new(
347 ErrorKind::Unexpected,
348 format!("Statistics {value} is not match with field type {ty}."),
349 ));
350 };
351
352 self.update_state_max(field_id, max_datum);
353 }
354
355 Ok(())
356 }
357
358 fn produce(self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) {
360 (self.lower_bounds, self.upper_bounds)
361 }
362}
363
364impl ParquetWriter {
365 #[allow(dead_code)]
367 pub(crate) async fn parquet_files_to_data_files(
368 file_io: &FileIO,
369 file_paths: Vec<String>,
370 table_metadata: &TableMetadata,
371 ) -> Result<Vec<DataFile>> {
372 let mut data_files: Vec<DataFile> = Vec::new();
374
375 for file_path in file_paths {
376 let input_file = file_io.new_input(&file_path)?;
377 let file_metadata = input_file.metadata().await?;
378 let file_size_in_bytes = file_metadata.size as usize;
379 let reader = input_file.reader().await?;
380
381 let mut parquet_reader = ArrowFileReader::new(file_metadata, reader);
382 let parquet_metadata = parquet_reader.get_metadata(None).await.map_err(|err| {
383 Error::new(
384 ErrorKind::DataInvalid,
385 format!("Error reading Parquet metadata: {err}"),
386 )
387 })?;
388 let mut builder = ParquetWriter::parquet_to_data_file_builder(
389 table_metadata.current_schema().clone(),
390 parquet_metadata,
391 file_size_in_bytes,
392 file_path,
393 HashMap::new(),
395 )?;
396 builder.partition_spec_id(table_metadata.default_partition_spec_id());
397 let data_file = builder.build().unwrap();
398 data_files.push(data_file);
399 }
400
401 Ok(data_files)
402 }
403
404 pub(crate) fn parquet_to_data_file_builder(
406 schema: SchemaRef,
407 metadata: Arc<ParquetMetaData>,
408 written_size: usize,
409 file_path: String,
410 nan_value_counts: HashMap<i32, u64>,
411 ) -> Result<DataFileBuilder> {
412 let index_by_parquet_path = {
413 let mut visitor = IndexByParquetPathName::new();
414 visit_schema(&schema, &mut visitor)?;
415 visitor
416 };
417
418 let (column_sizes, value_counts, null_value_counts, (lower_bounds, upper_bounds)) = {
419 let mut per_col_size: HashMap<i32, u64> = HashMap::new();
420 let mut per_col_val_num: HashMap<i32, u64> = HashMap::new();
421 let mut per_col_null_val_num: HashMap<i32, u64> = HashMap::new();
422 let mut min_max_agg = MinMaxColAggregator::new(schema);
423
424 for row_group in metadata.row_groups() {
425 for column_chunk_metadata in row_group.columns() {
426 let parquet_path = column_chunk_metadata.column_descr().path().string();
427
428 let Some(&field_id) = index_by_parquet_path.get(&parquet_path) else {
429 continue;
430 };
431
432 *per_col_size.entry(field_id).or_insert(0) +=
433 column_chunk_metadata.compressed_size() as u64;
434 *per_col_val_num.entry(field_id).or_insert(0) +=
435 column_chunk_metadata.num_values() as u64;
436
437 if let Some(statistics) = column_chunk_metadata.statistics() {
438 if let Some(null_count) = statistics.null_count_opt() {
439 *per_col_null_val_num.entry(field_id).or_insert(0) += null_count;
440 }
441
442 min_max_agg.update(field_id, statistics.clone())?;
443 }
444 }
445 }
446 (
447 per_col_size,
448 per_col_val_num,
449 per_col_null_val_num,
450 min_max_agg.produce(),
451 )
452 };
453
454 let mut builder = DataFileBuilder::default();
455 builder
456 .content(DataContentType::Data)
457 .file_path(file_path)
458 .file_format(DataFileFormat::Parquet)
459 .partition(Struct::empty())
460 .record_count(metadata.file_metadata().num_rows() as u64)
461 .file_size_in_bytes(written_size as u64)
462 .column_sizes(column_sizes)
463 .value_counts(value_counts)
464 .null_value_counts(null_value_counts)
465 .nan_value_counts(nan_value_counts)
466 .lower_bounds(lower_bounds)
469 .upper_bounds(upper_bounds)
470 .split_offsets(Some(
471 metadata
472 .row_groups()
473 .iter()
474 .filter_map(|group| group.file_offset())
475 .collect(),
476 ));
477
478 Ok(builder)
479 }
480
481 #[allow(dead_code)]
482 fn partition_value_from_bounds(
483 table_spec: Arc<PartitionSpec>,
484 lower_bounds: &HashMap<i32, Datum>,
485 upper_bounds: &HashMap<i32, Datum>,
486 ) -> Result<Struct> {
487 let mut partition_literals: Vec<Option<Literal>> = Vec::new();
488
489 for field in table_spec.fields() {
490 if let (Some(lower), Some(upper)) = (
491 lower_bounds.get(&field.source_id),
492 upper_bounds.get(&field.source_id),
493 ) {
494 if !field.transform.preserves_order() {
495 return Err(Error::new(
496 ErrorKind::DataInvalid,
497 format!(
498 "cannot infer partition value for non linear partition field (needs to preserve order): {} with transform {}",
499 field.name, field.transform
500 ),
501 ));
502 }
503
504 if lower != upper {
505 return Err(Error::new(
506 ErrorKind::DataInvalid,
507 format!(
508 "multiple partition values for field {}: lower: {:?}, upper: {:?}",
509 field.name, lower, upper
510 ),
511 ));
512 }
513
514 let transform_fn = create_transform_function(&field.transform)?;
515 let transform_literal =
516 Literal::from(transform_fn.transform_literal_result(lower)?);
517
518 partition_literals.push(Some(transform_literal));
519 } else {
520 partition_literals.push(None);
521 }
522 }
523
524 let partition_struct = Struct::from_iter(partition_literals);
525
526 Ok(partition_struct)
527 }
528}
529
530fn validate_arrow_schema_matches_iceberg(
533 arrow_schema: &ArrowSchema,
534 iceberg_schema: &crate::spec::Schema,
535) -> Result<()> {
536 let expected: ArrowSchema = iceberg_schema.try_into()?;
537 validate_schemas_match(&expected, arrow_schema, "")
538}
539
540fn validate_schemas_match(expected: &ArrowSchema, actual: &ArrowSchema, path: &str) -> Result<()> {
541 if expected.fields().len() != actual.fields().len() {
542 return Err(Error::new(
543 ErrorKind::DataInvalid,
544 format!(
545 "Schema mismatch at '{}': expected {} fields, got {}",
546 path,
547 expected.fields().len(),
548 actual.fields().len()
549 ),
550 ));
551 }
552 for (e, a) in expected.fields().iter().zip(actual.fields().iter()) {
553 validate_fields_match(e, a, path)?;
554 }
555 Ok(())
556}
557
558fn validate_fields_match(expected: &Field, actual: &Field, parent_path: &str) -> Result<()> {
559 let path = if parent_path.is_empty() {
560 expected.name().to_string()
561 } else {
562 format!("{}.{}", parent_path, expected.name())
563 };
564
565 if expected.name() != actual.name() {
566 return Err(Error::new(
567 ErrorKind::DataInvalid,
568 format!(
569 "Field name mismatch at '{}': expected '{}', got '{}'",
570 parent_path,
571 expected.name(),
572 actual.name()
573 ),
574 ));
575 }
576
577 if expected.is_nullable() != actual.is_nullable() {
578 return Err(Error::new(
579 ErrorKind::DataInvalid,
580 format!(
581 "Nullability mismatch at '{}': expected {}, got {}",
582 path,
583 expected.is_nullable(),
584 actual.is_nullable()
585 ),
586 ));
587 }
588
589 validate_datatypes_match(expected.data_type(), actual.data_type(), &path)
590}
591
592fn validate_datatypes_match(expected: &DataType, actual: &DataType, path: &str) -> Result<()> {
593 match (expected, actual) {
595 (DataType::Struct(e_fields), DataType::Struct(a_fields)) => {
596 if e_fields.len() != a_fields.len() {
597 return Err(Error::new(
598 ErrorKind::DataInvalid,
599 format!(
600 "Struct field count mismatch at '{}': expected {}, got {}",
601 path,
602 e_fields.len(),
603 a_fields.len()
604 ),
605 ));
606 }
607 for (e, a) in e_fields.iter().zip(a_fields.iter()) {
608 validate_fields_match(e, a, path)?;
609 }
610 Ok(())
611 }
612 (DataType::List(e_inner), DataType::List(a_inner))
613 | (DataType::LargeList(e_inner), DataType::LargeList(a_inner)) => {
614 validate_fields_match(e_inner, a_inner, path)
615 }
616 (DataType::Map(e_entries, _), DataType::Map(a_entries, _)) => {
617 validate_fields_match(e_entries, a_entries, path)
618 }
619 _ => {
620 if std::mem::discriminant(expected) != std::mem::discriminant(actual) {
622 return Err(Error::new(
623 ErrorKind::DataInvalid,
624 format!(
625 "Type mismatch at '{}': expected {:?}, got {:?}",
626 path, expected, actual
627 ),
628 ));
629 }
630 Ok(())
631 }
632 }
633}
634
635impl FileWriter for ParquetWriter {
636 async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
637 if batch.num_rows() == 0 {
639 return Ok(());
640 }
641
642 self.current_row_num += batch.num_rows();
643
644 self.nan_value_count_visitor
645 .compute(self.schema.clone(), batch.clone())?;
646
647 let writer = if let Some(writer) = &mut self.inner_writer {
648 writer
649 } else {
650 let arrow_schema: Arc<ArrowSchema> = match &self.arrow_schema {
653 Some(schema) => Arc::clone(schema),
654 None => Arc::new(self.schema.as_ref().try_into()?),
655 };
656 let inner_writer = self.output_file.writer().await?;
657 let async_writer = AsyncFileWriter::new(inner_writer);
658 let writer = AsyncArrowWriter::try_new(
659 async_writer,
660 arrow_schema,
661 Some(self.writer_properties.clone()),
662 )
663 .map_err(|err| {
664 Error::new(ErrorKind::Unexpected, "Failed to build parquet writer.")
665 .with_source(err)
666 })?;
667 self.inner_writer = Some(writer);
668 self.inner_writer.as_mut().unwrap()
669 };
670
671 writer.write(batch).await.map_err(|err| {
672 Error::new(
673 ErrorKind::Unexpected,
674 "Failed to write using parquet writer.",
675 )
676 .with_source(err)
677 })?;
678
679 Ok(())
680 }
681
682 async fn close(mut self) -> Result<Vec<DataFileBuilder>> {
683 let mut writer = match self.inner_writer.take() {
684 Some(writer) => writer,
685 None => return Ok(vec![]),
686 };
687
688 let metadata = writer.finish().await.map_err(|err| {
689 Error::new(ErrorKind::Unexpected, "Failed to finish parquet writer.").with_source(err)
690 })?;
691
692 let written_size = writer.bytes_written();
693
694 if self.current_row_num == 0 {
695 self.output_file.delete().await.map_err(|err| {
696 Error::new(
697 ErrorKind::Unexpected,
698 "Failed to delete empty parquet file.",
699 )
700 .with_source(err)
701 })?;
702 Ok(vec![])
703 } else {
704 let parquet_metadata = Arc::new(metadata);
705
706 Ok(vec![Self::parquet_to_data_file_builder(
707 self.schema,
708 parquet_metadata,
709 written_size,
710 self.output_file.location().to_string(),
711 self.nan_value_count_visitor.nan_value_counts,
712 )?])
713 }
714 }
715}
716
717impl CurrentFileStatus for ParquetWriter {
718 fn current_file_path(&self) -> String {
719 self.output_file.location().to_string()
720 }
721
722 fn current_row_num(&self) -> usize {
723 self.current_row_num
724 }
725
726 fn current_written_size(&self) -> usize {
727 if let Some(inner) = self.inner_writer.as_ref() {
728 inner.bytes_written() + inner.in_progress_size()
731 } else {
732 0
734 }
735 }
736
737 fn current_schema(&self) -> SchemaRef {
738 self.schema.clone()
739 }
740}
741
742struct AsyncFileWriter(Box<dyn FileWrite>);
748
749impl AsyncFileWriter {
750 pub fn new(writer: Box<dyn FileWrite>) -> Self {
752 Self(writer)
753 }
754}
755
756impl ArrowAsyncFileWriter for AsyncFileWriter {
757 fn write(&mut self, bs: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> {
758 Box::pin(async {
759 self.0
760 .write(bs)
761 .await
762 .map_err(|err| parquet::errors::ParquetError::External(Box::new(err)))
763 })
764 }
765
766 fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> {
767 Box::pin(async {
768 self.0
769 .close()
770 .await
771 .map_err(|err| parquet::errors::ParquetError::External(Box::new(err)))
772 })
773 }
774}
775
776#[cfg(test)]
777mod tests {
778 use std::collections::HashMap;
779 use std::sync::Arc;
780
781 use anyhow::Result;
782 use arrow_array::builder::{Float32Builder, Int32Builder, MapBuilder};
783 use arrow_array::types::{Float32Type, Int64Type};
784 use arrow_array::{
785 Array, ArrayRef, BooleanArray, Decimal128Array, Float32Array, Float64Array, Int32Array,
786 Int64Array, ListArray, MapArray, RecordBatch, StructArray,
787 };
788 use arrow_schema::{DataType, Field, Fields, SchemaRef as ArrowSchemaRef};
789 use arrow_select::concat::concat_batches;
790 use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
791 use parquet::file::statistics::ValueStatistics;
792 use tempfile::TempDir;
793 use uuid::Uuid;
794
795 use super::*;
796 use crate::arrow::schema_to_arrow_schema;
797 use crate::io::FileIO;
798 use crate::spec::decimal_utils::{decimal_mantissa, decimal_new, decimal_scale};
799 use crate::spec::{PrimitiveLiteral, Struct, *};
800 use crate::writer::file_writer::location_generator::{
801 DefaultFileNameGenerator, DefaultLocationGenerator, FileNameGenerator, LocationGenerator,
802 };
803 use crate::writer::tests::check_parquet_data_file;
804
805 fn schema_for_all_type() -> Schema {
806 Schema::builder()
807 .with_schema_id(1)
808 .with_fields(vec![
809 NestedField::optional(0, "boolean", Type::Primitive(PrimitiveType::Boolean)).into(),
810 NestedField::optional(1, "int", Type::Primitive(PrimitiveType::Int)).into(),
811 NestedField::optional(2, "long", Type::Primitive(PrimitiveType::Long)).into(),
812 NestedField::optional(3, "float", Type::Primitive(PrimitiveType::Float)).into(),
813 NestedField::optional(4, "double", Type::Primitive(PrimitiveType::Double)).into(),
814 NestedField::optional(5, "string", Type::Primitive(PrimitiveType::String)).into(),
815 NestedField::optional(6, "binary", Type::Primitive(PrimitiveType::Binary)).into(),
816 NestedField::optional(7, "date", Type::Primitive(PrimitiveType::Date)).into(),
817 NestedField::optional(8, "time", Type::Primitive(PrimitiveType::Time)).into(),
818 NestedField::optional(9, "timestamp", Type::Primitive(PrimitiveType::Timestamp))
819 .into(),
820 NestedField::optional(
821 10,
822 "timestamptz",
823 Type::Primitive(PrimitiveType::Timestamptz),
824 )
825 .into(),
826 NestedField::optional(
827 11,
828 "timestamp_ns",
829 Type::Primitive(PrimitiveType::TimestampNs),
830 )
831 .into(),
832 NestedField::optional(
833 12,
834 "timestamptz_ns",
835 Type::Primitive(PrimitiveType::TimestamptzNs),
836 )
837 .into(),
838 NestedField::optional(
839 13,
840 "decimal",
841 Type::Primitive(PrimitiveType::Decimal {
842 precision: 10,
843 scale: 5,
844 }),
845 )
846 .into(),
847 NestedField::optional(14, "uuid", Type::Primitive(PrimitiveType::Uuid)).into(),
848 NestedField::optional(15, "fixed", Type::Primitive(PrimitiveType::Fixed(10)))
849 .into(),
850 NestedField::optional(
853 16,
854 "decimal_38",
855 Type::Primitive(PrimitiveType::Decimal {
856 precision: 38,
857 scale: 5,
858 }),
859 )
860 .into(),
861 ])
862 .build()
863 .unwrap()
864 }
865
866 fn nested_schema_for_test() -> Schema {
867 Schema::builder()
869 .with_schema_id(1)
870 .with_fields(vec![
871 NestedField::required(0, "col0", Type::Primitive(PrimitiveType::Long)).into(),
872 NestedField::required(
873 1,
874 "col1",
875 Type::Struct(StructType::new(vec![
876 NestedField::required(5, "col_1_5", Type::Primitive(PrimitiveType::Long))
877 .into(),
878 NestedField::required(6, "col_1_6", Type::Primitive(PrimitiveType::Long))
879 .into(),
880 ])),
881 )
882 .into(),
883 NestedField::required(2, "col2", Type::Primitive(PrimitiveType::String)).into(),
884 NestedField::required(
885 3,
886 "col3",
887 Type::List(ListType::new(
888 NestedField::required(7, "element", Type::Primitive(PrimitiveType::Long))
889 .into(),
890 )),
891 )
892 .into(),
893 NestedField::required(
894 4,
895 "col4",
896 Type::Struct(StructType::new(vec![
897 NestedField::required(
898 8,
899 "col_4_8",
900 Type::Struct(StructType::new(vec![
901 NestedField::required(
902 9,
903 "col_4_8_9",
904 Type::Primitive(PrimitiveType::Long),
905 )
906 .into(),
907 ])),
908 )
909 .into(),
910 ])),
911 )
912 .into(),
913 NestedField::required(
914 10,
915 "col5",
916 Type::Map(MapType::new(
917 NestedField::required(11, "key", Type::Primitive(PrimitiveType::String))
918 .into(),
919 NestedField::required(
920 12,
921 "value",
922 Type::List(ListType::new(
923 NestedField::required(
924 13,
925 "item",
926 Type::Primitive(PrimitiveType::Long),
927 )
928 .into(),
929 )),
930 )
931 .into(),
932 )),
933 )
934 .into(),
935 ])
936 .build()
937 .unwrap()
938 }
939
940 #[tokio::test]
941 async fn test_index_by_parquet_path() {
942 let expect = HashMap::from([
943 ("col0".to_string(), 0),
944 ("col1.col_1_5".to_string(), 5),
945 ("col1.col_1_6".to_string(), 6),
946 ("col2".to_string(), 2),
947 ("col3.list.element".to_string(), 7),
948 ("col4.col_4_8.col_4_8_9".to_string(), 9),
949 ("col5.key_value.key".to_string(), 11),
950 ("col5.key_value.value.list.item".to_string(), 13),
951 ]);
952 let mut visitor = IndexByParquetPathName::new();
953 visit_schema(&nested_schema_for_test(), &mut visitor).unwrap();
954 assert_eq!(visitor.name_to_id, expect);
955 }
956
957 #[tokio::test]
958 async fn test_parquet_writer() -> Result<()> {
959 let temp_dir = TempDir::new().unwrap();
960 let file_io = FileIO::new_with_fs();
961 let location_gen = DefaultLocationGenerator::with_data_location(
962 temp_dir.path().to_str().unwrap().to_string(),
963 );
964 let file_name_gen =
965 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
966
967 let schema = {
969 let fields =
970 vec![
971 Field::new("col", DataType::Int64, true).with_metadata(HashMap::from([(
972 PARQUET_FIELD_ID_META_KEY.to_string(),
973 "0".to_string(),
974 )])),
975 ];
976 Arc::new(arrow_schema::Schema::new(fields))
977 };
978 let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
979 let null_col = Arc::new(Int64Array::new_null(1024)) as ArrayRef;
980 let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
981 let to_write_null = RecordBatch::try_new(schema.clone(), vec![null_col]).unwrap();
982
983 let output_file = file_io.new_output(
984 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
985 )?;
986
987 let mut pw = ParquetWriterBuilder::new(
989 WriterProperties::builder()
990 .set_max_row_group_row_count(Some(128))
991 .build(),
992 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
993 )
994 .build(output_file)
995 .await?;
996 pw.write(&to_write).await?;
997 pw.write(&to_write_null).await?;
998 let res = pw.close().await?;
999 assert_eq!(res.len(), 1);
1000 let data_file = res
1001 .into_iter()
1002 .next()
1003 .unwrap()
1004 .content(DataContentType::Data)
1006 .partition(Struct::empty())
1007 .partition_spec_id(0)
1008 .build()
1009 .unwrap();
1010
1011 assert_eq!(data_file.record_count(), 2048);
1013 assert_eq!(*data_file.value_counts(), HashMap::from([(0, 2048)]));
1014 assert_eq!(
1015 *data_file.lower_bounds(),
1016 HashMap::from([(0, Datum::long(0))])
1017 );
1018 assert_eq!(
1019 *data_file.upper_bounds(),
1020 HashMap::from([(0, Datum::long(1023))])
1021 );
1022 assert_eq!(*data_file.null_value_counts(), HashMap::from([(0, 1024)]));
1023
1024 let expect_batch = concat_batches(&schema, vec![&to_write, &to_write_null]).unwrap();
1026 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
1027
1028 Ok(())
1029 }
1030
1031 #[tokio::test]
1032 async fn test_parquet_writer_with_complex_schema() -> Result<()> {
1033 let temp_dir = TempDir::new().unwrap();
1034 let file_io = FileIO::new_with_fs();
1035 let location_gen = DefaultLocationGenerator::with_data_location(
1036 temp_dir.path().to_str().unwrap().to_string(),
1037 );
1038 let file_name_gen =
1039 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1040
1041 let schema = nested_schema_for_test();
1043 let arrow_schema: ArrowSchemaRef = Arc::new((&schema).try_into().unwrap());
1044 let col0 = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
1045 let col1 = Arc::new(StructArray::new(
1046 {
1047 if let DataType::Struct(fields) = arrow_schema.field(1).data_type() {
1048 fields.clone()
1049 } else {
1050 unreachable!()
1051 }
1052 },
1053 vec![
1054 Arc::new(Int64Array::from_iter_values(0..1024)),
1055 Arc::new(Int64Array::from_iter_values(0..1024)),
1056 ],
1057 None,
1058 ));
1059 let col2 = Arc::new(arrow_array::StringArray::from_iter_values(
1060 (0..1024).map(|n| n.to_string()),
1061 )) as ArrayRef;
1062 let col3 = Arc::new({
1063 let list_parts = arrow_array::ListArray::from_iter_primitive::<Int64Type, _, _>(
1064 (0..1024).map(|n| Some(vec![Some(n)])),
1065 )
1066 .into_parts();
1067 arrow_array::ListArray::new(
1068 {
1069 if let DataType::List(field) = arrow_schema.field(3).data_type() {
1070 field.clone()
1071 } else {
1072 unreachable!()
1073 }
1074 },
1075 list_parts.1,
1076 list_parts.2,
1077 list_parts.3,
1078 )
1079 }) as ArrayRef;
1080 let col4 = Arc::new(StructArray::new(
1081 {
1082 if let DataType::Struct(fields) = arrow_schema.field(4).data_type() {
1083 fields.clone()
1084 } else {
1085 unreachable!()
1086 }
1087 },
1088 vec![Arc::new(StructArray::new(
1089 {
1090 if let DataType::Struct(fields) = arrow_schema.field(4).data_type() {
1091 if let DataType::Struct(fields) = fields[0].data_type() {
1092 fields.clone()
1093 } else {
1094 unreachable!()
1095 }
1096 } else {
1097 unreachable!()
1098 }
1099 },
1100 vec![Arc::new(Int64Array::from_iter_values(0..1024))],
1101 None,
1102 ))],
1103 None,
1104 ));
1105 let col5 = Arc::new({
1106 let mut map_array_builder = MapBuilder::new(
1107 None,
1108 arrow_array::builder::StringBuilder::new(),
1109 arrow_array::builder::ListBuilder::new(arrow_array::builder::PrimitiveBuilder::<
1110 Int64Type,
1111 >::new()),
1112 );
1113 for i in 0..1024 {
1114 map_array_builder.keys().append_value(i.to_string());
1115 map_array_builder
1116 .values()
1117 .append_value(vec![Some(i as i64); i + 1]);
1118 map_array_builder.append(true)?;
1119 }
1120 let (_, offset_buffer, struct_array, null_buffer, ordered) =
1121 map_array_builder.finish().into_parts();
1122 let struct_array = {
1123 let (_, mut arrays, nulls) = struct_array.into_parts();
1124 let list_array = {
1125 let list_array = arrays[1]
1126 .as_any()
1127 .downcast_ref::<ListArray>()
1128 .unwrap()
1129 .clone();
1130 let (_, offsets, array, nulls) = list_array.into_parts();
1131 let list_field = {
1132 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1133 if let DataType::Struct(fields) = map_field.data_type() {
1134 if let DataType::List(list_field) = fields[1].data_type() {
1135 list_field.clone()
1136 } else {
1137 unreachable!()
1138 }
1139 } else {
1140 unreachable!()
1141 }
1142 } else {
1143 unreachable!()
1144 }
1145 };
1146 ListArray::new(list_field, offsets, array, nulls)
1147 };
1148 arrays[1] = Arc::new(list_array) as ArrayRef;
1149 StructArray::new(
1150 {
1151 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1152 if let DataType::Struct(fields) = map_field.data_type() {
1153 fields.clone()
1154 } else {
1155 unreachable!()
1156 }
1157 } else {
1158 unreachable!()
1159 }
1160 },
1161 arrays,
1162 nulls,
1163 )
1164 };
1165 arrow_array::MapArray::new(
1166 {
1167 if let DataType::Map(map_field, _) = arrow_schema.field(5).data_type() {
1168 map_field.clone()
1169 } else {
1170 unreachable!()
1171 }
1172 },
1173 offset_buffer,
1174 struct_array,
1175 null_buffer,
1176 ordered,
1177 )
1178 }) as ArrayRef;
1179 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
1180 col0, col1, col2, col3, col4, col5,
1181 ])
1182 .unwrap();
1183 let output_file = file_io.new_output(
1184 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1185 )?;
1186
1187 let mut pw =
1189 ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(schema))
1190 .build(output_file)
1191 .await?;
1192 pw.write(&to_write).await?;
1193 let res = pw.close().await?;
1194 assert_eq!(res.len(), 1);
1195 let data_file = res
1196 .into_iter()
1197 .next()
1198 .unwrap()
1199 .content(crate::spec::DataContentType::Data)
1201 .partition(Struct::empty())
1202 .partition_spec_id(0)
1203 .build()
1204 .unwrap();
1205
1206 assert_eq!(data_file.record_count(), 1024);
1208 assert_eq!(
1209 *data_file.value_counts(),
1210 HashMap::from([
1211 (0, 1024),
1212 (5, 1024),
1213 (6, 1024),
1214 (2, 1024),
1215 (7, 1024),
1216 (9, 1024),
1217 (11, 1024),
1218 (13, (1..1025).sum()),
1219 ])
1220 );
1221 assert_eq!(
1222 *data_file.lower_bounds(),
1223 HashMap::from([
1224 (0, Datum::long(0)),
1225 (5, Datum::long(0)),
1226 (6, Datum::long(0)),
1227 (2, Datum::string("0")),
1228 (7, Datum::long(0)),
1229 (9, Datum::long(0)),
1230 (11, Datum::string("0")),
1231 (13, Datum::long(0))
1232 ])
1233 );
1234 assert_eq!(
1235 *data_file.upper_bounds(),
1236 HashMap::from([
1237 (0, Datum::long(1023)),
1238 (5, Datum::long(1023)),
1239 (6, Datum::long(1023)),
1240 (2, Datum::string("999")),
1241 (7, Datum::long(1023)),
1242 (9, Datum::long(1023)),
1243 (11, Datum::string("999")),
1244 (13, Datum::long(1023))
1245 ])
1246 );
1247
1248 check_parquet_data_file(&file_io, &data_file, &to_write).await;
1250
1251 Ok(())
1252 }
1253
1254 #[tokio::test]
1255 async fn test_all_type_for_write() -> Result<()> {
1256 let temp_dir = TempDir::new().unwrap();
1257 let file_io = FileIO::new_with_fs();
1258 let location_gen = DefaultLocationGenerator::with_data_location(
1259 temp_dir.path().to_str().unwrap().to_string(),
1260 );
1261 let file_name_gen =
1262 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1263
1264 let schema = schema_for_all_type();
1267 let arrow_schema: ArrowSchemaRef = Arc::new((&schema).try_into().unwrap());
1268 let col0 = Arc::new(BooleanArray::from(vec![
1269 Some(true),
1270 Some(false),
1271 None,
1272 Some(true),
1273 ])) as ArrayRef;
1274 let col1 = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])) as ArrayRef;
1275 let col2 = Arc::new(Int64Array::from(vec![Some(1), Some(2), None, Some(4)])) as ArrayRef;
1276 let col3 = Arc::new(arrow_array::Float32Array::from(vec![
1277 Some(0.5),
1278 Some(2.0),
1279 None,
1280 Some(3.5),
1281 ])) as ArrayRef;
1282 let col4 = Arc::new(arrow_array::Float64Array::from(vec![
1283 Some(0.5),
1284 Some(2.0),
1285 None,
1286 Some(3.5),
1287 ])) as ArrayRef;
1288 let col5 = Arc::new(arrow_array::StringArray::from(vec![
1289 Some("a"),
1290 Some("b"),
1291 None,
1292 Some("d"),
1293 ])) as ArrayRef;
1294 let col6 = Arc::new(arrow_array::LargeBinaryArray::from_opt_vec(vec![
1295 Some(b"one"),
1296 None,
1297 Some(b""),
1298 Some(b"zzzz"),
1299 ])) as ArrayRef;
1300 let col7 = Arc::new(arrow_array::Date32Array::from(vec![
1301 Some(0),
1302 Some(1),
1303 None,
1304 Some(3),
1305 ])) as ArrayRef;
1306 let col8 = Arc::new(arrow_array::Time64MicrosecondArray::from(vec![
1307 Some(0),
1308 Some(1),
1309 None,
1310 Some(3),
1311 ])) as ArrayRef;
1312 let col9 = Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![
1313 Some(0),
1314 Some(1),
1315 None,
1316 Some(3),
1317 ])) as ArrayRef;
1318 let col10 = Arc::new(
1319 arrow_array::TimestampMicrosecondArray::from(vec![Some(0), Some(1), None, Some(3)])
1320 .with_timezone_utc(),
1321 ) as ArrayRef;
1322 let col11 = Arc::new(arrow_array::TimestampNanosecondArray::from(vec![
1323 Some(0),
1324 Some(1),
1325 None,
1326 Some(3),
1327 ])) as ArrayRef;
1328 let col12 = Arc::new(
1329 arrow_array::TimestampNanosecondArray::from(vec![Some(0), Some(1), None, Some(3)])
1330 .with_timezone_utc(),
1331 ) as ArrayRef;
1332 let col13 = Arc::new(
1333 arrow_array::Decimal128Array::from(vec![Some(1), Some(2), None, Some(100)])
1334 .with_precision_and_scale(10, 5)
1335 .unwrap(),
1336 ) as ArrayRef;
1337 let col14 = Arc::new(
1338 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1339 vec![
1340 Some(Uuid::from_u128(0).as_bytes().to_vec()),
1341 Some(Uuid::from_u128(1).as_bytes().to_vec()),
1342 None,
1343 Some(Uuid::from_u128(3).as_bytes().to_vec()),
1344 ]
1345 .into_iter(),
1346 16,
1347 )
1348 .unwrap(),
1349 ) as ArrayRef;
1350 let col15 = Arc::new(
1351 arrow_array::FixedSizeBinaryArray::try_from_sparse_iter_with_size(
1352 vec![
1353 Some(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]),
1354 Some(vec![11, 12, 13, 14, 15, 16, 17, 18, 19, 20]),
1355 None,
1356 Some(vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30]),
1357 ]
1358 .into_iter(),
1359 10,
1360 )
1361 .unwrap(),
1362 ) as ArrayRef;
1363 let col16 = Arc::new(
1364 arrow_array::Decimal128Array::from(vec![Some(1), Some(2), None, Some(100)])
1365 .with_precision_and_scale(38, 5)
1366 .unwrap(),
1367 ) as ArrayRef;
1368 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
1369 col0, col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11, col12, col13,
1370 col14, col15, col16,
1371 ])
1372 .unwrap();
1373 let output_file = file_io.new_output(
1374 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1375 )?;
1376
1377 let mut pw =
1379 ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(schema))
1380 .build(output_file)
1381 .await?;
1382 pw.write(&to_write).await?;
1383 let res = pw.close().await?;
1384 assert_eq!(res.len(), 1);
1385 let data_file = res
1386 .into_iter()
1387 .next()
1388 .unwrap()
1389 .content(crate::spec::DataContentType::Data)
1391 .partition(Struct::empty())
1392 .partition_spec_id(0)
1393 .build()
1394 .unwrap();
1395
1396 assert_eq!(data_file.record_count(), 4);
1398 assert!(data_file.value_counts().iter().all(|(_, &v)| { v == 4 }));
1399 assert!(
1400 data_file
1401 .null_value_counts()
1402 .iter()
1403 .all(|(_, &v)| { v == 1 })
1404 );
1405 assert_eq!(
1406 *data_file.lower_bounds(),
1407 HashMap::from([
1408 (0, Datum::bool(false)),
1409 (1, Datum::int(1)),
1410 (2, Datum::long(1)),
1411 (3, Datum::float(0.5)),
1412 (4, Datum::double(0.5)),
1413 (5, Datum::string("a")),
1414 (6, Datum::binary(vec![])),
1415 (7, Datum::date(0)),
1416 (8, Datum::time_micros(0).unwrap()),
1417 (9, Datum::timestamp_micros(0)),
1418 (10, Datum::timestamptz_micros(0)),
1419 (11, Datum::timestamp_nanos(0)),
1420 (12, Datum::timestamptz_nanos(0)),
1421 (
1422 13,
1423 Datum::new(
1424 PrimitiveType::Decimal {
1425 precision: 10,
1426 scale: 5
1427 },
1428 PrimitiveLiteral::Int128(1)
1429 )
1430 ),
1431 (14, Datum::uuid(Uuid::from_u128(0))),
1432 (15, Datum::fixed(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
1433 (
1434 16,
1435 Datum::new(
1436 PrimitiveType::Decimal {
1437 precision: 38,
1438 scale: 5
1439 },
1440 PrimitiveLiteral::Int128(1)
1441 )
1442 ),
1443 ])
1444 );
1445 assert_eq!(
1446 *data_file.upper_bounds(),
1447 HashMap::from([
1448 (0, Datum::bool(true)),
1449 (1, Datum::int(4)),
1450 (2, Datum::long(4)),
1451 (3, Datum::float(3.5)),
1452 (4, Datum::double(3.5)),
1453 (5, Datum::string("d")),
1454 (6, Datum::binary(vec![122, 122, 122, 122])),
1455 (7, Datum::date(3)),
1456 (8, Datum::time_micros(3).unwrap()),
1457 (9, Datum::timestamp_micros(3)),
1458 (10, Datum::timestamptz_micros(3)),
1459 (11, Datum::timestamp_nanos(3)),
1460 (12, Datum::timestamptz_nanos(3)),
1461 (
1462 13,
1463 Datum::new(
1464 PrimitiveType::Decimal {
1465 precision: 10,
1466 scale: 5
1467 },
1468 PrimitiveLiteral::Int128(100)
1469 )
1470 ),
1471 (14, Datum::uuid(Uuid::from_u128(3))),
1472 (
1473 15,
1474 Datum::fixed(vec![21, 22, 23, 24, 25, 26, 27, 28, 29, 30])
1475 ),
1476 (
1477 16,
1478 Datum::new(
1479 PrimitiveType::Decimal {
1480 precision: 38,
1481 scale: 5
1482 },
1483 PrimitiveLiteral::Int128(100)
1484 )
1485 ),
1486 ])
1487 );
1488
1489 check_parquet_data_file(&file_io, &data_file, &to_write).await;
1491
1492 Ok(())
1493 }
1494
1495 #[tokio::test]
1496 async fn test_decimal_bound() -> Result<()> {
1497 let temp_dir = TempDir::new().unwrap();
1498 let file_io = FileIO::new_with_fs();
1499 let location_gen = DefaultLocationGenerator::with_data_location(
1500 temp_dir.path().to_str().unwrap().to_string(),
1501 );
1502 let file_name_gen =
1503 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1504
1505 let schema = Arc::new(
1507 Schema::builder()
1508 .with_fields(vec![
1509 NestedField::optional(
1510 0,
1511 "decimal",
1512 Type::Primitive(PrimitiveType::Decimal {
1513 precision: 28,
1514 scale: 10,
1515 }),
1516 )
1517 .into(),
1518 ])
1519 .build()
1520 .unwrap(),
1521 );
1522 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1523 let output_file = file_io.new_output(
1524 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1525 )?;
1526 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone())
1527 .build(output_file)
1528 .await?;
1529 let col0 = Arc::new(
1530 Decimal128Array::from(vec![Some(22000000000), Some(11000000000)])
1531 .with_data_type(DataType::Decimal128(28, 10)),
1532 ) as ArrayRef;
1533 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1534 pw.write(&to_write).await?;
1535 let res = pw.close().await?;
1536 assert_eq!(res.len(), 1);
1537 let data_file = res
1538 .into_iter()
1539 .next()
1540 .unwrap()
1541 .content(crate::spec::DataContentType::Data)
1542 .partition(Struct::empty())
1543 .partition_spec_id(0)
1544 .build()
1545 .unwrap();
1546 assert_eq!(
1547 data_file.upper_bounds().get(&0),
1548 Some(Datum::decimal_with_precision(decimal_new(22000000000_i64, 10), 28).unwrap())
1549 .as_ref()
1550 );
1551 assert_eq!(
1552 data_file.lower_bounds().get(&0),
1553 Some(Datum::decimal_with_precision(decimal_new(11000000000_i64, 10), 28).unwrap())
1554 .as_ref()
1555 );
1556
1557 let schema = Arc::new(
1559 Schema::builder()
1560 .with_fields(vec![
1561 NestedField::optional(
1562 0,
1563 "decimal",
1564 Type::Primitive(PrimitiveType::Decimal {
1565 precision: 28,
1566 scale: 10,
1567 }),
1568 )
1569 .into(),
1570 ])
1571 .build()
1572 .unwrap(),
1573 );
1574 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1575 let output_file = file_io.new_output(
1576 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1577 )?;
1578 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema.clone())
1579 .build(output_file)
1580 .await?;
1581 let col0 = Arc::new(
1582 Decimal128Array::from(vec![Some(-22000000000), Some(-11000000000)])
1583 .with_data_type(DataType::Decimal128(28, 10)),
1584 ) as ArrayRef;
1585 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1586 pw.write(&to_write).await?;
1587 let res = pw.close().await?;
1588 assert_eq!(res.len(), 1);
1589 let data_file = res
1590 .into_iter()
1591 .next()
1592 .unwrap()
1593 .content(crate::spec::DataContentType::Data)
1594 .partition(Struct::empty())
1595 .partition_spec_id(0)
1596 .build()
1597 .unwrap();
1598 assert_eq!(
1599 data_file.upper_bounds().get(&0),
1600 Some(Datum::decimal_with_precision(decimal_new(-11000000000_i64, 10), 28).unwrap())
1601 .as_ref()
1602 );
1603 assert_eq!(
1604 data_file.lower_bounds().get(&0),
1605 Some(Datum::decimal_with_precision(decimal_new(-22000000000_i64, 10), 28).unwrap())
1606 .as_ref()
1607 );
1608
1609 use crate::spec::decimal_utils::decimal_from_str_exact;
1612 let decimal_max = decimal_from_str_exact("99999999999999999999999999999999999999").unwrap();
1613 let decimal_min =
1614 decimal_from_str_exact("-99999999999999999999999999999999999999").unwrap();
1615 assert_eq!(decimal_scale(&decimal_max), decimal_scale(&decimal_min));
1616 let schema = Arc::new(
1617 Schema::builder()
1618 .with_fields(vec![
1619 NestedField::optional(
1620 0,
1621 "decimal",
1622 Type::Primitive(PrimitiveType::Decimal {
1623 precision: 38,
1624 scale: decimal_scale(&decimal_max),
1625 }),
1626 )
1627 .into(),
1628 ])
1629 .build()
1630 .unwrap(),
1631 );
1632 let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1633 let output_file = file_io.new_output(
1634 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1635 )?;
1636 let mut pw = ParquetWriterBuilder::new(WriterProperties::builder().build(), schema)
1637 .build(output_file)
1638 .await?;
1639 let col0 = Arc::new(
1640 Decimal128Array::from(vec![
1641 Some(decimal_mantissa(&decimal_max)),
1642 Some(decimal_mantissa(&decimal_min)),
1643 ])
1644 .with_data_type(DataType::Decimal128(38, 0)),
1645 ) as ArrayRef;
1646 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1647 pw.write(&to_write).await?;
1648 let res = pw.close().await?;
1649 assert_eq!(res.len(), 1);
1650 let data_file = res
1651 .into_iter()
1652 .next()
1653 .unwrap()
1654 .content(crate::spec::DataContentType::Data)
1655 .partition(Struct::empty())
1656 .partition_spec_id(0)
1657 .build()
1658 .unwrap();
1659 assert_eq!(
1660 data_file.upper_bounds().get(&0),
1661 Some(Datum::decimal(decimal_max).unwrap()).as_ref()
1662 );
1663 assert_eq!(
1664 data_file.lower_bounds().get(&0),
1665 Some(Datum::decimal(decimal_min).unwrap()).as_ref()
1666 );
1667
1668 Ok(())
1738 }
1739
1740 #[tokio::test]
1741 async fn test_empty_write() -> Result<()> {
1742 let temp_dir = TempDir::new().unwrap();
1743 let file_io = FileIO::new_with_fs();
1744 let location_gen = DefaultLocationGenerator::with_data_location(
1745 temp_dir.path().to_str().unwrap().to_string(),
1746 );
1747 let file_name_gen =
1748 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1749
1750 let schema = {
1752 let fields = vec![
1753 arrow_schema::Field::new("col", arrow_schema::DataType::Int64, true).with_metadata(
1754 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())]),
1755 ),
1756 ];
1757 Arc::new(arrow_schema::Schema::new(fields))
1758 };
1759 let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
1760 let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
1761 let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
1762 let output_file = file_io.new_output(&file_path)?;
1763 let mut pw = ParquetWriterBuilder::new(
1764 WriterProperties::builder().build(),
1765 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1766 )
1767 .build(output_file)
1768 .await?;
1769 pw.write(&to_write).await?;
1770 pw.close().await.unwrap();
1771 assert!(file_io.exists(&file_path).await.unwrap());
1772
1773 let file_name_gen =
1775 DefaultFileNameGenerator::new("test_empty".to_string(), None, DataFileFormat::Parquet);
1776 let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
1777 let output_file = file_io.new_output(&file_path)?;
1778 let pw = ParquetWriterBuilder::new(
1779 WriterProperties::builder().build(),
1780 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1781 )
1782 .build(output_file)
1783 .await?;
1784 pw.close().await.unwrap();
1785 assert!(!file_io.exists(&file_path).await.unwrap());
1786
1787 Ok(())
1788 }
1789
1790 #[tokio::test]
1791 async fn test_nan_val_cnts_primitive_type() -> Result<()> {
1792 let temp_dir = TempDir::new().unwrap();
1793 let file_io = FileIO::new_with_fs();
1794 let location_gen = DefaultLocationGenerator::with_data_location(
1795 temp_dir.path().to_str().unwrap().to_string(),
1796 );
1797 let file_name_gen =
1798 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
1799
1800 let arrow_schema = {
1802 let fields = vec![
1803 Field::new("col", arrow_schema::DataType::Float32, false).with_metadata(
1804 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())]),
1805 ),
1806 Field::new("col2", arrow_schema::DataType::Float64, false).with_metadata(
1807 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1808 ),
1809 ];
1810 Arc::new(arrow_schema::Schema::new(fields))
1811 };
1812
1813 let float_32_col = Arc::new(Float32Array::from_iter_values_with_nulls(
1814 [1.0_f32, f32::NAN, 2.0, 2.0].into_iter(),
1815 None,
1816 )) as ArrayRef;
1817
1818 let float_64_col = Arc::new(Float64Array::from_iter_values_with_nulls(
1819 [1.0_f64, f64::NAN, 2.0, 2.0].into_iter(),
1820 None,
1821 )) as ArrayRef;
1822
1823 let to_write =
1824 RecordBatch::try_new(arrow_schema.clone(), vec![float_32_col, float_64_col]).unwrap();
1825 let output_file = file_io.new_output(
1826 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1827 )?;
1828
1829 let mut pw = ParquetWriterBuilder::new(
1831 WriterProperties::builder().build(),
1832 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1833 )
1834 .build(output_file)
1835 .await?;
1836
1837 pw.write(&to_write).await?;
1838 let res = pw.close().await?;
1839 assert_eq!(res.len(), 1);
1840 let data_file = res
1841 .into_iter()
1842 .next()
1843 .unwrap()
1844 .content(crate::spec::DataContentType::Data)
1846 .partition(Struct::empty())
1847 .partition_spec_id(0)
1848 .build()
1849 .unwrap();
1850
1851 assert_eq!(data_file.record_count(), 4);
1853 assert_eq!(*data_file.value_counts(), HashMap::from([(0, 4), (1, 4)]));
1854 assert_eq!(
1855 *data_file.lower_bounds(),
1856 HashMap::from([(0, Datum::float(1.0)), (1, Datum::double(1.0)),])
1857 );
1858 assert_eq!(
1859 *data_file.upper_bounds(),
1860 HashMap::from([(0, Datum::float(2.0)), (1, Datum::double(2.0)),])
1861 );
1862 assert_eq!(
1863 *data_file.null_value_counts(),
1864 HashMap::from([(0, 0), (1, 0)])
1865 );
1866 assert_eq!(
1867 *data_file.nan_value_counts(),
1868 HashMap::from([(0, 1), (1, 1)])
1869 );
1870
1871 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
1873 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
1874
1875 Ok(())
1876 }
1877
1878 #[tokio::test]
1879 async fn test_nan_val_cnts_struct_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 schema_struct_float_fields = Fields::from(vec![
1889 Field::new("col4", DataType::Float32, false).with_metadata(HashMap::from([(
1890 PARQUET_FIELD_ID_META_KEY.to_string(),
1891 "4".to_string(),
1892 )])),
1893 ]);
1894
1895 let schema_struct_nested_float_fields = Fields::from(vec![
1896 Field::new("col7", DataType::Float32, false).with_metadata(HashMap::from([(
1897 PARQUET_FIELD_ID_META_KEY.to_string(),
1898 "7".to_string(),
1899 )])),
1900 ]);
1901
1902 let schema_struct_nested_fields = Fields::from(vec![
1903 Field::new(
1904 "col6",
1905 arrow_schema::DataType::Struct(schema_struct_nested_float_fields.clone()),
1906 false,
1907 )
1908 .with_metadata(HashMap::from([(
1909 PARQUET_FIELD_ID_META_KEY.to_string(),
1910 "6".to_string(),
1911 )])),
1912 ]);
1913
1914 let arrow_schema = {
1916 let fields = vec![
1917 Field::new(
1918 "col3",
1919 arrow_schema::DataType::Struct(schema_struct_float_fields.clone()),
1920 false,
1921 )
1922 .with_metadata(HashMap::from([(
1923 PARQUET_FIELD_ID_META_KEY.to_string(),
1924 "3".to_string(),
1925 )])),
1926 Field::new(
1927 "col5",
1928 arrow_schema::DataType::Struct(schema_struct_nested_fields.clone()),
1929 false,
1930 )
1931 .with_metadata(HashMap::from([(
1932 PARQUET_FIELD_ID_META_KEY.to_string(),
1933 "5".to_string(),
1934 )])),
1935 ];
1936 Arc::new(arrow_schema::Schema::new(fields))
1937 };
1938
1939 let float_32_col = Arc::new(Float32Array::from_iter_values_with_nulls(
1940 [1.0_f32, f32::NAN, 2.0, 2.0].into_iter(),
1941 None,
1942 )) as ArrayRef;
1943
1944 let struct_float_field_col = Arc::new(StructArray::new(
1945 schema_struct_float_fields,
1946 vec![float_32_col.clone()],
1947 None,
1948 )) as ArrayRef;
1949
1950 let struct_nested_float_field_col = Arc::new(StructArray::new(
1951 schema_struct_nested_fields,
1952 vec![Arc::new(StructArray::new(
1953 schema_struct_nested_float_fields,
1954 vec![float_32_col.clone()],
1955 None,
1956 )) as ArrayRef],
1957 None,
1958 )) as ArrayRef;
1959
1960 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
1961 struct_float_field_col,
1962 struct_nested_float_field_col,
1963 ])
1964 .unwrap();
1965 let output_file = file_io.new_output(
1966 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
1967 )?;
1968
1969 let mut pw = ParquetWriterBuilder::new(
1971 WriterProperties::builder().build(),
1972 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
1973 )
1974 .build(output_file)
1975 .await?;
1976
1977 pw.write(&to_write).await?;
1978 let res = pw.close().await?;
1979 assert_eq!(res.len(), 1);
1980 let data_file = res
1981 .into_iter()
1982 .next()
1983 .unwrap()
1984 .content(crate::spec::DataContentType::Data)
1986 .partition(Struct::empty())
1987 .partition_spec_id(0)
1988 .build()
1989 .unwrap();
1990
1991 assert_eq!(data_file.record_count(), 4);
1993 assert_eq!(*data_file.value_counts(), HashMap::from([(4, 4), (7, 4)]));
1994 assert_eq!(
1995 *data_file.lower_bounds(),
1996 HashMap::from([(4, Datum::float(1.0)), (7, Datum::float(1.0)),])
1997 );
1998 assert_eq!(
1999 *data_file.upper_bounds(),
2000 HashMap::from([(4, Datum::float(2.0)), (7, Datum::float(2.0)),])
2001 );
2002 assert_eq!(
2003 *data_file.null_value_counts(),
2004 HashMap::from([(4, 0), (7, 0)])
2005 );
2006 assert_eq!(
2007 *data_file.nan_value_counts(),
2008 HashMap::from([(4, 1), (7, 1)])
2009 );
2010
2011 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2013 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2014
2015 Ok(())
2016 }
2017
2018 #[tokio::test]
2019 async fn test_nan_val_cnts_list_type() -> Result<()> {
2020 let temp_dir = TempDir::new().unwrap();
2021 let file_io = FileIO::new_with_fs();
2022 let location_gen = DefaultLocationGenerator::with_data_location(
2023 temp_dir.path().to_str().unwrap().to_string(),
2024 );
2025 let file_name_gen =
2026 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2027
2028 let schema_list_float_field = Field::new("element", DataType::Float32, true).with_metadata(
2029 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
2030 );
2031
2032 let schema_struct_list_float_field = Field::new("element", DataType::Float32, true)
2033 .with_metadata(HashMap::from([(
2034 PARQUET_FIELD_ID_META_KEY.to_string(),
2035 "4".to_string(),
2036 )]));
2037
2038 let schema_struct_list_field = Fields::from(vec![
2039 Field::new_list("col2", schema_struct_list_float_field.clone(), true).with_metadata(
2040 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())]),
2041 ),
2042 ]);
2043
2044 let arrow_schema = {
2045 let fields = vec![
2046 Field::new_list("col0", schema_list_float_field.clone(), true).with_metadata(
2047 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())]),
2048 ),
2049 Field::new_struct("col1", schema_struct_list_field.clone(), true)
2050 .with_metadata(HashMap::from([(
2051 PARQUET_FIELD_ID_META_KEY.to_string(),
2052 "2".to_string(),
2053 )]))
2054 .clone(),
2055 ];
2059 Arc::new(arrow_schema::Schema::new(fields))
2060 };
2061
2062 let list_parts = ListArray::from_iter_primitive::<Float32Type, _, _>(vec![Some(vec![
2063 Some(1.0_f32),
2064 Some(f32::NAN),
2065 Some(2.0),
2066 Some(2.0),
2067 ])])
2068 .into_parts();
2069
2070 let list_float_field_col = Arc::new({
2071 let list_parts = list_parts.clone();
2072 ListArray::new(
2073 {
2074 if let DataType::List(field) = arrow_schema.field(0).data_type() {
2075 field.clone()
2076 } else {
2077 unreachable!()
2078 }
2079 },
2080 list_parts.1,
2081 list_parts.2,
2082 list_parts.3,
2083 )
2084 }) as ArrayRef;
2085
2086 let struct_list_fields_schema =
2087 if let DataType::Struct(fields) = arrow_schema.field(1).data_type() {
2088 fields.clone()
2089 } else {
2090 unreachable!()
2091 };
2092
2093 let struct_list_float_field_col = Arc::new({
2094 ListArray::new(
2095 {
2096 if let DataType::List(field) = struct_list_fields_schema
2097 .first()
2098 .expect("could not find first list field")
2099 .data_type()
2100 {
2101 field.clone()
2102 } else {
2103 unreachable!()
2104 }
2105 },
2106 list_parts.1,
2107 list_parts.2,
2108 list_parts.3,
2109 )
2110 }) as ArrayRef;
2111
2112 let struct_list_float_field_col = Arc::new(StructArray::new(
2113 struct_list_fields_schema,
2114 vec![struct_list_float_field_col.clone()],
2115 None,
2116 )) as ArrayRef;
2117
2118 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
2119 list_float_field_col,
2120 struct_list_float_field_col,
2121 ])
2123 .expect("Could not form record batch");
2124 let output_file = file_io.new_output(
2125 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
2126 )?;
2127
2128 let mut pw = ParquetWriterBuilder::new(
2130 WriterProperties::builder().build(),
2131 Arc::new(
2132 to_write
2133 .schema()
2134 .as_ref()
2135 .try_into()
2136 .expect("Could not convert iceberg schema"),
2137 ),
2138 )
2139 .build(output_file)
2140 .await?;
2141
2142 pw.write(&to_write).await?;
2143 let res = pw.close().await?;
2144 assert_eq!(res.len(), 1);
2145 let data_file = res
2146 .into_iter()
2147 .next()
2148 .unwrap()
2149 .content(crate::spec::DataContentType::Data)
2150 .partition(Struct::empty())
2151 .partition_spec_id(0)
2152 .build()
2153 .unwrap();
2154
2155 assert_eq!(data_file.record_count(), 1);
2157 assert_eq!(*data_file.value_counts(), HashMap::from([(1, 4), (4, 4)]));
2158 assert_eq!(
2159 *data_file.lower_bounds(),
2160 HashMap::from([(1, Datum::float(1.0)), (4, Datum::float(1.0))])
2161 );
2162 assert_eq!(
2163 *data_file.upper_bounds(),
2164 HashMap::from([(1, Datum::float(2.0)), (4, Datum::float(2.0))])
2165 );
2166 assert_eq!(
2167 *data_file.null_value_counts(),
2168 HashMap::from([(1, 0), (4, 0)])
2169 );
2170 assert_eq!(
2171 *data_file.nan_value_counts(),
2172 HashMap::from([(1, 1), (4, 1)])
2173 );
2174
2175 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2177 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2178
2179 Ok(())
2180 }
2181
2182 macro_rules! construct_map_arr {
2183 ($map_key_field_schema:ident, $map_value_field_schema:ident) => {{
2184 let int_builder = Int32Builder::new();
2185 let float_builder = Float32Builder::with_capacity(4);
2186 let mut builder = MapBuilder::new(None, int_builder, float_builder);
2187 builder.keys().append_value(1);
2188 builder.values().append_value(1.0_f32);
2189 builder.append(true).unwrap();
2190 builder.keys().append_value(2);
2191 builder.values().append_value(f32::NAN);
2192 builder.append(true).unwrap();
2193 builder.keys().append_value(3);
2194 builder.values().append_value(2.0);
2195 builder.append(true).unwrap();
2196 builder.keys().append_value(4);
2197 builder.values().append_value(2.0);
2198 builder.append(true).unwrap();
2199 let array = builder.finish();
2200
2201 let (_field, offsets, entries, nulls, ordered) = array.into_parts();
2202 let new_struct_fields_schema =
2203 Fields::from(vec![$map_key_field_schema, $map_value_field_schema]);
2204
2205 let entries = {
2206 let (_, arrays, nulls) = entries.into_parts();
2207 StructArray::new(new_struct_fields_schema.clone(), arrays, nulls)
2208 };
2209
2210 let field = Arc::new(Field::new(
2211 DEFAULT_MAP_FIELD_NAME,
2212 DataType::Struct(new_struct_fields_schema),
2213 false,
2214 ));
2215
2216 Arc::new(MapArray::new(field, offsets, entries, nulls, ordered))
2217 }};
2218 }
2219
2220 #[tokio::test]
2221 async fn test_nan_val_cnts_map_type() -> Result<()> {
2222 let temp_dir = TempDir::new().unwrap();
2223 let file_io = FileIO::new_with_fs();
2224 let location_gen = DefaultLocationGenerator::with_data_location(
2225 temp_dir.path().to_str().unwrap().to_string(),
2226 );
2227 let file_name_gen =
2228 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2229
2230 let map_key_field_schema =
2231 Field::new(MAP_KEY_FIELD_NAME, DataType::Int32, false).with_metadata(HashMap::from([
2232 (PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string()),
2233 ]));
2234
2235 let map_value_field_schema =
2236 Field::new(MAP_VALUE_FIELD_NAME, DataType::Float32, true).with_metadata(HashMap::from(
2237 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
2238 ));
2239
2240 let struct_map_key_field_schema =
2241 Field::new(MAP_KEY_FIELD_NAME, DataType::Int32, false).with_metadata(HashMap::from([
2242 (PARQUET_FIELD_ID_META_KEY.to_string(), "6".to_string()),
2243 ]));
2244
2245 let struct_map_value_field_schema =
2246 Field::new(MAP_VALUE_FIELD_NAME, DataType::Float32, true).with_metadata(HashMap::from(
2247 [(PARQUET_FIELD_ID_META_KEY.to_string(), "7".to_string())],
2248 ));
2249
2250 let schema_struct_map_field = Fields::from(vec![
2251 Field::new_map(
2252 "col3",
2253 DEFAULT_MAP_FIELD_NAME,
2254 struct_map_key_field_schema.clone(),
2255 struct_map_value_field_schema.clone(),
2256 false,
2257 false,
2258 )
2259 .with_metadata(HashMap::from([(
2260 PARQUET_FIELD_ID_META_KEY.to_string(),
2261 "5".to_string(),
2262 )])),
2263 ]);
2264
2265 let arrow_schema = {
2266 let fields = vec![
2267 Field::new_map(
2268 "col0",
2269 DEFAULT_MAP_FIELD_NAME,
2270 map_key_field_schema.clone(),
2271 map_value_field_schema.clone(),
2272 false,
2273 false,
2274 )
2275 .with_metadata(HashMap::from([(
2276 PARQUET_FIELD_ID_META_KEY.to_string(),
2277 "0".to_string(),
2278 )])),
2279 Field::new_struct("col1", schema_struct_map_field.clone(), true)
2280 .with_metadata(HashMap::from([(
2281 PARQUET_FIELD_ID_META_KEY.to_string(),
2282 "3".to_string(),
2283 )]))
2284 .clone(),
2285 ];
2286 Arc::new(arrow_schema::Schema::new(fields))
2287 };
2288
2289 let map_array = construct_map_arr!(map_key_field_schema, map_value_field_schema);
2290
2291 let struct_map_arr =
2292 construct_map_arr!(struct_map_key_field_schema, struct_map_value_field_schema);
2293
2294 let struct_list_float_field_col = Arc::new(StructArray::new(
2295 schema_struct_map_field,
2296 vec![struct_map_arr],
2297 None,
2298 )) as ArrayRef;
2299
2300 let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![
2301 map_array,
2302 struct_list_float_field_col,
2303 ])
2304 .expect("Could not form record batch");
2305 let output_file = file_io.new_output(
2306 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
2307 )?;
2308
2309 let mut pw = ParquetWriterBuilder::new(
2311 WriterProperties::builder().build(),
2312 Arc::new(
2313 to_write
2314 .schema()
2315 .as_ref()
2316 .try_into()
2317 .expect("Could not convert iceberg schema"),
2318 ),
2319 )
2320 .build(output_file)
2321 .await?;
2322
2323 pw.write(&to_write).await?;
2324 let res = pw.close().await?;
2325 assert_eq!(res.len(), 1);
2326 let data_file = res
2327 .into_iter()
2328 .next()
2329 .unwrap()
2330 .content(crate::spec::DataContentType::Data)
2331 .partition(Struct::empty())
2332 .partition_spec_id(0)
2333 .build()
2334 .unwrap();
2335
2336 assert_eq!(data_file.record_count(), 4);
2338 assert_eq!(
2339 *data_file.value_counts(),
2340 HashMap::from([(1, 4), (2, 4), (6, 4), (7, 4)])
2341 );
2342 assert_eq!(
2343 *data_file.lower_bounds(),
2344 HashMap::from([
2345 (1, Datum::int(1)),
2346 (2, Datum::float(1.0)),
2347 (6, Datum::int(1)),
2348 (7, Datum::float(1.0))
2349 ])
2350 );
2351 assert_eq!(
2352 *data_file.upper_bounds(),
2353 HashMap::from([
2354 (1, Datum::int(4)),
2355 (2, Datum::float(2.0)),
2356 (6, Datum::int(4)),
2357 (7, Datum::float(2.0))
2358 ])
2359 );
2360 assert_eq!(
2361 *data_file.null_value_counts(),
2362 HashMap::from([(1, 0), (2, 0), (6, 0), (7, 0)])
2363 );
2364 assert_eq!(
2365 *data_file.nan_value_counts(),
2366 HashMap::from([(2, 1), (7, 1)])
2367 );
2368
2369 let expect_batch = concat_batches(&arrow_schema, vec![&to_write]).unwrap();
2371 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
2372
2373 Ok(())
2374 }
2375
2376 #[tokio::test]
2377 async fn test_write_empty_parquet_file() {
2378 let temp_dir = TempDir::new().unwrap();
2379 let file_io = FileIO::new_with_fs();
2380 let location_gen = DefaultLocationGenerator::with_data_location(
2381 temp_dir.path().to_str().unwrap().to_string(),
2382 );
2383 let file_name_gen =
2384 DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
2385 let output_file = file_io
2386 .new_output(location_gen.generate_location(None, &file_name_gen.generate_file_name()))
2387 .unwrap();
2388
2389 let pw = ParquetWriterBuilder::new(
2391 WriterProperties::builder().build(),
2392 Arc::new(
2393 Schema::builder()
2394 .with_schema_id(1)
2395 .with_fields(vec![
2396 NestedField::required(0, "col", Type::Primitive(PrimitiveType::Long))
2397 .with_id(0)
2398 .into(),
2399 ])
2400 .build()
2401 .expect("Failed to create schema"),
2402 ),
2403 )
2404 .build(output_file)
2405 .await
2406 .unwrap();
2407
2408 let res = pw.close().await.unwrap();
2409 assert_eq!(res.len(), 0);
2410
2411 assert_eq!(std::fs::read_dir(temp_dir.path()).unwrap().count(), 0);
2413 }
2414
2415 #[test]
2416 fn test_min_max_aggregator() {
2417 let schema = Arc::new(
2418 Schema::builder()
2419 .with_schema_id(1)
2420 .with_fields(vec![
2421 NestedField::required(0, "col", Type::Primitive(PrimitiveType::Int))
2422 .with_id(0)
2423 .into(),
2424 ])
2425 .build()
2426 .expect("Failed to create schema"),
2427 );
2428 let mut min_max_agg = MinMaxColAggregator::new(schema);
2429 let create_statistics =
2430 |min, max| Statistics::Int32(ValueStatistics::new(min, max, None, None, false));
2431 min_max_agg
2432 .update(0, create_statistics(None, Some(42)))
2433 .unwrap();
2434 min_max_agg
2435 .update(0, create_statistics(Some(0), Some(i32::MAX)))
2436 .unwrap();
2437 min_max_agg
2438 .update(0, create_statistics(Some(i32::MIN), None))
2439 .unwrap();
2440 min_max_agg
2441 .update(0, create_statistics(None, None))
2442 .unwrap();
2443
2444 let (lower_bounds, upper_bounds) = min_max_agg.produce();
2445
2446 assert_eq!(lower_bounds, HashMap::from([(0, Datum::int(i32::MIN))]));
2447 assert_eq!(upper_bounds, HashMap::from([(0, Datum::int(i32::MAX))]));
2448 }
2449
2450 fn cdc_test_schema() -> SchemaRef {
2455 Arc::new(
2456 Schema::builder()
2457 .with_schema_id(1)
2458 .with_fields(vec![
2459 NestedField::required(1, "id", Type::Primitive(PrimitiveType::Long)).into(),
2460 NestedField::required(2, "payload", Type::Primitive(PrimitiveType::String))
2461 .into(),
2462 ])
2463 .build()
2464 .unwrap(),
2465 )
2466 }
2467
2468 fn table_props(entries: HashMap<String, String>) -> TableProperties {
2469 TableProperties::try_from(&entries).unwrap()
2470 }
2471
2472 #[test]
2473 fn test_from_table_properties_no_cdc_by_default() {
2474 let tp = table_props(HashMap::new());
2475 let builder = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema());
2476 assert!(builder.props.content_defined_chunking().is_none());
2477 }
2478
2479 #[tokio::test]
2480 async fn test_from_table_properties_propagate_to_writer() {
2481 let tp = table_props(HashMap::from([
2490 (
2491 TableProperties::PROPERTY_PARQUET_CDC_ENABLED.to_string(),
2492 "true".to_string(),
2493 ),
2494 (
2495 TableProperties::PROPERTY_PARQUET_CDC_MIN_CHUNK_SIZE.to_string(),
2496 "4096".to_string(),
2497 ),
2498 (
2499 TableProperties::PROPERTY_PARQUET_CDC_MAX_CHUNK_SIZE.to_string(),
2500 "8192".to_string(),
2501 ),
2502 (
2503 TableProperties::PROPERTY_PARQUET_CDC_NORM_LEVEL.to_string(),
2504 "2".to_string(),
2505 ),
2506 ]));
2507
2508 let tmp = TempDir::new().unwrap();
2509 let output = FileIO::new_with_fs()
2510 .new_output(format!("{}/cdc.parquet", tmp.path().to_str().unwrap()))
2511 .unwrap();
2512 let writer = ParquetWriterBuilder::from_table_properties(&tp, cdc_test_schema())
2513 .build(output)
2514 .await
2515 .unwrap();
2516
2517 let cdc = writer
2518 .writer_properties
2519 .content_defined_chunking()
2520 .copied()
2521 .expect("CDC should be enabled on the built writer");
2522 assert_eq!(cdc.min_chunk_size, 4096);
2523 assert_eq!(cdc.max_chunk_size, 8192);
2524 assert_eq!(cdc.norm_level, 2);
2525 }
2526}