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