You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/28 05:55:23 UTC
(datafusion) 01/01: Add datafusion.execution.soft_max_bytes_per_output_file config parameter (#24924)
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-24924-ed43e69fdca8737c9d065b53b5fc5e4183302775
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit fb88eb529b41b64a1685d3044186c07f3768bbcd
Author: Thomas Tanon <th...@pellissier-tanon.fr>
AuthorDate: Mon Sep 28 05:54:54 2026 +0000
Add datafusion.execution.soft_max_bytes_per_output_file config parameter (#24924)
## Rationale for this change
Allows to roughly limit the size of parquet files when the number of
rows is a hard-to-use estimator.
This is especially useful when a column contains potentially large
blobs.
## What changes are included in this PR?
- Adds the config parameter
- Add a new `AsyncWrite` implementation created by
`ObjectWriterBuilder`. It sits between the optional compression and the
`BufWriter` implementation, counting the number of effectively flushed
bytes in a `Arc<AtomicUsize>`.
- The `Arc<AtomicU64>` is used in the demuxer logic to get an estimation
of the written size. This estimation lags because writes happens
asynchronously but this avoid having to block on encoding and
compression.
- Use `ObjectWriterBuilder` when creating the single file
`AsyncArrowWriter`
---
Cargo.lock | 1 +
datafusion/common/src/config.rs | 9 +-
datafusion/core/src/datasource/listing/table.rs | 214 ++++++++++++++++++++-
datafusion/datasource-arrow/src/file_format.rs | 5 +-
datafusion/datasource-avro/src/file_format.rs | 7 +-
datafusion/datasource-parquet/src/sink.rs | 38 ++--
datafusion/datasource/Cargo.toml | 1 +
datafusion/datasource/src/file_compression_type.rs | 5 +-
datafusion/datasource/src/write/demux.rs | 100 ++++++----
datafusion/datasource/src/write/mod.rs | 88 ++++++++-
datafusion/datasource/src/write/orchestration.rs | 31 +--
.../sqllogictest/test_files/information_schema.slt | 4 +-
.../sqllogictest/test_files/set_variable.slt | 8 +
docs/source/library-user-guide/upgrading/56.0.0.md | 35 ++++
docs/source/user-guide/configs.md | 3 +-
15 files changed, 462 insertions(+), 87 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index d63db0b271..6573ef5aa7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1987,6 +1987,7 @@ dependencies = [
"log",
"object_store",
"parking_lot",
+ "pin-project-lite",
"rand 0.9.4",
"tempfile",
"tokio",
diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs
index 5f45705301..4246dcc451 100644
--- a/datafusion/common/src/config.rs
+++ b/datafusion/common/src/config.rs
@@ -1059,7 +1059,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.
@@ -1068,6 +1068,13 @@ 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 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
/// 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 580302c882..14be8a18ce 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(),
@@ -1791,16 +1795,18 @@ 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),
];
- 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}"
);
@@ -1810,10 +1816,25 @@ 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(),
+ );
+ if format == "parquet" {
+ config_map.insert(
+ "datafusion.execution.parquet.compression".into(),
+ "uncompressed".into(),
+ );
+ }
let file_extension = match format {
"json" => JsonFormat::default().get_ext(),
@@ -1835,6 +1856,187 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_copy_respects_soft_max_bytes_per_output_file() -> Result<()> {
+ struct TestCase {
+ format: &'static str,
+ options: &'static str,
+ // Arrow, Avro, and the parallel Parquet writer buffer about 1 MB
+ // before forwarding bytes to the object writer (and therefore the
+ // demuxer). Their batches must be larger than that buffer for
+ // rotation to be observable.
+ payload_size: usize,
+ high_entropy: bool,
+ }
+
+ let test_cases = [
+ TestCase {
+ format: "csv",
+ options: "",
+ payload_size: 32 * 1024,
+ high_entropy: true,
+ },
+ TestCase {
+ format: "csv",
+ options: "OPTIONS ('format.compression' 'gzip')",
+ payload_size: 32 * 1024,
+ high_entropy: true,
+ },
+ TestCase {
+ format: "csv",
+ options: "OPTIONS ('format.compression' 'gzip')",
+ payload_size: 32 * 1024,
+ high_entropy: false,
+ },
+ TestCase {
+ format: "json",
+ options: "",
+ payload_size: 32 * 1024,
+ high_entropy: true,
+ },
+ TestCase {
+ format: "json",
+ options: "OPTIONS ('format.compression' 'zstd')",
+ payload_size: 32 * 1024,
+ high_entropy: true,
+ },
+ TestCase {
+ format: "json",
+ options: "OPTIONS ('format.compression' 'zstd')",
+ payload_size: 32 * 1024,
+ high_entropy: false,
+ },
+ TestCase {
+ format: "arrow",
+ options: "",
+ payload_size: 1536 * 1024,
+ high_entropy: true,
+ },
+ #[cfg(feature = "parquet")]
+ TestCase {
+ format: "parquet",
+ options: "OPTIONS ('format.compression' 'uncompressed', 'format.max_row_group_size' '1', 'format.allow_single_file_parallelism' 'false')",
+ payload_size: 32 * 1024,
+ high_entropy: true,
+ },
+ #[cfg(feature = "parquet")]
+ TestCase {
+ format: "parquet",
+ options: "OPTIONS ('format.compression' 'zstd(3)', 'format.max_row_group_size' '1')",
+ payload_size: 1536 * 1024,
+ high_entropy: true,
+ },
+ #[cfg(feature = "avro")]
+ TestCase {
+ format: "avro",
+ options: "",
+ payload_size: 1536 * 1024,
+ high_entropy: true,
+ },
+ ];
+
+ for test_case in test_cases {
+ 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(),
+ "1000000".into(),
+ );
+ config_map.insert(
+ "datafusion.execution.soft_max_bytes_per_output_file".into(),
+ "1024".into(),
+ );
+ // Use the smallest supported buffer so the writer cannot get far
+ // ahead of the demuxer. There is still an intentional delay, so the
+ // test only requires that rotation eventually occurs.
+ config_map.insert(
+ "datafusion.execution.max_buffered_batches_per_output_file".into(),
+ "2".into(),
+ );
+ let config = SessionConfig::from_string_hash_map(&config_map)?;
+ let ctx = SessionContext::new_with_config(config);
+
+ let payload = if test_case.high_entropy {
+ let alphabet =
+ b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
+ let mut state = 0x4d595df4d0f33173_u64;
+ (0..test_case.payload_size)
+ .map(|_| {
+ state ^= state << 13;
+ state ^= state >> 7;
+ state ^= state << 17;
+ alphabet[state as usize % alphabet.len()] as char
+ })
+ .collect::<String>()
+ } else {
+ "a".repeat(test_case.payload_size)
+ };
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("id", DataType::Int32, false),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+ let batches = (0..8)
+ .map(|id| {
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![
+ Arc::new(arrow::array::Int32Array::from(vec![id])),
+ Arc::new(arrow::array::StringArray::from(vec![
+ payload.as_str(),
+ ])),
+ ],
+ )
+ })
+ .collect::<std::result::Result<Vec<_>, _>>()?;
+ ctx.register_table(
+ "soft_limit_source",
+ Arc::new(MemTable::try_new(schema, vec![batches])?),
+ )?;
+
+ let output = TempDir::new()?;
+ let copy_sql = format!(
+ "COPY soft_limit_source TO '{}' STORED AS {} {}",
+ output.path().display(),
+ test_case.format,
+ test_case.options,
+ );
+ let result = ctx.sql(©_sql).await?.collect().await?;
+ assert_eq!(
+ ScalarValue::try_from_array(result[0].column(0), 0)?,
+ ScalarValue::UInt64(Some(8)),
+ "COPY row count for {}",
+ test_case.format,
+ );
+
+ let files = output
+ .path()
+ .read_dir()?
+ .collect::<std::io::Result<Vec<_>>>()?;
+ if test_case.high_entropy {
+ assert!(
+ files.len() > 1,
+ "{} {} did not rotate after reaching the byte limit; files: {:?}",
+ test_case.format,
+ test_case.options,
+ files.iter().map(|entry| entry.path()).collect::<Vec<_>>(),
+ );
+ } else {
+ assert_eq!(
+ files.len(),
+ 1,
+ "{} {} rotated despite compressed output remaining below the byte limit; files: {:?}",
+ test_case.format,
+ test_case.options,
+ files.iter().map(|entry| entry.path()).collect::<Vec<_>>(),
+ );
+ }
+ }
+ 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 1da354b863..b874466e2f 100644
--- a/datafusion/datasource-arrow/src/file_format.rs
+++ b/datafusion/datasource-arrow/src/file_format.rs
@@ -274,7 +274,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 +283,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(
@@ -293,6 +293,7 @@ impl FileSink for ArrowFileSink {
.execution
.objectstore_writer_buffer_size,
))
+ .with_bytes_written_counter(file_metadata.size)
.build()?;
file_write_tasks.spawn(async move {
let mut row_count = 0;
diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs
index 93dad800c6..0c081b0ee9 100644
--- a/datafusion/datasource-avro/src/file_format.rs
+++ b/datafusion/datasource-avro/src/file_format.rs
@@ -219,11 +219,11 @@ impl FileSink for AvroFileSink {
mut file_stream_rx: DemuxedStreamReceiver,
object_store: Arc<dyn ObjectStore>,
) -> Result<u64> {
- let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =
+ let mut file_write_tasks: JoinSet<Result<usize, DataFusionError>> =
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<SharedBuffer> =
WriterBuilder::new(writer_schema.as_ref().clone())
@@ -233,7 +233,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(
@@ -243,6 +243,7 @@ impl FileSink for AvroFileSink {
.execution
.objectstore_writer_buffer_size,
))
+ .with_bytes_written_counter(file_metadata.size)
.build()?;
file_write_tasks.spawn(async move {
let mut row_count = 0;
diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs
index 92286c37a6..725b4c74c9 100644
--- a/datafusion/datasource-parquet/src/sink.rs
+++ b/datafusion/datasource-parquet/src/sink.rs
@@ -35,7 +35,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, FileMetadata};
use datafusion_datasource::write::{
ObjectWriterBuilder, SharedBuffer, get_writer_schema,
};
@@ -50,7 +50,6 @@ use datafusion_physical_plan::metrics::{
};
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,
@@ -177,20 +176,25 @@ impl ParquetSink {
/// AsyncArrowWriters are used when individual parquet file serialization is not parallelized
fn create_async_arrow_writer(
&self,
- location: &Path,
+ file_metadata: &FileMetadata,
object_store: Arc<dyn ObjectStore>,
context: &Arc<TaskContext>,
parquet_props: WriterProperties,
- ) -> Result<AsyncArrowWriter<BufWriter>> {
- let buf_writer = BufWriter::with_capacity(
+ ) -> Result<AsyncArrowWriter<Box<dyn AsyncWrite + Send + Unpin>>> {
+ let buf_writer = ObjectWriterBuilder::new(
+ FileCompressionType::UNCOMPRESSED,
+ &file_metadata.path,
object_store,
- location.clone(),
+ )
+ .with_buffer_size(Some(
context
.session_config()
.options()
.execution
.objectstore_writer_buffer_size,
- );
+ ))
+ .with_bytes_written_counter(Arc::clone(&file_metadata.size))
+ .build()?;
let options = ArrowWriterOptions::new()
.with_properties(parquet_props)
.with_skip_arrow_metadata(self.parquet_options.global.skip_arrow_metadata);
@@ -293,21 +297,24 @@ 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,
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 {
@@ -318,7 +325,7 @@ impl FileSink for ParquetSink {
.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(
@@ -337,6 +344,7 @@ impl FileSink for ParquetSink {
.execution
.objectstore_writer_buffer_size,
))
+ .with_bytes_written_counter(file_metadata.size)
.build()?;
let ctx = ParquetFileWriteContext {
schema: get_writer_schema(&self.config),
@@ -354,7 +362,7 @@ impl FileSink for ParquetSink {
encoding_time,
)
.await?;
- Ok((path, parquet_meta_data))
+ Ok((file_metadata.path, parquet_meta_data))
});
}
}
diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml
index f09447b694..dd4111cf95 100644
--- a/datafusion/datasource/Cargo.toml
+++ b/datafusion/datasource/Cargo.toml
@@ -74,6 +74,7 @@ liblzma = { workspace = true, optional = true }
log = { workspace = true }
object_store = { workspace = true }
parking_lot = { workspace = true }
+pin-project-lite = { workspace = true }
rand = { workspace = true }
tempfile = { workspace = true, optional = true }
tokio = { workspace = true }
diff --git a/datafusion/datasource/src/file_compression_type.rs b/datafusion/datasource/src/file_compression_type.rs
index 89efb58065..8709706157 100644
--- a/datafusion/datasource/src/file_compression_type.rs
+++ b/datafusion/datasource/src/file_compression_type.rs
@@ -45,7 +45,6 @@ use futures::TryStreamExt;
use futures::stream::BoxStream;
#[cfg(feature = "compression")]
use liblzma::read::XzDecoder;
-use object_store::buffered::BufWriter;
use tokio::io::AsyncWrite;
#[cfg(feature = "compression")]
use tokio_util::io::{ReaderStream, StreamReader};
@@ -158,7 +157,7 @@ impl FileCompressionType {
/// according to this `FileCompressionType` using the default compression level.
pub fn convert_async_writer(
&self,
- w: BufWriter,
+ w: impl AsyncWrite + Send + Unpin + 'static,
) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
self.convert_async_writer_with_level(w, None)
}
@@ -170,7 +169,7 @@ impl FileCompressionType {
/// compression level. If `None`, the default level for each algorithm is used.
pub fn convert_async_writer_with_level(
&self,
- w: BufWriter,
+ w: impl AsyncWrite + Send + Unpin + 'static,
compression_level: Option<u32>,
) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
#[cfg(feature = "compression")]
diff --git a/datafusion/datasource/src/write/demux.rs b/datafusion/datasource/src/write/demux.rs
index 81e8962d74..cbd0a9747b 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::{AtomicU64, 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<AtomicU64>;
type RecordBatchReceiver = Receiver<RecordBatch>;
-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<RecordBatch>)>,
+ mut tx: UnboundedSender<(FileMetadata, RecordBatchReceiver)>,
mut input: SendableRecordBatchStream,
context: Arc<TaskContext>,
base_output_path: ListingTableUrl,
@@ -156,12 +159,14 @@ 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;
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);
@@ -179,9 +184,15 @@ async fn row_count_demuxer(
max_rows_per_file
};
+ let max_bytes_per_file = if single_file_output {
+ u64::MAX
+ } else {
+ max_bytes_per_file as u64
+ };
+
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,
@@ -189,7 +200,9 @@ 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);
part_idx += 1;
}
@@ -201,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,
@@ -209,12 +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);
part_idx += 1;
- } else if row_counts[next_send_steam] >= max_rows_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;
- 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,
@@ -223,6 +241,8 @@ 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();
@@ -277,26 +297,30 @@ fn create_new_file_stream(
file_extension: &str,
single_file_output: bool,
max_buffered_batches: usize,
- tx: &mut UnboundedSender<(Path, Receiver<RecordBatch>)>,
-) -> Result<Sender<RecordBatch>> {
- 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<RecordBatch>, FileSize)> {
+ let file_size = Arc::new(AtomicU64::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<RecordBatch>)>,
+ tx: UnboundedSender<(FileMetadata, RecordBatchReceiver)>,
mut input: SendableRecordBatchStream,
context: Arc<TaskContext>,
partition_by: Vec<(String, DataType)>,
@@ -337,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::<RecordBatch>(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(AtomicU64::new(0)),
+ };
+
+ tx.send((file_metadata, part_rx)).map_err(|_| {
exec_datafusion_err!("Error sending new file stream!")
})?;
@@ -372,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/mod.rs b/datafusion/datasource/src/write/mod.rs
index c8c85112f0..c2b177875b 100644
--- a/datafusion/datasource/src/write/mod.rs
+++ b/datafusion/datasource/src/write/mod.rs
@@ -18,19 +18,22 @@
//! Module containing helper methods/traits related to enabling
//! write support for the various file formats
-use std::io::Write;
-use std::sync::Arc;
-
use crate::file_compression_type::FileCompressionType;
use crate::file_sink_config::FileSinkConfig;
-use datafusion_common::error::Result;
-
+use crate::write::demux::FileSize;
use arrow::array::RecordBatch;
use arrow::datatypes::Schema;
use bytes::Bytes;
+use datafusion_common::error::Result;
use object_store::ObjectStore;
use object_store::buffered::BufWriter;
use object_store::path::Path;
+use pin_project_lite::pin_project;
+use std::io::{IoSlice, Write};
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::task::{Context, Poll, ready};
use tokio::io::AsyncWrite;
pub mod demux;
@@ -117,6 +120,8 @@ pub struct ObjectWriterBuilder {
buffer_size: Option<usize>,
/// The compression level for the object writer.
compression_level: Option<u32>,
+ /// Counts the number of bytes actually written.
+ bytes_written_counter: Option<Arc<AtomicU64>>,
}
impl ObjectWriterBuilder {
@@ -132,6 +137,7 @@ impl ObjectWriterBuilder {
object_store,
buffer_size: None,
compression_level: None,
+ bytes_written_counter: None,
}
}
@@ -205,6 +211,11 @@ impl ObjectWriterBuilder {
self.compression_level
}
+ pub fn with_bytes_written_counter(mut self, counter: FileSize) -> Self {
+ self.bytes_written_counter = Some(counter);
+ self
+ }
+
/// Return a writer object that writes to the object store location.
///
/// If a buffer size has not been set, the default buffer buffer size will
@@ -219,14 +230,77 @@ impl ObjectWriterBuilder {
object_store,
buffer_size,
compression_level,
+ bytes_written_counter,
} = self;
let buf_writer = match buffer_size {
Some(size) => BufWriter::with_capacity(object_store, location, size),
None => BufWriter::new(object_store, location),
};
+ if let Some(counter) = bytes_written_counter {
+ file_compression_type.convert_async_writer_with_level(
+ CountingWriter {
+ inner: buf_writer,
+ counter,
+ },
+ compression_level,
+ )
+ } else {
+ file_compression_type
+ .convert_async_writer_with_level(buf_writer, compression_level)
+ }
+ }
+}
+
+pin_project! {
+ /// [`AsyncWrite`] wrapper that counts the number of effectively written bytes
+ struct CountingWriter<W> {
+ #[pin]
+ inner: W,
+ counter: Arc<AtomicU64>,
+ }
+}
+
+impl<W: AsyncWrite> AsyncWrite for CountingWriter<W> {
+ fn poll_write(
+ self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ buf: &[u8],
+ ) -> Poll<std::io::Result<usize>> {
+ let this = self.project();
+ Poll::Ready(ready!(this.inner.poll_write(cx, buf)).inspect(|count| {
+ this.counter.fetch_add((*count) as u64, Ordering::Relaxed);
+ }))
+ }
+
+ fn poll_flush(
+ self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<std::io::Result<()>> {
+ self.project().inner.poll_flush(cx)
+ }
+
+ fn poll_shutdown(
+ self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ ) -> Poll<std::io::Result<()>> {
+ self.project().inner.poll_shutdown(cx)
+ }
+
+ fn poll_write_vectored(
+ self: Pin<&mut Self>,
+ cx: &mut Context<'_>,
+ bufs: &[IoSlice<'_>],
+ ) -> Poll<std::io::Result<usize>> {
+ let this = self.project();
+ Poll::Ready(
+ ready!(this.inner.poll_write_vectored(cx, bufs)).inspect(|count| {
+ this.counter.fetch_add((*count) as u64, Ordering::Relaxed);
+ }),
+ )
+ }
- file_compression_type
- .convert_async_writer_with_level(buf_writer, compression_level)
+ fn is_write_vectored(&self) -> bool {
+ self.inner.is_write_vectored()
}
}
diff --git a/datafusion/datasource/src/write/orchestration.rs b/datafusion/datasource/src/write/orchestration.rs
index cd821b3b87..8975d0f85e 100644
--- a/datafusion/datasource/src/write/orchestration.rs
+++ b/datafusion/datasource/src/write/orchestration.rs
@@ -25,14 +25,13 @@ use super::demux::DemuxedStreamReceiver;
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;
@@ -264,18 +263,22 @@ 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)
+ .with_bytes_written_counter(file_metadata.size)
+ .build()?;
if tx_file_bundle
.send((rb_stream, Arc::clone(&serializer), writer))
diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt
index 5305a76133..f6633a165d 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 [...]
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 smal [...]
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 casti [...]
@@ -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 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 respe [...]
diff --git a/datafusion/sqllogictest/test_files/set_variable.slt b/datafusion/sqllogictest/test_files/set_variable.slt
index bf9145fe44..8bb65519db 100644
--- a/datafusion/sqllogictest/test_files/set_variable.slt
+++ b/datafusion/sqllogictest/test_files/set_variable.slt
@@ -815,6 +815,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/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md
index 79067402a7..312668c485 100644
--- a/docs/source/library-user-guide/upgrading/56.0.0.md
+++ b/docs/source/library-user-guide/upgrading/56.0.0.md
@@ -263,6 +263,41 @@ let mut config = SessionConfig::new();
config.options_mut().execution.enable_nlj_coordinated_fallback = false;
```
+### `ExecutionOptions` has a new `soft_max_bytes_per_output_file` field
+
+`ExecutionOptions` gained a public `soft_max_bytes_per_output_file: ConfigNonZeroUsize`
+field (default 4 GiB minus one byte, or 4,294,967,295 bytes). When writing
+multiple output files, DataFusion starts a new file once a writer reports that
+the encoded size has reached this target or the existing
+`soft_max_rows_per_output_file` target has been reached. The byte limit is soft:
+buffered batches, in-flight batch serialization, individual batch sizes, and
+file metadata can make a completed file larger than the target. It does not
+apply to single-file output.
+
+**Who is affected:**
+
+- Users constructing `ExecutionOptions` (or `ConfigOptions`) with an exhaustive
+ struct literal must provide the new field.
+- Users writing multiple output files can set a smaller byte target to rotate
+ files before they reach the default size.
+
+**Migration guide:**
+
+Set the field explicitly, or fill it from `Default` as shown above. For example,
+to target 128 MiB output files in Rust:
+
+```rust,ignore
+let mut config = SessionConfig::new();
+config.options_mut().execution.soft_max_bytes_per_output_file =
+ ConfigNonZeroUsize::try_new(128 * 1024 * 1024)?;
+```
+
+The same setting can be changed with SQL:
+
+```sql
+SET datafusion.execution.soft_max_bytes_per_output_file = 134217728;
+```
+
### `datafusion.optimizer.use_statistics_registry` is deprecated and ignored
The `datafusion.optimizer.use_statistics_registry` config flag is deprecated and
diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md
index 1af6680bf6..ba8ea52bb0 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 creatio [...]
| 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 shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition [...]
| 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 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 m [...]
| 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). [...]
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org