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