This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23606-1a189157dba4c0ebd838824f3056969aef91e29a in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
