diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 52b78c844a73a..6d24eef0de04a 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -102,6 +102,7 @@ topk_tpch: Benchmark of top-k (sorting with limit) queries on TPC-H topk_sorted_tpch: Benchmark of top-k queries on TPC-H lineitem ordered by l_orderkey (SF=1) push_down_topk: Benchmark of ORDER BY ... LIMIT over outer joins on TPC-H dataset (SF=1) — exercises pushing TopK through a join external_aggr: External aggregation benchmark on TPC-H dataset (SF=1) +large_values: Queries over very large row values (multi-KiB strings; batches within batch_size rows but 100MB+ in memory); data is generated on first run wide_schema: Small-projection queries on a wide synthetic dataset (1024 cols × 256 files) — measures per-file metadata overhead (runs both 'wide' and 'narrow' subgroups: narrow is an internal baseline; the wide-vs-narrow ratio is the signal) predicate_eval: Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an @@ -554,6 +555,9 @@ main() { external_aggr) run_external_aggr ;; + large_values) + run_large_values + ;; sort_pushdown) run_sort_pushdown ;; @@ -1235,6 +1239,16 @@ run_h2o_window() { } # Runs the external aggregation benchmark +# Runs the large_values benchmark: queries over multi-KiB string values whose +# batches stay within batch_size rows but are 100MB+ in memory. The dfbench +# subcommand generates its dataset on first use. +run_large_values() { + RESULTS_FILE="${RESULTS_DIR}/large_values.json" + echo "RESULTS_FILE: ${RESULTS_FILE}" + echo "Running large_values benchmark..." + debug_run $CARGO_COMMAND --bin dfbench -- large-values --iterations 5 --path "${DATA_DIR}/large_values" -o "${RESULTS_FILE}" ${QUERY_ARG} +} + run_external_aggr() { # Use TPC-H SF1 dataset TPCH_DIR="${DATA_DIR}/tpch_sf1" diff --git a/benchmarks/src/bin/dfbench.rs b/benchmarks/src/bin/dfbench.rs index 50dd99368b7f0..3cd61e7478516 100644 --- a/benchmarks/src/bin/dfbench.rs +++ b/benchmarks/src/bin/dfbench.rs @@ -32,7 +32,8 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; use datafusion_benchmarks::{ - cancellation, clickbench, dict, h2o, hj, imdb, nlj, smj, sort_tpch, tpcds, tpch, + cancellation, clickbench, dict, h2o, hj, imdb, large_values, nlj, smj, sort_tpch, + tpcds, tpch, }; #[derive(Debug, Parser)] @@ -50,6 +51,7 @@ enum Options { H2o(h2o::RunOpt), HJ(hj::RunOpt), Imdb(imdb::RunOpt), + LargeValues(large_values::RunOpt), Nlj(nlj::RunOpt), Smj(smj::RunOpt), SortPushdown(sort_pushdown::RunOpt), @@ -71,6 +73,7 @@ pub async fn main() -> Result<()> { Options::H2o(opt) => opt.run().await, Options::HJ(opt) => opt.run().await, Options::Imdb(opt) => Box::pin(opt.run()).await, + Options::LargeValues(opt) => opt.run().await, Options::Nlj(opt) => opt.run().await, Options::Smj(opt) => opt.run().await, Options::SortPushdown(opt) => opt.run().await, diff --git a/benchmarks/src/large_values.rs b/benchmarks/src/large_values.rs new file mode 100644 index 0000000000000..594354a7ccebd --- /dev/null +++ b/benchmarks/src/large_values.rs @@ -0,0 +1,322 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmark for tables with very large row values (multi-KiB strings). +//! +//! With such data, a batch that respects the row-based `batch_size` limit +//! (8192 rows) can still be hundreds of MiB in memory, which stresses +//! byte-blind operators: sorts reserve ~2x the batch size before buffering, +//! aggregations materialize huge transient state, and joins move huge +//! batches through their probe side. Existing suites (TPC-H, ClickBench) +//! have ~1MiB batches and never hit this regime. +//! +//! The dataset is generated on first use: `--gen-rows` rows of +//! (id BIGINT, key BIGINT, val VARCHAR) split over `--gen-files` files, +//! where `val` is a random alphanumeric string of `--gen-value-size` bytes +//! (incompressible, so decoded batches are as large as they look). +//! +//! Intended usage is an A/B of `datafusion.execution.target_batch_size_bytes` +//! (e.g. via the `DATAFUSION_EXECUTION_TARGET_BATCH_SIZE_BYTES` environment +//! variable), optionally with `--memory-limit` to test whether queries that +//! exhaust memory on oversized batches can complete when batches are +//! normalized. + +use std::fs; +use std::fs::File; +use std::path::PathBuf; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use clap::Args; +use futures::StreamExt; +use parquet::arrow::ArrowWriter; +use parquet::basic::{Compression, ZstdLevel}; +use parquet::file::properties::WriterProperties; +use rand::SeedableRng; +use rand::distr::{Alphanumeric, SampleString}; +use rand::rngs::StdRng; + +use datafusion::error::Result; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_plan::display::DisplayableExecutionPlan; +use datafusion::physical_plan::{displayable, execute_stream}; +use datafusion::prelude::*; +use datafusion_common::instant::Instant; + +use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; + +/// Key cardinality for the `key` column (`id % KEY_CARDINALITY`) +const KEY_CARDINALITY: i64 = 1000; + +/// Rows generated per append to keep generation memory bounded +const GEN_CHUNK_ROWS: usize = 512; + +#[derive(Debug, Args)] +pub struct RunOpt { + /// Common options + #[command(flatten)] + common: CommonOpt, + + /// Query number. If not specified, runs all queries + #[arg(short, long)] + pub query: Option, + + /// Directory for the generated dataset + #[arg(required = true, short = 'p', long = "path")] + path: PathBuf, + + /// Path to JSON benchmark result to be compared using `compare.py` + #[arg(short = 'o', long = "output")] + output_path: Option, + + /// Total number of rows to generate + #[arg(long = "gen-rows", default_value = "65536")] + gen_rows: usize, + + /// Number of parquet files to split the dataset into + #[arg(long = "gen-files", default_value = "8")] + gen_files: usize, + + /// Size in bytes of each `val` string value + #[arg(long = "gen-value-size", default_value = "16384")] + gen_value_size: usize, + + /// Regenerate the dataset even if it already exists + #[arg(long = "regenerate")] + regenerate: bool, +} + +pub const LARGE_VALUES_QUERY_START_ID: usize = 1; +pub const LARGE_VALUES_QUERY_END_ID: usize = 5; + +/// Queries chosen to stress one byte-blind operator each. `count(*)` / +/// drained-stream shapes keep result sets small so the benchmark measures +/// the operators, not result materialization. +const LARGE_VALUES_QUERIES: [&str; 5] = [ + // Q1: scan + filter on the large values (baseline for the cost of + // re-chunking itself: scan-bound, no stateful operator) + r#" + SELECT count(*) + FROM large_values + WHERE val LIKE '%qq%' + "#, + // Q2: full sort, small key, large payload rows + r#" + SELECT key, val + FROM large_values + ORDER BY key + "#, + // Q3: full sort keyed on the large values themselves + r#" + SELECT val + FROM large_values + ORDER BY val + "#, + // Q4: aggregation with large-value accumulator state + r#" + SELECT key, min(val), max(val) + FROM large_values + GROUP BY key + "#, + // Q5: hash join with a modest build side; huge batches flow through + // the probe side and the join output + r#" + SELECT count(*) + FROM large_values a + JOIN (SELECT id, val FROM large_values WHERE key < 20) b + ON a.id = b.id + WHERE a.val != b.val + "#, +]; + +impl RunOpt { + pub async fn run(&self) -> Result<()> { + self.generate_data_if_needed()?; + + let query_range = match self.query { + Some(query_id) => query_id..=query_id, + None => LARGE_VALUES_QUERY_START_ID..=LARGE_VALUES_QUERY_END_ID, + }; + + let mut benchmark_run = BenchmarkRun::new(); + for query_id in query_range { + benchmark_run.start_new_case(&format!("{query_id}")); + + match self.benchmark_query(query_id).await { + Ok(query_results) => { + for iter in query_results { + benchmark_run.write_iter(iter.elapsed, iter.row_count); + } + } + Err(e) => { + benchmark_run.mark_failed(); + eprintln!("Query {query_id} failed: {e}"); + } + } + } + + benchmark_run.maybe_write_json(self.output_path.as_ref())?; + benchmark_run.maybe_print_failures(); + Ok(()) + } + + async fn benchmark_query(&self, query_id: usize) -> Result> { + let config = self.common.config()?; + let rt = self.common.build_runtime()?; + let state = SessionStateBuilder::new() + .with_config(config) + .with_runtime_env(rt) + .with_default_features() + .build(); + let ctx = SessionContext::from(state); + + let path = self.path.to_str().unwrap(); + ctx.register_parquet("large_values", path, ParquetReadOptions::default()) + .await?; + + let mut millis = vec![]; + let mut query_results = vec![]; + for i in 0..self.common.iterations { + let start = Instant::now(); + let sql = LARGE_VALUES_QUERIES[query_id - 1]; + let row_count = self.execute_query(&ctx, sql).await?; + let elapsed = start.elapsed(); + let ms = elapsed.as_secs_f64() * 1000.0; + millis.push(ms); + + println!( + "Query {query_id} iteration {i} took {ms:.1} ms and returned {row_count} rows" + ); + query_results.push(QueryResult { elapsed, row_count }); + } + + let avg = millis.iter().sum::() / millis.len() as f64; + println!("Query {query_id} avg time: {avg:.2} ms"); + print_memory_stats(); + + Ok(query_results) + } + + async fn execute_query(&self, ctx: &SessionContext, sql: &str) -> Result { + let debug = self.common.debug; + let plan = ctx.sql(sql).await?; + let (state, plan) = plan.into_parts(); + let plan = state.optimize(&plan)?; + let physical_plan = state.create_physical_plan(&plan).await?; + if debug { + println!( + "=== Physical plan ===\n{}\n", + displayable(physical_plan.as_ref()).indent(true) + ); + } + + // Drain the stream instead of collecting: sorted output of this + // dataset is ~1GiB and buffering it would dominate the measurement + let mut row_count = 0; + let mut stream = execute_stream(physical_plan.clone(), state.task_ctx())?; + while let Some(batch) = stream.next().await { + row_count += batch?.num_rows(); + } + + if debug { + println!( + "=== Physical plan with metrics ===\n{}\n", + DisplayableExecutionPlan::with_metrics(physical_plan.as_ref()) + .indent(true) + ); + } + Ok(row_count) + } + + fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("key", DataType::Int64, false), + Field::new("val", DataType::Utf8, false), + ])) + } + + fn generate_data_if_needed(&self) -> Result<()> { + let marker = self.path.join(format!( + "large_values_r{}_f{}_v{}.ok", + self.gen_rows, self.gen_files, self.gen_value_size + )); + if marker.exists() && !self.regenerate { + return Ok(()); + } + if self.path.exists() { + fs::remove_dir_all(&self.path)?; + } + fs::create_dir_all(&self.path)?; + + println!( + "Generating {} rows x {} byte values over {} files in {}...", + self.gen_rows, + self.gen_value_size, + self.gen_files, + self.path.display() + ); + let start = Instant::now(); + let schema = Self::schema(); + let rows_per_file = self.gen_rows.div_ceil(self.gen_files); + let props = WriterProperties::builder() + .set_compression(Compression::ZSTD(ZstdLevel::try_new(1)?)) + .build(); + + for file_idx in 0..self.gen_files { + let file_path = self.path.join(format!("part-{file_idx}.parquet")); + let file = File::create(&file_path)?; + let mut writer = + ArrowWriter::try_new(file, Arc::clone(&schema), Some(props.clone()))?; + // deterministic per-file seed for reproducibility + let mut rng = StdRng::seed_from_u64(42 + file_idx as u64); + + let file_start = file_idx * rows_per_file; + let file_rows = rows_per_file.min(self.gen_rows.saturating_sub(file_start)); + let mut written = 0; + while written < file_rows { + let n = GEN_CHUNK_ROWS.min(file_rows - written); + let base = (file_start + written) as i64; + let ids = Int64Array::from_iter_values(base..base + n as i64); + let keys = Int64Array::from_iter_values( + (base..base + n as i64).map(|id| id % KEY_CARDINALITY), + ); + let vals = + StringArray::from_iter_values((0..n).map(|_| { + Alphanumeric.sample_string(&mut rng, self.gen_value_size) + })); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(ids) as ArrayRef, + Arc::new(keys) as ArrayRef, + Arc::new(vals) as ArrayRef, + ], + )?; + writer.write(&batch)?; + written += n; + } + writer.close()?; + } + + fs::write(&marker, b"")?; + println!("Generated in {:.1}s", start.elapsed().as_secs_f64()); + Ok(()) + } +} diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index 8d24d44a174e3..96426f6047b6d 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -22,6 +22,7 @@ pub mod dict; pub mod h2o; pub mod hj; pub mod imdb; +pub mod large_values; pub mod nlj; pub mod smj; pub mod sort_pushdown; diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 454af28c14b4a..89b27c4094536 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -732,6 +732,29 @@ config_namespace! { /// metadata memory consumption pub batch_size: ConfigNonZeroUsize, default = non_zero_usize_default(8192) + /// Soft target, in bytes, for the in-memory size of batches, in addition to + /// the row-based `batch_size` target. When set, data source output batches + /// are re-chunked towards both targets: batches much larger than this (by + /// in-memory size) are split into compact copies of roughly this size even + /// when within `batch_size` rows -- protecting downstream operators from + /// very large batches (e.g. wide rows with large string values) -- and + /// small batches are coalesced until they reach either target. Batches near + /// the targets are passed through without copying. Sorts, sort-preserving + /// merges and aggregation spills also bound their output and spill batches + /// to roughly this size instead of emitting `batch_size`-row batches of + /// unbounded byte size. When `None` (the default), batches are only + /// chunked by row count. + pub target_batch_size_bytes: Option, default = None + + /// When `target_batch_size_bytes` is unset and the memory pool has a finite + /// limit, derive a byte-size target for batches automatically from the pool's + /// per-partition share, so queries whose batches are large relative to the + /// memory budget re-chunk them (and can complete) instead of failing with + /// ResourcesExhausted. Has no effect when there is no memory limit (nothing + /// to protect: batches pass through untouched) or when + /// `target_batch_size_bytes` is set explicitly. + pub adaptive_target_batch_size: bool, default = true + /// A perfect hash join (see `HashJoinExec` for more details) will be considered /// if the range of keys (max - min) on the build side is < this threshold. /// This provides a fast path for joins with very small key ranges, diff --git a/datafusion/core/tests/datasource/batch_normalizer.rs b/datafusion/core/tests/datasource/batch_normalizer.rs new file mode 100644 index 0000000000000..c96199c73f4ba --- /dev/null +++ b/datafusion/core/tests/datasource/batch_normalizer.rs @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end tests for `datafusion.execution.target_batch_size_bytes`: +//! byte-aware re-chunking of data source output batches + +use std::sync::Arc; + +use arrow::array::{ArrayRef, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::Result; +use datafusion_common::utils::memory::get_record_batch_memory_size; + +fn wide_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])) +} + +/// Batch with `rows` rows of `str_len`-byte strings +fn wide_batch(rows: usize, str_len: usize) -> RecordBatch { + let s = StringArray::from_iter_values((0..rows).map(|i| format!("{i:0>str_len$}"))); + RecordBatch::try_new(wide_schema(), vec![Arc::new(s) as ArrayRef]).unwrap() +} + +async fn run_scan( + config: SessionConfig, + batches: Vec, +) -> Result> { + let ctx = SessionContext::new_with_config(config.with_target_partitions(1)); + let table = MemTable::try_new(wide_schema(), vec![batches])?; + ctx.register_table("t", Arc::new(table))?; + ctx.sql("SELECT * FROM t").await?.collect().await +} + +#[tokio::test] +async fn target_batch_size_bytes_splits_wide_batches() -> Result<()> { + // One ~1MB batch of 100 wide rows; a 64KiB byte target must split it + let input = wide_batch(100, 10 * 1024); + let target_bytes = 64 * 1024; + + let mut config = SessionConfig::new(); + config.options_mut().execution.target_batch_size_bytes = Some(target_bytes); + + let results = run_scan(config, vec![input.clone()]).await?; + assert!( + results.len() > 1, + "expected wide batch to be split, got {} batch(es)", + results.len() + ); + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 100); + for batch in &results { + // each chunk must be a compact copy of roughly the target size, not + // a zero-copy slice pinning the ~1MB parent + assert!( + get_record_batch_memory_size(batch) <= 2 * target_bytes, + "output batch retains {} bytes, expected <= {}", + get_record_batch_memory_size(batch), + 2 * target_bytes + ); + } + Ok(()) +} + +#[tokio::test] +async fn wide_batches_not_split_by_default() -> Result<()> { + // Without a byte target, splitting is row-based only: a single wide + // batch within `batch_size` rows passes through whole + let input = wide_batch(100, 10 * 1024); + let results = run_scan(SessionConfig::new(), vec![input]).await?; + assert_eq!(results.len(), 1); + assert_eq!(results[0].num_rows(), 100); + Ok(()) +} + +#[tokio::test] +async fn target_batch_size_bytes_coalesces_small_batches() -> Result<()> { + // 50 batches of 10 tiny rows each: the normalizer coalesces them + let batches: Vec<_> = (0..50).map(|_| wide_batch(10, 8)).collect(); + + let mut config = SessionConfig::new(); + config.options_mut().execution.target_batch_size_bytes = Some(1024 * 1024); + + let results = run_scan(config, batches).await?; + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 500); + assert_eq!( + results.len(), + 1, + "expected small batches to be coalesced, got {} batches", + results.len() + ); + Ok(()) +} diff --git a/datafusion/core/tests/datasource/mod.rs b/datafusion/core/tests/datasource/mod.rs index 3785aa0766182..6de4712828b6d 100644 --- a/datafusion/core/tests/datasource/mod.rs +++ b/datafusion/core/tests/datasource/mod.rs @@ -20,5 +20,6 @@ //! Note tests for the Parquet format are in `parquet_integration` binary // Include tests in csv module +mod batch_normalizer; mod csv; mod object_store_access; diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 9eb92e7e3525d..fcaa5adcec497 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -23,6 +23,9 @@ use std::fmt::{Debug, Formatter}; use std::sync::{Arc, OnceLock}; use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_plan::batch_normalizer::{ + BatchNormalizerMetrics, BatchNormalizerStream, effective_target_batch_size_bytes, +}; use datafusion_physical_plan::execution_plan::{ Boundedness, EmissionType, SchedulingType, }; @@ -408,11 +411,29 @@ impl ExecutionPlan for DataSourceExec { .with_shared_state(shared_state); let stream = self.data_source.open_with_args(args)?; let batch_size = context.session_config().batch_size(); + let metrics = self.data_source.metrics(); + + // With a byte target configured, normalize output batches by both + // row count and in-memory size (splitting oversized batches and + // coalescing small ones); otherwise only split by row count. + let target_batch_size_bytes = effective_target_batch_size_bytes(&context); + if let Some(target_bytes) = target_batch_size_bytes { + log::debug!( + "Batch normalization enabled for partition {partition}: \ + batch_size={batch_size} target_batch_size_bytes={target_bytes}" + ); + let normalizer_metrics = BatchNormalizerMetrics::new(&metrics, partition); + return Ok(Box::pin(BatchNormalizerStream::new( + stream, + batch_size, + target_bytes, + normalizer_metrics, + ))); + } log::debug!( "Batch splitting enabled for partition {partition}: batch_size={batch_size}" ); - let metrics = self.data_source.metrics(); let split_metrics = SplitMetrics::new(&metrics, partition); Ok(Box::pin(BatchSplitStream::new( stream, diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 0d00e5c4d0d86..2be6870fa6d3d 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -306,6 +306,11 @@ pub(crate) struct GroupedHashAggregateStream { /// max rows in output RecordBatches batch_size: usize, + /// Soft byte-size target for spill and spill-merge batches (from + /// `datafusion.execution.target_batch_size_bytes`); `None` bounds + /// them by row count only + target_batch_size_bytes: Option, + /// Optional soft limit on the number of `group_values` in a batch /// If the number of `group_values` in a single batch exceeds this value, /// the `GroupedHashAggregateStream` operation immediately switches to @@ -390,6 +395,8 @@ impl GroupedHashAggregateStream { let agg_filter_expr = Arc::clone(&agg.filter_expr); let batch_size = context.session_config().batch_size(); + let target_batch_size_bytes = + crate::batch_normalizer::effective_target_batch_size_bytes(context); let input = agg.input.execute(partition, Arc::clone(context))?; let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); @@ -605,6 +612,7 @@ impl GroupedHashAggregateStream { baseline_metrics, group_by_metrics, batch_size, + target_batch_size_bytes, group_ordering, input_done: false, spill_state, @@ -1204,6 +1212,7 @@ impl GroupedHashAggregateStream { emit, self.spill_state.spill_expr.clone(), self.batch_size, + self.target_batch_size_bytes, ); let spillfile = self .spill_state @@ -1297,6 +1306,7 @@ impl GroupedHashAggregateStream { .with_expressions(&self.spill_state.spill_expr) .with_metrics(self.baseline_metrics.clone()) .with_batch_size(self.batch_size) + .with_batch_size_bytes(self.target_batch_size_bytes) .with_reservation(self.reservation.new_empty()) .build()?; self.input_done = false; diff --git a/datafusion/physical-plan/src/batch_normalizer.rs b/datafusion/physical-plan/src/batch_normalizer.rs new file mode 100644 index 0000000000000..0b64300ce9428 --- /dev/null +++ b/datafusion/physical-plan/src/batch_normalizer.rs @@ -0,0 +1,1032 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`BatchNormalizer`]: re-chunk a stream of [`RecordBatch`]es towards a +//! target row count *and* byte size. +//! +//! Operators generally bound their output batches by row count +//! (`datafusion.execution.batch_size`) only. With wide rows (large strings, +//! many columns) a batch within the row limit can still be arbitrarily large +//! in bytes (multiple GB), which breaks memory accounting assumptions and +//! causes OOMs. Conversely, streams of tiny batches waste per-batch overhead. +//! +//! [`BatchNormalizer`] addresses both with a single component that, per input +//! batch, takes one of four actions: +//! +//! 1. **Pass through** (zero copy): the batch is already acceptably sized. +//! The acceptance band is deliberately wide so that near-target batches +//! are never copied. +//! 2. **Coalesce**: small batches are buffered (copied) and emitted once the +//! buffer reaches the target row count *or* the target byte size, +//! whichever comes first. +//! 3. **Split**: oversized batches (by rows or by bytes) are re-emitted as +//! compact copies of roughly the target size. Note that copying is +//! load-bearing: `RecordBatch::slice` is zero-copy, so a plain slice of a +//! 4GB batch still pins the entire 4GB of buffers. Splitting must +//! materialize fresh buffers to actually release memory. +//! 4. **Compact**: batches whose logical (sliced) size is small but which +//! pin much larger buffers (e.g. a small slice of a huge batch, or a +//! StringView array referencing mostly-garbage data buffers) are copied +//! so the underlying buffers can be freed. +//! +//! The copy machinery is arrow's [`BatchCoalescer`], which compacts all +//! buffered data (including StringView garbage collection). + +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, ready}; + +use arrow::array::{ + Array, BinaryViewArray, GenericByteViewArray, RecordBatch, StringViewArray, + UInt64Array, +}; +use arrow::compute::{BatchCoalescer, take_record_batch}; +use arrow::datatypes::{ByteViewType, SchemaRef}; +use datafusion_common::utils::memory::get_record_batch_memory_size; +use datafusion_common::{Result, assert_or_internal_err}; +use datafusion_execution::memory_pool::MemoryLimit; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use futures::{Stream, StreamExt}; + +use crate::metrics::{Count, ExecutionPlanMetricsSet, MetricBuilder}; +use crate::spill::gc_view_arrays; +use crate::stream::EmptyRecordBatchStream; + +/// Multiple of `target_bytes` above which a batch is split. +/// +/// The generous slop avoids wasting work copying batches that are only +/// somewhat over target (splitting copies data; see module docs). +const SPLIT_SLOP_FACTOR: usize = 2; + +/// A batch is "wasteful" (and gets compacted) when the memory it retains is +/// more than this multiple of its logical size... +const WASTE_RATIO: usize = 2; + +/// ...and the absolute waste exceeds `max(MIN_WASTE_BYTES, target_bytes)`. +/// Compaction copies data, which only pays off when it frees memory that is +/// large on the scale the byte target cares about. In particular, +/// parquet-decoded string/view batches routinely retain 2-4x their logical +/// size because decode buffers are shared across consecutive batches; +/// compacting every such batch costs a copy per batch for no benefit (the +/// batches die immediately in streaming pipelines). +const MIN_WASTE_BYTES: usize = 1024 * 1024; + +/// Divisor applied to the per-partition share of a finite memory pool to +/// derive an adaptive byte target: `pool_size / target_partitions / +/// ADAPTIVE_TARGET_DIVISOR`. +/// +/// The margin covers what each partition's operators need per batch beyond +/// the batch itself: an external sort reserves ~2x a batch's size to buffer +/// it, its spill merge reserves ~4x the largest spilled batch per stream +/// (with a read-ahead of 2), and the final sort-preserving merge buffers +/// batches from every stream concurrently. /16 keeps several such batches +/// per partition comfortably inside the partition's fair share, empirically +/// turning "ResourcesExhausted on oversized batches" into completed queries +/// without shrinking batches so far that per-batch overhead dominates. +pub const ADAPTIVE_TARGET_DIVISOR: usize = 16; + +/// Adaptive targeting only activates when the derived target reaches this +/// minimum (i.e. the pool holds at least 16MiB per partition). Below it the +/// pool is too small for re-chunking to change the outcome, so behavior is +/// left exactly as with no target -- this also keeps carefully-constructed +/// small-pool OOM scenarios (e.g. in tests) unchanged. +pub const MIN_ADAPTIVE_TARGET_BYTES: usize = 1024 * 1024; + +/// Ceiling for the adaptive byte target. Batches beyond ~16MiB buy no +/// per-batch-overhead amortization, and empirically the sort/merge machinery +/// under a tight memory limit is only reliably stable when spilled batches +/// stay at or below this size; a larger pool share is headroom, not an +/// invitation for bigger batches. +pub const MAX_ADAPTIVE_TARGET_BYTES: usize = 16 * 1024 * 1024; + +/// Resolve the effective byte-size target for batches under this task +/// context: +/// +/// 1. An explicit `datafusion.execution.target_batch_size_bytes` always wins. +/// 2. Otherwise, if `datafusion.execution.adaptive_target_batch_size` is on +/// and the memory pool reports a finite limit, the target is derived from +/// the per-partition share of the pool (see [`ADAPTIVE_TARGET_DIVISOR`]). +/// With no memory limit there is nothing to protect, so no target is set +/// and batches are never split or copied. +/// 3. `None` otherwise: batches are chunked by row count only. +pub fn effective_target_batch_size_bytes(context: &TaskContext) -> Option { + let options = context.session_config().options(); + if let Some(target) = options.execution.target_batch_size_bytes { + return Some(target); + } + if !options.execution.adaptive_target_batch_size { + return None; + } + match context.memory_pool().memory_limit() { + MemoryLimit::Finite(pool_size) => { + let partitions = options.execution.target_partitions.max(1); + let target = pool_size / partitions / ADAPTIVE_TARGET_DIVISOR; + (target >= MIN_ADAPTIVE_TARGET_BYTES) + .then(|| target.min(MAX_ADAPTIVE_TARGET_BYTES)) + } + _ => None, + } +} + +/// Metrics for [`BatchNormalizer`] +#[derive(Debug, Clone, Default)] +pub struct BatchNormalizerMetrics { + /// Input batches emitted unchanged (zero copy) + pub batches_passed_through: Count, + /// Input batches buffered for coalescing + pub batches_coalesced: Count, + /// Input batches split because they exceeded the row or byte target + pub batches_split: Count, + /// Input batches copied to release pinned memory + pub batches_compacted: Count, +} + +impl BatchNormalizerMetrics { + pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + Self { + batches_passed_through: MetricBuilder::new(metrics) + .counter("batches_passed_through", partition), + batches_coalesced: MetricBuilder::new(metrics) + .counter("batches_coalesced", partition), + batches_split: MetricBuilder::new(metrics) + .counter("batches_split", partition), + batches_compacted: MetricBuilder::new(metrics) + .counter("batches_compacted", partition), + } + } +} + +/// Logical (materialized) size of a batch in bytes: the number of bytes +/// needed if exactly the rows in this batch were copied into fresh, +/// minimally-sized buffers. +/// +/// This is intentionally different from +/// [`get_record_batch_memory_size`], which returns the memory the batch +/// *retains* (full backing buffers, including the parts a slice does not +/// reference). The ratio between the two is used to detect wasteful batches. +fn logical_batch_size(batch: &RecordBatch) -> Result { + let mut total = 0; + for array in batch.columns() { + let data = array.to_data(); + total += data.get_slice_memory_size()?; + // `get_slice_memory_size` does not count the variable-length data + // buffers of view arrays (arrow-rs#8230): add the bytes actually + // referenced by the (non-inlined) views. Nested view arrays are not + // corrected, matching the spill code's `GetSlicedSize`. + if let Some(a) = array.as_any().downcast_ref::() { + total += referenced_view_bytes(a); + } else if let Some(a) = array.as_any().downcast_ref::() { + total += referenced_view_bytes(a); + } + } + Ok(total) +} + +/// Total bytes in the data buffers referenced by non-null, non-inlined views +fn referenced_view_bytes(array: &GenericByteViewArray) -> usize { + let nulls = array.nulls(); + array + .views() + .iter() + .enumerate() + .map(|(i, v)| { + if nulls.map(|n| n.is_null(i)).unwrap_or(false) { + return 0; + } + // low 32 bits of a view are the value's length; values longer + // than 12 bytes live in the data buffers + let len = (*v as u32) as usize; + if len > 12 { len } else { 0 } + }) + .sum() +} + +/// Copy `length` rows of `batch` starting at `offset` into fresh, +/// minimally-sized buffers. +/// +/// This must NOT be implemented with `slice`/`concat`: `RecordBatch::slice` +/// is zero-copy and arrow's `concat` short-circuits a single input array to +/// a zero-copy slice, so neither releases the parent batch's buffers. `take` +/// materializes fresh buffers for all types; view arrays additionally need a +/// garbage-collection pass because `take` copies their (small) views while +/// still sharing the underlying data buffers. +fn compact_slice( + batch: &RecordBatch, + offset: usize, + length: usize, +) -> Result { + let indices = UInt64Array::from_iter_values(offset as u64..(offset + length) as u64); + let taken = take_record_batch(batch, &indices)?; + gc_view_arrays(&taken) +} + +/// An oversized batch being split incrementally: one chunk is copied out per +/// [`BatchNormalizer::produce`] call so that peak memory stays at the parent +/// batch plus ~one chunk, rather than the parent plus all its chunks. +#[derive(Debug)] +struct PendingSplit { + batch: RecordBatch, + offset: usize, + chunk_rows: usize, + /// If true the chunks are emitted as zero-copy slices instead of compact + /// copies. Only used for batches with no columns (row count only), where + /// a slice retains nothing. + zero_copy: bool, +} + +/// See [module docs](self) for details. +/// +/// This is the sans-IO core; [`BatchNormalizerStream`] adapts it to a +/// [`SendableRecordBatchStream`]. +/// +/// Call [`push_batch`](Self::push_batch) when [`produce`](Self::produce) +/// returns `Ok(None)`, and [`finish`](Self::finish) at end of input to flush +/// any buffered rows. +#[derive(Debug)] +pub struct BatchNormalizer { + schema: SchemaRef, + /// Copy/compaction machinery: buffers and compacts pushed batches + /// (including StringView garbage collection), auto-completing an output + /// batch whenever `target_rows` rows are buffered. + inner: BatchCoalescer, + target_rows: usize, + target_bytes: usize, + /// Logical bytes currently buffered in `inner` + buffered_bytes: usize, + /// Output batches ready to be returned by `produce` + completed: VecDeque, + /// Oversized batch currently being split, if any + pending_split: Option, + finished: bool, + metrics: BatchNormalizerMetrics, +} + +impl BatchNormalizer { + /// Create a new normalizer targeting `target_rows` rows and + /// `target_bytes` (logical) bytes per output batch. + pub fn new( + schema: SchemaRef, + target_rows: usize, + target_bytes: usize, + metrics: BatchNormalizerMetrics, + ) -> Self { + Self { + inner: BatchCoalescer::new(Arc::clone(&schema), target_rows), + schema, + target_rows, + target_bytes: target_bytes.max(1), + buffered_bytes: 0, + completed: VecDeque::new(), + pending_split: None, + finished: false, + metrics, + } + } + + pub fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + /// Push an input batch. + /// + /// Must only be called when [`produce`](Self::produce) returned + /// `Ok(None)` (i.e. all pending output has been drained) and before + /// [`finish`](Self::finish). + pub fn push_batch(&mut self, batch: RecordBatch) -> Result<()> { + assert_or_internal_err!( + !self.finished, + "BatchNormalizer: push_batch after finish" + ); + assert_or_internal_err!( + self.pending_split.is_none(), + "BatchNormalizer: push_batch while a split is pending; drain produce() first" + ); + + let rows = batch.num_rows(); + if rows == 0 { + return Ok(()); + } + + // Row-count-only batches: nothing to copy, split by rows if needed + if batch.num_columns() == 0 { + if rows > self.target_rows { + self.metrics.batches_split.add(1); + self.pending_split = Some(PendingSplit { + batch, + offset: 0, + chunk_rows: self.target_rows, + zero_copy: true, + }); + } else { + self.metrics.batches_passed_through.add(1); + self.completed.push_back(batch); + } + return Ok(()); + } + + let logical = logical_batch_size(&batch)?; + + // Oversized (by rows or bytes): split into ~target-sized chunks + if rows > self.target_rows || logical > SPLIT_SLOP_FACTOR * self.target_bytes { + let bytes_per_row = (logical / rows).max(1); + let chunk_rows = + (self.target_bytes / bytes_per_row).clamp(1, self.target_rows); + if chunk_rows < rows { + self.metrics.batches_split.add(1); + // Chunks are emitted directly (not via the coalescing + // buffer), so flush buffered rows first to preserve order + self.flush_buffer()?; + self.pending_split = Some(PendingSplit { + batch, + offset: 0, + chunk_rows, + zero_copy: false, + }); + return Ok(()); + } + // A single row wider than the split threshold cannot be split + // further: emit as-is. (`chunk_rows >= rows` is only reachable + // here for `rows == 1`.) + return self.pass_through(batch); + } + + // Acceptable size, but does it pin much more memory than it uses? + // (e.g. a small slice of a huge batch): copy it so the backing + // buffers can be freed. + let retained = get_record_batch_memory_size(&batch); + let min_waste = MIN_WASTE_BYTES.max(self.target_bytes); + let wasteful = retained > WASTE_RATIO * logical + && retained.saturating_sub(logical) >= min_waste; + let batch = if wasteful { + self.metrics.batches_compacted.add(1); + compact_slice(&batch, 0, rows)? + } else { + batch + }; + + // In the acceptance band: emit as-is + if rows >= self.target_rows / 2 || logical >= self.target_bytes / 2 { + if !wasteful { + self.metrics.batches_passed_through.add(1); + } + return self.emit(batch); + } + + // Small: buffer for coalescing + if !wasteful { + self.metrics.batches_coalesced.add(1); + } + self.buffer(batch, logical) + } + + /// Signal end of input, flushing any buffered rows. + pub fn finish(&mut self) -> Result<()> { + assert_or_internal_err!( + self.pending_split.is_none(), + "BatchNormalizer: finish while a split is pending; drain produce() first" + ); + self.flush_buffer()?; + self.finished = true; + Ok(()) + } + + /// Return the next output batch, doing incremental split work if an + /// oversized batch is pending. Returns `Ok(None)` when more input is + /// needed (or, after [`finish`](Self::finish), when fully drained). + pub fn produce(&mut self) -> Result> { + loop { + if let Some(batch) = self.completed.pop_front() { + return Ok(Some(batch)); + } + let Some(mut pending) = self.pending_split.take() else { + return Ok(None); + }; + // Copy out the next chunk of the pending oversized batch + let remaining = pending.batch.num_rows() - pending.offset; + let take = remaining.min(pending.chunk_rows); + let chunk = if pending.zero_copy { + pending.batch.slice(pending.offset, take) + } else { + compact_slice(&pending.batch, pending.offset, take)? + }; + pending.offset += take; + if pending.offset < pending.batch.num_rows() { + self.pending_split = Some(pending); + } + self.completed.push_back(chunk); + } + } + + /// Push `batch` (of logical size `logical`) into the copying buffer, + /// flushing if the byte target is reached. Arrow's coalescer flushes by + /// itself when the row target is reached. + fn buffer(&mut self, batch: RecordBatch, logical: usize) -> Result<()> { + self.inner.push_batch(batch)?; + self.buffered_bytes += logical; + self.drain_inner()?; + if self.buffered_bytes >= self.target_bytes { + self.flush_buffer()?; + } + Ok(()) + } + + /// Emit `batch` unchanged, zero copy + fn pass_through(&mut self, batch: RecordBatch) -> Result<()> { + self.metrics.batches_passed_through.add(1); + self.emit(batch) + } + + /// Emit `batch`, flushing any buffered rows first to preserve input order + fn emit(&mut self, batch: RecordBatch) -> Result<()> { + self.flush_buffer()?; + self.completed.push_back(batch); + Ok(()) + } + + /// Complete the currently buffered rows (if any) as an output batch + fn flush_buffer(&mut self) -> Result<()> { + if self.inner.get_buffered_rows() > 0 { + self.inner.finish_buffered_batch()?; + self.drain_inner()?; + } + self.buffered_bytes = 0; + Ok(()) + } + + /// Move batches completed by the inner coalescer into the output queue + fn drain_inner(&mut self) -> Result<()> { + while let Some(batch) = self.inner.next_completed_batch() { + // completed batches are freshly compacted: logical == retained + self.buffered_bytes = self + .buffered_bytes + .saturating_sub(logical_batch_size(&batch)?); + self.completed.push_back(batch); + } + Ok(()) + } +} + +/// Stream adapter for [`BatchNormalizer`]. +/// +/// Applies the normalizer to `input`, forwarding errors and flushing +/// buffered rows when the input is exhausted. +pub struct BatchNormalizerStream { + input: SendableRecordBatchStream, + normalizer: BatchNormalizer, + input_done: bool, +} + +impl BatchNormalizerStream { + pub fn new( + input: SendableRecordBatchStream, + target_rows: usize, + target_bytes: usize, + metrics: BatchNormalizerMetrics, + ) -> Self { + let normalizer = + BatchNormalizer::new(input.schema(), target_rows, target_bytes, metrics); + Self { + input, + normalizer, + input_done: false, + } + } +} + +impl Stream for BatchNormalizerStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + match self.normalizer.produce() { + Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Ok(None) => {} + Err(e) => return Poll::Ready(Some(Err(e))), + } + if self.input_done { + return Poll::Ready(None); + } + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + if let Err(e) = self.normalizer.push_batch(batch) { + return Poll::Ready(Some(Err(e))); + } + } + Some(Err(e)) => return Poll::Ready(Some(Err(e))), + None => { + if let Err(e) = self.normalizer.finish() { + return Poll::Ready(Some(Err(e))); + } + // Release the input pipeline's resources + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + self.input_done = true; + } + } + } + } +} + +impl RecordBatchStream for BatchNormalizerStream { + fn schema(&self) -> SchemaRef { + self.normalizer.schema() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int64Array, StringArray}; + use arrow::compute::concat_batches; + use arrow::datatypes::{DataType, Field, Schema}; + use futures::TryStreamExt; + + const TARGET_ROWS: usize = 100; + const TARGET_BYTES: usize = 10 * 1024; // 10 KiB + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("s", DataType::Utf8, false), + ])) + } + + /// Batch with `rows` rows where each string value is `str_len` bytes + fn string_batch(rows: usize, str_len: usize) -> RecordBatch { + let ids = Int64Array::from_iter_values(0..rows as i64); + let s = + StringArray::from_iter_values((0..rows).map(|i| format!("{i:0>str_len$}"))); + RecordBatch::try_new( + test_schema(), + vec![Arc::new(ids) as ArrayRef, Arc::new(s) as ArrayRef], + ) + .unwrap() + } + + fn normalizer() -> BatchNormalizer { + BatchNormalizer::new( + test_schema(), + TARGET_ROWS, + TARGET_BYTES, + BatchNormalizerMetrics::default(), + ) + } + + /// Push all batches, then finish, collecting every produced batch. + fn run( + normalizer: &mut BatchNormalizer, + batches: impl IntoIterator, + ) -> Vec { + let mut out = vec![]; + for batch in batches { + normalizer.push_batch(batch).unwrap(); + while let Some(b) = normalizer.produce().unwrap() { + out.push(b); + } + } + normalizer.finish().unwrap(); + while let Some(b) = normalizer.produce().unwrap() { + out.push(b); + } + out + } + + /// Assert that the concatenation of `outputs` equals concatenation of `inputs` + fn assert_same_data(inputs: &[RecordBatch], outputs: &[RecordBatch]) { + let schema = inputs[0].schema(); + let expected = concat_batches(&schema, inputs).unwrap(); + let actual = concat_batches(&schema, outputs).unwrap(); + assert_eq!(expected, actual); + } + + fn logical_size(batch: &RecordBatch) -> usize { + logical_batch_size(batch).unwrap() + } + + // === adaptive target === + + #[test] + fn adaptive_target_from_memory_limit() { + use datafusion_execution::TaskContext; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::memory_pool::{FairSpillPool, GreedyMemoryPool}; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + let ctx_with = + |pool: Option>, + target: Option, + adaptive: bool, + partitions: usize| { + let mut config = SessionConfig::new().with_target_partitions(partitions); + config.options_mut().execution.target_batch_size_bytes = target; + config.options_mut().execution.adaptive_target_batch_size = adaptive; + let mut rt = RuntimeEnvBuilder::new(); + if let Some(pool) = pool { + rt = rt.with_memory_pool(pool); + } + TaskContext::default() + .with_session_config(config) + .with_runtime(Arc::new(rt.build().unwrap())) + }; + + // no limit, no explicit target -> disabled + let ctx = ctx_with(None, None, true, 8); + assert_eq!(effective_target_batch_size_bytes(&ctx), None); + + // finite pool -> pool_size / partitions / ADAPTIVE_TARGET_DIVISOR + let pool: Arc = + Arc::new(FairSpillPool::new(2 * 1024 * 1024 * 1024)); + let ctx = ctx_with(Some(Arc::clone(&pool)), None, true, 8); + assert_eq!( + effective_target_batch_size_bytes(&ctx), + Some(2 * 1024 * 1024 * 1024 / 8 / ADAPTIVE_TARGET_DIVISOR) + ); + + // explicit target always wins + let ctx = ctx_with(Some(Arc::clone(&pool)), Some(123456), true, 8); + assert_eq!(effective_target_batch_size_bytes(&ctx), Some(123456)); + + // adaptive disabled -> None even with a finite pool + let ctx = ctx_with(Some(Arc::clone(&pool)), None, false, 8); + assert_eq!(effective_target_batch_size_bytes(&ctx), None); + + // huge pool / few partitions -> clamped to the 16MiB ceiling + let pool: Arc = + Arc::new(GreedyMemoryPool::new(64 * 1024 * 1024 * 1024)); + let ctx = ctx_with(Some(pool), None, true, 4); + assert_eq!( + effective_target_batch_size_bytes(&ctx), + Some(MAX_ADAPTIVE_TARGET_BYTES) + ); + + // pool too small for the minimum target -> adaptive stays off so + // small-pool behavior (and crafted OOM scenarios) is unchanged + let pool: Arc = + Arc::new(GreedyMemoryPool::new(64 * 1024 * 1024)); + let ctx = ctx_with(Some(pool), None, true, 512); + assert_eq!(effective_target_batch_size_bytes(&ctx), None); + } + + // === pass through === + + #[test] + fn passthrough_normal_batch_is_zero_copy() { + // rows == target, bytes within band -> exact same arrays out + let batch = string_batch(TARGET_ROWS, 20); + let out = run(&mut normalizer(), [batch.clone()]); + assert_eq!(out.len(), 1); + for (in_col, out_col) in batch.columns().iter().zip(out[0].columns()) { + assert!( + Arc::ptr_eq(in_col, out_col), + "expected zero-copy pass through" + ); + } + } + + #[test] + fn passthrough_by_bytes_with_few_rows() { + // Only 10 rows (< target_rows/2) but ~8KiB (> target_bytes/2): + // "few wide rows" is an acceptable batch, not one to coalesce + let batch = string_batch(10, 800); + assert!(logical_size(&batch) >= TARGET_BYTES / 2); + let out = run(&mut normalizer(), [batch.clone()]); + assert_eq!(out.len(), 1); + assert!(Arc::ptr_eq(batch.column(1), out[0].column(1))); + } + + #[test] + fn passthrough_slightly_oversized_batch() { + // ~1.6x target bytes: within the slop band, must NOT be split/copied + let batch = string_batch(TARGET_ROWS, 160); + let size = logical_size(&batch); + assert!(size > TARGET_BYTES && size < SPLIT_SLOP_FACTOR * TARGET_BYTES); + let out = run(&mut normalizer(), [batch.clone()]); + assert_eq!(out.len(), 1); + assert!(Arc::ptr_eq(batch.column(1), out[0].column(1))); + } + + // === coalesce === + + #[test] + fn coalesce_small_batches_by_rows() { + // 10-row tiny batches accumulate to target_rows + let inputs: Vec<_> = (0..25).map(|_| string_batch(10, 5)).collect(); + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert_eq!( + out.iter().map(|b| b.num_rows()).collect::>(), + vec![100, 100, 50], + ); + } + + #[test] + fn coalesce_flushes_on_bytes_before_rows() { + // 10 rows x ~400B = ~4KiB logical per batch; the byte target (10KiB) + // is reached after ~3 batches, well before 100 buffered rows + let inputs: Vec<_> = (0..6).map(|_| string_batch(10, 400)).collect(); + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert!(out.len() >= 2, "byte target should have forced a flush"); + for b in &out { + assert!(b.num_rows() < TARGET_ROWS); + assert!(logical_size(b) <= SPLIT_SLOP_FACTOR * TARGET_BYTES); + } + } + + // === split === + + #[test] + fn split_wide_batch_by_bytes() { + // 50 rows x ~1KiB = ~51KiB >> 2x target (20KiB): split into ~10-row + // chunks of ~target_bytes each + let batch = string_batch(50, 1024); + let inputs = [batch.clone()]; + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert!(out.len() > 1, "expected batch to be split"); + for b in &out { + assert!( + logical_size(b) <= SPLIT_SLOP_FACTOR * TARGET_BYTES, + "chunk of {} bytes exceeds split threshold", + logical_size(b) + ); + // chunks must be compact copies: they must not retain the + // original ~51KiB of buffers + assert!( + get_record_batch_memory_size(b) < logical_size(&batch) / 2, + "chunk retains the parent batch's buffers" + ); + } + } + + #[test] + fn split_by_rows() { + // 1000 tiny rows > target_rows: split into 100-row chunks + let batch = string_batch(1000, 5); + let inputs = [batch.clone()]; + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert_eq!(out.len(), 10); + for b in &out { + assert_eq!(b.num_rows(), TARGET_ROWS); + } + } + + #[test] + fn split_produces_output_incrementally() { + // Oversized batches must not materialize all chunks eagerly: + // after push, produce() yields chunks one at a time + let batch = string_batch(50, 1024); + let mut n = normalizer(); + n.push_batch(batch).unwrap(); + let first = n.produce().unwrap(); + assert!(first.is_some()); + // there must be more chunks pending + let second = n.produce().unwrap(); + assert!(second.is_some()); + } + + #[test] + fn giant_single_row_passes_through() { + // One row larger than the split threshold cannot be split further; + // it must be emitted as-is (and not loop or error) + let batch = string_batch(1, 5 * TARGET_BYTES); + let inputs = [batch.clone()]; + let out = run(&mut normalizer(), inputs.clone()); + assert_eq!(out.len(), 1); + assert_eq!(out[0].num_rows(), 1); + assert_same_data(&inputs, &out); + } + + #[test] + fn split_string_view_batch_releases_data_buffers() { + // View arrays need explicit GC when split: a copied chunk's views + // would otherwise still reference the parent's large data buffers + use arrow::array::StringViewArray; + let rows = 50; + let s = + StringViewArray::from_iter_values((0..rows).map(|i| format!("{i:0>1024}"))); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s) as ArrayRef]) + .unwrap(); + let parent_size = get_record_batch_memory_size(&batch); + let mut n = BatchNormalizer::new( + schema, + TARGET_ROWS, + TARGET_BYTES, + BatchNormalizerMetrics::default(), + ); + let inputs = [batch]; + let out = run(&mut n, inputs.clone()); + assert_same_data(&inputs, &out); + assert!(out.len() > 1, "expected batch to be split"); + for b in &out { + assert!( + get_record_batch_memory_size(b) < parent_size / 2, + "chunk retains the parent's view data buffers" + ); + } + } + + // === compact === + + #[test] + fn wasteful_slice_is_compacted() { + // A small slice of a big batch pins the whole parent buffer. + // logical size is in the pass-through band, but retained size is + // ~50x larger -> must be copied so the parent can be freed. + let parent = string_batch(10_000, 300); // ~3MB + let slice = parent.slice(0, 20); // ~6KiB logical + let logical = logical_size(&slice); + assert!(logical >= TARGET_BYTES / 2, "slice must be in pass band"); + assert!(get_record_batch_memory_size(&slice) > WASTE_RATIO * logical); + + let inputs = [slice.clone()]; + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert_eq!(out.len(), 1); + assert!( + get_record_batch_memory_size(&out[0]) < logical * 2, + "output still retains the parent's buffers" + ); + } + + #[test] + fn small_waste_is_not_compacted() { + // Waste below MIN_WASTE_BYTES is not worth a copy: a slice of a + // small parent should pass through zero-copy + let parent = string_batch(200, 60); + let slice = parent.slice(0, TARGET_ROWS); + let out = run(&mut normalizer(), [slice.clone()]); + assert_eq!(out.len(), 1); + assert!(Arc::ptr_eq(slice.column(1), out[0].column(1))); + } + + #[test] + fn moderately_wasteful_batch_below_target_passes_through() { + // ClickBench regression shape: parquet-decoded view/string batches + // routinely retain 2-4x their logical size because decode buffers are + // shared across consecutive batches. Waste that is large in ratio but + // small relative to target_bytes must NOT trigger a per-batch copy -- + // in a streaming pipeline these batches die immediately and + // compacting them is pure overhead. + let target_bytes = 16 * 1024 * 1024; + let parent = string_batch(16 * 1024, 1024); // ~17MB retained + let slice = parent.slice(0, 5000); // ~5MB logical, ~12MB waste, ratio >2 + let logical = logical_size(&slice); + let retained = get_record_batch_memory_size(&slice); + assert!(retained > WASTE_RATIO * logical); + assert!(retained - logical > MIN_WASTE_BYTES); + assert!(retained - logical < target_bytes); + + let mut n = BatchNormalizer::new( + test_schema(), + 8192, + target_bytes, + BatchNormalizerMetrics::default(), + ); + let out = run(&mut n, [slice.clone()]); + assert_eq!(out.len(), 1); + assert!( + Arc::ptr_eq(slice.column(1), out[0].column(1)), + "moderately wasteful batch should pass through zero-copy" + ); + } + + // === ordering / flush === + + #[test] + fn ordering_preserved_when_passthrough_interrupts_buffer() { + // buffered small batches must be flushed (as a runt batch) before a + // pass-through batch is emitted, preserving input order + let small1 = string_batch(10, 5); + let small2 = string_batch(10, 5); + let big = string_batch(TARGET_ROWS, 20); + let inputs = vec![small1, small2, big]; + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert_eq!( + out.iter().map(|b| b.num_rows()).collect::>(), + vec![20, TARGET_ROWS], + ); + } + + #[test] + fn finish_flushes_buffered_rows() { + let inputs = vec![string_batch(10, 5), string_batch(10, 5)]; + let out = run(&mut normalizer(), inputs.clone()); + assert_same_data(&inputs, &out); + assert_eq!(out.len(), 1); + assert_eq!(out[0].num_rows(), 20); + } + + #[test] + fn empty_batches_are_dropped() { + let inputs = vec![string_batch(0, 5), string_batch(10, 5)]; + let out = run(&mut normalizer(), inputs); + assert_eq!(out.len(), 1); + assert_eq!(out[0].num_rows(), 10); + } + + #[test] + fn no_column_batches_split_by_rows() { + use arrow::array::RecordBatchOptions; + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + vec![], + &RecordBatchOptions::new().with_row_count(Some(250)), + ) + .unwrap(); + let mut n = BatchNormalizer::new( + schema, + TARGET_ROWS, + TARGET_BYTES, + BatchNormalizerMetrics::default(), + ); + let out = run(&mut n, [batch]); + assert_eq!( + out.iter().map(|b| b.num_rows()).collect::>(), + vec![100, 100, 50], + ); + } + + #[test] + fn metrics_are_recorded() { + let metrics = ExecutionPlanMetricsSet::new(); + let m = BatchNormalizerMetrics::new(&metrics, 0); + let mut n = + BatchNormalizer::new(test_schema(), TARGET_ROWS, TARGET_BYTES, m.clone()); + run( + &mut n, + [ + string_batch(TARGET_ROWS, 20), // pass through + string_batch(10, 5), // coalesce + string_batch(50, 1024), // split + ], + ); + assert_eq!(m.batches_passed_through.value(), 1); + assert_eq!(m.batches_coalesced.value(), 1); + assert_eq!(m.batches_split.value(), 1); + } + + // === stream adapter === + + #[tokio::test] + async fn stream_end_to_end() { + use crate::ExecutionPlan; + use crate::test::TestMemoryExec; + use datafusion_execution::TaskContext; + + let inputs: Vec<_> = (0..5) + .map(|_| string_batch(10, 5)) + .chain([string_batch(50, 1024)]) + .collect(); + let schema = test_schema(); + let exec = TestMemoryExec::try_new_exec( + std::slice::from_ref(&inputs), + Arc::clone(&schema), + None, + ) + .unwrap(); + let input_stream = exec.execute(0, Arc::new(TaskContext::default())).unwrap(); + + let stream = BatchNormalizerStream::new( + input_stream, + TARGET_ROWS, + TARGET_BYTES, + BatchNormalizerMetrics::default(), + ); + assert_eq!(stream.schema(), schema); + let out: Vec = Box::pin(stream).try_collect().await.unwrap(); + assert_same_data(&inputs, &out); + assert!(out.len() > 1); + for b in &out { + assert!(b.num_rows() <= TARGET_ROWS); + assert!(logical_size(b) <= SPLIT_SLOP_FACTOR * TARGET_BYTES); + } + } +} diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 6cc6e44c32cc3..03b8335916551 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -64,6 +64,7 @@ mod visitor; pub mod aggregates; pub mod analyze; pub mod async_func; +pub mod batch_normalizer; pub mod buffer; pub mod coalesce; pub mod coalesce_batches; diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 75eb2ff980325..7f2885be9b8af 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::spill::get_record_batch_memory_size; +use crate::spill::{gc_view_arrays, get_record_batch_memory_size}; use arrow::array::ArrayRef; use arrow::compute::interleave; use arrow::datatypes::SchemaRef; @@ -32,6 +32,9 @@ struct BatchCursor { batch_idx: usize, /// The row index within the given batch row_idx: usize, + /// Estimated in-memory bytes per row of the current batch, used to + /// track the byte size of in-progress output rows + bytes_per_row: usize, } /// Provides an API to incrementally build a [`RecordBatch`] from partitioned [`RecordBatch`] @@ -67,14 +70,26 @@ pub struct BatchBuilder { /// The accumulated stream indexes from which to pull rows /// Consists of a tuple of `(batch_idx, row_idx)` indices: Vec<(usize, usize)>, + + /// Emit an output batch once the estimated bytes of in-progress rows + /// reach this target, even if fewer than `batch_size` rows are buffered. + /// `None` disables byte-based emission (row-count only). + target_bytes: Option, + + /// Estimated in-memory bytes of the in-progress rows + in_progress_bytes: usize, } impl BatchBuilder { - /// Create a new [`BatchBuilder`] with the provided `stream_count` and `batch_size` + /// Create a new [`BatchBuilder`] with the provided `stream_count` and + /// `batch_size`. If `target_bytes` is set, [`Self::byte_target_reached`] + /// reports when the in-progress rows reach that estimated size so the + /// caller can emit byte-bounded batches. pub fn new( schema: SchemaRef, stream_count: usize, batch_size: usize, + target_bytes: Option, reservation: MemoryReservation, ) -> Self { let initial_reservation = reservation.size(); @@ -86,6 +101,8 @@ impl BatchBuilder { reservation, batches_mem_used: 0, initial_reservation, + target_bytes, + in_progress_bytes: 0, } } @@ -98,10 +115,12 @@ impl BatchBuilder { // pre-reserved bytes from sort_spill_reservation_bytes). try_grow_reservation_to_at_least(&mut self.reservation, self.batches_mem_used)?; let batch_idx = self.batches.len(); + let bytes_per_row = size / batch.num_rows().max(1); self.batches.push((stream_idx, batch)); self.cursors[stream_idx] = BatchCursor { batch_idx, row_idx: 0, + bytes_per_row, }; Ok(()) } @@ -112,6 +131,14 @@ impl BatchBuilder { let row_idx = cursor.row_idx; cursor.row_idx += 1; self.indices.push((cursor.batch_idx, row_idx)); + self.in_progress_bytes += cursor.bytes_per_row; + } + + /// Returns `true` if a byte target is configured and the estimated size + /// of the in-progress rows has reached it + pub fn byte_target_reached(&self) -> bool { + self.target_bytes + .is_some_and(|target| self.in_progress_bytes >= target) } /// Returns the number of in-progress rows in this [`BatchBuilder`] @@ -154,6 +181,17 @@ impl BatchBuilder { rows_to_emit: usize, columns: Vec, ) -> Result { + // Scale the in-progress byte estimate by the fraction of rows left + // behind (only the offset-overflow partial-emit path keeps any); + // exact per-row attribution is not worth tracking here. + let total_rows = self.indices.len(); + let remaining_rows = total_rows - rows_to_emit; + self.in_progress_bytes = if remaining_rows == 0 { + 0 + } else { + (self.in_progress_bytes / total_rows.max(1)) * remaining_rows + }; + // Remove consumed indices, keeping any remaining for the next call. self.indices.drain(..rows_to_emit); @@ -194,7 +232,17 @@ impl BatchBuilder { self.reservation.shrink(self.reservation.size() - target); } - RecordBatch::try_new(Arc::clone(&self.schema), columns).map_err(Into::into) + let batch = RecordBatch::try_new(Arc::clone(&self.schema), columns)?; + if self.target_bytes.is_some() { + // `interleave` copies view-array views but shares their data + // buffers, so a byte-targeted output batch would otherwise pin + // every source batch it draws from -- inflating its retained + // size (and downstream memory reservations) far past the + // target. GC when the waste is real; no-op otherwise. Row-only + // mode (target unset) keeps today's behavior byte-for-byte. + return gc_view_arrays(&batch); + } + Ok(batch) } /// Drains the in_progress row indexes, and builds a new RecordBatch from them @@ -341,13 +389,53 @@ mod tests { assert!(matches!(error, DataFusionError::Execution(msg) if msg == "boom")); } + #[test] + fn test_byte_target_reached() { + use arrow::array::StringArray; + // 100 rows of ~1KiB strings + let s = StringArray::from_iter_values((0..100).map(|i| format!("{i:0>1024}"))); + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)]).unwrap(); + + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + let mut builder = + BatchBuilder::new(Arc::clone(&schema), 1, 8192, Some(10 * 1024), reservation); + builder.push_batch(0, batch.clone()).unwrap(); + + let mut pushed = 0; + while !builder.byte_target_reached() && pushed < 100 { + builder.push_row(0); + pushed += 1; + } + // ~10 rows of ~1KiB reach the 10KiB byte target, far below the row target + assert!(builder.byte_target_reached()); + assert!((5..=20).contains(&pushed), "pushed {pushed} rows"); + + // building the pending rows resets byte progress + let out = builder.build_record_batch().unwrap().unwrap(); + assert_eq!(out.num_rows(), pushed); + assert!(!builder.byte_target_reached()); + + // without a byte target the check never fires + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let reservation = MemoryConsumer::new("test").register(&pool); + let mut builder = + BatchBuilder::new(Arc::clone(&schema), 1, 8192, None, reservation); + builder.push_batch(0, batch).unwrap(); + for _ in 0..100 { + builder.push_row(0); + } + assert!(!builder.byte_target_reached()); + } + #[test] fn test_try_interleave_columns_surfaces_arrow_offset_overflow() { let batch = overflow_list_batch(); let schema = batch.schema(); let pool: Arc = Arc::new(UnboundedMemoryPool::default()); let reservation = MemoryConsumer::new("test").register(&pool); - let mut builder = BatchBuilder::new(schema, 1, 2, reservation); + let mut builder = BatchBuilder::new(schema, 1, 2, None, reservation); builder.push_batch(0, batch).unwrap(); let error = builder diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 4583d19e91061..241e864a0dcda 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -156,11 +156,13 @@ pub(crate) struct SortPreservingMergeStream { } impl SortPreservingMergeStream { + #[expect(clippy::too_many_arguments)] pub(crate) fn new( streams: CursorStream, schema: SchemaRef, metrics: BaselineMetrics, batch_size: usize, + batch_size_bytes: Option, fetch: Option, reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, @@ -168,7 +170,13 @@ impl SortPreservingMergeStream { let stream_count = streams.partitions(); Self { - in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), + in_progress: BatchBuilder::new( + schema, + stream_count, + batch_size, + batch_size_bytes, + reservation, + ), streams, metrics, done: false, @@ -318,7 +326,9 @@ impl SortPreservingMergeStream { if self.fetch_reached() { self.done = true; self.drain_in_progress_on_done = true; - } else if self.in_progress.len() < self.batch_size { + } else if self.in_progress.len() < self.batch_size + && !self.in_progress.byte_target_reached() + { continue; } } @@ -652,6 +662,7 @@ mod tests { Arc::clone(&schema), BaselineMetrics::new(&metrics, 0), 16, + None, Some(1), reservation, true, diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 4d108ac046eb0..54eff2a042e62 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -143,6 +143,7 @@ pub(crate) struct MultiLevelMergeBuilder { expr: LexOrdering, metrics: BaselineMetrics, batch_size: usize, + batch_size_bytes: Option, reservation: MemoryReservation, fetch: Option, enable_round_robin_tie_breaker: bool, @@ -164,6 +165,7 @@ impl MultiLevelMergeBuilder { expr: LexOrdering, metrics: BaselineMetrics, batch_size: usize, + batch_size_bytes: Option, reservation: MemoryReservation, fetch: Option, enable_round_robin_tie_breaker: bool, @@ -176,6 +178,7 @@ impl MultiLevelMergeBuilder { expr, metrics, batch_size, + batch_size_bytes, reservation, enable_round_robin_tie_breaker, fetch, @@ -376,6 +379,7 @@ impl MultiLevelMergeBuilder { .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) .with_batch_size(self.batch_size) + .with_batch_size_bytes(self.batch_size_bytes) .with_fetch(self.fetch) .with_metrics(if is_output { // Only add the metrics to the last run @@ -759,6 +763,7 @@ mod tests { expr, BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), batch_size, + None, reservation, None, false, @@ -938,6 +943,7 @@ mod tests { expr, metrics, 1024, + None, reservation, None, false, diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 792c432155a8b..7f74bac289e2e 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use parking_lot::RwLock; +use crate::batch_normalizer::effective_target_batch_size_bytes; use crate::common::spawn_buffered; use crate::execution_plan::{ Boundedness, CardinalityEffect, EmissionType, has_same_children_properties, @@ -218,6 +219,12 @@ struct ExternalSorter { expr: LexOrdering, /// The target number of rows for output batches batch_size: usize, + /// Soft byte-size target for output batches (from + /// `datafusion.execution.target_batch_size_bytes`): sorted output chunks + /// and merged batches are additionally bounded to about this many bytes, + /// so wide rows do not produce arbitrarily large batches. `None` bounds + /// output by row count only. + target_batch_size_bytes: Option, /// If the in size of buffered memory batches is below this size, /// the data will be concatenated and sorted in place rather than /// sort/merged. @@ -273,6 +280,7 @@ impl ExternalSorter { schema: SchemaRef, expr: LexOrdering, batch_size: usize, + target_batch_size_bytes: Option, sort_spill_reservation_bytes: usize, sort_in_place_threshold_bytes: usize, // Configured via `datafusion.execution.spill_compression`. @@ -308,6 +316,7 @@ impl ExternalSorter { merge_reservation, runtime, batch_size, + target_batch_size_bytes, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, }) @@ -365,6 +374,7 @@ impl ExternalSorter { .with_expressions(&self.expr.clone()) .with_metrics(self.metrics.baseline.clone()) .with_batch_size(self.batch_size) + .with_batch_size_bytes(self.target_batch_size_bytes) .with_fetch(None) .with_reservation(self.merge_reservation.take()) .build() @@ -664,6 +674,7 @@ impl ExternalSorter { self.metrics.baseline.intermediate() }) .with_batch_size(self.batch_size) + .with_batch_size_bytes(self.target_batch_size_bytes) .with_fetch(None) .with_reservation(self.merge_reservation.new_empty()) .build() @@ -743,12 +754,14 @@ impl ExternalSorter { let schema = batch.schema(); let expressions = self.expr.clone(); let batch_size = self.batch_size; + let target_bytes = self.target_batch_size_bytes; let stream = futures::stream::once(async move { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; + let sorted_batches = + sort_batch_chunked(&batch, &expressions, batch_size, target_bytes)?; // Resize the reservation to match the actual sorted output size. // Using try_resize avoids a release-then-reacquire cycle, which @@ -916,12 +929,22 @@ pub fn sort_batch( /// Sort a batch and return the result as multiple batches of size `batch_size`. /// This is useful when you want to avoid creating one large sorted batch in memory, /// and instead want to process the sorted data in smaller chunks. +/// +/// If `target_bytes` is set, chunks are additionally bounded to about that +/// many bytes each (wide rows would otherwise produce huge chunks). pub fn sort_batch_chunked( batch: &RecordBatch, expressions: &LexOrdering, batch_size: usize, + target_bytes: Option, ) -> Result> { - IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size).collect() + IncrementalSortIterator::new( + batch.clone(), + expressions.clone(), + batch_size, + target_bytes, + ) + .collect() } /// Sort execution plan. @@ -1381,6 +1404,7 @@ impl ExecutionPlan for SortExec { input.schema(), self.expr.clone(), context.session_config().batch_size(), + effective_target_batch_size_bytes(&context), execution_options.sort_spill_reservation_bytes, execution_options.sort_in_place_threshold_bytes, context.session_config().spill_compression(), @@ -2889,6 +2913,124 @@ mod tests { Ok((sorted_batches, metrics)) } + #[tokio::test] + async fn test_sort_batch_chunked_byte_target() -> Result<()> { + use arrow::array::StringArray; + // 100 rows x ~1KiB strings (~100KiB): an 8KiB byte target must + // produce many small chunks even though 100 rows < batch_size + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let s = + StringArray::from_iter_values((0..100).rev().map(|i| format!("{i:0>1024}"))); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)])?; + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("s", 0)))].into(); + + let target_bytes = 8 * 1024; + let result_batches = + sort_batch_chunked(&batch, &expressions, 1000, Some(target_bytes))?; + + assert!( + result_batches.len() > 1, + "expected byte target to split the sorted output" + ); + let mut total_rows = 0; + for chunk in &result_batches { + assert!( + get_record_batch_memory_size(chunk) <= 2 * target_bytes, + "chunk of {} bytes exceeds byte target", + get_record_batch_memory_size(chunk) + ); + total_rows += chunk.num_rows(); + } + assert_eq!(total_rows, 100); + + // Concatenated output must equal a full sort of the input + let expected = sort_batch(&batch, &expressions, None)?; + let actual = concat_batches(&schema, &result_batches)?; + assert_eq!(expected, actual); + Ok(()) + } + + #[tokio::test] + async fn test_sort_batch_chunked_byte_target_releases_view_buffers() -> Result<()> { + use arrow::array::StringViewArray; + // take() on view arrays shares data buffers; chunked sort output + // must GC them or each chunk pins the whole parent batch + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + let s = StringViewArray::from_iter_values( + (0..100).rev().map(|i| format!("{i:0>1024}")), + ); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)])?; + let parent_size = get_record_batch_memory_size(&batch); + + let expressions: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("s", 0)))].into(); + let result_batches = + sort_batch_chunked(&batch, &expressions, 1000, Some(8 * 1024))?; + assert!(result_batches.len() > 1); + for chunk in &result_batches { + assert!( + get_record_batch_memory_size(chunk) < parent_size / 2, + "chunk retains {} bytes of the parent's view data buffers", + get_record_batch_memory_size(chunk) + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_sort_exec_byte_bounded_output() -> Result<()> { + use arrow::array::StringArray; + // Multiple wide-row input batches; with target_batch_size_bytes set, + // the sorted output batches must be bounded by bytes, not just rows + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batches: Vec = (0..4) + .map(|b| { + let s = StringArray::from_iter_values( + (0..50).map(|i| format!("{:0>1024}", (b * 50 + i * 7919) % 200)), + ); + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)]).unwrap() + }) + .collect(); + + let target_bytes = 16 * 1024; + let mut config = SessionConfig::new(); + config.options_mut().execution.target_batch_size_bytes = Some(target_bytes); + let task_ctx = Arc::new(TaskContext::default().with_session_config(config)); + + let plan = TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?; + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("s", 0)))].into(), + plan, + )); + let result = collect(sort_exec, task_ctx).await?; + + assert!( + result.len() > 1, + "expected byte-bounded output batches, got {}", + result.len() + ); + let total_rows: usize = result.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 200); + for batch in &result { + assert!( + get_record_batch_memory_size(batch) <= 2 * target_bytes, + "output batch of {} bytes exceeds byte target", + get_record_batch_memory_size(batch) + ); + } + Ok(()) + } + #[tokio::test] async fn test_sort_batch_chunked_basic() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); @@ -2907,7 +3049,7 @@ mod tests { [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); // Sort with batch_size = 250 - let result_batches = sort_batch_chunked(&batch, &expressions, 250)?; + let result_batches = sort_batch_chunked(&batch, &expressions, 250, None)?; // Verify 4 batches are returned assert_eq!(result_batches.len(), 4); @@ -2960,7 +3102,7 @@ mod tests { [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); // Sort with batch_size = 100 - let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + let result_batches = sort_batch_chunked(&batch, &expressions, 100, None)?; // Should return exactly 1 batch assert_eq!(result_batches.len(), 1); @@ -2992,7 +3134,7 @@ mod tests { [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); // Sort with batch_size = 100 - let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + let result_batches = sort_batch_chunked(&batch, &expressions, 100, None)?; // Should return exactly 10 batches of 100 rows each assert_eq!(result_batches.len(), 10); @@ -3019,7 +3161,7 @@ mod tests { let expressions: LexOrdering = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); - let result_batches = sort_batch_chunked(&batch, &expressions, 100)?; + let result_batches = sort_batch_chunked(&batch, &expressions, 100, None)?; // Empty input produces no output batches (0 chunks) assert_eq!(result_batches.len(), 0); @@ -3273,7 +3415,8 @@ mod tests { 0, Arc::clone(&schema), [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(), - 128, // batch_size + 128, // batch_size + None, // target_batch_size_bytes sort_spill_reservation_bytes, usize::MAX, // sort_in_place_threshold_bytes (high to avoid concat path) SpillCompression::Uncompressed, diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 77a7d8f8f2e11..ef9d16941cf3f 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -19,6 +19,7 @@ use std::sync::Arc; +use crate::batch_normalizer::effective_target_batch_size_bytes; use crate::common::spawn_buffered; use crate::limit::LimitStream; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; @@ -363,6 +364,7 @@ impl ExecutionPlan for SortPreservingMergeExec { .with_expressions(&self.expr) .with_metrics(BaselineMetrics::new(&self.metrics, partition)) .with_batch_size(context.session_config().batch_size()) + .with_batch_size_bytes(effective_target_batch_size_bytes(&context)) .with_fetch(self.fetch) .with_reservation(reservation) .with_round_robin_tie_breaker(self.enable_round_robin_repartition) @@ -838,6 +840,124 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_merge_byte_bounded_output() { + use crate::spill::get_record_batch_memory_size; + use crate::stream::RecordBatchStreamAdapter; + use arrow::array::StringArray; + use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; + + // Two sorted streams of wide rows (~1KiB strings). With a byte + // target set, merged output batches must be bounded by bytes even + // though the row target (8192) is never reached. + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let make_stream = |offset: usize| -> SendableRecordBatchStream { + let s = StringArray::from_iter_values( + (0..100).map(|i| format!("{:0>1024}", 2 * i + offset)), + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)]).unwrap(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )) + }; + + let exprs: LexOrdering = + [PhysicalSortExpr::new_default(col("s", &schema).unwrap())].into(); + let target_bytes = 16 * 1024; + + let merged = StreamingMergeBuilder::new() + .with_streams(vec![make_stream(0), make_stream(1)]) + .with_schema(Arc::clone(&schema)) + .with_expressions(&exprs) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(8192) + .with_batch_size_bytes(Some(target_bytes)) + .with_reservation({ + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("test").register(&pool) + }) + .build() + .unwrap(); + + let batches = common::collect(merged).await.unwrap(); + assert!( + batches.len() > 1, + "expected byte target to bound merged batches, got {}", + batches.len() + ); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 200); + for batch in &batches { + assert!( + get_record_batch_memory_size(batch) <= 2 * target_bytes, + "merged batch of {} bytes exceeds byte target", + get_record_batch_memory_size(batch) + ); + } + } + + #[tokio::test] + async fn test_merge_byte_bounded_output_releases_view_buffers() { + use crate::spill::get_record_batch_memory_size; + use crate::stream::RecordBatchStreamAdapter; + use arrow::array::StringViewArray; + use datafusion_execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; + + // Utf8View inputs: arrow's interleave copies views but shares the + // underlying data buffers, so without a GC pass every merged output + // batch would pin all source batches -- inflating retained size (and + // therefore downstream memory reservations) by orders of magnitude. + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Utf8View, + false, + )])); + let source_size = 200 * 1024; // ~200KiB per source batch + let make_stream = |offset: usize| -> SendableRecordBatchStream { + let s = StringViewArray::from_iter_values( + (0..200).map(|i| format!("{:0>1024}", 2 * i + offset)), + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s)]).unwrap(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch)]), + )) + }; + + let exprs: LexOrdering = + [PhysicalSortExpr::new_default(col("s", &schema).unwrap())].into(); + let target_bytes = 32 * 1024; + + let merged = StreamingMergeBuilder::new() + .with_streams(vec![make_stream(0), make_stream(1)]) + .with_schema(Arc::clone(&schema)) + .with_expressions(&exprs) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(8192) + .with_batch_size_bytes(Some(target_bytes)) + .with_reservation({ + let pool: Arc = Arc::new(UnboundedMemoryPool::default()); + MemoryConsumer::new("test").register(&pool) + }) + .build() + .unwrap(); + + let batches = common::collect(merged).await.unwrap(); + assert!(batches.len() > 1); + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 400); + for batch in &batches { + assert!( + get_record_batch_memory_size(batch) < source_size, + "merged batch retains {} bytes -- it is pinning the source \ + batches' view data buffers", + get_record_batch_memory_size(batch) + ); + } + } + // Split the provided record batch into multiple batch_size record batches fn split_batch(sorted: &RecordBatch, batch_size: usize) -> Vec { let batches = sorted.num_rows().div_ceil(batch_size); diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index 107631074ed3d..b3be7fe492454 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -16,6 +16,7 @@ // under the License. use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues}; +use crate::spill::{gc_view_arrays, get_record_batch_memory_size}; use crate::{EmptyRecordBatchStream, SendableRecordBatchStream}; use crate::{PhysicalExpr, PhysicalSortExpr}; use arrow::array::{Array, UInt32Array}; @@ -311,20 +312,38 @@ pub(crate) struct IncrementalSortIterator { batch: RecordBatch, expressions: LexOrdering, batch_size: usize, + /// When byte-targeted, GC view arrays in each chunk: `take` shares view + /// data buffers, so an un-GC'd chunk would pin the whole parent batch + gc_views: bool, indices: Option, cursor: usize, } impl IncrementalSortIterator { + /// Create a new iterator yielding sorted chunks of at most `batch_size` + /// rows. If `target_bytes` is set, the chunk row count is further reduced + /// so each chunk's estimated in-memory size is about `target_bytes`, + /// bounding output batches by bytes as well as rows (wide rows would + /// otherwise produce arbitrarily large chunks). pub(crate) fn new( batch: RecordBatch, expressions: LexOrdering, batch_size: usize, + target_bytes: Option, ) -> Self { + let batch_size = match target_bytes { + Some(target) if batch.num_rows() > 0 => { + let bytes_per_row = + (get_record_batch_memory_size(&batch) / batch.num_rows()).max(1); + (target / bytes_per_row).clamp(1, batch_size) + } + _ => batch_size, + }; Self { batch, expressions, batch_size, + gc_views: target_bytes.is_some(), cursor: 0, indices: None, } @@ -369,6 +388,14 @@ impl Iterator for IncrementalSortIterator { Ok(batch) => batch, Err(e) => return Some(Err(e.into())), }; + let new_batch = if self.gc_views { + match gc_view_arrays(&new_batch) { + Ok(batch) => batch, + Err(e) => return Some(Err(e)), + } + } else { + new_batch + }; self.cursor += batch_size; @@ -430,7 +457,7 @@ mod tests { .unwrap(); let mut total_rows = 0; - IncrementalSortIterator::new(batch.clone(), expressions, batch_size).try_for_each( + IncrementalSortIterator::new(batch.clone(), expressions, batch_size, None).try_for_each( |result| { let chunk = result?; total_rows += chunk.num_rows(); diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index ade24ff0534ff..461728c8480c3 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -43,7 +43,7 @@ macro_rules! primitive_merge_helper { } macro_rules! merge_helper { - ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ + ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $batch_size_bytes:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); return Ok(Box::pin(SortPreservingMergeStream::new( @@ -51,6 +51,7 @@ macro_rules! merge_helper { $schema, $tracking_metrics, $batch_size, + $batch_size_bytes, $fetch, $reservation, $enable_round_robin_tie_breaker, @@ -92,6 +93,7 @@ pub struct StreamingMergeBuilder<'a> { expressions: Option<&'a LexOrdering>, metrics: Option, batch_size: Option, + batch_size_bytes: Option, fetch: Option, reservation: Option, enable_round_robin_tie_breaker: bool, @@ -143,6 +145,15 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Also emit an output batch whenever the estimated size of the + /// in-progress rows reaches `batch_size_bytes`, so that wide rows do not + /// produce arbitrarily large merged batches. `None` (the default) emits + /// on row count only. + pub fn with_batch_size_bytes(mut self, batch_size_bytes: Option) -> Self { + self.batch_size_bytes = batch_size_bytes; + self + } + pub fn with_fetch(mut self, fetch: Option) -> Self { self.fetch = fetch; self @@ -184,6 +195,7 @@ impl<'a> StreamingMergeBuilder<'a> { schema, metrics, batch_size, + batch_size_bytes, reservation, fetch, expressions, @@ -212,6 +224,7 @@ impl<'a> StreamingMergeBuilder<'a> { expressions.clone(), metrics, batch_size, + batch_size_bytes, reservation, fetch, enable_round_robin_tie_breaker, @@ -238,12 +251,12 @@ impl<'a> StreamingMergeBuilder<'a> { let sort = expressions[0].clone(); let data_type = sort.expr.data_type(schema.as_ref())?; downcast_primitive! { - data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker), - DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker), + DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker) _ => {} } } @@ -259,6 +272,7 @@ impl<'a> StreamingMergeBuilder<'a> { schema, metrics, batch_size, + batch_size_bytes, fetch, reservation, enable_round_robin_tie_breaker, diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index bf45564e26333..32fde1a76e214 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -214,6 +214,7 @@ datafusion.catalog.has_header true datafusion.catalog.information_schema true datafusion.catalog.location NULL datafusion.catalog.newlines_in_values false +datafusion.execution.adaptive_target_batch_size true datafusion.execution.batch_size 8192 datafusion.execution.coalesce_batches true datafusion.execution.collect_statistics true @@ -281,6 +282,7 @@ datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed datafusion.execution.split_file_groups_by_statistics false +datafusion.execution.target_batch_size_bytes NULL datafusion.execution.target_partitions 7 datafusion.execution.time_zone NULL datafusion.execution.use_row_number_estimates_to_optimize_partitioning false @@ -373,6 +375,7 @@ datafusion.catalog.has_header true Default value for `format.has_header` for `CR datafusion.catalog.information_schema true Should DataFusion provide access to `information_schema` virtual tables for displaying schema information datafusion.catalog.location NULL Location scanned to load tables for `default` schema datafusion.catalog.newlines_in_values false Specifies whether newlines in (quoted) CSV values are supported. This is the default value for `format.newlines_in_values` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. Parsing newlines in quoted values may be affected by execution behaviour such as parallel file scanning. Setting this to `true` ensures that newlines in values are parsed successfully, which may reduce performance. +datafusion.execution.adaptive_target_batch_size true When `target_batch_size_bytes` is unset and the memory pool has a finite limit, derive a byte-size target for batches automatically from the pool's per-partition share, so queries whose batches are large relative to the memory budget re-chunk them (and can complete) instead of failing with ResourcesExhausted. Has no effect when there is no memory limit (nothing to protect: batches pass through untouched) or when `target_batch_size_bytes` is set explicitly. datafusion.execution.batch_size 8192 Default batch size while creating new batches, it's especially useful for buffer-in-memory batches since creating tiny batches would result in too much metadata memory consumption datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. @@ -440,6 +443,7 @@ datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer cap datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). datafusion.execution.spill_compression uncompressed Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. datafusion.execution.split_file_groups_by_statistics false Attempt to eliminate sorts by packing & sorting files with non-overlapping statistics into the same file groups. Currently experimental +datafusion.execution.target_batch_size_bytes NULL Soft target, in bytes, for the in-memory size of batches, in addition to the row-based `batch_size` target. When set, data source output batches are re-chunked towards both targets: batches much larger than this (by in-memory size) are split into compact copies of roughly this size even when within `batch_size` rows -- protecting downstream operators from very large batches (e.g. wide rows with large string values) -- and small batches are coalesced until they reach either target. Batches near the targets are passed through without copying. Sorts, sort-preserving merges and aggregation spills also bound their output and spill batches to roughly this size instead of emitting `batch_size`-row batches of unbounded byte size. When `None` (the default), batches are only chunked by row count. datafusion.execution.target_partitions 7 Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system datafusion.execution.time_zone NULL The default time zone Some functions, e.g. `now` return timestamps in this time zone datafusion.execution.use_row_number_estimates_to_optimize_partitioning false Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 03340c366d70f..aa4cf0268ef1f 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -74,6 +74,8 @@ The following configuration settings are available: | datafusion.catalog.has_header | true | Default value for `format.has_header` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. | | datafusion.catalog.newlines_in_values | false | Specifies whether newlines in (quoted) CSV values are supported. This is the default value for `format.newlines_in_values` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. Parsing newlines in quoted values may be affected by execution behaviour such as parallel file scanning. Setting this to `true` ensures that newlines in values are parsed successfully, which may reduce performance. | | datafusion.execution.batch_size | 8192 | Default batch size while creating new batches, it's especially useful for buffer-in-memory batches since creating tiny batches would result in too much metadata memory consumption | +| datafusion.execution.target_batch_size_bytes | NULL | Soft target, in bytes, for the in-memory size of batches, in addition to the row-based `batch_size` target. When set, data source output batches are re-chunked towards both targets: batches much larger than this (by in-memory size) are split into compact copies of roughly this size even when within `batch_size` rows -- protecting downstream operators from very large batches (e.g. wide rows with large string values) -- and small batches are coalesced until they reach either target. Batches near the targets are passed through without copying. Sorts, sort-preserving merges and aggregation spills also bound their output and spill batches to roughly this size instead of emitting `batch_size`-row batches of unbounded byte size. When `None` (the default), batches are only chunked by row count. | +| datafusion.execution.adaptive_target_batch_size | true | When `target_batch_size_bytes` is unset and the memory pool has a finite limit, derive a byte-size target for batches automatically from the pool's per-partition share, so queries whose batches are large relative to the memory budget re-chunk them (and can complete) instead of failing with ResourcesExhausted. Has no effect when there is no memory limit (nothing to protect: batches pass through untouched) or when `target_batch_size_bytes` is set explicitly. | | datafusion.execution.perfect_hash_join_small_build_threshold | 1024 | A perfect hash join (see `HashJoinExec` for more details) will be considered if the range of keys (max - min) on the build side is < this threshold. This provides a fast path for joins with very small key ranges, bypassing the density check. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.perfect_hash_join_min_key_density | 0.15 | The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.coalesce_batches | true | When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting |