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-22347-b4739e5031f16e1f0c151e52c976973f99ab8642 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit c8b784a01f5d0bcbe0dac806730fb61afc0be8ef Author: Adrian Garcia Badaracco <[email protected]> AuthorDate: Mon May 18 21:15:51 2026 -0700 refactor(parquet-datasource): split sink and schema_coercion out of file_format.rs (#22347) ## Which issue does this PR close? Relates to the discussion in #22024 about the Parquet datasource crate becoming hard to navigate. Split out of #22156, which bundled several code-motion moves into one PR — this is one of three smaller, independently-reviewable PRs that replace it. ## Rationale for this change `file_format.rs` had grown to ~2,000 LOC, bundling several distinct responsibilities into one file. That makes it hard to read and hard to review changes in isolation. This PR is **pure code motion**: no behavior change and no public API change. ## What changes are included in this PR? Extracts two responsibilities from `file_format.rs` into focused modules (`file_format.rs` drops to ~660 LOC): - `sink.rs` — `ParquetSink` and the parallel-write machinery (`column_serializer_task`, `spawn_column_parallel_row_group_writer`, `output_single_parquet_file_parallelized`, `concatenate_parallel_row_groups`, etc.). - `schema_coercion.rs` — the Arrow-schema coercion utilities (`apply_file_schema_type_coercions`, `coerce_int96_to_resolution`, `coerce_file_schema_to_view_type`, `coerce_file_schema_to_string_type`, `transform_schema_to_view`, `transform_binary_to_string`, `field_with_new_type`) and their tests. Every previously-public item is still reachable at the same path: the crate root re-exports `sink::ParquetSink` and the `schema_coercion::*` functions, and the historical `file_format::ParquetSink` path is preserved via `pub use` (datafusion-proto depends on it). ## Are these changes tested? Yes, covered by existing tests (the `coerce_int96_to_resolution_*` tests moved with the function to `schema_coercion.rs`). `cargo test -p datafusion-datasource-parquet --all-features` (122 passing) and `cargo clippy -p datafusion-datasource-parquet --all-targets --all-features -- -D warnings` both pass. `datafusion-proto` (a downstream `ParquetSink` consumer) builds clean. ## Are there any user-facing changes? No. Public API is unchanged — every previously-public item is still reachable at the same crate-root path. The only difference is the file organization inside the crate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.7 (1M context) <[email protected]> --- datafusion/datasource-parquet/src/file_format.rs | 1410 +------------------- datafusion/datasource-parquet/src/mod.rs | 9 + .../datasource-parquet/src/schema_coercion.rs | 692 ++++++++++ datafusion/datasource-parquet/src/sink.rs | 754 +++++++++++ 4 files changed, 1473 insertions(+), 1392 deletions(-) diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 5676ee940c..3a15793ff1 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -17,48 +17,41 @@ //! [`ParquetFormat`]: Parquet [`FileFormat`] abstractions -use std::cell::RefCell; +use std::fmt; use std::fmt::Debug; use std::ops::Range; -use std::rc::Rc; use std::sync::Arc; -use std::{fmt, vec}; -use arrow::array::RecordBatch; -use arrow::datatypes::{Fields, Schema, SchemaRef, TimeUnit}; +// Re-export so the historical `file_format::*` paths still resolve. +#[expect(deprecated)] +pub use crate::schema_coercion::{ + apply_file_schema_type_coercions, coerce_file_schema_to_string_type, + coerce_file_schema_to_view_type, coerce_int96_to_resolution, + transform_binary_to_string, transform_schema_to_view, +}; +pub use crate::sink::ParquetSink; + +use arrow::datatypes::{Fields, Schema, SchemaRef}; use datafusion_datasource::TableSchema; use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; -use datafusion_datasource::write::{ - ObjectWriterBuilder, SharedBuffer, get_writer_schema, -}; +use datafusion_datasource::file_sink_config::FileSinkConfig; use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; -use datafusion_datasource::write::demux::DemuxedStreamReceiver; -use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::Statistics; use datafusion_common::config::{ConfigField, ConfigFileType, TableParquetOptions}; use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{ - DEFAULT_PARQUET_EXTENSION, DataFusionError, GetExt, HashSet, Result, - internal_datafusion_err, internal_err, not_impl_err, + DEFAULT_PARQUET_EXTENSION, DataFusionError, GetExt, Result, internal_datafusion_err, + internal_err, not_impl_err, }; -use datafusion_common::{HashMap, Statistics}; -use datafusion_common_runtime::{JoinSet, SpawnedTask}; -use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::sink::{DataSink, DataSinkExec}; -use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; -use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_datasource::sink::DataSinkExec; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement}; -use datafusion_physical_plan::metrics::{ - ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, - MetricsSet, Time, -}; -use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; +use datafusion_physical_plan::ExecutionPlan; use datafusion_session::Session; use crate::metadata::{DFParquetMetadata, lex_ordering_to_sorting_columns}; @@ -68,38 +61,13 @@ use async_trait::async_trait; use bytes::Bytes; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::cache::cache_manager::FileMetadataCache; -use datafusion_execution::runtime_env::RuntimeEnv; use futures::future::BoxFuture; use futures::{FutureExt, StreamExt, TryStreamExt}; -use object_store::buffered::BufWriter; use object_store::path::Path; use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt}; -use parquet::arrow::arrow_writer::{ - ArrowColumnChunk, ArrowColumnWriter, ArrowLeafColumn, ArrowRowGroupWriterFactory, - ArrowWriterOptions, compute_leaves, -}; use parquet::arrow::async_reader::MetadataFetch; -use parquet::arrow::{ArrowWriter, AsyncArrowWriter}; -use parquet::basic::Type; -#[cfg(feature = "parquet_encryption")] -use parquet::encryption::encrypt::FileEncryptionProperties; use parquet::errors::ParquetError; -use parquet::file::metadata::{ParquetMetaData, SortingColumn}; -use parquet::file::properties::{ - DEFAULT_MAX_ROW_GROUP_ROW_COUNT, WriterProperties, WriterPropertiesBuilder, -}; -use parquet::file::writer::SerializedFileWriter; -use parquet::schema::types::SchemaDescriptor; -use tokio::io::{AsyncWrite, AsyncWriteExt}; -use tokio::sync::mpsc::{self, Receiver, Sender}; - -/// Initial writing buffer size. Note this is just a size hint for efficiency. It -/// will grow beyond the set value if needed. -const INITIAL_BUFFER_BYTES: usize = 1048576; - -/// When writing parquet files in parallel, if the buffered Parquet data exceeds -/// this size, it is flushed to object store -const BUFFER_FLUSH_BYTES: usize = 1024000; +use parquet::file::metadata::ParquetMetaData; #[derive(Default)] /// Factory struct used to create [ParquetFormat] @@ -609,450 +577,6 @@ impl ParquetFormat { } } -/// Apply necessary schema type coercions to make file schema match table schema. -/// -/// This function performs two main types of transformations in a single pass: -/// 1. Binary types to string types conversion - Converts binary data types to their -/// corresponding string types when the table schema expects string data -/// 2. Regular to view types conversion - Converts standard string/binary types to -/// view types when the table schema uses view types -/// -/// # Arguments -/// * `table_schema` - The table schema containing the desired types -/// * `file_schema` - The file schema to be transformed -/// -/// # Returns -/// * `Some(Schema)` - If any transformations were applied, returns the transformed schema -/// * `None` - If no transformations were needed -pub fn apply_file_schema_type_coercions( - table_schema: &Schema, - file_schema: &Schema, -) -> Option<Schema> { - let mut needs_view_transform = false; - let mut needs_string_transform = false; - - // Create a mapping of table field names to their data types for fast lookup - // and simultaneously check if we need any transformations - let table_fields: HashMap<_, _> = table_schema - .fields() - .iter() - .map(|f| { - let dt = f.data_type(); - // Check if we need view type transformation - if matches!(dt, &DataType::Utf8View | &DataType::BinaryView) { - needs_view_transform = true; - } - // Check if we need string type transformation - if matches!( - dt, - &DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View - ) { - needs_string_transform = true; - } - - (f.name(), dt) - }) - .collect(); - - // Early return if no transformation needed - if !needs_view_transform && !needs_string_transform { - return None; - } - - let transformed_fields: Vec<Arc<Field>> = file_schema - .fields() - .iter() - .map(|field| { - let field_name = field.name(); - let field_type = field.data_type(); - - // Look up the corresponding field type in the table schema - if let Some(table_type) = table_fields.get(field_name) { - match (table_type, field_type) { - // table schema uses string type, coerce the file schema to use string type - ( - &DataType::Utf8, - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - return field_with_new_type(field, DataType::Utf8); - } - // table schema uses large string type, coerce the file schema to use large string type - ( - &DataType::LargeUtf8, - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - return field_with_new_type(field, DataType::LargeUtf8); - } - // table schema uses string view type, coerce the file schema to use view type - ( - &DataType::Utf8View, - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - return field_with_new_type(field, DataType::Utf8View); - } - // Handle view type conversions - (&DataType::Utf8View, DataType::Utf8 | DataType::LargeUtf8) => { - return field_with_new_type(field, DataType::Utf8View); - } - (&DataType::BinaryView, DataType::Binary | DataType::LargeBinary) => { - return field_with_new_type(field, DataType::BinaryView); - } - _ => {} - } - } - - // If no transformation is needed, keep the original field - Arc::clone(field) - }) - .collect(); - - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) -} - -/// Coerces the file schema's Timestamps to the provided TimeUnit if Parquet schema contains INT96. -pub fn coerce_int96_to_resolution( - parquet_schema: &SchemaDescriptor, - file_schema: &Schema, - time_unit: &TimeUnit, -) -> Option<Schema> { - // Traverse the parquet_schema columns looking for int96 physical types. If encountered, insert - // the field's full path into a set. - let int96_fields: HashSet<_> = parquet_schema - .columns() - .iter() - .filter(|f| f.physical_type() == Type::INT96) - .map(|f| f.path().string()) - .collect(); - - if int96_fields.is_empty() { - // The schema doesn't contain any int96 fields, so skip the remaining logic. - return None; - } - - // Do a DFS into the schema using a stack, looking for timestamp(nanos) fields that originated - // as int96 to coerce to the provided time_unit. - - type NestedFields = Rc<RefCell<Vec<FieldRef>>>; - type StackContext<'a> = ( - Vec<&'a str>, // The Parquet column path (e.g., "c0.list.element.c1") for the current field. - &'a FieldRef, // The current field to be processed. - NestedFields, // The parent's fields that this field will be (possibly) type-coerced and - // inserted into. All fields have a parent, so this is not an Option type. - Option<NestedFields>, // Nested types need to create their own vector of fields for their - // children. For primitive types this will remain None. For nested - // types it is None the first time they are processed. Then, we - // instantiate a vector for its children, push the field back onto the - // stack to be processed again, and DFS into its children. The next - // time we process the field, we know we have DFS'd into the children - // because this field is Some. - ); - - // This is our top-level fields from which we will construct our schema. We pass this into our - // initial stack context as the parent fields, and the DFS populates it. - let fields = Rc::new(RefCell::new(Vec::with_capacity(file_schema.fields.len()))); - - // TODO: It might be possible to only DFS into nested fields that we know contain an int96 if we - // use some sort of LPM data structure to check if we're currently DFS'ing nested types that are - // in a column path that contains an int96. That can be a future optimization for large schemas. - let transformed_schema = { - // Populate the stack with our top-level fields. - let mut stack: Vec<StackContext> = file_schema - .fields() - .iter() - .rev() - .map(|f| (vec![f.name().as_str()], f, Rc::clone(&fields), None)) - .collect(); - - // Pop fields to DFS into until we have exhausted the stack. - while let Some((parquet_path, current_field, parent_fields, child_fields)) = - stack.pop() - { - match (current_field.data_type(), child_fields) { - (DataType::Struct(unprocessed_children), None) => { - // This is the first time popping off this struct. We don't yet know the - // correct types of its children (i.e., if they need coercing) so we create - // a vector for child_fields, push the struct node back onto the stack to be - // processed again (see below) after processing all its children. - let child_fields = Rc::new(RefCell::new(Vec::with_capacity( - unprocessed_children.len(), - ))); - // Note that here we push the struct back onto the stack with its - // parent_fields in the same position, now with Some(child_fields). - stack.push(( - parquet_path.clone(), - current_field, - parent_fields, - Some(Rc::clone(&child_fields)), - )); - // Push all the children in reverse to maintain original schema order due to - // stack processing. - for child in unprocessed_children.into_iter().rev() { - let mut child_path = parquet_path.clone(); - // Build up a normalized path that we'll use as a key into the original - // int96_fields set above to test if this originated as int96. - child_path.push("."); - child_path.push(child.name()); - // Note that here we push the field onto the stack using the struct's - // new child_fields vector as the field's parent_fields. - stack.push((child_path, child, Rc::clone(&child_fields), None)); - } - } - (DataType::Struct(unprocessed_children), Some(processed_children)) => { - // This is the second time popping off this struct. The child_fields vector - // now contains each field that has been DFS'd into, and we can construct - // the resulting struct with correct child types. - let processed_children = processed_children.borrow(); - assert_eq!(processed_children.len(), unprocessed_children.len()); - let processed_struct = Field::new_struct( - current_field.name(), - processed_children.as_slice(), - current_field.is_nullable(), - ); - parent_fields.borrow_mut().push(Arc::new(processed_struct)); - } - (DataType::List(unprocessed_child), None) => { - // This is the first time popping off this list. See struct docs above. - let child_fields = Rc::new(RefCell::new(Vec::with_capacity(1))); - stack.push(( - parquet_path.clone(), - current_field, - parent_fields, - Some(Rc::clone(&child_fields)), - )); - let mut child_path = parquet_path.clone(); - // Spark uses a definition for arrays/lists that results in a group - // named "list" that is not maintained when parsing to Arrow. We just push - // this name into the path. - child_path.push(".list."); - child_path.push(unprocessed_child.name()); - stack.push(( - child_path.clone(), - unprocessed_child, - Rc::clone(&child_fields), - None, - )); - } - (DataType::List(_), Some(processed_children)) => { - // This is the second time popping off this list. See struct docs above. - let processed_children = processed_children.borrow(); - assert_eq!(processed_children.len(), 1); - let processed_list = Field::new_list( - current_field.name(), - Arc::clone(&processed_children[0]), - current_field.is_nullable(), - ); - parent_fields.borrow_mut().push(Arc::new(processed_list)); - } - (DataType::Map(unprocessed_child, _), None) => { - // This is the first time popping off this map. See struct docs above. - let child_fields = Rc::new(RefCell::new(Vec::with_capacity(1))); - stack.push(( - parquet_path.clone(), - current_field, - parent_fields, - Some(Rc::clone(&child_fields)), - )); - let mut child_path = parquet_path.clone(); - child_path.push("."); - child_path.push(unprocessed_child.name()); - stack.push(( - child_path.clone(), - unprocessed_child, - Rc::clone(&child_fields), - None, - )); - } - (DataType::Map(_, sorted), Some(processed_children)) => { - // This is the second time popping off this map. See struct docs above. - let processed_children = processed_children.borrow(); - assert_eq!(processed_children.len(), 1); - let processed_map = Field::new( - current_field.name(), - DataType::Map(Arc::clone(&processed_children[0]), *sorted), - current_field.is_nullable(), - ); - parent_fields.borrow_mut().push(Arc::new(processed_map)); - } - (DataType::Timestamp(TimeUnit::Nanosecond, None), None) - if int96_fields.contains(parquet_path.concat().as_str()) => - // We found a timestamp(nanos) and it originated as int96. Coerce it to the correct - // time_unit. - { - parent_fields.borrow_mut().push(field_with_new_type( - current_field, - DataType::Timestamp(*time_unit, None), - )); - } - // Other types can be cloned as they are. - _ => parent_fields.borrow_mut().push(Arc::clone(current_field)), - } - } - assert_eq!(fields.borrow().len(), file_schema.fields.len()); - Schema::new_with_metadata( - fields.borrow_mut().clone(), - file_schema.metadata.clone(), - ) - }; - - Some(transformed_schema) -} - -/// Coerces the file schema if the table schema uses a view type. -#[deprecated( - since = "47.0.0", - note = "Use `apply_file_schema_type_coercions` instead" -)] -pub fn coerce_file_schema_to_view_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option<Schema> { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| { - let dt = f.data_type(); - if dt.equals_datatype(&DataType::Utf8View) - || dt.equals_datatype(&DataType::BinaryView) - { - transform = true; - } - (f.name(), dt) - }) - .collect(); - - if !transform { - return None; - } - - let transformed_fields: Vec<Arc<Field>> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - (Some(DataType::Utf8View), DataType::Utf8 | DataType::LargeUtf8) => { - field_with_new_type(field, DataType::Utf8View) - } - ( - Some(DataType::BinaryView), - DataType::Binary | DataType::LargeBinary, - ) => field_with_new_type(field, DataType::BinaryView), - _ => Arc::clone(field), - }, - ) - .collect(); - - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) -} - -/// If the table schema uses a string type, coerce the file schema to use a string type. -/// -/// See [ParquetFormat::binary_as_string] for details -#[deprecated( - since = "47.0.0", - note = "Use `apply_file_schema_type_coercions` instead" -)] -pub fn coerce_file_schema_to_string_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option<Schema> { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| (f.name(), f.data_type())) - .collect(); - let transformed_fields: Vec<Arc<Field>> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - // table schema uses string type, coerce the file schema to use string type - ( - Some(DataType::Utf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8) - } - // table schema uses large string type, coerce the file schema to use large string type - ( - Some(DataType::LargeUtf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::LargeUtf8) - } - // table schema uses string view type, coerce the file schema to use view type - ( - Some(DataType::Utf8View), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8View) - } - _ => Arc::clone(field), - }, - ) - .collect(); - - if !transform { - None - } else { - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) - } -} - -/// Create a new field with the specified data type, copying the other -/// properties from the input field -fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { - Arc::new(field.as_ref().clone().with_data_type(new_type)) -} - -/// Transform a schema to use view types for Utf8 and Binary -/// -/// See [ParquetFormat::force_view_types] for details -pub fn transform_schema_to_view(schema: &Schema) -> Schema { - let transformed_fields: Vec<Arc<Field>> = schema - .fields - .iter() - .map(|field| match field.data_type() { - DataType::Utf8 | DataType::LargeUtf8 => { - field_with_new_type(field, DataType::Utf8View) - } - DataType::Binary | DataType::LargeBinary => { - field_with_new_type(field, DataType::BinaryView) - } - _ => Arc::clone(field), - }) - .collect(); - Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) -} - -/// Transform a schema so that any binary types are strings -pub fn transform_binary_to_string(schema: &Schema) -> Schema { - let transformed_fields: Vec<Arc<Field>> = schema - .fields - .iter() - .map(|field| match field.data_type() { - DataType::Binary => field_with_new_type(field, DataType::Utf8), - DataType::LargeBinary => field_with_new_type(field, DataType::LargeUtf8), - DataType::BinaryView => field_with_new_type(field, DataType::Utf8View), - _ => Arc::clone(field), - }) - .collect(); - Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) -} - /// [`MetadataFetch`] adapter for reading bytes from an [`ObjectStore`] pub struct ObjectStoreFetch<'a> { store: &'a dyn ObjectStore, @@ -1138,901 +662,3 @@ pub fn statistics_from_parquet_meta_calc( ) -> Result<Statistics> { DFParquetMetadata::statistics_from_parquet_metadata(metadata, &table_schema) } - -/// Implements [`DataSink`] for writing to a parquet file. -pub struct ParquetSink { - /// Config options for writing data - config: FileSinkConfig, - /// Underlying parquet options - parquet_options: TableParquetOptions, - /// File metadata from successfully produced parquet files. The Mutex is only used - /// to allow inserting to HashMap from behind borrowed reference in DataSink::write_all. - written: Arc<parking_lot::Mutex<HashMap<Path, ParquetMetaData>>>, - /// Optional sorting columns to write to Parquet metadata - sorting_columns: Option<Vec<SortingColumn>>, - /// Metrics for tracking write operations - metrics: ExecutionPlanMetricsSet, -} - -impl Debug for ParquetSink { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("ParquetSink").finish() - } -} - -impl DisplayAs for ParquetSink { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!(f, "ParquetSink(file_groups=",)?; - FileGroupDisplay(&self.config.file_group).fmt_as(t, f)?; - write!(f, ")") - } - DisplayFormatType::TreeRender => { - // TODO: collect info - write!(f, "") - } - } - } -} - -impl ParquetSink { - /// Create from config. - pub fn new(config: FileSinkConfig, parquet_options: TableParquetOptions) -> Self { - Self { - config, - parquet_options, - written: Default::default(), - sorting_columns: None, - metrics: ExecutionPlanMetricsSet::new(), - } - } - - /// Set sorting columns for the Parquet file metadata. - pub fn with_sorting_columns( - mut self, - sorting_columns: Option<Vec<SortingColumn>>, - ) -> Self { - self.sorting_columns = sorting_columns; - self - } - - /// Retrieve the file metadata for the written files, keyed to the path - /// which may be partitioned (in the case of hive style partitioning). - pub fn written(&self) -> HashMap<Path, ParquetMetaData> { - self.written.lock().clone() - } - - /// Create writer properties based upon configuration settings, - /// including partitioning and the inclusion of arrow schema metadata. - async fn create_writer_props( - &self, - runtime: &Arc<RuntimeEnv>, - path: &Path, - ) -> Result<WriterProperties> { - let schema = self.config.output_schema(); - - // TODO: avoid this clone in follow up PR, where the writer properties & schema - // are calculated once on `ParquetSink::new` - let mut parquet_opts = self.parquet_options.clone(); - if !self.parquet_options.global.skip_arrow_metadata { - parquet_opts.arrow_schema(schema); - } - - let mut builder = WriterPropertiesBuilder::try_from(&parquet_opts)?; - - // Set sorting columns if configured - if let Some(ref sorting_columns) = self.sorting_columns { - builder = builder.set_sorting_columns(Some(sorting_columns.clone())); - } - - builder = set_writer_encryption_properties( - builder, - runtime, - parquet_opts, - schema, - path, - ) - .await?; - Ok(builder.build()) - } - - /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore - /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized - async fn create_async_arrow_writer( - &self, - location: &Path, - object_store: Arc<dyn ObjectStore>, - context: &Arc<TaskContext>, - parquet_props: WriterProperties, - ) -> Result<AsyncArrowWriter<BufWriter>> { - let buf_writer = BufWriter::with_capacity( - object_store, - location.clone(), - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - ); - let options = ArrowWriterOptions::new() - .with_properties(parquet_props) - .with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata); - - let writer = AsyncArrowWriter::try_new_with_options( - buf_writer, - get_writer_schema(&self.config), - options, - )?; - Ok(writer) - } - - /// Parquet options - pub fn parquet_options(&self) -> &TableParquetOptions { - &self.parquet_options - } -} - -#[cfg(feature = "parquet_encryption")] -async fn set_writer_encryption_properties( - builder: WriterPropertiesBuilder, - runtime: &Arc<RuntimeEnv>, - parquet_opts: TableParquetOptions, - schema: &Arc<Schema>, - path: &Path, -) -> Result<WriterPropertiesBuilder> { - if let Some(file_encryption_properties) = parquet_opts.crypto.file_encryption { - // Encryption properties have been specified directly - return Ok(builder.with_file_encryption_properties(Arc::new( - FileEncryptionProperties::try_from(file_encryption_properties)?, - ))); - } else if let Some(encryption_factory_id) = &parquet_opts.crypto.factory_id.as_ref() { - // Encryption properties will be generated by an encryption factory - let encryption_factory = - runtime.parquet_encryption_factory(encryption_factory_id)?; - let file_encryption_properties = encryption_factory - .get_file_encryption_properties( - &parquet_opts.crypto.factory_options, - schema, - path, - ) - .await?; - if let Some(file_encryption_properties) = file_encryption_properties { - return Ok( - builder.with_file_encryption_properties(file_encryption_properties) - ); - } - } - Ok(builder) -} - -#[cfg(not(feature = "parquet_encryption"))] -async fn set_writer_encryption_properties( - builder: WriterPropertiesBuilder, - _runtime: &Arc<RuntimeEnv>, - _parquet_opts: TableParquetOptions, - _schema: &Arc<Schema>, - _path: &Path, -) -> Result<WriterPropertiesBuilder> { - Ok(builder) -} - -#[async_trait] -impl FileSink for ParquetSink { - fn config(&self) -> &FileSinkConfig { - &self.config - } - - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc<TaskContext>, - demux_task: SpawnedTask<Result<()>>, - mut file_stream_rx: DemuxedStreamReceiver, - object_store: Arc<dyn ObjectStore>, - ) -> Result<u64> { - let rows_written_counter = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Rows) - .global_counter("rows_written"); - // Note: bytes_written is the sum of compressed row group sizes, which - // may differ slightly from the actual on-disk file size (excludes footer, - // page indexes, and other Parquet metadata overhead). - let bytes_written_counter = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Bytes) - .global_counter("bytes_written"); - let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0); - - let parquet_opts = &self.parquet_options; - - let mut file_write_tasks: JoinSet< - std::result::Result<(Path, ParquetMetaData), DataFusionError>, - > = JoinSet::new(); - - let runtime = context.runtime_env(); - let parallel_options = ParallelParquetWriterOptions { - max_parallel_row_groups: parquet_opts - .global - .maximum_parallel_row_group_writers, - max_buffered_record_batches_per_stream: parquet_opts - .global - .maximum_buffered_record_batches_per_stream, - }; - - while let Some((path, mut rx)) = file_stream_rx.recv().await { - let parquet_props = self.create_writer_props(&runtime, &path).await?; - // CDC requires the sequential writer: the chunker state lives in ArrowWriter - // and persists across row groups. The parallel path bypasses ArrowWriter entirely. - if !parquet_opts.global.allow_single_file_parallelism - || parquet_opts.global.use_content_defined_chunking.is_some() - { - let mut writer = self - .create_async_arrow_writer( - &path, - Arc::clone(&object_store), - context, - parquet_props.clone(), - ) - .await?; - let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) - .register(context.memory_pool()); - file_write_tasks.spawn( - async move { - while let Some(batch) = rx.recv().await { - writer.write(&batch).await?; - reservation.try_resize(writer.memory_size())?; - } - let parquet_meta_data = writer - .close() - .await - .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; - Ok((path, parquet_meta_data)) - } - .with_elapsed_compute(elapsed_compute.clone()), - ); - } else { - let writer = ObjectWriterBuilder::new( - // Parquet files as a whole are never compressed, since they - // manage compressed blocks themselves. - FileCompressionType::UNCOMPRESSED, - &path, - Arc::clone(&object_store), - ) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .build()?; - let ctx = ParquetFileWriteContext { - schema: get_writer_schema(&self.config), - props: Arc::new(parquet_props), - skip_arrow_metadata: self.parquet_options.global.skip_arrow_metadata, - parallel_options: Arc::new(parallel_options.clone()), - pool: Arc::clone(context.memory_pool()), - }; - let encoding_time = elapsed_compute.clone(); - file_write_tasks.spawn(async move { - let parquet_meta_data = output_single_parquet_file_parallelized( - writer, - rx, - ctx, - encoding_time, - ) - .await?; - Ok((path, parquet_meta_data)) - }); - } - } - - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - let (path, parquet_meta_data) = r?; - let file_rows = parquet_meta_data.file_metadata().num_rows() as usize; - let file_bytes: usize = parquet_meta_data - .row_groups() - .iter() - .map(|rg| rg.compressed_size() as usize) - .sum(); - rows_written_counter.add(file_rows); - bytes_written_counter.add(file_bytes); - let mut written_files = self.written.lock(); - written_files - .try_insert(path.clone(), parquet_meta_data) - .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; - drop(written_files); - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } - - demux_task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - - Ok(rows_written_counter.value() as u64) - } -} - -#[async_trait] -impl DataSink for ParquetSink { - fn metrics(&self) -> Option<MetricsSet> { - Some(self.metrics.clone_inner()) - } - - fn schema(&self) -> &SchemaRef { - self.config.output_schema() - } - - async fn write_all( - &self, - data: SendableRecordBatchStream, - context: &Arc<TaskContext>, - ) -> Result<u64> { - FileSink::write_all(self, data, context).await - } -} - -/// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] -/// Once the channel is exhausted, returns the ArrowColumnWriter. -async fn column_serializer_task( - mut rx: Receiver<ArrowLeafColumn>, - mut writer: ArrowColumnWriter, - reservation: MemoryReservation, - encoding_time: Time, -) -> Result<(ArrowColumnWriter, MemoryReservation)> { - while let Some(col) = rx.recv().await { - let _timer = encoding_time.timer(); - writer.write(&col)?; - reservation.try_resize(writer.memory_size())?; - } - Ok((writer, reservation)) -} - -type ColumnWriterTask = SpawnedTask<Result<(ArrowColumnWriter, MemoryReservation)>>; -type ColSender = Sender<ArrowLeafColumn>; - -/// Spawns a parallel serialization task for each column -/// Returns join handles for each columns serialization task along with a send channel -/// to send arrow arrays to each serialization task. -fn spawn_column_parallel_row_group_writer( - col_writers: Vec<ArrowColumnWriter>, - max_buffer_size: usize, - pool: &Arc<dyn MemoryPool>, - encoding_time: &Time, -) -> Result<(Vec<ColumnWriterTask>, Vec<ColSender>)> { - let num_columns = col_writers.len(); - - let mut col_writer_tasks = Vec::with_capacity(num_columns); - let mut col_array_channels = Vec::with_capacity(num_columns); - for writer in col_writers.into_iter() { - // Buffer size of this channel limits the number of arrays queued up for column level serialization - let (send_array, receive_array) = - mpsc::channel::<ArrowLeafColumn>(max_buffer_size); - col_array_channels.push(send_array); - - let reservation = - MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(pool); - let task = SpawnedTask::spawn(column_serializer_task( - receive_array, - writer, - reservation, - encoding_time.clone(), - )); - col_writer_tasks.push(task); - } - - Ok((col_writer_tasks, col_array_channels)) -} - -/// Settings related to writing parquet files in parallel -#[derive(Clone)] -struct ParallelParquetWriterOptions { - max_parallel_row_groups: usize, - max_buffered_record_batches_per_stream: usize, -} - -/// Write configuration inputs shared across all parallel tasks that encode a -/// single Parquet file. These values are invariant for the duration of one file -/// write and do not change per row-group or per column. -/// -/// Separating these from per-call parameters (`object_store_writer`, `data`, -/// `encoding_time`) keeps the deep parallel call chain below the argument-count -/// limit without mixing configuration with runtime state. -#[derive(Clone)] -struct ParquetFileWriteContext { - schema: Arc<Schema>, - props: Arc<WriterProperties>, - skip_arrow_metadata: bool, - parallel_options: Arc<ParallelParquetWriterOptions>, - pool: Arc<dyn MemoryPool>, -} - -/// This is the return type of calling [ArrowColumnWriter].close() on each column -/// i.e. the Vec of encoded columns which can be appended to a row group -type RBStreamSerializeResult = Result<(Vec<ArrowColumnChunk>, MemoryReservation, usize)>; - -/// Sends the ArrowArrays in passed [RecordBatch] through the channels to their respective -/// parallel column serializers. -async fn send_arrays_to_col_writers( - col_array_channels: &[ColSender], - rb: &RecordBatch, - schema: Arc<Schema>, -) -> Result<()> { - // Each leaf column has its own channel, increment next_channel for each leaf column sent. - let mut next_channel = 0; - for (array, field) in rb.columns().iter().zip(schema.fields()) { - for c in compute_leaves(field, array)? { - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if col_array_channels[next_channel].send(c).await.is_err() { - return Ok(()); - } - - next_channel += 1; - } - } - - Ok(()) -} - -/// Spawns a tokio task which joins the parallel column writer tasks, -/// and finalizes the row group -fn spawn_rg_join_and_finalize_task( - column_writer_tasks: Vec<ColumnWriterTask>, - rg_rows: usize, - pool: &Arc<dyn MemoryPool>, - encoding_time: Time, -) -> SpawnedTask<RBStreamSerializeResult> { - let rg_reservation = - MemoryConsumer::new("ParquetSink(SerializedRowGroupWriter)").register(pool); - - SpawnedTask::spawn(async move { - let num_cols = column_writer_tasks.len(); - let mut finalized_rg = Vec::with_capacity(num_cols); - for task in column_writer_tasks.into_iter() { - let (writer, _col_reservation) = task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - let encoded_size = writer.get_estimated_total_bytes(); - rg_reservation.grow(encoded_size); - let _timer = encoding_time.timer(); - finalized_rg.push(writer.close()?); - } - - Ok((finalized_rg, rg_reservation, rg_rows)) - }) -} - -/// This task coordinates the serialization of a parquet file in parallel. -/// As the query produces RecordBatches, these are written to a RowGroup -/// via parallel [ArrowColumnWriter] tasks. Once the desired max rows per -/// row group is reached, the parallel tasks are joined on another separate task -/// and sent to a concatenation task. This task immediately continues to work -/// on the next row group in parallel. So, parquet serialization is parallelized -/// across both columns and row_groups, with a theoretical max number of parallel tasks -/// given by n_columns * num_row_groups. -fn spawn_parquet_parallel_serialization_task( - row_group_writer_factory: ArrowRowGroupWriterFactory, - mut data: Receiver<RecordBatch>, - serialize_tx: Sender<SpawnedTask<RBStreamSerializeResult>>, - ctx: ParquetFileWriteContext, - encoding_time: Time, -) -> SpawnedTask<Result<(), DataFusionError>> { - SpawnedTask::spawn(async move { - let max_buffer_rb = ctx.parallel_options.max_buffered_record_batches_per_stream; - let max_row_group_rows = ctx - .props - .max_row_group_row_count() - .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT); - let mut row_group_index = 0; - let col_writers = - row_group_writer_factory.create_column_writers(row_group_index)?; - let (mut column_writer_handles, mut col_array_channels) = - spawn_column_parallel_row_group_writer( - col_writers, - max_buffer_rb, - &ctx.pool, - &encoding_time, - )?; - let mut current_rg_rows = 0; - - while let Some(mut rb) = data.recv().await { - // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case - // when max_row_group_rows < execution.batch_size as an alternative to a recursive async - // function. - loop { - if current_rg_rows + rb.num_rows() < max_row_group_rows { - send_arrays_to_col_writers( - &col_array_channels, - &rb, - Arc::clone(&ctx.schema), - ) - .await?; - current_rg_rows += rb.num_rows(); - break; - } else { - let rows_left = max_row_group_rows - current_rg_rows; - let a = rb.slice(0, rows_left); - send_arrays_to_col_writers( - &col_array_channels, - &a, - Arc::clone(&ctx.schema), - ) - .await?; - - // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup - // on a separate task, so that we can immediately start on the next RG before waiting - // for the current one to finish. - drop(col_array_channels); - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - max_row_group_rows, - &ctx.pool, - encoding_time.clone(), - ); - - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { - return Ok(()); - } - - current_rg_rows = 0; - rb = rb.slice(rows_left, rb.num_rows() - rows_left); - - row_group_index += 1; - let col_writers = row_group_writer_factory - .create_column_writers(row_group_index)?; - (column_writer_handles, col_array_channels) = - spawn_column_parallel_row_group_writer( - col_writers, - max_buffer_rb, - &ctx.pool, - &encoding_time, - )?; - } - } - } - - drop(col_array_channels); - // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows - if current_rg_rows > 0 { - let finalize_rg_task = spawn_rg_join_and_finalize_task( - column_writer_handles, - current_rg_rows, - &ctx.pool, - encoding_time.clone(), - ); - - // Do not surface error from closed channel (means something - // else hit an error, and the plan is shutting down). - if serialize_tx.send(finalize_rg_task).await.is_err() { - return Ok(()); - } - } - - Ok(()) - }) -} - -/// Consume RowGroups serialized by other parallel tasks and concatenate them in -/// to the final parquet file, while flushing finalized bytes to an [ObjectStore] -async fn concatenate_parallel_row_groups( - mut parquet_writer: SerializedFileWriter<SharedBuffer>, - merged_buff: SharedBuffer, - mut serialize_rx: Receiver<SpawnedTask<RBStreamSerializeResult>>, - mut object_store_writer: Box<dyn AsyncWrite + Send + Unpin>, - pool: Arc<dyn MemoryPool>, -) -> Result<ParquetMetaData> { - let file_reservation = - MemoryConsumer::new("ParquetSink(SerializedFileWriter)").register(&pool); - - while let Some(task) = serialize_rx.recv().await { - let result = task.join_unwind().await; - let (serialized_columns, rg_reservation, _cnt) = - result.map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - - let mut rg_out = parquet_writer.next_row_group()?; - for chunk in serialized_columns { - chunk.append_to_row_group(&mut rg_out)?; - rg_reservation.free(); - - let mut buff_to_flush = merged_buff.buffer.try_lock().unwrap(); - file_reservation.try_resize(buff_to_flush.len())?; - - if buff_to_flush.len() > BUFFER_FLUSH_BYTES { - object_store_writer - .write_all(buff_to_flush.as_slice()) - .await?; - buff_to_flush.clear(); - file_reservation.try_resize(buff_to_flush.len())?; // will set to zero - } - } - rg_out.close()?; - } - - let parquet_meta_data = parquet_writer.close()?; - let final_buff = merged_buff.buffer.try_lock().unwrap(); - - object_store_writer.write_all(final_buff.as_slice()).await?; - object_store_writer.shutdown().await?; - file_reservation.free(); - - Ok(parquet_meta_data) -} - -/// Parallelizes the serialization of a single parquet file, by first serializing N -/// independent RecordBatch streams in parallel to RowGroups in memory. Another -/// task then stitches these independent RowGroups together and streams this large -/// single parquet file to an ObjectStore in multiple parts. -async fn output_single_parquet_file_parallelized( - object_store_writer: Box<dyn AsyncWrite + Send + Unpin>, - data: Receiver<RecordBatch>, - ctx: ParquetFileWriteContext, - encoding_time: Time, -) -> Result<ParquetMetaData> { - let max_rowgroups = ctx.parallel_options.max_parallel_row_groups; - // Buffer size of this channel limits maximum number of RowGroups being worked on in parallel - let (serialize_tx, serialize_rx) = - mpsc::channel::<SpawnedTask<RBStreamSerializeResult>>(max_rowgroups); - - let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES); - let options = ArrowWriterOptions::new() - .with_properties((*ctx.props).clone()) - .with_skip_arrow_metadata(ctx.skip_arrow_metadata); - let writer = ArrowWriter::try_new_with_options( - merged_buff.clone(), - Arc::clone(&ctx.schema), - options, - )?; - let (writer, row_group_writer_factory) = writer.into_serialized_writer()?; - - let pool = Arc::clone(&ctx.pool); - let launch_serialization_task = spawn_parquet_parallel_serialization_task( - row_group_writer_factory, - data, - serialize_tx, - ctx, - encoding_time, - ); - let parquet_meta_data = concatenate_parallel_row_groups( - writer, - merged_buff, - serialize_rx, - object_store_writer, - pool, - ) - .await?; - - launch_serialization_task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - Ok(parquet_meta_data) -} - -#[cfg(test)] -mod tests { - use parquet::arrow::parquet_to_arrow_schema; - - use super::*; - - use parquet::schema::parser::parse_message_type; - - #[test] - fn coerce_int96_to_resolution_with_mixed_timestamps() { - // Unclear if Spark (or other writer) could generate a file with mixed timestamps like this, - // but we want to test the scenario just in case since it's at least a valid schema as far - // as the Parquet spec is concerned. - let spark_schema = " - message spark_schema { - optional int96 c0; - optional int64 c1 (TIMESTAMP(NANOS,true)); - optional int64 c2 (TIMESTAMP(NANOS,false)); - optional int64 c3 (TIMESTAMP(MILLIS,true)); - optional int64 c4 (TIMESTAMP(MILLIS,false)); - optional int64 c5 (TIMESTAMP(MICROS,true)); - optional int64 c6 (TIMESTAMP(MICROS,false)); - } - "; - - let schema = parse_message_type(spark_schema).expect("should parse schema"); - let descr = SchemaDescriptor::new(Arc::new(schema)); - - let arrow_schema = parquet_to_arrow_schema(&descr, None).unwrap(); - - let result = - coerce_int96_to_resolution(&descr, &arrow_schema, &TimeUnit::Microsecond) - .unwrap(); - - // Only the first field (c0) should be converted to a microsecond timestamp because it's the - // only timestamp that originated from an INT96. - let expected_schema = Schema::new(vec![ - Field::new("c0", DataType::Timestamp(TimeUnit::Microsecond, None), true), - Field::new( - "c1", - DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), - true, - ), - Field::new("c2", DataType::Timestamp(TimeUnit::Nanosecond, None), true), - Field::new( - "c3", - DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), - true, - ), - Field::new("c4", DataType::Timestamp(TimeUnit::Millisecond, None), true), - Field::new( - "c5", - DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), - true, - ), - Field::new("c6", DataType::Timestamp(TimeUnit::Microsecond, None), true), - ]); - - assert_eq!(result, expected_schema); - } - - #[test] - fn coerce_int96_to_resolution_with_nested_types() { - // This schema is derived from Comet's CometFuzzTestSuite ParquetGenerator only using int96 - // primitive types with generateStruct, generateArray, and generateMap set to true, with one - // additional field added to c4's struct to make sure all fields in a struct get modified. - // https://github.com/apache/datafusion-comet/blob/main/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala - let spark_schema = " - message spark_schema { - optional int96 c0; - optional group c1 { - optional int96 c0; - } - optional group c2 { - optional group c0 (LIST) { - repeated group list { - optional int96 element; - } - } - } - optional group c3 (LIST) { - repeated group list { - optional int96 element; - } - } - optional group c4 (LIST) { - repeated group list { - optional group element { - optional int96 c0; - optional int96 c1; - } - } - } - optional group c5 (MAP) { - repeated group key_value { - required int96 key; - optional int96 value; - } - } - optional group c6 (LIST) { - repeated group list { - optional group element (MAP) { - repeated group key_value { - required int96 key; - optional int96 value; - } - } - } - } - } - "; - - let schema = parse_message_type(spark_schema).expect("should parse schema"); - let descr = SchemaDescriptor::new(Arc::new(schema)); - - let arrow_schema = parquet_to_arrow_schema(&descr, None).unwrap(); - - let result = - coerce_int96_to_resolution(&descr, &arrow_schema, &TimeUnit::Microsecond) - .unwrap(); - - let expected_schema = Schema::new(vec![ - Field::new("c0", DataType::Timestamp(TimeUnit::Microsecond, None), true), - Field::new_struct( - "c1", - vec![Field::new( - "c0", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - )], - true, - ), - Field::new_struct( - "c2", - vec![Field::new_list( - "c0", - Field::new( - "element", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - true, - )], - true, - ), - Field::new_list( - "c3", - Field::new( - "element", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - true, - ), - Field::new_list( - "c4", - Field::new_struct( - "element", - vec![ - Field::new( - "c0", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - Field::new( - "c1", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - ], - true, - ), - true, - ), - Field::new_map( - "c5", - "key_value", - Field::new( - "key", - DataType::Timestamp(TimeUnit::Microsecond, None), - false, - ), - Field::new( - "value", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - false, - true, - ), - Field::new_list( - "c6", - Field::new_map( - "element", - "key_value", - Field::new( - "key", - DataType::Timestamp(TimeUnit::Microsecond, None), - false, - ), - Field::new( - "value", - DataType::Timestamp(TimeUnit::Microsecond, None), - true, - ), - false, - true, - ), - true, - ), - ]); - - assert_eq!(result, expected_schema); - } -} diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index f9329cf74d..af2e103e6a 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -34,6 +34,8 @@ mod push_decoder; mod reader; mod row_filter; mod row_group_filter; +mod schema_coercion; +mod sink; mod sort; pub mod source; mod supported_predicates; @@ -47,4 +49,11 @@ pub use reader::*; // Expose so downstream crates can use it pub use row_filter::build_row_filter; pub use row_filter::can_expr_be_pushed_down_with_schemas; pub use row_group_filter::RowGroupAccessPlanFilter; +#[expect(deprecated)] +pub use schema_coercion::{ + apply_file_schema_type_coercions, coerce_file_schema_to_string_type, + coerce_file_schema_to_view_type, coerce_int96_to_resolution, + transform_binary_to_string, transform_schema_to_view, +}; +pub use sink::ParquetSink; pub use writer::plan_to_parquet; diff --git a/datafusion/datasource-parquet/src/schema_coercion.rs b/datafusion/datasource-parquet/src/schema_coercion.rs new file mode 100644 index 0000000000..e38ef5332e --- /dev/null +++ b/datafusion/datasource-parquet/src/schema_coercion.rs @@ -0,0 +1,692 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Arrow-schema coercion utilities used by the Parquet reader to make a +//! file schema match the table schema (binary→string, regular→view, +//! INT96→Timestamp). +//! +//! These helpers are independent of the [`ParquetFormat`](crate::file_format::ParquetFormat) +//! type and several have been re-exported at the crate root for use by +//! callers outside the format implementation. + +use std::cell::RefCell; +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef, Schema, TimeUnit}; +use parquet::basic::Type; +use parquet::schema::types::SchemaDescriptor; + +/// Apply necessary schema type coercions to make file schema match table schema. +/// +/// This function performs two main types of transformations in a single pass: +/// 1. Binary types to string types conversion - Converts binary data types to their +/// corresponding string types when the table schema expects string data +/// 2. Regular to view types conversion - Converts standard string/binary types to +/// view types when the table schema uses view types +/// +/// # Arguments +/// * `table_schema` - The table schema containing the desired types +/// * `file_schema` - The file schema to be transformed +/// +/// # Returns +/// * `Some(Schema)` - If any transformations were applied, returns the transformed schema +/// * `None` - If no transformations were needed +pub fn apply_file_schema_type_coercions( + table_schema: &Schema, + file_schema: &Schema, +) -> Option<Schema> { + let mut needs_view_transform = false; + let mut needs_string_transform = false; + + // Create a mapping of table field names to their data types for fast lookup + // and simultaneously check if we need any transformations + let table_fields: HashMap<_, _> = table_schema + .fields() + .iter() + .map(|f| { + let dt = f.data_type(); + // Check if we need view type transformation + if matches!(dt, &DataType::Utf8View | &DataType::BinaryView) { + needs_view_transform = true; + } + // Check if we need string type transformation + if matches!( + dt, + &DataType::Utf8 | &DataType::LargeUtf8 | &DataType::Utf8View + ) { + needs_string_transform = true; + } + + (f.name(), dt) + }) + .collect(); + + // Early return if no transformation needed + if !needs_view_transform && !needs_string_transform { + return None; + } + + let transformed_fields: Vec<Arc<Field>> = file_schema + .fields() + .iter() + .map(|field| { + let field_name = field.name(); + let field_type = field.data_type(); + + // Look up the corresponding field type in the table schema + if let Some(table_type) = table_fields.get(field_name) { + match (table_type, field_type) { + // table schema uses string type, coerce the file schema to use string type + ( + &DataType::Utf8, + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + return field_with_new_type(field, DataType::Utf8); + } + // table schema uses large string type, coerce the file schema to use large string type + ( + &DataType::LargeUtf8, + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + return field_with_new_type(field, DataType::LargeUtf8); + } + // table schema uses string view type, coerce the file schema to use view type + ( + &DataType::Utf8View, + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + return field_with_new_type(field, DataType::Utf8View); + } + // Handle view type conversions + (&DataType::Utf8View, DataType::Utf8 | DataType::LargeUtf8) => { + return field_with_new_type(field, DataType::Utf8View); + } + (&DataType::BinaryView, DataType::Binary | DataType::LargeBinary) => { + return field_with_new_type(field, DataType::BinaryView); + } + _ => {} + } + } + + // If no transformation is needed, keep the original field + Arc::clone(field) + }) + .collect(); + + Some(Schema::new_with_metadata( + transformed_fields, + file_schema.metadata.clone(), + )) +} + +/// Coerces the file schema's Timestamps to the provided TimeUnit if Parquet schema contains INT96. +pub fn coerce_int96_to_resolution( + parquet_schema: &SchemaDescriptor, + file_schema: &Schema, + time_unit: &TimeUnit, +) -> Option<Schema> { + // Traverse the parquet_schema columns looking for int96 physical types. If encountered, insert + // the field's full path into a set. + let int96_fields: HashSet<_> = parquet_schema + .columns() + .iter() + .filter(|f| f.physical_type() == Type::INT96) + .map(|f| f.path().string()) + .collect(); + + if int96_fields.is_empty() { + // The schema doesn't contain any int96 fields, so skip the remaining logic. + return None; + } + + // Do a DFS into the schema using a stack, looking for timestamp(nanos) fields that originated + // as int96 to coerce to the provided time_unit. + + type NestedFields = Rc<RefCell<Vec<FieldRef>>>; + type StackContext<'a> = ( + Vec<&'a str>, // The Parquet column path (e.g., "c0.list.element.c1") for the current field. + &'a FieldRef, // The current field to be processed. + NestedFields, // The parent's fields that this field will be (possibly) type-coerced and + // inserted into. All fields have a parent, so this is not an Option type. + Option<NestedFields>, // Nested types need to create their own vector of fields for their + // children. For primitive types this will remain None. For nested + // types it is None the first time they are processed. Then, we + // instantiate a vector for its children, push the field back onto the + // stack to be processed again, and DFS into its children. The next + // time we process the field, we know we have DFS'd into the children + // because this field is Some. + ); + + // This is our top-level fields from which we will construct our schema. We pass this into our + // initial stack context as the parent fields, and the DFS populates it. + let fields = Rc::new(RefCell::new(Vec::with_capacity(file_schema.fields.len()))); + + // TODO: It might be possible to only DFS into nested fields that we know contain an int96 if we + // use some sort of LPM data structure to check if we're currently DFS'ing nested types that are + // in a column path that contains an int96. That can be a future optimization for large schemas. + let transformed_schema = { + // Populate the stack with our top-level fields. + let mut stack: Vec<StackContext> = file_schema + .fields() + .iter() + .rev() + .map(|f| (vec![f.name().as_str()], f, Rc::clone(&fields), None)) + .collect(); + + // Pop fields to DFS into until we have exhausted the stack. + while let Some((parquet_path, current_field, parent_fields, child_fields)) = + stack.pop() + { + match (current_field.data_type(), child_fields) { + (DataType::Struct(unprocessed_children), None) => { + // This is the first time popping off this struct. We don't yet know the + // correct types of its children (i.e., if they need coercing) so we create + // a vector for child_fields, push the struct node back onto the stack to be + // processed again (see below) after processing all its children. + let child_fields = Rc::new(RefCell::new(Vec::with_capacity( + unprocessed_children.len(), + ))); + // Note that here we push the struct back onto the stack with its + // parent_fields in the same position, now with Some(child_fields). + stack.push(( + parquet_path.clone(), + current_field, + parent_fields, + Some(Rc::clone(&child_fields)), + )); + // Push all the children in reverse to maintain original schema order due to + // stack processing. + for child in unprocessed_children.into_iter().rev() { + let mut child_path = parquet_path.clone(); + // Build up a normalized path that we'll use as a key into the original + // int96_fields set above to test if this originated as int96. + child_path.push("."); + child_path.push(child.name()); + // Note that here we push the field onto the stack using the struct's + // new child_fields vector as the field's parent_fields. + stack.push((child_path, child, Rc::clone(&child_fields), None)); + } + } + (DataType::Struct(unprocessed_children), Some(processed_children)) => { + // This is the second time popping off this struct. The child_fields vector + // now contains each field that has been DFS'd into, and we can construct + // the resulting struct with correct child types. + let processed_children = processed_children.borrow(); + assert_eq!(processed_children.len(), unprocessed_children.len()); + let processed_struct = Field::new_struct( + current_field.name(), + processed_children.as_slice(), + current_field.is_nullable(), + ); + parent_fields.borrow_mut().push(Arc::new(processed_struct)); + } + (DataType::List(unprocessed_child), None) => { + // This is the first time popping off this list. See struct docs above. + let child_fields = Rc::new(RefCell::new(Vec::with_capacity(1))); + stack.push(( + parquet_path.clone(), + current_field, + parent_fields, + Some(Rc::clone(&child_fields)), + )); + let mut child_path = parquet_path.clone(); + // Spark uses a definition for arrays/lists that results in a group + // named "list" that is not maintained when parsing to Arrow. We just push + // this name into the path. + child_path.push(".list."); + child_path.push(unprocessed_child.name()); + stack.push(( + child_path.clone(), + unprocessed_child, + Rc::clone(&child_fields), + None, + )); + } + (DataType::List(_), Some(processed_children)) => { + // This is the second time popping off this list. See struct docs above. + let processed_children = processed_children.borrow(); + assert_eq!(processed_children.len(), 1); + let processed_list = Field::new_list( + current_field.name(), + Arc::clone(&processed_children[0]), + current_field.is_nullable(), + ); + parent_fields.borrow_mut().push(Arc::new(processed_list)); + } + (DataType::Map(unprocessed_child, _), None) => { + // This is the first time popping off this map. See struct docs above. + let child_fields = Rc::new(RefCell::new(Vec::with_capacity(1))); + stack.push(( + parquet_path.clone(), + current_field, + parent_fields, + Some(Rc::clone(&child_fields)), + )); + let mut child_path = parquet_path.clone(); + child_path.push("."); + child_path.push(unprocessed_child.name()); + stack.push(( + child_path.clone(), + unprocessed_child, + Rc::clone(&child_fields), + None, + )); + } + (DataType::Map(_, sorted), Some(processed_children)) => { + // This is the second time popping off this map. See struct docs above. + let processed_children = processed_children.borrow(); + assert_eq!(processed_children.len(), 1); + let processed_map = Field::new( + current_field.name(), + DataType::Map(Arc::clone(&processed_children[0]), *sorted), + current_field.is_nullable(), + ); + parent_fields.borrow_mut().push(Arc::new(processed_map)); + } + (DataType::Timestamp(TimeUnit::Nanosecond, None), None) + if int96_fields.contains(parquet_path.concat().as_str()) => + // We found a timestamp(nanos) and it originated as int96. Coerce it to the correct + // time_unit. + { + parent_fields.borrow_mut().push(field_with_new_type( + current_field, + DataType::Timestamp(*time_unit, None), + )); + } + // Other types can be cloned as they are. + _ => parent_fields.borrow_mut().push(Arc::clone(current_field)), + } + } + assert_eq!(fields.borrow().len(), file_schema.fields.len()); + Schema::new_with_metadata( + fields.borrow_mut().clone(), + file_schema.metadata.clone(), + ) + }; + + Some(transformed_schema) +} + +/// Coerces the file schema if the table schema uses a view type. +#[deprecated( + since = "47.0.0", + note = "Use `apply_file_schema_type_coercions` instead" +)] +pub fn coerce_file_schema_to_view_type( + table_schema: &Schema, + file_schema: &Schema, +) -> Option<Schema> { + let mut transform = false; + let table_fields: HashMap<_, _> = table_schema + .fields + .iter() + .map(|f| { + let dt = f.data_type(); + if dt.equals_datatype(&DataType::Utf8View) + || dt.equals_datatype(&DataType::BinaryView) + { + transform = true; + } + (f.name(), dt) + }) + .collect(); + + if !transform { + return None; + } + + let transformed_fields: Vec<Arc<Field>> = file_schema + .fields + .iter() + .map( + |field| match (table_fields.get(field.name()), field.data_type()) { + (Some(DataType::Utf8View), DataType::Utf8 | DataType::LargeUtf8) => { + field_with_new_type(field, DataType::Utf8View) + } + ( + Some(DataType::BinaryView), + DataType::Binary | DataType::LargeBinary, + ) => field_with_new_type(field, DataType::BinaryView), + _ => Arc::clone(field), + }, + ) + .collect(); + + Some(Schema::new_with_metadata( + transformed_fields, + file_schema.metadata.clone(), + )) +} + +/// If the table schema uses a string type, coerce the file schema to use a string type. +/// +/// See [`ParquetFormat::binary_as_string`](crate::file_format::ParquetFormat::binary_as_string) for details +#[deprecated( + since = "47.0.0", + note = "Use `apply_file_schema_type_coercions` instead" +)] +pub fn coerce_file_schema_to_string_type( + table_schema: &Schema, + file_schema: &Schema, +) -> Option<Schema> { + let mut transform = false; + let table_fields: HashMap<_, _> = table_schema + .fields + .iter() + .map(|f| (f.name(), f.data_type())) + .collect(); + let transformed_fields: Vec<Arc<Field>> = file_schema + .fields + .iter() + .map( + |field| match (table_fields.get(field.name()), field.data_type()) { + // table schema uses string type, coerce the file schema to use string type + ( + Some(DataType::Utf8), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::Utf8) + } + // table schema uses large string type, coerce the file schema to use large string type + ( + Some(DataType::LargeUtf8), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::LargeUtf8) + } + // table schema uses string view type, coerce the file schema to use view type + ( + Some(DataType::Utf8View), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView, + ) => { + transform = true; + field_with_new_type(field, DataType::Utf8View) + } + _ => Arc::clone(field), + }, + ) + .collect(); + + if !transform { + None + } else { + Some(Schema::new_with_metadata( + transformed_fields, + file_schema.metadata.clone(), + )) + } +} + +/// Create a new field with the specified data type, copying the other +/// properties from the input field +fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { + Arc::new(field.as_ref().clone().with_data_type(new_type)) +} + +/// Transform a schema to use view types for Utf8 and Binary +/// +/// See [`ParquetFormat::force_view_types`](crate::file_format::ParquetFormat::force_view_types) for details +pub fn transform_schema_to_view(schema: &Schema) -> Schema { + let transformed_fields: Vec<Arc<Field>> = schema + .fields + .iter() + .map(|field| match field.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + field_with_new_type(field, DataType::Utf8View) + } + DataType::Binary | DataType::LargeBinary => { + field_with_new_type(field, DataType::BinaryView) + } + _ => Arc::clone(field), + }) + .collect(); + Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) +} + +/// Transform a schema so that any binary types are strings +pub fn transform_binary_to_string(schema: &Schema) -> Schema { + let transformed_fields: Vec<Arc<Field>> = schema + .fields + .iter() + .map(|field| match field.data_type() { + DataType::Binary => field_with_new_type(field, DataType::Utf8), + DataType::LargeBinary => field_with_new_type(field, DataType::LargeUtf8), + DataType::BinaryView => field_with_new_type(field, DataType::Utf8View), + _ => Arc::clone(field), + }) + .collect(); + Schema::new_with_metadata(transformed_fields, schema.metadata.clone()) +} +#[cfg(test)] +mod tests { + use parquet::arrow::parquet_to_arrow_schema; + + use super::*; + + use parquet::schema::parser::parse_message_type; + + #[test] + fn coerce_int96_to_resolution_with_mixed_timestamps() { + // Unclear if Spark (or other writer) could generate a file with mixed timestamps like this, + // but we want to test the scenario just in case since it's at least a valid schema as far + // as the Parquet spec is concerned. + let spark_schema = " + message spark_schema { + optional int96 c0; + optional int64 c1 (TIMESTAMP(NANOS,true)); + optional int64 c2 (TIMESTAMP(NANOS,false)); + optional int64 c3 (TIMESTAMP(MILLIS,true)); + optional int64 c4 (TIMESTAMP(MILLIS,false)); + optional int64 c5 (TIMESTAMP(MICROS,true)); + optional int64 c6 (TIMESTAMP(MICROS,false)); + } + "; + + let schema = parse_message_type(spark_schema).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + + let arrow_schema = parquet_to_arrow_schema(&descr, None).unwrap(); + + let result = + coerce_int96_to_resolution(&descr, &arrow_schema, &TimeUnit::Microsecond) + .unwrap(); + + // Only the first field (c0) should be converted to a microsecond timestamp because it's the + // only timestamp that originated from an INT96. + let expected_schema = Schema::new(vec![ + Field::new("c0", DataType::Timestamp(TimeUnit::Microsecond, None), true), + Field::new( + "c1", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + true, + ), + Field::new("c2", DataType::Timestamp(TimeUnit::Nanosecond, None), true), + Field::new( + "c3", + DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())), + true, + ), + Field::new("c4", DataType::Timestamp(TimeUnit::Millisecond, None), true), + Field::new( + "c5", + DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + Field::new("c6", DataType::Timestamp(TimeUnit::Microsecond, None), true), + ]); + + assert_eq!(result, expected_schema); + } + + #[test] + fn coerce_int96_to_resolution_with_nested_types() { + // This schema is derived from Comet's CometFuzzTestSuite ParquetGenerator only using int96 + // primitive types with generateStruct, generateArray, and generateMap set to true, with one + // additional field added to c4's struct to make sure all fields in a struct get modified. + // https://github.com/apache/datafusion-comet/blob/main/spark/src/main/scala/org/apache/comet/testing/ParquetGenerator.scala + let spark_schema = " + message spark_schema { + optional int96 c0; + optional group c1 { + optional int96 c0; + } + optional group c2 { + optional group c0 (LIST) { + repeated group list { + optional int96 element; + } + } + } + optional group c3 (LIST) { + repeated group list { + optional int96 element; + } + } + optional group c4 (LIST) { + repeated group list { + optional group element { + optional int96 c0; + optional int96 c1; + } + } + } + optional group c5 (MAP) { + repeated group key_value { + required int96 key; + optional int96 value; + } + } + optional group c6 (LIST) { + repeated group list { + optional group element (MAP) { + repeated group key_value { + required int96 key; + optional int96 value; + } + } + } + } + } + "; + + let schema = parse_message_type(spark_schema).expect("should parse schema"); + let descr = SchemaDescriptor::new(Arc::new(schema)); + + let arrow_schema = parquet_to_arrow_schema(&descr, None).unwrap(); + + let result = + coerce_int96_to_resolution(&descr, &arrow_schema, &TimeUnit::Microsecond) + .unwrap(); + + let expected_schema = Schema::new(vec![ + Field::new("c0", DataType::Timestamp(TimeUnit::Microsecond, None), true), + Field::new_struct( + "c1", + vec![Field::new( + "c0", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + )], + true, + ), + Field::new_struct( + "c2", + vec![Field::new_list( + "c0", + Field::new( + "element", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + true, + )], + true, + ), + Field::new_list( + "c3", + Field::new( + "element", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + true, + ), + Field::new_list( + "c4", + Field::new_struct( + "element", + vec![ + Field::new( + "c0", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + Field::new( + "c1", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + ], + true, + ), + true, + ), + Field::new_map( + "c5", + "key_value", + Field::new( + "key", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new( + "value", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + false, + true, + ), + Field::new_list( + "c6", + Field::new_map( + "element", + "key_value", + Field::new( + "key", + DataType::Timestamp(TimeUnit::Microsecond, None), + false, + ), + Field::new( + "value", + DataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + false, + true, + ), + true, + ), + ]); + + assert_eq!(result, expected_schema); + } +} diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs new file mode 100644 index 0000000000..a73be8d2e6 --- /dev/null +++ b/datafusion/datasource-parquet/src/sink.rs @@ -0,0 +1,754 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ParquetSink`] — DataFusion `DataSink` implementation that writes one +//! or more Parquet files to an [`ObjectStore`], optionally with parallel +//! per-column and per-row-group serialization. + +use std::fmt; +use std::fmt::Debug; +use std::sync::Arc; + +use arrow::array::RecordBatch; +use arrow::datatypes::{Schema, SchemaRef}; +use async_trait::async_trait; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::{DataFusionError, HashMap, Result, internal_datafusion_err}; +use datafusion_common_runtime::{JoinSet, SpawnedTask}; +use datafusion_datasource::display::FileGroupDisplay; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; +use datafusion_datasource::sink::DataSink; +use datafusion_datasource::write::demux::DemuxedStreamReceiver; +use datafusion_datasource::write::{ + ObjectWriterBuilder, SharedBuffer, get_writer_schema, +}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +use datafusion_physical_plan::metrics::{ + ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, + MetricsSet, Time, +}; +use datafusion_physical_plan::{DisplayAs, DisplayFormatType}; +use object_store::ObjectStore; +use object_store::buffered::BufWriter; +use object_store::path::Path; +use parquet::arrow::arrow_writer::{ + ArrowColumnChunk, ArrowColumnWriter, ArrowLeafColumn, ArrowRowGroupWriterFactory, + ArrowWriterOptions, compute_leaves, +}; +use parquet::arrow::{ArrowWriter, AsyncArrowWriter}; +#[cfg(feature = "parquet_encryption")] +use parquet::encryption::encrypt::FileEncryptionProperties; +use parquet::file::metadata::{ParquetMetaData, SortingColumn}; +use parquet::file::properties::{ + DEFAULT_MAX_ROW_GROUP_ROW_COUNT, WriterProperties, WriterPropertiesBuilder, +}; +use parquet::file::writer::SerializedFileWriter; +use tokio::io::{AsyncWrite, AsyncWriteExt}; +use tokio::sync::mpsc::{self, Receiver, Sender}; + +/// Initial writing buffer size. Note this is just a size hint for efficiency. It +/// will grow beyond the set value if needed. +const INITIAL_BUFFER_BYTES: usize = 1048576; + +/// When writing parquet files in parallel, if the buffered Parquet data exceeds +/// this size, it is flushed to object store +const BUFFER_FLUSH_BYTES: usize = 1024000; + +/// Implements [`DataSink`] for writing to a parquet file. +pub struct ParquetSink { + /// Config options for writing data + config: FileSinkConfig, + /// Underlying parquet options + parquet_options: TableParquetOptions, + /// File metadata from successfully produced parquet files. The Mutex is only used + /// to allow inserting to HashMap from behind borrowed reference in DataSink::write_all. + written: Arc<parking_lot::Mutex<HashMap<Path, ParquetMetaData>>>, + /// Optional sorting columns to write to Parquet metadata + sorting_columns: Option<Vec<SortingColumn>>, + /// Metrics for tracking write operations + metrics: ExecutionPlanMetricsSet, +} + +impl Debug for ParquetSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ParquetSink").finish() + } +} + +impl DisplayAs for ParquetSink { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ParquetSink(file_groups=",)?; + FileGroupDisplay(&self.config.file_group).fmt_as(t, f)?; + write!(f, ")") + } + DisplayFormatType::TreeRender => { + // TODO: collect info + write!(f, "") + } + } + } +} + +impl ParquetSink { + /// Create from config. + pub fn new(config: FileSinkConfig, parquet_options: TableParquetOptions) -> Self { + Self { + config, + parquet_options, + written: Default::default(), + sorting_columns: None, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + /// Set sorting columns for the Parquet file metadata. + pub fn with_sorting_columns( + mut self, + sorting_columns: Option<Vec<SortingColumn>>, + ) -> Self { + self.sorting_columns = sorting_columns; + self + } + + /// Retrieve the file metadata for the written files, keyed to the path + /// which may be partitioned (in the case of hive style partitioning). + pub fn written(&self) -> HashMap<Path, ParquetMetaData> { + self.written.lock().clone() + } + + /// Create writer properties based upon configuration settings, + /// including partitioning and the inclusion of arrow schema metadata. + async fn create_writer_props( + &self, + runtime: &Arc<RuntimeEnv>, + path: &Path, + ) -> Result<WriterProperties> { + let schema = self.config.output_schema(); + + // TODO: avoid this clone in follow up PR, where the writer properties & schema + // are calculated once on `ParquetSink::new` + let mut parquet_opts = self.parquet_options.clone(); + if !self.parquet_options.global.skip_arrow_metadata { + parquet_opts.arrow_schema(schema); + } + + let mut builder = WriterPropertiesBuilder::try_from(&parquet_opts)?; + + // Set sorting columns if configured + if let Some(ref sorting_columns) = self.sorting_columns { + builder = builder.set_sorting_columns(Some(sorting_columns.clone())); + } + + builder = set_writer_encryption_properties( + builder, + runtime, + parquet_opts, + schema, + path, + ) + .await?; + Ok(builder.build()) + } + + /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore + /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized + async fn create_async_arrow_writer( + &self, + location: &Path, + object_store: Arc<dyn ObjectStore>, + context: &Arc<TaskContext>, + parquet_props: WriterProperties, + ) -> Result<AsyncArrowWriter<BufWriter>> { + let buf_writer = BufWriter::with_capacity( + object_store, + location.clone(), + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + ); + let options = ArrowWriterOptions::new() + .with_properties(parquet_props) + .with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata); + + let writer = AsyncArrowWriter::try_new_with_options( + buf_writer, + get_writer_schema(&self.config), + options, + )?; + Ok(writer) + } + + /// Parquet options + pub fn parquet_options(&self) -> &TableParquetOptions { + &self.parquet_options + } +} + +#[cfg(feature = "parquet_encryption")] +async fn set_writer_encryption_properties( + builder: WriterPropertiesBuilder, + runtime: &Arc<RuntimeEnv>, + parquet_opts: TableParquetOptions, + schema: &Arc<Schema>, + path: &Path, +) -> Result<WriterPropertiesBuilder> { + if let Some(file_encryption_properties) = parquet_opts.crypto.file_encryption { + // Encryption properties have been specified directly + return Ok(builder.with_file_encryption_properties(Arc::new( + FileEncryptionProperties::try_from(file_encryption_properties)?, + ))); + } else if let Some(encryption_factory_id) = &parquet_opts.crypto.factory_id.as_ref() { + // Encryption properties will be generated by an encryption factory + let encryption_factory = + runtime.parquet_encryption_factory(encryption_factory_id)?; + let file_encryption_properties = encryption_factory + .get_file_encryption_properties( + &parquet_opts.crypto.factory_options, + schema, + path, + ) + .await?; + if let Some(file_encryption_properties) = file_encryption_properties { + return Ok( + builder.with_file_encryption_properties(file_encryption_properties) + ); + } + } + Ok(builder) +} + +#[cfg(not(feature = "parquet_encryption"))] +async fn set_writer_encryption_properties( + builder: WriterPropertiesBuilder, + _runtime: &Arc<RuntimeEnv>, + _parquet_opts: TableParquetOptions, + _schema: &Arc<Schema>, + _path: &Path, +) -> Result<WriterPropertiesBuilder> { + Ok(builder) +} + +#[async_trait] +impl FileSink for ParquetSink { + fn config(&self) -> &FileSinkConfig { + &self.config + } + + async fn spawn_writer_tasks_and_join( + &self, + context: &Arc<TaskContext>, + demux_task: SpawnedTask<Result<()>>, + mut file_stream_rx: DemuxedStreamReceiver, + object_store: Arc<dyn ObjectStore>, + ) -> Result<u64> { + let rows_written_counter = MetricBuilder::new(&self.metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + // Note: bytes_written is the sum of compressed row group sizes, which + // may differ slightly from the actual on-disk file size (excludes footer, + // page indexes, and other Parquet metadata overhead). + let bytes_written_counter = MetricBuilder::new(&self.metrics) + .with_category(MetricCategory::Bytes) + .global_counter("bytes_written"); + let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0); + + let parquet_opts = &self.parquet_options; + + let mut file_write_tasks: JoinSet< + std::result::Result<(Path, ParquetMetaData), DataFusionError>, + > = JoinSet::new(); + + let runtime = context.runtime_env(); + let parallel_options = ParallelParquetWriterOptions { + max_parallel_row_groups: parquet_opts + .global + .maximum_parallel_row_group_writers, + max_buffered_record_batches_per_stream: parquet_opts + .global + .maximum_buffered_record_batches_per_stream, + }; + + while let Some((path, mut rx)) = file_stream_rx.recv().await { + let parquet_props = self.create_writer_props(&runtime, &path).await?; + // CDC requires the sequential writer: the chunker state lives in ArrowWriter + // and persists across row groups. The parallel path bypasses ArrowWriter entirely. + if !parquet_opts.global.allow_single_file_parallelism + || parquet_opts.global.use_content_defined_chunking.is_some() + { + let mut writer = self + .create_async_arrow_writer( + &path, + Arc::clone(&object_store), + context, + parquet_props.clone(), + ) + .await?; + let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) + .register(context.memory_pool()); + file_write_tasks.spawn( + async move { + while let Some(batch) = rx.recv().await { + writer.write(&batch).await?; + reservation.try_resize(writer.memory_size())?; + } + let parquet_meta_data = writer + .close() + .await + .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; + Ok((path, parquet_meta_data)) + } + .with_elapsed_compute(elapsed_compute.clone()), + ); + } else { + let writer = ObjectWriterBuilder::new( + // Parquet files as a whole are never compressed, since they + // manage compressed blocks themselves. + FileCompressionType::UNCOMPRESSED, + &path, + Arc::clone(&object_store), + ) + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .build()?; + let ctx = ParquetFileWriteContext { + schema: get_writer_schema(&self.config), + props: Arc::new(parquet_props), + skip_arrow_metadata: self.parquet_options.global.skip_arrow_metadata, + parallel_options: Arc::new(parallel_options.clone()), + pool: Arc::clone(context.memory_pool()), + }; + let encoding_time = elapsed_compute.clone(); + file_write_tasks.spawn(async move { + let parquet_meta_data = output_single_parquet_file_parallelized( + writer, + rx, + ctx, + encoding_time, + ) + .await?; + Ok((path, parquet_meta_data)) + }); + } + } + + while let Some(result) = file_write_tasks.join_next().await { + match result { + Ok(r) => { + let (path, parquet_meta_data) = r?; + let file_rows = parquet_meta_data.file_metadata().num_rows() as usize; + let file_bytes: usize = parquet_meta_data + .row_groups() + .iter() + .map(|rg| rg.compressed_size() as usize) + .sum(); + rows_written_counter.add(file_rows); + bytes_written_counter.add(file_bytes); + let mut written_files = self.written.lock(); + written_files + .try_insert(path.clone(), parquet_meta_data) + .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; + drop(written_files); + } + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } + } + } + } + + demux_task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + + Ok(rows_written_counter.value() as u64) + } +} + +#[async_trait] +impl DataSink for ParquetSink { + fn metrics(&self) -> Option<MetricsSet> { + Some(self.metrics.clone_inner()) + } + + fn schema(&self) -> &SchemaRef { + self.config.output_schema() + } + + async fn write_all( + &self, + data: SendableRecordBatchStream, + context: &Arc<TaskContext>, + ) -> Result<u64> { + FileSink::write_all(self, data, context).await + } +} + +/// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] +/// Once the channel is exhausted, returns the ArrowColumnWriter. +async fn column_serializer_task( + mut rx: Receiver<ArrowLeafColumn>, + mut writer: ArrowColumnWriter, + reservation: MemoryReservation, + encoding_time: Time, +) -> Result<(ArrowColumnWriter, MemoryReservation)> { + while let Some(col) = rx.recv().await { + let _timer = encoding_time.timer(); + writer.write(&col)?; + reservation.try_resize(writer.memory_size())?; + } + Ok((writer, reservation)) +} + +type ColumnWriterTask = SpawnedTask<Result<(ArrowColumnWriter, MemoryReservation)>>; +type ColSender = Sender<ArrowLeafColumn>; + +/// Spawns a parallel serialization task for each column +/// Returns join handles for each columns serialization task along with a send channel +/// to send arrow arrays to each serialization task. +fn spawn_column_parallel_row_group_writer( + col_writers: Vec<ArrowColumnWriter>, + max_buffer_size: usize, + pool: &Arc<dyn MemoryPool>, + encoding_time: &Time, +) -> Result<(Vec<ColumnWriterTask>, Vec<ColSender>)> { + let num_columns = col_writers.len(); + + let mut col_writer_tasks = Vec::with_capacity(num_columns); + let mut col_array_channels = Vec::with_capacity(num_columns); + for writer in col_writers.into_iter() { + // Buffer size of this channel limits the number of arrays queued up for column level serialization + let (send_array, receive_array) = + mpsc::channel::<ArrowLeafColumn>(max_buffer_size); + col_array_channels.push(send_array); + + let reservation = + MemoryConsumer::new("ParquetSink(ArrowColumnWriter)").register(pool); + let task = SpawnedTask::spawn(column_serializer_task( + receive_array, + writer, + reservation, + encoding_time.clone(), + )); + col_writer_tasks.push(task); + } + + Ok((col_writer_tasks, col_array_channels)) +} + +/// Settings related to writing parquet files in parallel +#[derive(Clone)] +struct ParallelParquetWriterOptions { + max_parallel_row_groups: usize, + max_buffered_record_batches_per_stream: usize, +} + +/// Write configuration inputs shared across all parallel tasks that encode a +/// single Parquet file. These values are invariant for the duration of one file +/// write and do not change per row-group or per column. +/// +/// Separating these from per-call parameters (`object_store_writer`, `data`, +/// `encoding_time`) keeps the deep parallel call chain below the argument-count +/// limit without mixing configuration with runtime state. +#[derive(Clone)] +struct ParquetFileWriteContext { + schema: Arc<Schema>, + props: Arc<WriterProperties>, + skip_arrow_metadata: bool, + parallel_options: Arc<ParallelParquetWriterOptions>, + pool: Arc<dyn MemoryPool>, +} + +/// This is the return type of calling [ArrowColumnWriter].close() on each column +/// i.e. the Vec of encoded columns which can be appended to a row group +type RBStreamSerializeResult = Result<(Vec<ArrowColumnChunk>, MemoryReservation, usize)>; + +/// Sends the ArrowArrays in passed [RecordBatch] through the channels to their respective +/// parallel column serializers. +async fn send_arrays_to_col_writers( + col_array_channels: &[ColSender], + rb: &RecordBatch, + schema: Arc<Schema>, +) -> Result<()> { + // Each leaf column has its own channel, increment next_channel for each leaf column sent. + let mut next_channel = 0; + for (array, field) in rb.columns().iter().zip(schema.fields()) { + for c in compute_leaves(field, array)? { + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if col_array_channels[next_channel].send(c).await.is_err() { + return Ok(()); + } + + next_channel += 1; + } + } + + Ok(()) +} + +/// Spawns a tokio task which joins the parallel column writer tasks, +/// and finalizes the row group +fn spawn_rg_join_and_finalize_task( + column_writer_tasks: Vec<ColumnWriterTask>, + rg_rows: usize, + pool: &Arc<dyn MemoryPool>, + encoding_time: Time, +) -> SpawnedTask<RBStreamSerializeResult> { + let rg_reservation = + MemoryConsumer::new("ParquetSink(SerializedRowGroupWriter)").register(pool); + + SpawnedTask::spawn(async move { + let num_cols = column_writer_tasks.len(); + let mut finalized_rg = Vec::with_capacity(num_cols); + for task in column_writer_tasks.into_iter() { + let (writer, _col_reservation) = task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + let encoded_size = writer.get_estimated_total_bytes(); + rg_reservation.grow(encoded_size); + let _timer = encoding_time.timer(); + finalized_rg.push(writer.close()?); + } + + Ok((finalized_rg, rg_reservation, rg_rows)) + }) +} + +/// This task coordinates the serialization of a parquet file in parallel. +/// As the query produces RecordBatches, these are written to a RowGroup +/// via parallel [ArrowColumnWriter] tasks. Once the desired max rows per +/// row group is reached, the parallel tasks are joined on another separate task +/// and sent to a concatenation task. This task immediately continues to work +/// on the next row group in parallel. So, parquet serialization is parallelized +/// across both columns and row_groups, with a theoretical max number of parallel tasks +/// given by n_columns * num_row_groups. +fn spawn_parquet_parallel_serialization_task( + row_group_writer_factory: ArrowRowGroupWriterFactory, + mut data: Receiver<RecordBatch>, + serialize_tx: Sender<SpawnedTask<RBStreamSerializeResult>>, + ctx: ParquetFileWriteContext, + encoding_time: Time, +) -> SpawnedTask<Result<(), DataFusionError>> { + SpawnedTask::spawn(async move { + let max_buffer_rb = ctx.parallel_options.max_buffered_record_batches_per_stream; + let max_row_group_rows = ctx + .props + .max_row_group_row_count() + .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT); + let mut row_group_index = 0; + let col_writers = + row_group_writer_factory.create_column_writers(row_group_index)?; + let (mut column_writer_handles, mut col_array_channels) = + spawn_column_parallel_row_group_writer( + col_writers, + max_buffer_rb, + &ctx.pool, + &encoding_time, + )?; + let mut current_rg_rows = 0; + + while let Some(mut rb) = data.recv().await { + // This loop allows the "else" block to repeatedly split the RecordBatch to handle the case + // when max_row_group_rows < execution.batch_size as an alternative to a recursive async + // function. + loop { + if current_rg_rows + rb.num_rows() < max_row_group_rows { + send_arrays_to_col_writers( + &col_array_channels, + &rb, + Arc::clone(&ctx.schema), + ) + .await?; + current_rg_rows += rb.num_rows(); + break; + } else { + let rows_left = max_row_group_rows - current_rg_rows; + let a = rb.slice(0, rows_left); + send_arrays_to_col_writers( + &col_array_channels, + &a, + Arc::clone(&ctx.schema), + ) + .await?; + + // Signal the parallel column writers that the RowGroup is done, join and finalize RowGroup + // on a separate task, so that we can immediately start on the next RG before waiting + // for the current one to finish. + drop(col_array_channels); + let finalize_rg_task = spawn_rg_join_and_finalize_task( + column_writer_handles, + max_row_group_rows, + &ctx.pool, + encoding_time.clone(), + ); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + + current_rg_rows = 0; + rb = rb.slice(rows_left, rb.num_rows() - rows_left); + + row_group_index += 1; + let col_writers = row_group_writer_factory + .create_column_writers(row_group_index)?; + (column_writer_handles, col_array_channels) = + spawn_column_parallel_row_group_writer( + col_writers, + max_buffer_rb, + &ctx.pool, + &encoding_time, + )?; + } + } + } + + drop(col_array_channels); + // Handle leftover rows as final rowgroup, which may be smaller than max_row_group_rows + if current_rg_rows > 0 { + let finalize_rg_task = spawn_rg_join_and_finalize_task( + column_writer_handles, + current_rg_rows, + &ctx.pool, + encoding_time.clone(), + ); + + // Do not surface error from closed channel (means something + // else hit an error, and the plan is shutting down). + if serialize_tx.send(finalize_rg_task).await.is_err() { + return Ok(()); + } + } + + Ok(()) + }) +} + +/// Consume RowGroups serialized by other parallel tasks and concatenate them in +/// to the final parquet file, while flushing finalized bytes to an [ObjectStore] +async fn concatenate_parallel_row_groups( + mut parquet_writer: SerializedFileWriter<SharedBuffer>, + merged_buff: SharedBuffer, + mut serialize_rx: Receiver<SpawnedTask<RBStreamSerializeResult>>, + mut object_store_writer: Box<dyn AsyncWrite + Send + Unpin>, + pool: Arc<dyn MemoryPool>, +) -> Result<ParquetMetaData> { + let file_reservation = + MemoryConsumer::new("ParquetSink(SerializedFileWriter)").register(&pool); + + while let Some(task) = serialize_rx.recv().await { + let result = task.join_unwind().await; + let (serialized_columns, rg_reservation, _cnt) = + result.map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + + let mut rg_out = parquet_writer.next_row_group()?; + for chunk in serialized_columns { + chunk.append_to_row_group(&mut rg_out)?; + rg_reservation.free(); + + let mut buff_to_flush = merged_buff.buffer.try_lock().unwrap(); + file_reservation.try_resize(buff_to_flush.len())?; + + if buff_to_flush.len() > BUFFER_FLUSH_BYTES { + object_store_writer + .write_all(buff_to_flush.as_slice()) + .await?; + buff_to_flush.clear(); + file_reservation.try_resize(buff_to_flush.len())?; // will set to zero + } + } + rg_out.close()?; + } + + let parquet_meta_data = parquet_writer.close()?; + let final_buff = merged_buff.buffer.try_lock().unwrap(); + + object_store_writer.write_all(final_buff.as_slice()).await?; + object_store_writer.shutdown().await?; + file_reservation.free(); + + Ok(parquet_meta_data) +} + +/// Parallelizes the serialization of a single parquet file, by first serializing N +/// independent RecordBatch streams in parallel to RowGroups in memory. Another +/// task then stitches these independent RowGroups together and streams this large +/// single parquet file to an ObjectStore in multiple parts. +async fn output_single_parquet_file_parallelized( + object_store_writer: Box<dyn AsyncWrite + Send + Unpin>, + data: Receiver<RecordBatch>, + ctx: ParquetFileWriteContext, + encoding_time: Time, +) -> Result<ParquetMetaData> { + let max_rowgroups = ctx.parallel_options.max_parallel_row_groups; + // Buffer size of this channel limits maximum number of RowGroups being worked on in parallel + let (serialize_tx, serialize_rx) = + mpsc::channel::<SpawnedTask<RBStreamSerializeResult>>(max_rowgroups); + + let merged_buff = SharedBuffer::new(INITIAL_BUFFER_BYTES); + let options = ArrowWriterOptions::new() + .with_properties((*ctx.props).clone()) + .with_skip_arrow_metadata(ctx.skip_arrow_metadata); + let writer = ArrowWriter::try_new_with_options( + merged_buff.clone(), + Arc::clone(&ctx.schema), + options, + )?; + let (writer, row_group_writer_factory) = writer.into_serialized_writer()?; + + let pool = Arc::clone(&ctx.pool); + let launch_serialization_task = spawn_parquet_parallel_serialization_task( + row_group_writer_factory, + data, + serialize_tx, + ctx, + encoding_time, + ); + let parquet_meta_data = concatenate_parallel_row_groups( + writer, + merged_buff, + serialize_rx, + object_store_writer, + pool, + ) + .await?; + + launch_serialization_task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + Ok(parquet_meta_data) +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
