Skip to content

Add datafusion.execution.soft_max_bytes_per_output_file config parameter #24924

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Tpt wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Add datafusion.execution.soft_max_bytes_per_output_file config parameter #24924

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d67a0ce
Add datafusion.execution.soft_max_bytes_per_output_file config parameter
Tpt Sep 3, 2026
cb03a02
Use the actual written file size
Tpt Sep 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion datafusion/common/src/config.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -1065,6 +1065,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
Expand Down
107 changes: 101 additions & 6 deletions datafusion/core/src/datasource/listing/table.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -721,6 +725,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(),
Expand Down Expand Up @@ -1719,16 +1727,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}"
);
Expand All @@ -1738,10 +1748,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(),
Expand All @@ -1763,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::<std::io::Result<Vec<_>>>()?;
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(
Expand Down
10 changes: 8 additions & 2 deletions datafusion/datasource-arrow/src/file_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(),
Expand All @@ -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(
Expand All @@ -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();
}
}
Expand Down
12 changes: 9 additions & 3 deletions datafusion/datasource-avro/src/file_format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -219,11 +220,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())
Expand All @@ -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(
Expand All @@ -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();
}
}
Expand Down
Loading