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