This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11971-bad29dde10217915e808f364d5aee92759e33ea4 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit ed478d7b795d5e03b916213e40b7a03e860e5366 Author: Hippolyte Barraud <[email protected]> AuthorDate: Sat Oct 31 11:30:33 2026 +0000 perf(parquet): share Arrow write-schema plans across row groups (#11971) # Which issue does this PR close? <!-- It is not necessary to raise an issue, but it is encouraged if it is a non-trivial enhancement/bug. If this PR relates to an existing issue, but does not fully resolve it, please change the wording below from "Closes #NNN" to "Relates to #NNN" or "Part of #NNN" to avoid false closing an issue. --> - Contributes to #9731 - Contributes to #9722 - Part of #11273 # Rationale for this change <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. Please try to explain the problem you are trying to solve in terms of the user-visible behavior, rather than the implementation. --> The eager Arrow writer repeatedly derives the same schema-only field-to-leaf mapping when creating and validating row groups, even though that metadata is invariant for the lifetime of the writer factory. # What changes are included in this PR? <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> Build an immutable write-schema plan once per factory and share it across row groups and column writers. Reuse the cached target fields, leaf ranges, and Parquet descriptors for both ordinary and CDC paths, while continuing to derive leaves and levels from each incoming batch. This removes repeated schema traversal without retaining input arrays or changing dense, dictionary, REE, view, required-null, or leaf-ownership semantics. # Are these changes tested? <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests, if so identify these tests)? If this PR claims a performance improvement, please include evidence such as benchmark results. DO NOT simply paste a list of commands ran to verify the code; this is what our CI is for and it pollutes the PR. ONLY explain if testing was included with this PR, or why it wasn't. --> All tests passing. # Are there any significant user-facing changes? <!-- Significant user-facing changes include: - Breaking public API change - Significant change to existing functionality (e.g. function returns something different) These need to be explicitly called out to ensure we have proper documentation if required and to ensure we follow proper semantic versioning. DO NOT restate the changes of this PR; if a bug was fixed and already explained above, we don't need this to be restated as a user-facing change. --> None. --------- Signed-off-by: Hippolyte Barraud <[email protected]> Co-authored-by: Ed Seidl <[email protected]> Co-authored-by: Andrew Lamb <[email protected]> --- parquet/benches/bit_packing.rs | 2 +- parquet/src/arrow/arrow_writer/mod.rs | 460 ++++++++++++++++++++++----- parquet/tests/arrow_writer/mod.rs | 2 +- parquet/tests/encryption/encryption_async.rs | 6 +- 4 files changed, 393 insertions(+), 77 deletions(-) diff --git a/parquet/benches/bit_packing.rs b/parquet/benches/bit_packing.rs index 52ffb6dfad..c11e48240b 100644 --- a/parquet/benches/bit_packing.rs +++ b/parquet/benches/bit_packing.rs @@ -156,7 +156,7 @@ fn bench_arrow_dictionary_indices(c: &mut Criterion, values: &[i32]) { .unwrap(); let file_writer = SerializedFileWriter::new(Vec::new(), parquet_schema.root_schema_ptr(), props).unwrap(); - let factory = ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema); + let factory = ArrowRowGroupWriterFactory::try_new(&file_writer, arrow_schema).unwrap(); let mut leaves = compute_leaves(&field, &array).unwrap(); let leaf = leaves.pop().unwrap(); diff --git a/parquet/src/arrow/arrow_writer/mod.rs b/parquet/src/arrow/arrow_writer/mod.rs index 7e6264c398..e901e3132d 100644 --- a/parquet/src/arrow/arrow_writer/mod.rs +++ b/parquet/src/arrow/arrow_writer/mod.rs @@ -27,10 +27,10 @@ use std::sync::{Arc, Mutex}; use std::vec::IntoIter; use arrow_array::cast::AsArray; -use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter}; +use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter, new_empty_array}; use arrow_array::{PrimitiveArray, types::*}; use arrow_schema::{ - ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit, + ArrowError, DataType as ArrowDataType, Field, FieldRef, IntervalUnit, SchemaRef, TimeUnit, }; use super::schema::{ @@ -56,7 +56,7 @@ use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData}; use crate::file::properties::{WriterProperties, WriterPropertiesPtr}; use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter}; use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift}; -use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor}; +use crate::schema::types::{ColumnDescPtr, SchemaDescriptor}; use levels::{ArrayLevels, calculate_array_levels}; mod byte_array; @@ -277,7 +277,7 @@ impl<W: Write + Send> ArrowWriter<W> { SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?; let mut row_group_writer_factory = - ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone()); + ArrowRowGroupWriterFactory::try_new(&file_writer, arrow_schema.clone())?; if let Some(page_store_factory) = options.page_store_factory { row_group_writer_factory = row_group_writer_factory.with_page_store_factory(page_store_factory); @@ -1037,7 +1037,7 @@ impl ArrowColumnChunk { /// .unwrap(); /// /// // Create a factory for building Arrow column writers -/// let row_group_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); +/// let row_group_factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap(); /// // Create column writers for the 0th row group /// let col_writers = row_group_factory.create_column_writers(0).unwrap(); /// @@ -1259,18 +1259,85 @@ impl ArrowColumnWriter { } } -fn arrow_leaf_ranges(schema: &SchemaRef) -> Result<Vec<Range<usize>>> { - let mut count = 0; - schema - .fields - .iter() - .map(|field| { +/// Associates a top-level Arrow field with its range of Parquet leaf writers. +#[derive(Debug)] +struct ArrowTopLevelWriterSpec { + /// The writer schema's field supplies the target nullability and nesting, + /// while each batch can use a compatible physical layout (dictionary, + /// run-end, string/binary offset width, or view). + field: FieldRef, + leaf_range: Range<usize>, +} + +/// Replayable schema-planning errors kept only for the deprecated +/// [`ArrowRowGroupWriterFactory::new`] constructor. It cannot return errors, +/// and repeated `create_column_writers(&self)` calls need fresh owned errors. +/// [`ArrowRowGroupWriterFactory::try_new`] returns the original error directly. +#[derive(Debug)] +enum ArrowWriteSchemaPlanError { + General(String), + Nyi(String), + ArrowError(String), +} + +impl ArrowWriteSchemaPlanError { + fn to_parquet_error(&self) -> ParquetError { + match self { + Self::General(message) => ParquetError::General(message.clone()), + Self::Nyi(message) => ParquetError::NYI(message.clone()), + Self::ArrowError(message) => ParquetError::ArrowError(message.clone()), + } + } + + fn from_parquet_error(error: ParquetError) -> Self { + match error { + ParquetError::NYI(message) => Self::Nyi(message), + ParquetError::General(message) => Self::General(message), + ParquetError::ArrowError(message) => Self::ArrowError(message), + error => Self::General(error.to_string()), + } + } +} + +/// Immutable schema glue shared by every row group. +/// +/// This deliberately caches only schema fields and Parquet descriptors, never +/// arrays or per-batch level/value plans. +#[derive(Debug)] +struct ArrowWriteSchemaPlan { + fields: Box<[ArrowTopLevelWriterSpec]>, + leaves: Box<[ColumnDescPtr]>, +} + +impl ArrowWriteSchemaPlan { + fn try_new(parquet: &SchemaDescriptor, arrow: &SchemaRef) -> Result<Self> { + let parquet_leaves = parquet.columns(); + let mut leaf_count = 0; + let mut fields = Vec::with_capacity(arrow.fields.len()); + + for field in &arrow.fields { validate_map_key_type(field.data_type())?; - let start = count; - count += compute_leaves(field, &arrow_array::new_empty_array(field.data_type()))?.len(); - Ok(start..count) + let start = leaf_count; + leaf_count += compute_leaves(field, &new_empty_array(field.data_type()))?.len(); + fields.push(ArrowTopLevelWriterSpec { + field: Arc::clone(field), + leaf_range: start..leaf_count, + }); + } + + if leaf_count != parquet_leaves.len() { + return Err(ParquetError::General(format!( + "Arrow schema maps to {} leaf columns but Parquet schema has {}", + leaf_count, + parquet_leaves.len() + ))); + } + + Ok(Self { + fields: fields.into_boxed_slice(), + leaves: parquet_leaves.to_vec().into_boxed_slice(), }) - .collect() + } } /// Encodes [`RecordBatch`] to a parquet row group @@ -1282,27 +1349,29 @@ fn arrow_leaf_ranges(schema: &SchemaRef) -> Result<Vec<Range<usize>>> { #[derive(Debug)] struct ArrowRowGroupWriter { writers: Vec<ArrowColumnWriter>, - schema: SchemaRef, + schema_plan: Arc<ArrowWriteSchemaPlan>, buffered_rows: usize, } impl ArrowRowGroupWriter { - fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self { + fn new(writers: Vec<ArrowColumnWriter>, schema_plan: Arc<ArrowWriteSchemaPlan>) -> Self { + debug_assert_eq!(writers.len(), schema_plan.leaves.len()); Self { writers, - schema: arrow.clone(), + schema_plan, buffered_rows: 0, } } fn write(&mut self, batch: &RecordBatch) -> Result<()> { - let ranges = self.validate_batch_shape(batch, None)?; + self.validate_batch_shape(batch, None)?; self.buffered_rows += batch.num_rows(); - for (column_idx, (field, range)) in self.schema.fields.iter().zip(ranges).enumerate() { - let leaves = compute_leaves(field, batch.column(column_idx))?; - Self::validate_leaf_count(field, &range, leaves.len())?; + + for (column_idx, field) in self.schema_plan.fields.iter().enumerate() { + let leaves = compute_leaves(field.field.as_ref(), batch.column(column_idx))?; + self.validate_leaf_count(field, leaves.len())?; for (offset, leaf) in leaves.into_iter().enumerate() { - self.writers[range.start + offset].write(&leaf)?; + self.writers[field.leaf_range.start + offset].write(&leaf)?; } } Ok(()) @@ -1313,14 +1382,15 @@ impl ArrowRowGroupWriter { batch: &RecordBatch, chunkers: &mut [ContentDefinedChunker], ) -> Result<()> { - let ranges = self.validate_batch_shape(batch, Some(chunkers.len()))?; + self.validate_batch_shape(batch, Some(chunkers.len()))?; self.buffered_rows += batch.num_rows(); - for (column_idx, (field, range)) in self.schema.fields.iter().zip(ranges).enumerate() { - let leaves = compute_leaves(field, batch.column(column_idx))?; - Self::validate_leaf_count(field, &range, leaves.len())?; + + for (column_idx, field) in self.schema_plan.fields.iter().enumerate() { + let leaves = compute_leaves(field.field.as_ref(), batch.column(column_idx))?; + self.validate_leaf_count(field, leaves.len())?; for (offset, leaf) in leaves.into_iter().enumerate() { - let index = range.start + offset; - self.writers[index].write_with_chunker(&leaf, &mut chunkers[index])?; + let leaf_idx = field.leaf_range.start + offset; + self.writers[leaf_idx].write_with_chunker(&leaf, &mut chunkers[leaf_idx])?; } } Ok(()) @@ -1330,39 +1400,38 @@ impl ArrowRowGroupWriter { &self, batch: &RecordBatch, chunker_count: Option<usize>, - ) -> Result<Vec<Range<usize>>> { - if batch.num_columns() != self.schema.fields.len() { + ) -> Result<()> { + if batch.num_columns() != self.schema_plan.fields.len() { return Err(ParquetError::ArrowError(format!( "Incompatible schema: writer has {} top-level fields but batch has {} columns", - self.schema.fields.len(), + self.schema_plan.fields.len(), batch.num_columns() ))); } - let ranges = arrow_leaf_ranges(&self.schema)?; - let leaf_count = ranges.last().map_or(0, |r| r.end); - if self.writers.len() != leaf_count { + if self.writers.len() != self.schema_plan.leaves.len() { return Err(ParquetError::General(format!( "Arrow row-group writer has {} column writers for {} planned leaves", self.writers.len(), - leaf_count + self.schema_plan.leaves.len() ))); } if let Some(actual) = chunker_count - && actual != leaf_count + && actual != self.schema_plan.leaves.len() { return Err(ParquetError::General(format!( - "content-defined chunking has {actual} chunkers for {leaf_count} planned leaves" + "content-defined chunking has {actual} chunkers for {} planned leaves", + self.schema_plan.leaves.len() ))); } - Ok(ranges) + Ok(()) } - fn validate_leaf_count(field: &Field, range: &Range<usize>, actual: usize) -> Result<()> { - let expected = range.len(); + fn validate_leaf_count(&self, field: &ArrowTopLevelWriterSpec, actual: usize) -> Result<()> { + let expected = field.leaf_range.end - field.leaf_range.start; if actual != expected { return Err(ParquetError::ArrowError(format!( "Incompatible schema: field '{}' produced {actual} leaf columns but writer expects {expected}", - field.name() + field.field.name() ))); } Ok(()) @@ -1386,12 +1455,15 @@ impl ArrowRowGroupWriter { /// Factory that creates new column writers for each row group in the Parquet file. /// -/// You can create this structure via an [`ArrowWriter::into_serialized_writer`]. +/// Create this structure with [`Self::try_new`] or [`ArrowWriter::into_serialized_writer`]. /// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel #[derive(Debug)] pub struct ArrowRowGroupWriterFactory { - schema: SchemaDescPtr, - arrow_schema: SchemaRef, + // Only deprecated `new` can store an Err; `try_new` always stores an Ok. + // TODO: When removing `new`, store Arc<ArrowWriteSchemaPlan> directly, + // remove ArrowWriteSchemaPlanError and schema_plan(), and simplify + // from_schema_plan to accept the plan without a Result. + schema_plan: std::result::Result<Arc<ArrowWriteSchemaPlan>, ArrowWriteSchemaPlanError>, props: WriterPropertiesPtr, page_store_factory: Arc<dyn PageStoreFactory>, #[cfg(feature = "encryption")] @@ -1399,16 +1471,52 @@ pub struct ArrowRowGroupWriterFactory { } impl ArrowRowGroupWriterFactory { - /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema + /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema. + /// + /// Schema-planning errors are deferred until [`Self::create_column_writers`] is called. + /// Use [`Self::try_new`] to return these errors during construction instead. + #[deprecated( + since = "60.2.0", + note = "Use `ArrowRowGroupWriterFactory::try_new` instead" + )] pub fn new<W: Write + Send>( file_writer: &SerializedFileWriter<W>, arrow_schema: SchemaRef, ) -> Self { - let schema = Arc::clone(file_writer.schema_descr_ptr()); + let schema_plan = + ArrowWriteSchemaPlan::try_new(file_writer.schema_descr_ptr().as_ref(), &arrow_schema) + .map(Arc::new) + .map_err(ArrowWriteSchemaPlanError::from_parquet_error); + Self::from_schema_plan(file_writer, schema_plan) + } + + /// Try to create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema. + /// + /// Validates and caches the schema-only field-to-leaf mapping for reuse across row groups. + /// + /// # Errors + /// + /// Returns schema-planning errors immediately, including invalid map-key schemas and + /// mismatched Arrow and Parquet leaf counts. + pub fn try_new<W: Write + Send>( + file_writer: &SerializedFileWriter<W>, + arrow_schema: SchemaRef, + ) -> Result<Self> { + let schema_plan = + ArrowWriteSchemaPlan::try_new(file_writer.schema_descr_ptr().as_ref(), &arrow_schema)?; + Ok(Self::from_schema_plan( + file_writer, + Ok(Arc::new(schema_plan)), + )) + } + + fn from_schema_plan<W: Write + Send>( + file_writer: &SerializedFileWriter<W>, + schema_plan: std::result::Result<Arc<ArrowWriteSchemaPlan>, ArrowWriteSchemaPlanError>, + ) -> Self { let props = Arc::clone(file_writer.properties()); Self { - schema, - arrow_schema, + schema_plan, props, page_store_factory: Arc::new(InMemoryPageStoreFactory), #[cfg(feature = "encryption")] @@ -1428,27 +1536,34 @@ impl ArrowRowGroupWriterFactory { } fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> { - let writers = self.create_column_writers(row_group_index)?; - Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema)) + let schema_plan = Arc::clone(self.schema_plan()?); + let writers = self.create_column_writers_from_plan(row_group_index, &schema_plan)?; + Ok(ArrowRowGroupWriter::new(writers, schema_plan)) } /// Create column writers for a new row group, with the given row group index pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> { - let ranges = arrow_leaf_ranges(&self.arrow_schema)?; - let leaf_count = ranges.last().map_or(0, |r| r.end); - if leaf_count != self.schema.num_columns() { - return Err(ParquetError::General(format!( - "Arrow schema maps to {} leaf columns but Parquet schema has {}", - leaf_count, - self.schema.num_columns() - ))); + self.create_column_writers_from_plan(row_group_index, self.schema_plan()?) + } + + fn schema_plan(&self) -> Result<&Arc<ArrowWriteSchemaPlan>> { + match &self.schema_plan { + Ok(plan) => Ok(plan), + Err(error) => Err(error.to_parquet_error()), } - let mut writers = Vec::with_capacity(self.arrow_schema.fields.len()); - let mut leaves = self.schema.columns().iter(); + } + + fn create_column_writers_from_plan( + &self, + row_group_index: usize, + schema_plan: &ArrowWriteSchemaPlan, + ) -> Result<Vec<ArrowColumnWriter>> { + let mut writers = Vec::with_capacity(schema_plan.leaves.len()); + let mut leaves = schema_plan.leaves.iter(); let column_factory = self.column_writer_factory(row_group_index); - for field in &self.arrow_schema.fields { + for field in &schema_plan.fields { column_factory.get_arrow_column_writer( - field.data_type(), + field.field.data_type(), &self.props, &mut leaves, &mut writers, @@ -3279,10 +3394,15 @@ mod tests { SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), Arc::clone(&props)) .unwrap(); + let eager_error = + ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&two_fields)).unwrap_err(); + assert!(matches!(&eager_error, ParquetError::General(_))); + #[expect(deprecated)] // Verify that the legacy constructor still defers errors. let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&two_fields)); let Err(err) = factory.create_column_writers(0) else { panic!("mismatched schemas should not create column writers") }; + assert_eq!(err.to_string(), eager_error.to_string()); assert!( err.to_string() .contains("Arrow schema maps to 2 leaf columns but Parquet schema has 1"), @@ -3292,10 +3412,15 @@ mod tests { let parquet = ArrowSchemaConverter::new().convert(&two_fields).unwrap(); let file_writer = SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), props).unwrap(); + let eager_error = + ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&one_field)).unwrap_err(); + assert!(matches!(&eager_error, ParquetError::General(_))); + #[expect(deprecated)] // Verify that the legacy constructor still defers errors. let factory = ArrowRowGroupWriterFactory::new(&file_writer, one_field); let Err(err) = factory.create_column_writers(0) else { panic!("mismatched schemas should not create column writers") }; + assert_eq!(err.to_string(), eager_error.to_string()); assert!( err.to_string() .contains("Arrow schema maps to 1 leaf columns but Parquet schema has 2"), @@ -3479,15 +3604,25 @@ mod tests { Arc::new(WriterProperties::default()), ) .unwrap(); - let factory = ArrowRowGroupWriterFactory::new(&file, schema); - let Err(err) = factory.create_column_writers(0) else { - panic!("nullable map-key schema should not create column writers") - }; - assert!( - err.to_string() - .contains("Map key field 'key' must be non-nullable"), - "{err}" + let eager_error = + ArrowRowGroupWriterFactory::try_new(&file, Arc::clone(&schema)).unwrap_err(); + assert!(matches!(&eager_error, ParquetError::ArrowError(_))); + assert_eq!( + eager_error.to_string(), + "Arrow: Map key field 'key' must be non-nullable, including its encoded values" ); + #[expect(deprecated)] // Verify that the legacy constructor still defers errors. + let factory = ArrowRowGroupWriterFactory::new(&file, schema); + for row_group_index in 0..2 { + let Err(err) = factory.create_column_writers(row_group_index) else { + panic!("nullable map-key schema should not create column writers") + }; + assert!(matches!(&err, ParquetError::ArrowError(_)), "{err:?}"); + assert_eq!( + err.to_string(), + "Arrow: Map key field 'key' must be non-nullable, including its encoded values" + ); + } } } @@ -3518,6 +3653,185 @@ mod tests { assert_eq!(group.buffered_rows, 0); } + #[test] + fn arrow_write_schema_plan_preserves_error_variants() { + for error in [ + ParquetError::General("schema mismatch".to_owned()), + ParquetError::NYI("unsupported type".to_owned()), + ParquetError::ArrowError("invalid schema".to_owned()), + ] { + let expected_variant = std::mem::discriminant(&error); + let expected_message = error.to_string(); + let cached = ArrowWriteSchemaPlanError::from_parquet_error(error); + for _ in 0..2 { + let error = cached.to_parquet_error(); + assert_eq!(std::mem::discriminant(&error), expected_variant); + assert_eq!(error.to_string(), expected_message); + } + } + } + + #[test] + fn arrow_write_schema_plan_caches_nested_leaf_ranges() { + let nested = ArrowDataType::Struct(Fields::from(vec![ + Field::new("i", ArrowDataType::Int32, false), + Field::new("s", ArrowDataType::Utf8, true), + ])); + let list = ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Boolean, true))); + let ree = ArrowDataType::RunEndEncoded( + Arc::new(Field::new("run_ends", ArrowDataType::Int32, false)), + Arc::new(Field::new("values", ArrowDataType::Int64, false)), + ); + let dictionary_fsb = ArrowDataType::Dictionary( + Box::new(ArrowDataType::Int8), + Box::new(ArrowDataType::FixedSizeBinary(2)), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("nested", nested, true), + Field::new("list", list, true), + Field::new("ree", ree, false), + Field::new("dictionary_fsb", dictionary_fsb, true), + ])); + let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap(); + + let plan = ArrowWriteSchemaPlan::try_new(&parquet, &schema).unwrap(); + let ranges = plan + .fields + .iter() + .map(|field| field.leaf_range.clone()) + .collect::<Vec<_>>(); + assert_eq!(ranges, vec![0..2, 2..3, 3..4, 4..5]); + assert_eq!(plan.leaves.len(), parquet.num_columns()); + assert_eq!( + plan.leaves[4].physical_type(), + crate::basic::Type::FIXED_LEN_BYTE_ARRAY + ); + } + + #[test] + fn cached_schema_plan_preserves_compatible_physical_layout_alternation() { + let writer_schema = Arc::new(Schema::new(vec![ + Field::new("number", ArrowDataType::Int32, false), + Field::new("text", ArrowDataType::Utf8, false), + Field::new("bytes", ArrowDataType::Binary, false), + ])); + + let dense = RecordBatch::try_new( + Arc::clone(&writer_schema), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec!["a"])), + Arc::new(BinaryArray::from_iter_values([b"a".as_slice()])), + ], + ) + .unwrap(); + + let number_dictionary = DictionaryArray::<Int8Type>::try_new( + Int8Array::from(vec![0]), + Arc::new(Int32Array::from(vec![2])), + ) + .unwrap(); + let text_dictionary = DictionaryArray::<Int8Type>::try_new( + Int8Array::from(vec![0]), + Arc::new(StringArray::from(vec!["b"])), + ) + .unwrap(); + let bytes_dictionary = DictionaryArray::<Int8Type>::try_new( + Int8Array::from(vec![0]), + Arc::new(BinaryArray::from_iter_values([b"b".as_slice()])), + ) + .unwrap(); + let dictionary = RecordBatch::try_from_iter(vec![ + ("number", Arc::new(number_dictionary) as ArrayRef), + ("text", Arc::new(text_dictionary) as ArrayRef), + ("bytes", Arc::new(bytes_dictionary) as ArrayRef), + ]) + .unwrap(); + + let run_ends = Int32Array::from(vec![1]); + let number_ree: ArrayRef = + Arc::new(Int32RunArray::try_new(&run_ends, &Int32Array::from(vec![3])).unwrap()); + let text_ree: ArrayRef = + Arc::new(Int32RunArray::try_new(&run_ends, &StringArray::from(vec!["c"])).unwrap()); + let bytes_ree: ArrayRef = Arc::new( + Int32RunArray::try_new(&run_ends, &BinaryArray::from_iter_values([b"c".as_slice()])) + .unwrap(), + ); + let ree = RecordBatch::try_from_iter(vec![ + ("number", number_ree), + ("text", text_ree), + ("bytes", bytes_ree), + ]) + .unwrap(); + + let alternate = RecordBatch::try_from_iter(vec![ + ("number", Arc::new(Int32Array::from(vec![4])) as ArrayRef), + ( + "text", + Arc::new(LargeStringArray::from(vec!["d"])) as ArrayRef, + ), + ( + "bytes", + Arc::new(BinaryViewArray::from_iter_values([b"d".as_slice()])) as ArrayRef, + ), + ]) + .unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(1)) + .build(); + let mut out = Vec::new(); + let mut writer = + ArrowWriter::try_new(&mut out, Arc::clone(&writer_schema), Some(props)).unwrap(); + for batch in [&dense, &dictionary, &ree, &alternate] { + writer.write(batch).unwrap(); + } + writer.close().unwrap(); + + let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(out)).unwrap(); + assert_eq!(builder.metadata().num_row_groups(), 4); + let mut reader = builder.build().unwrap(); + let actual = reader.next().unwrap().unwrap(); + assert_eq!( + actual.column(0).as_primitive::<Int32Type>(), + &Int32Array::from(vec![1, 2, 3, 4]) + ); + assert_eq!( + actual.column(1).as_string::<i32>(), + &StringArray::from(vec!["a", "b", "c", "d"]) + ); + assert_eq!( + actual.column(2).as_binary::<i32>(), + &BinaryArray::from_iter_values([ + b"a".as_slice(), + b"b".as_slice(), + b"c".as_slice(), + b"d".as_slice(), + ]) + ); + } + + #[test] + fn cached_schema_plan_does_not_retain_batch_arrays() { + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + ArrowDataType::Int32, + false, + )])); + let array = Arc::new(Int32Array::from(vec![1, 2, 3])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap(); + let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap(); + + writer.write(&batch).unwrap(); + drop(batch); + assert_eq!( + Arc::strong_count(&array), + 1, + "schema and row-group caches must not retain input ArrayRefs" + ); + writer.close().unwrap(); + } + #[test] fn arrow_writer_page_size() { let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); @@ -5455,7 +5769,7 @@ mod tests { SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone()) .unwrap(); - let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + let factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap(); let mut col_writers = factory.create_column_writers(0).unwrap(); let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64)); for leaves in compute_leaves(schema.field(0), &arr).unwrap() { @@ -5805,7 +6119,7 @@ mod tests { fn arrow_bounds(mut file: SerializedFileWriter<Vec<u8>>, array: ArrayRef) -> Vec<u8> { let field = Field::new("value", array.data_type().clone(), false); let schema = Arc::new(Schema::new(vec![field.clone()])); - let factory = ArrowRowGroupWriterFactory::new(&file, schema); + let factory = ArrowRowGroupWriterFactory::try_new(&file, schema).unwrap(); let mut writer: ArrowColumnWriter = factory.create_column_writers(0).unwrap().remove(0); let leaf = compute_leaves(&field, &array).unwrap().remove(0); writer.write(&leaf).unwrap(); diff --git a/parquet/tests/arrow_writer/mod.rs b/parquet/tests/arrow_writer/mod.rs index 791205dfe5..8d5543f744 100644 --- a/parquet/tests/arrow_writer/mod.rs +++ b/parquet/tests/arrow_writer/mod.rs @@ -497,7 +497,7 @@ fn write_dict_columns(page_store_factory: Option<Arc<dyn PageStoreFactory>>, poo SerializedFileWriter::new(std::io::sink(), parquet_schema.root_schema_ptr(), props) .unwrap(); - let mut factory = ArrowRowGroupWriterFactory::new(&file, schema.clone()); + let mut factory = ArrowRowGroupWriterFactory::try_new(&file, schema.clone()).unwrap(); if let Some(f) = page_store_factory { factory = factory.with_page_store_factory(f); } diff --git a/parquet/tests/encryption/encryption_async.rs b/parquet/tests/encryption/encryption_async.rs index 189eaccb8b..299c1d986c 100644 --- a/parquet/tests/encryption/encryption_async.rs +++ b/parquet/tests/encryption/encryption_async.rs @@ -947,7 +947,8 @@ async fn test_concurrent_encrypted_writing_over_multiple_row_groups() { let writer = SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(schema)); + let row_group_writer_factory = + ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(schema)).unwrap(); let max_row_groups = 1; let (serialize_tx, serialize_rx) = @@ -1012,7 +1013,8 @@ async fn test_multi_threaded_encrypted_writing() { let temp_file = tempfile::tempfile().unwrap(); let mut writer = SerializedFileWriter::new(&temp_file, parquet_schema.root_schema_ptr(), props).unwrap(); - let row_group_writer_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema)); + let row_group_writer_factory = + ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap(); let (serialize_tx, mut serialize_rx) = tokio::sync::mpsc::channel::<JoinHandle<RBStreamSerializeResult>>(1);
