Skip to main content

iceberg/writer/file_writer/
parquet_writer.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! The module contains the file writer for parquet file format.
19
20use 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/// ParquetWriterBuilder is used to builder a [`ParquetWriter`]
51#[derive(Clone, Debug)]
52pub struct ParquetWriterBuilder {
53    props: WriterProperties,
54    schema: SchemaRef,
55    match_mode: FieldMatchMode,
56    /// Optional Arrow schema to use for writing. If provided, this schema is used
57    /// directly for the Parquet writer instead of converting from the Iceberg schema.
58    /// This allows preserving Arrow extension metadata in the output.
59    arrow_schema: Option<Arc<ArrowSchema>>,
60}
61
62impl ParquetWriterBuilder {
63    /// Create a new `ParquetWriterBuilder`
64    /// To construct the write result, the schema should contain the `PARQUET_FIELD_ID_META_KEY` metadata for each field.
65    ///
66    /// When writing into an existing Iceberg table, prefer
67    /// [`Self::from_table_properties`], which derives `WriterProperties` from
68    /// the table's `write.parquet.*` properties.
69    pub fn new(props: WriterProperties, schema: SchemaRef) -> Self {
70        Self::new_with_match_mode(props, schema, FieldMatchMode::Id)
71    }
72
73    /// Create a new `ParquetWriterBuilder` with custom match mode
74    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    /// Build a `ParquetWriterBuilder` from Iceberg table properties and a
88    /// schema, translating `write.parquet.*` settings into `WriterProperties`
89    /// instead of using parquet-rs defaults.
90    ///
91    /// Currently translates the content-defined-chunking keys
92    /// (`write.parquet.content-defined-chunking.*`); other keys fall back to
93    /// parquet-rs defaults.
94    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        // TODO: translate the remaining write.parquet.* keys (e.g. compression-codec,
101        // row-group-size-bytes, page-size-bytes).
102        // This constructor is intended to be the single place that maps them.
103        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    /// Set the field match mode used to map Arrow fields to Iceberg fields.
110    ///
111    /// Defaults to [`FieldMatchMode::Id`]. Use [`FieldMatchMode::Name`] when the
112    /// incoming Arrow schema does not carry Iceberg field-id metadata.
113    pub fn with_match_mode(mut self, match_mode: FieldMatchMode) -> Self {
114        self.match_mode = match_mode;
115        self
116    }
117
118    /// Set an Arrow schema to use for writing instead of deriving one from the Iceberg schema.
119    /// The Arrow schema's structure must match the Iceberg schema (field names, types, nullability),
120    /// but may contain additional metadata (e.g., Arrow extension types).
121    /// Returns an error if the schemas don't match structurally.
122    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
145/// A mapping from Parquet column path names to internal field id
146struct IndexByParquetPathName {
147    name_to_id: HashMap<String, i32>,
148
149    field_names: Vec<String>,
150
151    field_id: i32,
152}
153
154impl IndexByParquetPathName {
155    /// Creates a new, empty `IndexByParquetPathName`
156    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    /// Retrieves the internal field ID
165    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
263/// `ParquetWriter`` is used to write arrow data into parquet file on storage.
264pub struct ParquetWriter {
265    schema: SchemaRef,
266    /// Optional Arrow schema to use for writing. If provided, this is used directly
267    /// instead of converting from the Iceberg schema, preserving any extension metadata.
268    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
276/// Used to aggregate min and max value of each column.
277struct MinMaxColAggregator {
278    lower_bounds: HashMap<i32, Datum>,
279    upper_bounds: HashMap<i32, Datum>,
280    schema: SchemaRef,
281}
282
283impl MinMaxColAggregator {
284    /// Creates new and empty `MinMaxColAggregator`
285    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    /// Update statistics
316    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            // Following java implementation: https://github.com/apache/iceberg/blob/29a2c456353a6120b8c882ed2ab544975b168d7b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetUtil.java#L163
323            // Ignore the field if it is not in schema.
324            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    /// Returns lower and upper bounds
359    fn produce(self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) {
360        (self.lower_bounds, self.upper_bounds)
361    }
362}
363
364impl ParquetWriter {
365    /// Converts parquet files to data files
366    #[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        // TODO: support adding to partitioned table
373        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                // TODO: Implement nan_value_counts here
394                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    /// `ParquetMetadata` to data file builder
405    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            // # NOTE:
467            // - We can ignore implementing distinct_counts due to this: https://lists.apache.org/thread/j52tsojv0x4bopxyzsp7m7bqt23n5fnd
468            .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
530/// Validate that an Arrow schema's structure matches an Iceberg schema.
531/// This checks field names, types, and nullability, but ignores metadata differences.
532fn 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    // Compare types structurally, ignoring metadata
594    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            // For primitive types, use direct comparison (metadata not relevant)
621            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        // Skip empty batch
638        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            // Use the provided Arrow schema if available, otherwise derive from Iceberg schema.
651            // Using a provided Arrow schema allows preserving extension metadata in the output.
652            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/AsyncArrowWriter contains sync and async writers
729            // written size = bytes flushed to inner's async writer + bytes buffered in the inner's sync writer
730            inner.bytes_written() + inner.in_progress_size()
731        } else {
732            // inner writer is not initialized yet
733            0
734        }
735    }
736
737    fn current_schema(&self) -> SchemaRef {
738        self.schema.clone()
739    }
740}
741
742/// AsyncFileWriter is a wrapper of FileWrite to make it compatible with tokio::io::AsyncWrite.
743///
744/// # NOTES
745///
746/// We keep this wrapper been used inside only.
747struct AsyncFileWriter(Box<dyn FileWrite>);
748
749impl AsyncFileWriter {
750    /// Create a new `AsyncFileWriter` with the given writer.
751    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                // Parquet Statistics will use different representation for Decimal with precision 38 and scale 5,
851                // so we need to add a new field for it.
852                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        // Int, Struct(Int,Int), String, List(Int), Struct(Struct(Int)), Map(String, List(Int))
868        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        // prepare data
968        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        // write data
988        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            // Put dummy field for build successfully.
1005            .content(DataContentType::Data)
1006            .partition(Struct::empty())
1007            .partition_spec_id(0)
1008            .build()
1009            .unwrap();
1010
1011        // check data file
1012        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        // check the written file
1025        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        // prepare data
1042        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        // write data
1188        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            // Put dummy field for build successfully.
1200            .content(crate::spec::DataContentType::Data)
1201            .partition(Struct::empty())
1202            .partition_spec_id(0)
1203            .build()
1204            .unwrap();
1205
1206        // check data file
1207        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 the written file
1249        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        // prepare data
1265        // generate iceberg schema for all type
1266        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        // write data
1378        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            // Put dummy field for build successfully.
1390            .content(crate::spec::DataContentType::Data)
1391            .partition(Struct::empty())
1392            .partition_spec_id(0)
1393            .build()
1394            .unwrap();
1395
1396        // check data file
1397        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 the written file
1490        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        // test 1.1 and 2.2
1506        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        // test -1.1 and -2.2
1558        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        // test 38-digit precision decimal values (Iceberg spec max)
1610        // Note: fastnum D128::MAX/MIN have impractical exponents, so we use meaningful values
1611        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        // test max and min for scale 38
1669        // # TODO
1670        // Readd this case after resolve https://github.com/apache/iceberg-rust/issues/669
1671        // let schema = Arc::new(
1672        //     Schema::builder()
1673        //         .with_fields(vec![NestedField::optional(
1674        //             0,
1675        //             "decimal",
1676        //             Type::Primitive(PrimitiveType::Decimal {
1677        //                 precision: 38,
1678        //                 scale: 0,
1679        //             }),
1680        //         )
1681        //         .into()])
1682        //         .build()
1683        //         .unwrap(),
1684        // );
1685        // let arrow_schema: ArrowSchemaRef = Arc::new(schema_to_arrow_schema(&schema).unwrap());
1686        // let mut pw = ParquetWriterBuilder::new(
1687        //     WriterProperties::builder().build(),
1688        //     schema,
1689        //     file_io.clone(),
1690        //     loccation_gen,
1691        //     file_name_gen,
1692        // )
1693        // .build()
1694        // .await?;
1695        // let col0 = Arc::new(
1696        //     Decimal128Array::from(vec![
1697        //         Some(99999999999999999999999999999999999999_i128),
1698        //         Some(-99999999999999999999999999999999999999_i128),
1699        //     ])
1700        //     .with_data_type(DataType::Decimal128(38, 0)),
1701        // ) as ArrayRef;
1702        // let to_write = RecordBatch::try_new(arrow_schema.clone(), vec![col0]).unwrap();
1703        // pw.write(&to_write).await?;
1704        // let res = pw.close().await?;
1705        // assert_eq!(res.len(), 1);
1706        // let data_file = res
1707        //     .into_iter()
1708        //     .next()
1709        //     .unwrap()
1710        //     .content(crate::spec::DataContentType::Data)
1711        //     .partition(Struct::empty())
1712        //     .build()
1713        //     .unwrap();
1714        // assert_eq!(
1715        //     data_file.upper_bounds().get(&0),
1716        //     Some(Datum::new(
1717        //         PrimitiveType::Decimal {
1718        //             precision: 38,
1719        //             scale: 0
1720        //         },
1721        //         PrimitiveLiteral::Int128(99999999999999999999999999999999999999_i128)
1722        //     ))
1723        //     .as_ref()
1724        // );
1725        // assert_eq!(
1726        //     data_file.lower_bounds().get(&0),
1727        //     Some(Datum::new(
1728        //         PrimitiveType::Decimal {
1729        //             precision: 38,
1730        //             scale: 0
1731        //         },
1732        //         PrimitiveLiteral::Int128(-99999999999999999999999999999999999999_i128)
1733        //     ))
1734        //     .as_ref()
1735        // );
1736
1737        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        // Test that file will create if data to write
1751        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        // Test that file will not create if no data to write
1774        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        // prepare data
1801        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        // write data
1830        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            // Put dummy field for build successfully.
1845            .content(crate::spec::DataContentType::Data)
1846            .partition(Struct::empty())
1847            .partition_spec_id(0)
1848            .build()
1849            .unwrap();
1850
1851        // check data file
1852        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        // check the written file
1872        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        // prepare data
1915        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        // write data
1970        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            // Put dummy field for build successfully.
1985            .content(crate::spec::DataContentType::Data)
1986            .partition(Struct::empty())
1987            .partition_spec_id(0)
1988            .build()
1989            .unwrap();
1990
1991        // check data file
1992        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        // check the written file
2012        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                // Field::new_large_list("col3", schema_large_list_float_field.clone(), true).with_metadata(
2056                //     HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "5".to_string())]),
2057                // ).clone(),
2058            ];
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            // large_list_float_field_col,
2122        ])
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        // write data
2129        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        // check data file
2156        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        // check the written file
2176        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        // write data
2310        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        // check data file
2337        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        // check the written file
2370        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        // write data
2390        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        // Check that file should have been deleted.
2412        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    // -----------------------------------------------------------------
2451    // ParquetWriterBuilder::from_table_properties
2452    // -----------------------------------------------------------------
2453
2454    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        // `build()` must carry the translated `WriterProperties` through to the
2482        // `ParquetWriter` unchanged — otherwise the `write.parquet.*` settings
2483        // derived in `from_table_properties` would never reach parquet-rs.
2484        //
2485        // Asserting on the writer's `WriterProperties` (rather than re-reading a
2486        // written file) keeps this a direct propagation check: every future
2487        // `write.parquet.*` option just adds an assertion on its corresponding
2488        // `WriterProperties` getter here.
2489        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}