From d67a0cee64738859b5c9a5be60f387c298f40e0d Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Thu, 3 Sep 2026 17:44:50 +0200 Subject: [PATCH 1/2] Add datafusion.execution.soft_max_bytes_per_output_file config parameter Allows to roughly limit the size of parquet files when the number of rows is a hard-to-use estimator --- datafusion/common/src/config.rs | 8 ++++- .../core/src/datasource/listing/table.rs | 29 +++++++++++++++---- datafusion/datasource/src/write/demux.rs | 16 +++++++++- .../test_files/information_schema.slt | 4 ++- docs/source/user-guide/configs.md | 3 +- 5 files changed, 50 insertions(+), 10 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 360586b0e9bae..69f997d63e443 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1056,7 +1056,7 @@ config_namespace! { /// Guarantees a minimum level of output files running in parallel. /// RecordBatches will be distributed in round robin fashion to each /// parallel writer. Each writer is closed and a new file opened once - /// soft_max_rows_per_output_file is reached. + /// soft_max_rows_per_output_file or soft_max_bytes_per_output_file is reached. pub minimum_parallel_output_files: ConfigNonZeroUsize, default = non_zero_usize_default(4) /// Target number of rows in output files when writing multiple. @@ -1065,6 +1065,12 @@ config_namespace! { /// number of rows written is not roughly divisible by the soft max pub soft_max_rows_per_output_file: ConfigNonZeroUsize, default = non_zero_usize_default(50000000) + /// Target number of bytes in output files when writing multiple. + /// This is a soft max, so it can be exceeded slightly. There also + /// will be one file smaller than the limit if the total + /// number of rows written is not roughly divisible by the soft max + pub soft_max_bytes_per_output_file: ConfigNonZeroUsize, default = non_zero_usize_default(4294967295) + /// This is the maximum number of RecordBatches buffered /// for each output file being worked. Higher values can potentially /// give faster write performance at the cost of higher peak diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 982766dc88519..836cc1d88bace 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -721,6 +721,10 @@ mod tests { "datafusion.execution.soft_max_rows_per_output_file".into(), "10".into(), ); + config_map.insert( + "datafusion.execution.soft_max_bytes_per_output_file".into(), + "10".into(), + ); config_map.insert( "datafusion.execution.parquet.compression".into(), "zstd(5)".into(), @@ -1719,16 +1723,20 @@ mod tests { #[tokio::test] async fn test_insert_into_parameterized() -> Result<()> { let test_cases = vec![ - // (file_format, batch_size, soft_max_rows, expected_files) - ("json", 10, 10, 2), - ("csv", 10, 10, 2), + // (file_format, batch_size, soft_max_rows, soft_max_bytes, expected_files) + ("json", 10, 10, 1000, 2), + ("csv", 10, 10, 1000, 2), #[cfg(feature = "parquet")] - ("parquet", 10, 10, 2), + ("parquet", 10, 20, 1000, 1), #[cfg(feature = "parquet")] - ("parquet", 20, 20, 1), + ("parquet", 10, 10, 1000, 2), + #[cfg(feature = "parquet")] + ("parquet", 10, 20, 100, 2), ]; - for (format, batch_size, soft_max_rows, expected_files) in test_cases { + for (format, batch_size, soft_max_rows, soft_max_bytes, expected_files) in + test_cases + { println!( "Testing insert with format: {format}, batch_size: {batch_size}, expected files: {expected_files}" ); @@ -1738,10 +1746,19 @@ mod tests { "datafusion.execution.batch_size".into(), batch_size.to_string(), ); + // Isolate soft-limit rotation from the initial parallel-writer fan-out. + config_map.insert( + "datafusion.execution.minimum_parallel_output_files".into(), + "1".into(), + ); config_map.insert( "datafusion.execution.soft_max_rows_per_output_file".into(), soft_max_rows.to_string(), ); + config_map.insert( + "datafusion.execution.soft_max_bytes_per_output_file".into(), + soft_max_bytes.to_string(), + ); let file_extension = match format { "json" => JsonFormat::default().get_ext(), diff --git a/datafusion/datasource/src/write/demux.rs b/datafusion/datasource/src/write/demux.rs index 81e8962d740ef..6aa50b027bd4a 100644 --- a/datafusion/datasource/src/write/demux.rs +++ b/datafusion/datasource/src/write/demux.rs @@ -156,6 +156,7 @@ async fn row_count_demuxer( let exec_options = &context.session_config().options().execution; let max_rows_per_file = exec_options.soft_max_rows_per_output_file.get(); + let max_bytes_per_file = exec_options.soft_max_bytes_per_output_file.get(); let max_buffered_batches = exec_options.max_buffered_batches_per_output_file.get(); let minimum_parallel_files = exec_options.minimum_parallel_output_files.get(); let mut part_idx = 0; @@ -165,6 +166,7 @@ async fn row_count_demuxer( let mut next_send_steam = 0; let mut row_counts = Vec::with_capacity(minimum_parallel_files); + let mut bytes_counts = Vec::with_capacity(minimum_parallel_files); // Overrides if single_file_output is set let minimum_parallel_files = if single_file_output { @@ -179,6 +181,12 @@ async fn row_count_demuxer( max_rows_per_file }; + let max_bytes_per_file = if single_file_output { + usize::MAX + } else { + max_bytes_per_file + }; + if single_file_output { // ensure we have one file open, even when the input stream is empty open_file_streams.push(create_new_file_stream( @@ -191,6 +199,7 @@ async fn row_count_demuxer( &mut tx, )?); row_counts.push(0); + bytes_counts.push(0); part_idx += 1; } @@ -211,9 +220,13 @@ async fn row_count_demuxer( &mut tx, )?); row_counts.push(0); + bytes_counts.push(0); part_idx += 1; - } else if row_counts[next_send_steam] >= max_rows_per_file { + } else if row_counts[next_send_steam] >= max_rows_per_file + || bytes_counts[next_send_steam] >= max_bytes_per_file + { row_counts[next_send_steam] = 0; + bytes_counts[next_send_steam] = 0; open_file_streams[next_send_steam] = create_new_file_stream( &base_output_path, &write_id, @@ -226,6 +239,7 @@ async fn row_count_demuxer( part_idx += 1; } row_counts[next_send_steam] += rb.num_rows(); + bytes_counts[next_send_steam] += rb.get_array_memory_size(); open_file_streams[next_send_steam] .send(rb) .await diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b270eba99d7b0..0469a0197b656 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -276,6 +276,7 @@ datafusion.execution.planning_concurrency 13 datafusion.execution.skip_partial_aggregation_probe_ratio_threshold 0.8 datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false +datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 @@ -390,7 +391,7 @@ datafusion.execution.listing_table_ignore_subdirectory true Should sub directori datafusion.execution.max_buffered_batches_per_output_file 2 This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption. This budget is split evenly between two independent points in the write pipeline (see the demuxer diagram in #7791): how many files can be in flight from the demuxer to a writer task, and how many RecordBatches are buffered for a single file's writer. Must be at least 2 so each half gets at least 1 unit of buffering - 0 or 1 would leave one side with a zero-capacity channel and panic at write time. datafusion.execution.max_spill_file_size_bytes 134217728 Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB datafusion.execution.meta_fetch_concurrency 32 Number of files to read in parallel when inferring schema and statistics -datafusion.execution.minimum_parallel_output_files 4 Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file is reached. +datafusion.execution.minimum_parallel_output_files 4 Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file or soft_max_bytes_per_output_file is reached. datafusion.execution.objectstore_writer_buffer_size 10485760 Size (bytes) of data buffer DataFusion uses when writing output files. This affects the size of the data chunks that are uploaded to remote object stores (e.g. AWS S3). If very large (>= 100 GiB) output files are being written, it may be necessary to increase this size to avoid errors from the remote end point. datafusion.execution.parquet.allow_single_file_parallelism true (writing) Controls whether DataFusion will attempt to speed up writing parquet files by serializing them in parallel. Each column in each row group in each output file are serialized in parallel leveraging a maximum possible core count of n_files*n_row_groups*n_columns. datafusion.execution.parquet.binary_as_string false (reading) If true, parquet reader will read columns of `Binary/LargeBinary` with `Utf8`, and `BinaryView` with `Utf8View`. Parquet files generated by some legacy writers do not correctly set the UTF8 flag for strings, causing string columns to be loaded as BLOB instead. The parquet reader has special optimizations for `Utf8` validation, so reading such columns as strings is significantly faster than reading them as binary and then casting to string. @@ -437,6 +438,7 @@ datafusion.execution.planning_concurrency 13 Fan-out during initial physical pla datafusion.execution.skip_partial_aggregation_probe_ratio_threshold 0.8 Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode datafusion.execution.skip_physical_aggregate_schema_check false When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. +datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target number of bytes in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 0085d4ac7c1fa..d6239d6e1eb69 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -129,8 +129,9 @@ The following configuration settings are available: | datafusion.execution.max_spill_file_size_bytes | 134217728 | Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB | | datafusion.execution.enable_nlj_coordinated_fallback | true | Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback coordinates per-chunk left state (visited bitmap and probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) build a separate coordinator per task and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. | | datafusion.execution.meta_fetch_concurrency | 32 | Number of files to read in parallel when inferring schema and statistics | -| datafusion.execution.minimum_parallel_output_files | 4 | Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file is reached. | +| datafusion.execution.minimum_parallel_output_files | 4 | Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file or soft_max_bytes_per_output_file is reached. | | datafusion.execution.soft_max_rows_per_output_file | 50000000 | Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max | +| datafusion.execution.soft_max_bytes_per_output_file | 4294967295 | Target number of bytes in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max | | datafusion.execution.max_buffered_batches_per_output_file | 2 | This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption. This budget is split evenly between two independent points in the write pipeline (see the demuxer diagram in #7791): how many files can be in flight from the demuxer to a writer task, and how many RecordBatches are buffered for a single file's writer. Must be at least 2 so each half gets at least 1 unit of buffering - 0 or 1 would leave one side with a zero-capacity channel and panic at write time. | | datafusion.execution.listing_table_ignore_subdirectory | true | Should sub directories be ignored when scanning directories for data files. Defaults to true (ignores subdirectories), consistent with Hive. Note that this setting does not affect reading partitioned tables (e.g. `/table/year=2021/month=01/data.parquet`). | | datafusion.execution.listing_table_factory_infer_partitions | true | Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). | From cb03a02189c566a9508831ce523f59044bda53c7 Mon Sep 17 00:00:00 2001 From: Thomas Tanon Date: Wed, 16 Sep 2026 19:50:47 +0200 Subject: [PATCH 2/2] Use the actual written file size --- datafusion/common/src/config.rs | 9 +- .../core/src/datasource/listing/table.rs | 82 +++++++++++++++- .../datasource-arrow/src/file_format.rs | 10 +- datafusion/datasource-avro/src/file_format.rs | 12 ++- datafusion/datasource-parquet/src/sink.rs | 41 ++++++-- datafusion/datasource/src/write/demux.rs | 98 +++++++++++-------- .../datasource/src/write/orchestration.rs | 51 ++++++---- .../test_files/information_schema.slt | 2 +- .../sqllogictest/test_files/set_variable.slt | 8 ++ docs/source/user-guide/configs.md | 2 +- 10 files changed, 234 insertions(+), 81 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 69f997d63e443..b5397441e874c 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1065,10 +1065,11 @@ config_namespace! { /// number of rows written is not roughly divisible by the soft max pub soft_max_rows_per_output_file: ConfigNonZeroUsize, default = non_zero_usize_default(50000000) - /// Target number of bytes in output files when writing multiple. - /// This is a soft max, so it can be exceeded slightly. There also - /// will be one file smaller than the limit if the total - /// number of rows written is not roughly divisible by the soft max + /// Target encoded size in bytes of output files when writing multiple. + /// Writers asynchronously report the cumulative encoded size as they + /// process RecordBatches. The final file size may exceed this limit due + /// to batches buffered before the limit is observed, the size of a batch, + /// and file metadata written when the file is finalized. pub soft_max_bytes_per_output_file: ConfigNonZeroUsize, default = non_zero_usize_default(4294967295) /// This is the maximum number of RecordBatches buffered diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 836cc1d88bace..bbb4e33c37902 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -140,6 +140,8 @@ fn infer_options_boxed( #[cfg(test)] mod tests { + #[cfg(feature = "parquet")] + use crate::dataframe::DataFrameWriteOptions; #[cfg(feature = "parquet")] use crate::datasource::file_format::parquet::ParquetFormat; use crate::datasource::listing::table::ListingTableConfigExt; @@ -156,6 +158,8 @@ mod tests { object_store::make_test_store_and_state, object_store::register_test_store, }, }; + #[cfg(feature = "parquet")] + use arrow::array::BinaryArray; use arrow::{compute::SortOptions, record_batch::RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_catalog::TableProvider; @@ -1730,8 +1734,6 @@ mod tests { ("parquet", 10, 20, 1000, 1), #[cfg(feature = "parquet")] ("parquet", 10, 10, 1000, 2), - #[cfg(feature = "parquet")] - ("parquet", 10, 20, 100, 2), ]; for (format, batch_size, soft_max_rows, soft_max_bytes, expected_files) in @@ -1759,6 +1761,12 @@ mod tests { "datafusion.execution.soft_max_bytes_per_output_file".into(), soft_max_bytes.to_string(), ); + if format == "parquet" { + config_map.insert( + "datafusion.execution.parquet.compression".into(), + "uncompressed".into(), + ); + } let file_extension = match format { "json" => JsonFormat::default().get_ext(), @@ -1780,6 +1788,76 @@ mod tests { Ok(()) } + #[cfg(feature = "parquet")] + #[tokio::test] + async fn test_soft_max_bytes_uses_compressed_parquet_size() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "blob", + DataType::Binary, + false, + )])); + let blob = vec![0_u8; 2 * 1024 * 1024]; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(BinaryArray::from(vec![blob.as_slice()]))], + )?; + + for (compression, expect_single_file) in + [("uncompressed", false), ("zstd(3)", true)] + { + let mut config_map = HashMap::new(); + config_map.insert( + "datafusion.execution.minimum_parallel_output_files".into(), + "1".into(), + ); + config_map.insert( + "datafusion.execution.soft_max_rows_per_output_file".into(), + "100".into(), + ); + config_map.insert( + "datafusion.execution.soft_max_bytes_per_output_file".into(), + (64 * 1024).to_string(), + ); + config_map.insert( + "datafusion.execution.parquet.compression".into(), + compression.into(), + ); + config_map.insert( + "datafusion.execution.parquet.dictionary_enabled".into(), + "false".into(), + ); + let ctx = SessionContext::new_with_config( + SessionConfig::from_string_hash_map(&config_map)?, + ); + let source = Arc::new(MemTable::try_new( + Arc::clone(&schema), + vec![vec![batch.clone(); 20]], + )?); + let output_dir = TempDir::new()?; + + ctx.read_table(source)? + .write_parquet( + output_dir.path().to_str().unwrap(), + DataFrameWriteOptions::new(), + None, + ) + .await?; + + let output_files = output_dir + .path() + .read_dir()? + .collect::>>()?; + if expect_single_file { + assert_eq!(output_files.len(), 1); + assert!(output_files[0].metadata()?.len() < 64 * 1024); + } else { + assert!(output_files.len() > 1); + } + } + + Ok(()) + } + #[tokio::test] async fn test_basic_table_scan() -> Result<()> { let ctx = SessionContext::new_with_config( diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index 2bee57ef17581..e819e304b9380 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -23,6 +23,7 @@ use std::collections::HashMap; use std::fmt::{self, Debug}; use std::io::{Seek, SeekFrom}; use std::sync::Arc; +use std::sync::atomic::Ordering; use arrow::datatypes::{Schema, SchemaRef}; use arrow::error::ArrowError; @@ -274,7 +275,7 @@ impl FileSink for ArrowFileSink { let ipc_options = IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)? .try_with_compression(Some(CompressionType::LZ4_FRAME))?; - while let Some((path, mut rx)) = file_stream_rx.recv().await { + while let Some((file_metadata, mut rx)) = file_stream_rx.recv().await { let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options( shared_buffer.clone(), @@ -283,7 +284,7 @@ impl FileSink for ArrowFileSink { )?; let mut object_store_writer = ObjectWriterBuilder::new( FileCompressionType::UNCOMPRESSED, - &path, + &file_metadata.path, Arc::clone(&object_store), ) .with_buffer_size(Some( @@ -296,14 +297,19 @@ impl FileSink for ArrowFileSink { .build()?; file_write_tasks.spawn(async move { let mut row_count = 0; + let mut flushed_bytes = 0; while let Some(batch) = rx.recv().await { row_count += batch.num_rows(); arrow_writer.write(&batch)?; let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); + file_metadata + .size + .store(flushed_bytes + buff_to_flush.len(), Ordering::Relaxed); if buff_to_flush.len() > BUFFER_FLUSH_BYTES { object_store_writer .write_all(buff_to_flush.as_slice()) .await?; + flushed_bytes += buff_to_flush.len(); buff_to_flush.clear(); } } diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 93dad800c6f0d..a7e7b7766c9e6 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/datafusion/datasource-avro/src/file_format.rs @@ -19,6 +19,7 @@ use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; +use std::sync::atomic::Ordering; use std::{fmt, io}; use crate::read_avro_schema_from_reader; @@ -219,11 +220,11 @@ impl FileSink for AvroFileSink { mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, ) -> Result { - let mut file_write_tasks: JoinSet> = + let mut file_write_tasks: JoinSet> = JoinSet::new(); let writer_schema = get_writer_schema(&self.config); - while let Some((path, mut rx)) = file_stream_rx.recv().await { + while let Some((file_metadata, mut rx)) = file_stream_rx.recv().await { let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); let mut avro_writer: AvroWriter = WriterBuilder::new(writer_schema.as_ref().clone()) @@ -233,7 +234,7 @@ impl FileSink for AvroFileSink { })?; let mut object_store_writer = ObjectWriterBuilder::new( FileCompressionType::UNCOMPRESSED, - &path, + &file_metadata.path, Arc::clone(&object_store), ) .with_buffer_size(Some( @@ -246,16 +247,21 @@ impl FileSink for AvroFileSink { .build()?; file_write_tasks.spawn(async move { let mut row_count = 0; + let mut flushed_bytes = 0; while let Some(batch) = rx.recv().await { row_count += batch.num_rows(); avro_writer .write(&batch) .map_err(|e| internal_datafusion_err!("{e}"))?; let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); + file_metadata + .size + .store(flushed_bytes + buff_to_flush.len(), Ordering::Relaxed); if buff_to_flush.len() > BUFFER_FLUSH_BYTES { object_store_writer .write_all(buff_to_flush.as_slice()) .await?; + flushed_bytes += buff_to_flush.len(); buff_to_flush.clear(); } } diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 3c66d4dcd74fb..f25199b13013b 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -22,6 +22,7 @@ use std::fmt; use std::fmt::Debug; use std::sync::Arc; +use std::sync::atomic::Ordering; use arrow::array::RecordBatch; use arrow::datatypes::{Schema, SchemaRef}; @@ -35,7 +36,7 @@ use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; use datafusion_datasource::sink::DataSink; #[cfg(feature = "proto")] use datafusion_datasource::sink::DataSinkExec; -use datafusion_datasource::write::demux::DemuxedStreamReceiver; +use datafusion_datasource::write::demux::{DemuxedStreamReceiver, FileSize}; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, }; @@ -293,32 +294,38 @@ impl FileSink for ParquetSink { .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?; + while let Some((file_metadata, mut rx)) = file_stream_rx.recv().await { + let parquet_props = self + .create_writer_props(&runtime, &file_metadata.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.content_defined_chunking.enabled { let mut writer = self.create_async_arrow_writer( - &path, + &file_metadata.path, Arc::clone(&object_store), context, parquet_props.clone(), )?; - let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) - .register(context.memory_pool()); + let reservation = + MemoryConsumer::new(format!("ParquetSink[{}]", file_metadata.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 encoded_size = + writer.bytes_written() + writer.in_progress_size(); + file_metadata.size.store(encoded_size, Ordering::Relaxed); } let parquet_meta_data = writer .close() .await .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; - Ok((path, parquet_meta_data)) + Ok((file_metadata.path, parquet_meta_data)) } .with_elapsed_compute(elapsed_compute.clone()), ); @@ -327,7 +334,7 @@ impl FileSink for ParquetSink { // Parquet files as a whole are never compressed, since they // manage compressed blocks themselves. FileCompressionType::UNCOMPRESSED, - &path, + &file_metadata.path, Arc::clone(&object_store), ) .with_buffer_size(Some( @@ -352,9 +359,10 @@ impl FileSink for ParquetSink { rx, ctx, encoding_time, + file_metadata.size, ) .await?; - Ok((path, parquet_meta_data)) + Ok((file_metadata.path, parquet_meta_data)) }); } } @@ -521,11 +529,19 @@ async fn column_serializer_task( mut writer: ArrowColumnWriter, reservation: MemoryReservation, encoding_time: Time, + file_size: FileSize, ) -> Result<(ArrowColumnWriter, MemoryReservation)> { + let mut encoded_size = 0; while let Some(col) = rx.recv().await { let _timer = encoding_time.timer(); writer.write(&col)?; reservation.try_resize(writer.memory_size())?; + let new_encoded_size = writer.get_estimated_total_bytes(); + file_size.fetch_add( + new_encoded_size.saturating_sub(encoded_size), + Ordering::Relaxed, + ); + encoded_size = new_encoded_size; } Ok((writer, reservation)) } @@ -541,6 +557,7 @@ fn spawn_column_parallel_row_group_writer( max_buffer_size: usize, pool: &Arc, encoding_time: &Time, + file_size: &FileSize, ) -> Result<(Vec, Vec)> { let num_columns = col_writers.len(); @@ -559,6 +576,7 @@ fn spawn_column_parallel_row_group_writer( writer, reservation, encoding_time.clone(), + Arc::clone(file_size), )); col_writer_tasks.push(task); } @@ -660,6 +678,7 @@ fn spawn_parquet_parallel_serialization_task( serialize_tx: Sender>, ctx: ParquetFileWriteContext, encoding_time: Time, + file_size: FileSize, ) -> SpawnedTask> { SpawnedTask::spawn(async move { let max_buffer_rb = ctx.parallel_options.max_buffered_record_batches_per_stream; @@ -676,6 +695,7 @@ fn spawn_parquet_parallel_serialization_task( max_buffer_rb, &ctx.pool, &encoding_time, + &file_size, )?; let mut current_rg_rows = 0; @@ -732,6 +752,7 @@ fn spawn_parquet_parallel_serialization_task( max_buffer_rb, &ctx.pool, &encoding_time, + &file_size, )?; } } @@ -813,6 +834,7 @@ async fn output_single_parquet_file_parallelized( data: Receiver, ctx: ParquetFileWriteContext, encoding_time: Time, + file_size: FileSize, ) -> Result { 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 @@ -837,6 +859,7 @@ async fn output_single_parquet_file_parallelized( serialize_tx, ctx, encoding_time, + file_size, ); let parquet_meta_data = concatenate_parallel_row_groups( writer, diff --git a/datafusion/datasource/src/write/demux.rs b/datafusion/datasource/src/write/demux.rs index 6aa50b027bd4a..5d52842cab278 100644 --- a/datafusion/datasource/src/write/demux.rs +++ b/datafusion/datasource/src/write/demux.rs @@ -18,14 +18,14 @@ //! Module containing helper methods/traits related to enabling //! dividing input stream into multiple output files at execution time -use std::borrow::Cow; -use std::collections::HashMap; -use std::sync::Arc; - use crate::url::ListingTableUrl; use crate::write::FileSinkConfig; use datafusion_common::error::Result; use datafusion_physical_plan::SendableRecordBatchStream; +use std::borrow::Cow; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::array::{ ArrayAccessor, RecordBatch, StringArray, StructArray, builder::UInt64Builder, @@ -50,8 +50,10 @@ use object_store::path::Path; use rand::distr::SampleString; use tokio::sync::mpsc::{self, Receiver, Sender, UnboundedReceiver, UnboundedSender}; +/// Cumulative encoded bytes reported by an output file writer. +pub type FileSize = Arc; type RecordBatchReceiver = Receiver; -pub type DemuxedStreamReceiver = UnboundedReceiver<(Path, RecordBatchReceiver)>; +pub type DemuxedStreamReceiver = UnboundedReceiver<(FileMetadata, RecordBatchReceiver)>; /// Splits a single [SendableRecordBatchStream] into a dynamically determined /// number of partitions at execution time. @@ -144,9 +146,10 @@ pub(crate) fn start_demuxer_task( (task, rx) } -/// Dynamically partitions input stream to achieve desired maximum rows per file +/// Dynamically partitions the input stream to achieve the desired maximum rows +/// and encoded bytes per file. async fn row_count_demuxer( - mut tx: UnboundedSender<(Path, Receiver)>, + mut tx: UnboundedSender<(FileMetadata, RecordBatchReceiver)>, mut input: SendableRecordBatchStream, context: Arc, base_output_path: ListingTableUrl, @@ -163,10 +166,10 @@ async fn row_count_demuxer( let write_id = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 16); let mut open_file_streams = Vec::with_capacity(minimum_parallel_files); + let mut file_sizes = Vec::with_capacity(minimum_parallel_files); let mut next_send_steam = 0; let mut row_counts = Vec::with_capacity(minimum_parallel_files); - let mut bytes_counts = Vec::with_capacity(minimum_parallel_files); // Overrides if single_file_output is set let minimum_parallel_files = if single_file_output { @@ -189,7 +192,7 @@ async fn row_count_demuxer( if single_file_output { // ensure we have one file open, even when the input stream is empty - open_file_streams.push(create_new_file_stream( + let (file_stream, file_size) = create_new_file_stream( &base_output_path, &write_id, part_idx, @@ -197,9 +200,10 @@ async fn row_count_demuxer( single_file_output, max_buffered_batches, &mut tx, - )?); + )?; + open_file_streams.push(file_stream); + file_sizes.push(file_size); row_counts.push(0); - bytes_counts.push(0); part_idx += 1; } @@ -210,7 +214,7 @@ async fn row_count_demuxer( is_batch_received = true; // ensure we have at least minimum_parallel_files open if open_file_streams.len() < minimum_parallel_files { - open_file_streams.push(create_new_file_stream( + let (file_stream, file_size) = create_new_file_stream( &base_output_path, &write_id, part_idx, @@ -218,16 +222,17 @@ async fn row_count_demuxer( single_file_output, max_buffered_batches, &mut tx, - )?); + )?; + open_file_streams.push(file_stream); + file_sizes.push(file_size); row_counts.push(0); - bytes_counts.push(0); part_idx += 1; - } else if row_counts[next_send_steam] >= max_rows_per_file - || bytes_counts[next_send_steam] >= max_bytes_per_file + } + if row_counts[next_send_steam] >= max_rows_per_file + || file_sizes[next_send_steam].load(Ordering::Relaxed) >= max_bytes_per_file { row_counts[next_send_steam] = 0; - bytes_counts[next_send_steam] = 0; - open_file_streams[next_send_steam] = create_new_file_stream( + let (file_stream, file_size) = create_new_file_stream( &base_output_path, &write_id, part_idx, @@ -236,10 +241,11 @@ async fn row_count_demuxer( max_buffered_batches, &mut tx, )?; + open_file_streams[next_send_steam] = file_stream; + file_sizes[next_send_steam] = file_size; part_idx += 1; } row_counts[next_send_steam] += rb.num_rows(); - bytes_counts[next_send_steam] += rb.get_array_memory_size(); open_file_streams[next_send_steam] .send(rb) .await @@ -291,26 +297,30 @@ fn create_new_file_stream( file_extension: &str, single_file_output: bool, max_buffered_batches: usize, - tx: &mut UnboundedSender<(Path, Receiver)>, -) -> Result> { - let file_path = generate_file_path( - base_output_path, - write_id, - part_idx, - file_extension, - single_file_output, - ); + tx: &mut UnboundedSender<(FileMetadata, RecordBatchReceiver)>, +) -> Result<(Sender, FileSize)> { + let file_size = Arc::new(AtomicUsize::new(0)); + let file_metadata = FileMetadata { + path: generate_file_path( + base_output_path, + write_id, + part_idx, + file_extension, + single_file_output, + ), + size: Arc::clone(&file_size), + }; let (tx_file, rx_file) = mpsc::channel(max_buffered_batches / 2); - tx.send((file_path, rx_file)) + tx.send((file_metadata, rx_file)) .map_err(|_| exec_datafusion_err!("Error sending RecordBatch to file stream!"))?; - Ok(tx_file) + Ok((tx_file, file_size)) } /// Splits an input stream based on the distinct values of a set of columns /// Assumes standard hive style partition paths such as /// /col1=val1/col2=val2/outputfile.parquet async fn hive_style_partitions_demuxer( - tx: UnboundedSender<(Path, Receiver)>, + tx: UnboundedSender<(FileMetadata, RecordBatchReceiver)>, mut input: SendableRecordBatchStream, context: Arc, partition_by: Vec<(String, DataType)>, @@ -351,15 +361,18 @@ async fn hive_style_partitions_demuxer( // Create channel for previously unseen distinct partition key and notify consumer of new file let (part_tx, part_rx) = mpsc::channel::(max_buffered_recordbatches); - let file_path = compute_hive_style_file_path( - &part_key, - &partition_by, - &write_id, - &file_extension, - &base_output_path, - ); - - tx.send((file_path, part_rx)).map_err(|_| { + let file_metadata = FileMetadata { + path: compute_hive_style_file_path( + &part_key, + &partition_by, + &write_id, + &file_extension, + &base_output_path, + ), + size: Arc::new(AtomicUsize::new(0)), + }; + + tx.send((file_metadata, part_rx)).map_err(|_| { exec_datafusion_err!("Error sending new file stream!") })?; @@ -386,6 +399,11 @@ async fn hive_style_partitions_demuxer( Ok(()) } +pub struct FileMetadata { + pub path: Path, + pub size: FileSize, +} + fn compute_partition_keys_by_row<'a>( rb: &'a RecordBatch, partition_by: &'a [(String, DataType)], diff --git a/datafusion/datasource/src/write/orchestration.rs b/datafusion/datasource/src/write/orchestration.rs index cd821b3b87897..d53f4683192be 100644 --- a/datafusion/datasource/src/write/orchestration.rs +++ b/datafusion/datasource/src/write/orchestration.rs @@ -20,19 +20,19 @@ //! parallelization, and abort handling use std::sync::Arc; +use std::sync::atomic::Ordering; -use super::demux::DemuxedStreamReceiver; +use super::demux::{DemuxedStreamReceiver, FileSize}; use super::{BatchSerializer, ObjectWriterBuilder}; use crate::file_compression_type::FileCompressionType; use datafusion_common::error::Result; - -use arrow::array::RecordBatch; use datafusion_common::{ DataFusionError, exec_datafusion_err, internal_datafusion_err, internal_err, }; use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_execution::TaskContext; +use arrow::array::RecordBatch; use bytes::Bytes; use futures::join; use object_store::ObjectStore; @@ -86,6 +86,7 @@ pub(crate) async fn serialize_rb_stream_to_object_store( mut data_rx: Receiver, serializer: Arc, mut writer: WriterType, + file_size: FileSize, ) -> SerializedRecordBatchResult { let (tx, mut rx) = mpsc::channel::>>(100); @@ -111,9 +112,11 @@ pub(crate) async fn serialize_rb_stream_to_object_store( }); let mut row_count = 0; + let mut serialized_bytes = 0; while let Some(task) = rx.recv().await { match task.join().await { Ok(Ok((cnt, bytes))) => { + serialized_bytes += bytes.len(); match writer.write_all(&bytes).await { Ok(_) => (), Err(e) => { @@ -124,6 +127,7 @@ pub(crate) async fn serialize_rb_stream_to_object_store( } } row_count += cnt; + file_size.store(serialized_bytes, Ordering::Relaxed); } Ok(Err(e)) => { // Return the writer along with the error @@ -154,7 +158,7 @@ pub(crate) async fn serialize_rb_stream_to_object_store( SerializedRecordBatchResult::success(writer, row_count) } -type FileWriteBundle = (Receiver, SerializerType, WriterType); +type FileWriteBundle = (Receiver, SerializerType, WriterType, FileSize); /// Contains the common logic for serializing RecordBatches and /// writing the resulting bytes to an ObjectStore. /// Serialization is assumed to be stateless, i.e. @@ -173,9 +177,10 @@ pub(crate) async fn stateless_serialize_and_write_files( // if true, we may not have a guarantee that all written data was cleaned up. let mut any_abort_errors = false; let mut join_set = JoinSet::new(); - while let Some((data_rx, serializer, writer)) = rx.recv().await { + while let Some((data_rx, serializer, writer, file_size)) = rx.recv().await { join_set.spawn(async move { - serialize_rb_stream_to_object_store(data_rx, serializer, writer).await + serialize_rb_stream_to_object_store(data_rx, serializer, writer, file_size) + .await }); } let mut finished_writers = Vec::new(); @@ -264,21 +269,29 @@ pub async fn spawn_writer_tasks_and_join( let write_coordinator_task = SpawnedTask::spawn(async move { stateless_serialize_and_write_files(rx_file_bundle, tx_row_cnt).await }); - while let Some((location, rb_stream)) = file_stream_rx.recv().await { - let writer = - ObjectWriterBuilder::new(compression, &location, Arc::clone(&object_store)) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .with_compression_level(compression_level) - .build()?; + while let Some((file_metadata, rb_stream)) = file_stream_rx.recv().await { + let writer = ObjectWriterBuilder::new( + compression, + &file_metadata.path, + Arc::clone(&object_store), + ) + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .with_compression_level(compression_level) + .build()?; if tx_file_bundle - .send((rb_stream, Arc::clone(&serializer), writer)) + .send(( + rb_stream, + Arc::clone(&serializer), + writer, + file_metadata.size, + )) .await .is_err() { diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 0469a0197b656..7a6547f423c36 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -438,7 +438,7 @@ datafusion.execution.planning_concurrency 13 Fan-out during initial physical pla datafusion.execution.skip_partial_aggregation_probe_ratio_threshold 0.8 Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode datafusion.execution.skip_physical_aggregate_schema_check false When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. -datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target number of bytes in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max +datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. diff --git a/datafusion/sqllogictest/test_files/set_variable.slt b/datafusion/sqllogictest/test_files/set_variable.slt index 5d9cc3bd4ade1..26e9017553fb3 100644 --- a/datafusion/sqllogictest/test_files/set_variable.slt +++ b/datafusion/sqllogictest/test_files/set_variable.slt @@ -773,6 +773,14 @@ caused by Invalid or Unsupported Configuration: value must be greater than 0 +statement error +SET datafusion.execution.soft_max_bytes_per_output_file = 0 +---- +DataFusion error: Error setting config datafusion.execution.soft_max_bytes_per_output_file +caused by +Invalid or Unsupported Configuration: value must be greater than 0 + + statement error SET datafusion.execution.max_spill_file_size_bytes = 0 ---- diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index d6239d6e1eb69..159154b6fe925 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -131,7 +131,7 @@ The following configuration settings are available: | datafusion.execution.meta_fetch_concurrency | 32 | Number of files to read in parallel when inferring schema and statistics | | datafusion.execution.minimum_parallel_output_files | 4 | Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file or soft_max_bytes_per_output_file is reached. | | datafusion.execution.soft_max_rows_per_output_file | 50000000 | Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max | -| datafusion.execution.soft_max_bytes_per_output_file | 4294967295 | Target number of bytes in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max | +| datafusion.execution.soft_max_bytes_per_output_file | 4294967295 | Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. | | datafusion.execution.max_buffered_batches_per_output_file | 2 | This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption. This budget is split evenly between two independent points in the write pipeline (see the demuxer diagram in #7791): how many files can be in flight from the demuxer to a writer task, and how many RecordBatches are buffered for a single file's writer. Must be at least 2 so each half gets at least 1 unit of buffering - 0 or 1 would leave one side with a zero-capacity channel and panic at write time. | | datafusion.execution.listing_table_ignore_subdirectory | true | Should sub directories be ignored when scanning directories for data files. Defaults to true (ignores subdirectories), consistent with Hive. Note that this setting does not affect reading partitioned tables (e.g. `/table/year=2021/month=01/data.parquet`). | | datafusion.execution.listing_table_factory_infer_partitions | true | Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). |