1use std::collections::HashMap;
21use std::sync::Arc;
22
23use arrow_schema::SchemaRef as ArrowSchemaRef;
24use bytes::Bytes;
25use futures::future::BoxFuture;
26use itertools::Itertools;
27use parquet::arrow::AsyncArrowWriter;
28use parquet::arrow::async_reader::AsyncFileReader;
29use parquet::arrow::async_writer::AsyncFileWriter as ArrowAsyncFileWriter;
30use parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
31use parquet::encryption::encrypt::FileEncryptionProperties;
32use parquet::file::metadata::ParquetMetaData;
33use parquet::file::properties::{CdcOptions, WriterProperties};
34use parquet::file::statistics::Statistics;
35
36use super::{FileWriter, FileWriterBuilder};
37use crate::arrow::{
38 ArrowFileReader, DEFAULT_MAP_FIELD_NAME, FieldMatchMode, NanValueCountVisitor,
39 get_parquet_stat_max_as_datum, get_parquet_stat_min_as_datum,
40};
41use crate::compression::CompressionCodec;
42use crate::encryption::{EncryptionManager, StandardKeyMetadata};
43use crate::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#[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 pub fn new(props: WriterProperties, schema: SchemaRef) -> Self {
70 Self::new_with_match_mode(props, schema, FieldMatchMode::Id)
71 }
72
73 pub fn new_with_match_mode(
75 props: WriterProperties,
76 schema: SchemaRef,
77 match_mode: FieldMatchMode,
78 ) -> Self {
79 Self {
80 props,
81 schema,
82 match_mode,
83 encryption_manager: None,
84 }
85 }
86
87 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 pub fn with_match_mode(mut self, match_mode: FieldMatchMode) -> Self {
120 self.match_mode = match_mode;
121 self
122 }
123
124 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
187struct IndexByParquetPathName {
189 name_to_id: HashMap<String, i32>,
190
191 field_names: Vec<String>,
192
193 field_id: i32,
194}
195
196impl IndexByParquetPathName {
197 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 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
312pub 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
323struct MinMaxColAggregator {
325 lower_bounds: HashMap<i32, Datum>,
326 upper_bounds: HashMap<i32, Datum>,
327 schema: SchemaRef,
328}
329
330impl MinMaxColAggregator {
331 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 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 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 fn produce(self) -> (HashMap<i32, Datum>, HashMap<i32, Datum>) {
407 (self.lower_bounds, self.upper_bounds)
408 }
409}
410
411impl ParquetWriter {
412 #[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 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 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 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 .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 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 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.bytes_written() + inner.in_progress_size()
713 } else {
714 0
716 }
717 }
718}
719
720struct AsyncFileWriter(Box<dyn FileWrite>);
726
727impl AsyncFileWriter {
728 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 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 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 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 let schema = {
969 let fields =
970 vec![
971 Field::new("col", DataType::Int64, true).with_metadata(HashMap::from([(
972 PARQUET_FIELD_ID_META_KEY.to_string(),
973 "0".to_string(),
974 )])),
975 ];
976 Arc::new(arrow_schema::Schema::new(fields))
977 };
978 let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
979 let null_col = Arc::new(Int64Array::new_null(1024)) as ArrayRef;
980 let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
981 let to_write_null = RecordBatch::try_new(schema.clone(), vec![null_col]).unwrap();
982
983 let output_file = file_io.new_output(
984 location_gen.generate_location(None, &file_name_gen.generate_file_name()),
985 )?;
986
987 let mut pw = ParquetWriterBuilder::new(
989 WriterProperties::builder()
990 .set_max_row_group_row_count(Some(128))
991 .build(),
992 Arc::new(to_write.schema().as_ref().try_into().unwrap()),
993 )
994 .build(output_file)
995 .await?;
996 pw.write(&to_write).await?;
997 pw.write(&to_write_null).await?;
998 let res = pw.close().await?;
999 assert_eq!(res.len(), 1);
1000 let data_file = res
1001 .into_iter()
1002 .next()
1003 .unwrap()
1004 .content(DataContentType::Data)
1006 .partition(Struct::empty())
1007 .partition_spec_id(0)
1008 .build()
1009 .unwrap();
1010
1011 assert_eq!(data_file.record_count(), 2048);
1013 assert_eq!(*data_file.value_counts(), HashMap::from([(0, 2048)]));
1014 assert_eq!(
1015 *data_file.lower_bounds(),
1016 HashMap::from([(0, Datum::long(0))])
1017 );
1018 assert_eq!(
1019 *data_file.upper_bounds(),
1020 HashMap::from([(0, Datum::long(1023))])
1021 );
1022 assert_eq!(*data_file.null_value_counts(), HashMap::from([(0, 1024)]));
1023
1024 let expect_batch = concat_batches(&schema, vec![&to_write, &to_write_null]).unwrap();
1026 check_parquet_data_file(&file_io, &data_file, &expect_batch).await;
1027
1028 Ok(())
1029 }
1030
1031 #[tokio::test]
1032 async fn test_parquet_writer_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 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 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 .content(DataContentType::Data)
1291 .partition(Struct::empty())
1292 .partition_spec_id(0)
1293 .build()
1294 .unwrap();
1295
1296 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_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 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 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 .content(DataContentType::Data)
1481 .partition(Struct::empty())
1482 .partition_spec_id(0)
1483 .build()
1484 .unwrap();
1485
1486 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_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 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 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 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 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 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 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 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 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 .content(DataContentType::Data)
1940 .partition(Struct::empty())
1941 .partition_spec_id(0)
1942 .build()
1943 .unwrap();
1944
1945 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 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 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 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 .content(DataContentType::Data)
2080 .partition(Struct::empty())
2081 .partition_spec_id(0)
2082 .build()
2083 .unwrap();
2084
2085 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 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 ];
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 ])
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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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}