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