This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 0646a310cd fix: reduce peak memory usage when round robin tiebreaker
is disabled (#23606)
0646a310cd is described below
commit 0646a310cdeb25ba6a091e643dddd06da7ea181e
Author: Ariel Miculas-Trif <[email protected]>
AuthorDate: Fri Aug 7 11:22:12 2026 +0200
fix: reduce peak memory usage when round robin tiebreaker is disabled
(#23606)
## Which issue does this PR close?
- First part of https://github.com/apache/datafusion/issues/23604
## Rationale for this change
Don't pay the memory price when the tie breaking feature is disabled.
## What changes are included in this PR?
Mostly tests to show the issue
## Are these changes tested?
Yes
## Are there any user-facing changes?
No
---
Cargo.lock | 2 -
benchmarks/Cargo.toml | 5 -
benchmarks/src/bin/external_aggr.rs | 6 +-
benchmarks/src/util/memory.rs | 4 +-
benchmarks/src/util/mod.rs | 2 -
benchmarks/src/util/options.rs | 7 +-
benchmarks/src/util/run.rs | 3 +-
.../spilling_fuzz_in_memory_constrained_env.rs | 255 +++++++++++++++++++++
datafusion/execution/src/memory_pool/mod.rs | 2 +
.../execution/src/memory_pool/peak_recording.rs | 28 +--
datafusion/physical-plan/src/sorts/merge.rs | 75 +++---
11 files changed, 323 insertions(+), 66 deletions(-)
diff --git a/Cargo.lock b/Cargo.lock
index 7ace154139..93627f290f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1780,7 +1780,6 @@ name = "datafusion-benchmarks"
version = "54.1.0"
dependencies = [
"arrow",
- "arrow-buffer",
"async-trait",
"bytes",
"clap",
@@ -1788,7 +1787,6 @@ dependencies = [
"datafusion",
"datafusion-common",
"datafusion-common-runtime",
- "datafusion-execution",
"datafusion-proto",
"env_logger",
"futures",
diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 11f83cef5e..a8dcd704ef 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -65,11 +65,6 @@ tokio-util = { version = "0.7.17" }
toml = "0.9.8"
[dev-dependencies]
-# `pool`/`arrow_buffer_pool` are enabled only for tests, so the benchmark
-# binaries are built exactly as before. They let `memory_pool`'s tests cover
-# Arrow-side reservations reaching the pool via `ArrowMemoryPool`.
-arrow-buffer = { workspace = true, features = ["pool"] }
-datafusion-execution = { workspace = true, features = ["arrow_buffer_pool"] }
datafusion-proto = { workspace = true, features = ["parquet"] }
tempfile = { workspace = true }
diff --git a/benchmarks/src/bin/external_aggr.rs
b/benchmarks/src/bin/external_aggr.rs
index 226a619192..c19554eb33 100644
--- a/benchmarks/src/bin/external_aggr.rs
+++ b/benchmarks/src/bin/external_aggr.rs
@@ -34,14 +34,12 @@ use datafusion::datasource::listing::{
use datafusion::datasource::{MemTable, TableProvider};
use datafusion::error::Result;
use datafusion::execution::SessionStateBuilder;
-use datafusion::execution::memory_pool::FairSpillPool;
+use datafusion::execution::memory_pool::{FairSpillPool, PeakRecordingPool};
use datafusion::execution::runtime_env::RuntimeEnvBuilder;
use datafusion::physical_plan::display::DisplayableExecutionPlan;
use datafusion::physical_plan::{collect, displayable};
use datafusion::prelude::*;
-use datafusion_benchmarks::util::{
- BenchmarkRun, CommonOpt, PeakRecordingPool, QueryResult,
-};
+use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, QueryResult};
use datafusion_common::instant::Instant;
use datafusion_common::utils::get_available_parallelism;
use datafusion_common::{DEFAULT_PARQUET_EXTENSION, exec_err};
diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs
index 2b186c79c3..f0339c9cf0 100644
--- a/benchmarks/src/util/memory.rs
+++ b/benchmarks/src/util/memory.rs
@@ -15,9 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use datafusion::execution::memory_pool::MemoryPool;
-
-use super::PeakRecordingPool;
+use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool};
/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by
/// the peak reservation of `memory_pool` when a memory limit was configured.
diff --git a/benchmarks/src/util/mod.rs b/benchmarks/src/util/mod.rs
index 43855ea468..6dc11c0f42 100644
--- a/benchmarks/src/util/mod.rs
+++ b/benchmarks/src/util/mod.rs
@@ -18,11 +18,9 @@
//! Shared benchmark utilities
pub mod latency_object_store;
mod memory;
-mod memory_pool;
mod options;
mod run;
pub use memory::print_memory_stats;
-pub use memory_pool::PeakRecordingPool;
pub use options::CommonOpt;
pub use run::{BenchQuery, BenchmarkRun, QueryResult};
diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs
index c744d0bf31..4a1c14674a 100644
--- a/benchmarks/src/util/options.rs
+++ b/benchmarks/src/util/options.rs
@@ -21,7 +21,10 @@ use clap::Args;
use datafusion::{
execution::{
disk_manager::DiskManagerBuilder,
- memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool,
TrackConsumersPool},
+ memory_pool::{
+ FairSpillPool, GreedyMemoryPool, MemoryPool, PeakRecordingPool,
+ TrackConsumersPool,
+ },
object_store::ObjectStoreUrl,
runtime_env::{RuntimeEnv, RuntimeEnvBuilder},
},
@@ -30,7 +33,7 @@ use datafusion::{
use datafusion_common::{DataFusionError, Result};
use object_store::local::LocalFileSystem;
-use super::{latency_object_store::LatencyObjectStore,
memory_pool::PeakRecordingPool};
+use super::latency_object_store::LatencyObjectStore;
// Common benchmark options (don't use doc comments otherwise this doc
// shows up in help files)
diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs
index 6c63ceec64..772d421bc7 100644
--- a/benchmarks/src/util/run.rs
+++ b/benchmarks/src/util/run.rs
@@ -15,8 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use super::memory_pool::PeakRecordingPool;
-use datafusion::execution::memory_pool::MemoryPool;
+use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool};
use datafusion::{DATAFUSION_VERSION, error::Result};
use datafusion_common::utils::get_available_parallelism;
use serde::{Serialize, Serializer};
diff --git
a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
index 103c3e03c0..f82d0165f2 100644
---
a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
+++
b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs
@@ -23,6 +23,7 @@ use std::sync::Arc;
use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric;
use crate::fuzz_cases::once_exec::OnceExec;
use arrow::array::UInt64Array;
+use arrow::row::{RowConverter, SortField};
use arrow::{array::StringArray, compute::SortOptions,
record_batch::RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use datafusion::common::Result;
@@ -45,9 +46,19 @@ use datafusion_physical_plan::aggregates::{
AggregateExec, AggregateMode, PhysicalGroupBy,
};
use datafusion_physical_plan::metrics::MetricValue;
+use datafusion_physical_plan::spill::get_record_batch_memory_size;
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use futures::StreamExt;
+use arrow::array::Int32Array;
+use datafusion::datasource::memory::MemorySourceConfig;
+use
datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+use datafusion_execution::memory_pool::{
+ MemoryPool, PeakRecordingPool, UnboundedMemoryPool,
+};
+use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
+use datafusion_physical_plan::spill::SpillManager;
+
#[tokio::test]
async fn test_sort_with_limited_memory() -> Result<()> {
let record_batch_size = 8192;
@@ -290,6 +301,250 @@ async fn
test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
Ok(())
}
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, false,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false, false,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, true,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false, true,
false).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(true, false,
true).await
+}
+
+#[tokio::test]
+async fn
test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values()
+-> Result<()> {
+ run_sort_preserving_merge_peak_memory_with_spilled_input(false, false,
true).await
+}
+
+/// Intended to measure the maximum number of record batches held in memory by
+/// the SortPreservingMergeStream in a convoluted way by measuring the peak
+/// memory reservation. Relevant for merging spilled streams, where the
produced
+/// record batches suffer from the following issue:
+/// https://github.com/apache/arrow-rs/issues/6363
+///
+/// After an IPC roundtrip, all columns in a [`RecordBatch`] share a single
+/// parent buffer. It causes the memory reservation to be inflated, but the
+/// bigger issue is the increase in the peak allocated memory caused by
+/// prev_cursors in SortPreservingMergeExec. The increase is caused by the fact
+/// that the FieldCursor inside prev_cursors holds a reference for the entire
+/// Buffer allocated for the input record batch, preventing it from being
+/// dropped and thus increasing the number of concomitent input record batches
+/// living during the merging phase
+async fn run_sort_preserving_merge_peak_memory_with_spilled_input(
+ round_robin: bool,
+ multi_column_sort: bool,
+ tied_values: bool,
+) -> Result<()> {
+ let num_batches = 10usize;
+ let num_rows_per_batch = 100usize;
+ // payload is ~100x larger than the sort key (i32 = 4 bytes, string ≈ 400
bytes)
+ let large_string = "x".repeat(400);
+
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("sort_key", DataType::Int32, false),
+ Field::new("payload", DataType::Utf8, false),
+ ]));
+
+ // Unbounded env used only for spilling the input; the merge runs under its
+ // own pool below.
+ let spill_env = Arc::new(RuntimeEnvBuilder::new().build()?);
+
+ let mut partition_batches: Vec<Vec<RecordBatch>> = Vec::new();
+
+ for stream_idx in 0..2usize {
+ // Each stream covers a non-overlapping key range so both are
individually
+ // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000). When
+ // `tied_values` is set, every row of every batch in both streams
+ // instead carries the same sort key, so every comparison between the
+ // two streams is a tie.
+ let batches: Vec<RecordBatch> = (0..num_batches)
+ .map(|b| {
+ // Interleave streams: stream 0 → even slots [0,200,400,...],
+ // stream 1 → odd slots [100,300,500,...] so the merge
+ // alternates between them on every batch.
+ let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32;
+ let sort_col: Int32Array = if tied_values {
+ std::iter::repeat_n(0, num_rows_per_batch).collect()
+ } else {
+ (base..base + num_rows_per_batch as i32).collect()
+ };
+ let payload_col: StringArray =
+ std::iter::repeat_n(large_string.as_str(),
num_rows_per_batch)
+ .map(Some)
+ .collect();
+ RecordBatch::try_new(
+ Arc::clone(&schema),
+ vec![Arc::new(sort_col), Arc::new(payload_col)],
+ )
+ .unwrap()
+ })
+ .collect();
+
+ // Spill to disk then read back: each RecordBatch is now IPC-backed,
+ // meaning all columns share a single parent buffer. As a result,
+ // get_buffer_memory_size() on the sort_key column returns the full
+ // parent-buffer capacity (≈ batch size of both columns combined)
rather
+ // than just the key data (num_rows * 4 bytes).
+ let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+ let manager =
+ SpillManager::new(Arc::clone(&spill_env), metrics,
Arc::clone(&schema));
+ let spill_file = manager
+ .spill_record_batch_and_finish(&batches, "stream")?
+ .expect("non-empty input should produce a spill file");
+
+ let mut stream = manager.read_spill_as_stream(spill_file, None)?;
+ let mut ipc_batches: Vec<RecordBatch> = Vec::new();
+ while let Some(batch) = stream.next().await {
+ ipc_batches.push(batch?);
+ }
+ partition_batches.push(ipc_batches);
+ }
+
+ let ipc_batch_size =
get_record_batch_memory_size(&partition_batches[0][0]);
+
+ // Build a 2-partition plan from the IPC-recovered batches.
+ let input =
+ MemorySourceConfig::try_new_exec(&partition_batches,
Arc::clone(&schema), None)?;
+
+ let sort_key_expr = PhysicalSortExpr {
+ expr: col("sort_key", &schema)?,
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ };
+ // `payload` has the same value in every row, so adding it as a secondary
+ // sort key doesn't change the resulting order — it only forces the merge
+ // onto the row-oriented (`RowValues`/`RowCursorStream`) comparison path
+ // used whenever more than one sort expression is present.
+ let mut sort_exprs = vec![sort_key_expr];
+ if multi_column_sort {
+ sort_exprs.push(PhysicalSortExpr {
+ expr: col("payload", &schema)?,
+ options: SortOptions {
+ descending: false,
+ nulls_first: true,
+ },
+ });
+ }
+
+ // When sorting by more than one column, the merge switches to the
+ // row-oriented `RowValues`/`RowCursorStream` path
+ //
+ // `RowCursorStream` also tracks one *shared* (not per-partition)
+ // reservation sized to `converter.size()` (`stream.rs`:
+ // `self.reservation.try_resize(self.converter.size())`) — the
+ // `RowConverter`'s own fixed internal state, separate from the `Rows`
+ // it produces per batch.
+ let (row_batch_size, converter_size) = if multi_column_sort {
+ let sort_fields = sort_exprs
+ .iter()
+ .map(|s| {
+ let data_type = s.expr.data_type(&schema)?;
+ Ok(SortField::new_with_options(data_type, s.options))
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let converter = RowConverter::new(sort_fields)?;
+ let cols = sort_exprs
+ .iter()
+ .map(|s| {
+ s.expr
+ .evaluate(&partition_batches[0][0])?
+ .into_array(partition_batches[0][0].num_rows())
+ })
+ .collect::<Result<Vec<_>>>()?;
+ let rows = converter.convert_columns(&cols)?;
+ (rows.size(), converter.size())
+ } else {
+ (0, 0)
+ };
+
+ let merge = Arc::new(
+ SortPreservingMergeExec::new(LexOrdering::new(sort_exprs).unwrap(),
input)
+ .with_round_robin_repartition(round_robin),
+ );
+
+ // PeakRecordingPool records peak reserved bytes as a running high-water
mark
+ // (via grow/shrink deltas), independent of any per-consumer registration
+ // bookkeeping - unlike TrackConsumersPool, whose tracked-consumer entry
(and
+ // its peak) gets discarded the moment the consumer unregisters, which now
+ // happens mid-poll (inside the drain loop below) rather than when the
+ // caller eventually drops the returned stream.
+ let tracking_pool = Arc::new(PeakRecordingPool::new(Arc::new(
+ UnboundedMemoryPool::default(),
+ )));
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&tracking_pool) as Arc<dyn MemoryPool>)
+ .build()?;
+ let task_ctx = Arc::new(
+ TaskContext::default()
+
.with_session_config(SessionConfig::new().with_batch_size(num_rows_per_batch))
+ .with_runtime(Arc::new(runtime)),
+ );
+
+ let mut output = merge.execute(0, task_ctx)?;
+ let mut total_rows = 0usize;
+ while let Some(batch) = output.next().await {
+ total_rows += batch?.num_rows();
+ }
+ assert_eq!(total_rows, 2 * num_batches * num_rows_per_batch);
+
+ let peak_bytes = tracking_pool.peak_reserved();
+
+ // in the single column case, the cursor takes up an ipc_batch_size worth
of memory due to the
+ // IPC roundtrip issue
+ // for the multi-column case, we've calculated row_batch_size above
+ let cursor_unit = if multi_column_sort {
+ row_batch_size
+ } else {
+ ipc_batch_size
+ };
+
+ // BatchBuilder needs to hold 3 Record batches simultaneously to merge two
+ // streams (because a stream can cross a record batch boundary)
+ // there is also one cursor needed per stream
+ let mut max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size;
+
+ // with round robin enabled, 2 extra cursors live in memory
+ // see https://github.com/apache/datafusion/issues/23604
+ if round_robin {
+ max_peak += 2 * cursor_unit;
+ };
+
+ assert!(
+ peak_bytes > 0,
+ "peak reservation {peak_bytes} should be greater than 0"
+ );
+ assert!(
+ peak_bytes <= max_peak,
+ "peak reservation {peak_bytes} bytes exceeds max_peak ({max_peak}
bytes); \
+ round_robin={round_robin}, multi_column_sort={multi_column_sort}",
+ );
+
+ Ok(())
+}
+
struct RunTestWithLimitedMemoryArgs {
pool_size: usize,
task_ctx: Arc<TaskContext>,
diff --git a/datafusion/execution/src/memory_pool/mod.rs
b/datafusion/execution/src/memory_pool/mod.rs
index 2b36ee7f40..40a79d136b 100644
--- a/datafusion/execution/src/memory_pool/mod.rs
+++ b/datafusion/execution/src/memory_pool/mod.rs
@@ -24,6 +24,7 @@ use std::fmt::Display;
use std::hash::{Hash, Hasher};
use std::{cmp::Ordering, sync::Arc, sync::atomic};
+mod peak_recording;
mod pool;
#[cfg(feature = "arrow_buffer_pool")]
@@ -36,6 +37,7 @@ pub mod proxy {
pub use datafusion_common::{
human_readable_count, human_readable_duration, human_readable_size, units,
};
+pub use peak_recording::*;
pub use pool::*;
/// Tracks and potentially limits memory use across operators during execution.
diff --git a/benchmarks/src/util/memory_pool.rs
b/datafusion/execution/src/memory_pool/peak_recording.rs
similarity index 94%
rename from benchmarks/src/util/memory_pool.rs
rename to datafusion/execution/src/memory_pool/peak_recording.rs
index a3606ca0a7..b407cc0eaf 100644
--- a/benchmarks/src/util/memory_pool.rs
+++ b/datafusion/execution/src/memory_pool/peak_recording.rs
@@ -27,7 +27,7 @@
//! itself is never recorded — [`MemoryPool::reserved`] is a live value that
has
//! usually fallen back to zero by the time a query finishes. This module
records
//! the high-water mark so benchmarks can emit it alongside the peak RSS that
-//! [`print_memory_stats`] already prints, making the gap between the two
+//! `print_memory_stats` already prints, making the gap between the two
//! measurable.
//!
//! This is measurement only: nothing here enforces a relationship between the
@@ -38,8 +38,6 @@
//! through `ArrowMemoryPool` are included, because that adapter grows a
//! DataFusion reservation against the pool it wraps; nothing claims buffers
//! today, but the peak picks it up when something does.
-//!
-//! [`print_memory_stats`]: super::print_memory_stats
use std::{
fmt::{Debug, Display, Formatter},
@@ -49,9 +47,7 @@ use std::{
},
};
-use datafusion::execution::memory_pool::{
- MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation,
-};
+use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation};
use datafusion_common::Result;
/// Wraps a [`MemoryPool`], recording the high-water mark of
@@ -71,8 +67,7 @@ use datafusion_common::Result;
///
/// ```
/// # use std::sync::Arc;
-/// # use datafusion::execution::memory_pool::{GreedyMemoryPool,
MemoryConsumer, MemoryPool};
-/// # use datafusion_benchmarks::util::PeakRecordingPool;
+/// # use datafusion_execution::memory_pool::{GreedyMemoryPool,
MemoryConsumer, MemoryPool, PeakRecordingPool};
/// let recording =
Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024))));
/// let pool: Arc<dyn MemoryPool> = Arc::clone(&recording) as _;
///
@@ -116,10 +111,8 @@ impl PeakRecordingPool {
/// The recorder installed as `pool`, if there is one.
///
/// Returns `None` whenever a benchmark runs without a memory limit, since
- /// [`CommonOpt::runtime_env_builder`] only installs the wrapper alongside
a
+ /// `CommonOpt::runtime_env_builder` only installs the wrapper alongside a
/// pool it has a limit for.
- ///
- /// [`CommonOpt::runtime_env_builder`]:
super::CommonOpt::runtime_env_builder
pub fn from_pool(pool: &dyn MemoryPool) -> Option<&Self> {
pool.downcast_ref::<Self>()
}
@@ -140,12 +133,10 @@ impl PeakRecordingPool {
/// Reset the value returned by [`Self::peak_reserved`] to what is reserved
/// right now, so the next reading covers only what follows.
///
- /// [`BenchmarkRun::start_new_case`] calls this, giving each benchmark
query
+ /// `BenchmarkRun::start_new_case` calls this, giving each benchmark query
/// its own reading. Anything still held when a query starts — data the
/// benchmark loaded up front, say — stays in the reading, since the query
/// runs with those bytes reserved.
- ///
- /// [`BenchmarkRun::start_new_case`]: super::BenchmarkRun::start_new_case
pub fn reset_peak(&self) {
self.peak
.store(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed);
@@ -227,7 +218,7 @@ impl MemoryPool for PeakRecordingPool {
#[cfg(test)]
mod tests {
- use datafusion::execution::memory_pool::GreedyMemoryPool;
+ use crate::memory_pool::GreedyMemoryPool;
use super::*;
@@ -358,10 +349,15 @@ mod tests {
/// bytes show up in this peak without further changes — as long as the
/// adapter is built from the `RuntimeEnv`'s pool, which is the wrapped
one.
/// This test pins that.
+ ///
+ /// Only compiled with `--features arrow_buffer_pool`, since that's what
+ /// gates `crate::memory_pool::arrow` and `arrow_buffer::MemoryPool` in the
+ /// first place; not part of this crate's default feature set.
+ #[cfg(feature = "arrow_buffer_pool")]
#[test]
fn records_reservations_arriving_through_the_arrow_adapter() {
+ use crate::memory_pool::arrow::ArrowMemoryPool;
use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait;
- use datafusion_execution::memory_pool::arrow::ArrowMemoryPool;
let (recording, pool) = pool(4096);
diff --git a/datafusion/physical-plan/src/sorts/merge.rs
b/datafusion/physical-plan/src/sorts/merge.rs
index 310416c22d..6476490387 100644
--- a/datafusion/physical-plan/src/sorts/merge.rs
+++ b/datafusion/physical-plan/src/sorts/merge.rs
@@ -89,29 +89,6 @@ pub(crate) struct SortPreservingMergeStream<C: CursorValues>
{
/// Cursors for each input partition. `None` means the input is exhausted
cursors: Vec<Option<Cursor<C>>>,
- /// Configuration parameter to enable round-robin selection of tied
winners of loser tree.
- ///
- /// This option controls the tie-breaker strategy and attempts to avoid the
- /// issue of unbalanced polling between partitions
- ///
- /// If `true`, when multiple partitions have the same value, the partition
- /// that has the fewest poll counts is selected. This strategy ensures that
- /// multiple partitions with the same value are chosen equally,
distributing
- /// the polling load in a round-robin fashion. This approach balances the
- /// workload more effectively across partitions and avoids excessive buffer
- /// growth.
- ///
- /// if `false`, partitions with smaller indices are consistently chosen as
- /// the winners, which can lead to an uneven distribution of polling and
potentially
- /// causing upstream operator buffers for the other partitions to grow
- /// excessively, as they continued receiving data without consuming it.
- ///
- /// For example, an upstream operator like `RepartitionExec` execution
would
- /// keep sending data to certain partitions, but those partitions wouldn't
- /// consume the data if they weren't selected as winners. This resulted in
- /// inefficient buffer usage.
- enable_round_robin_tie_breaker: bool,
-
/// Flag indicating whether we are in the mode of round-robin
/// tie breaker for the loser tree winners.
round_robin_tie_breaker_mode: bool,
@@ -126,8 +103,9 @@ pub(crate) struct SortPreservingMergeStream<C:
CursorValues> {
/// Current reset count
current_reset_epoch: usize,
- /// Stores the previous value of each partitions for tracking the poll
counts on the same value.
- prev_cursors: Vec<Option<Cursor<C>>>,
+ /// Stores the previous value of each partitions for tracking the poll
counts on the same value
+ /// Used if and only if round robin tie breaker is enabled, otherwise None
+ prev_cursors: Option<Vec<Option<Cursor<C>>>>,
/// Optional number of rows to fetch
fetch: Option<usize>,
@@ -156,7 +134,11 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
streams,
metrics,
cursors: (0..stream_count).map(|_| None).collect(),
- prev_cursors: (0..stream_count).map(|_| None).collect(),
+ prev_cursors: if enable_round_robin_tie_breaker {
+ Some((0..stream_count).map(|_| None).collect())
+ } else {
+ None
+ },
round_robin_tie_breaker_mode: false,
num_of_polled_with_same_value: vec![0; stream_count],
current_reset_epoch: 0,
@@ -165,7 +147,6 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
batch_size,
fetch,
produced: 0,
- enable_round_robin_tie_breaker,
}
}
@@ -396,7 +377,13 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
if let Some(c) = cursor.as_mut() {
// Compare with the last row in the previous batch
- let prev_cursor = &self.prev_cursors[partition_idx];
+ let prev_cursor = self
+ .prev_cursors
+ .as_ref()
+ .map(|v| &v[partition_idx])
+ .expect(
+ "prev_cursor should be set when round robin tie breaker is
enabled",
+ );
if c.is_eq_to_prev_one(prev_cursor.as_ref()) {
self.num_of_polled_with_same_value[partition_idx] += 1;
} else {
@@ -405,6 +392,31 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
}
}
+ /// Whether round-robin selection of tied winners of loser tree is enabled.
+ ///
+ /// This option controls the tie-breaker strategy and attempts to avoid the
+ /// issue of unbalanced polling between partitions
+ ///
+ /// If `true`, when multiple partitions have the same value, the partition
+ /// that has the fewest poll counts is selected. This strategy ensures that
+ /// multiple partitions with the same value are chosen equally,
distributing
+ /// the polling load in a round-robin fashion. This approach balances the
+ /// workload more effectively across partitions and avoids excessive buffer
+ /// growth.
+ ///
+ /// if `false`, partitions with smaller indices are consistently chosen as
+ /// the winners, which can lead to an uneven distribution of polling and
potentially
+ /// causing upstream operator buffers for the other partitions to grow
+ /// excessively, as they continued receiving data without consuming it.
+ ///
+ /// For example, an upstream operator like `RepartitionExec` execution
would
+ /// keep sending data to certain partitions, but those partitions wouldn't
+ /// consume the data if they weren't selected as winners. This resulted in
+ /// inefficient buffer usage.
+ fn round_robin_tie_breaker_enabled(&self) -> bool {
+ self.prev_cursors.is_some()
+ }
+
fn fetch_reached(&mut self) -> bool {
self.fetch
.map(|fetch| self.produced + self.in_progress.len() >= fetch)
@@ -421,7 +433,10 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
let finished = cursor.is_finished();
if finished {
// Take the current cursor, leaving `None` in its place
- self.prev_cursors[stream_idx] =
self.cursors[stream_idx].take();
+ let taken = self.cursors[stream_idx].take();
+ if let Some(prev_cursors) = &mut self.prev_cursors {
+ prev_cursors[stream_idx] = taken;
+ }
}
return finished;
}
@@ -588,7 +603,7 @@ impl<C: CursorValues> SortPreservingMergeStream<C> {
if cmp_node == 1 {
let challenger = self.loser_tree[1];
// If round-robin tie-breaker is enabled and we're at the final
comparison (cmp_node == 1)
- if self.enable_round_robin_tie_breaker {
+ if self.round_robin_tie_breaker_enabled() {
match (&self.cursors[winner], &self.cursors[challenger]) {
(Some(ac), Some(bc)) => match ac.cmp(bc) {
std::cmp::Ordering::Equal => {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
