You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/16 00:33:21 UTC
(datafusion) branch gh-readonly-queue/main/pr-25250-8bd6629db6891640ad8c2993183131de627877e2 created (now 767d1422d9)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25250-8bd6629db6891640ad8c2993183131de627877e2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 767d1422d9 fix: register the sort-merge join stream as a spillable memory consumer (#25250)
This branch includes the following new commits:
new 767d1422d9 fix: register the sort-merge join stream as a spillable memory consumer (#25250)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: register the sort-merge join stream as a spillable memory consumer (#25250)
Posted by gi...@apache.org.
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-25250-8bd6629db6891640ad8c2993183131de627877e2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 767d1422d9a6459e6251c8b6cc0515abf07d08a2
Author: Jay Zhan <ja...@gmail.com>
AuthorDate: Wed Sep 16 00:32:12 2026 +0000
fix: register the sort-merge join stream as a spillable memory consumer (#25250)
## Which issue does this PR close?
- No existing issue. Found while preparing the hash join sort-merge
fallback, which runs this stream next to two external sorts inside one
partition.
## Rationale for this change
`SMJStream` spills: when `try_grow` for a buffered batch fails it writes
the buffered side to disk through its `SpillManager` and carries on. But
it registers its memory consumer with the default `can_spill = false`,
so the memory pool is told the opposite.
Only `FairSpillPool` reads that flag, and it uses it to run two
different budgets:
- A consumer that **can spill** is limited to an even share, `(pool_size
- unspillable) / num_spillable`, and is expected to spill when it
reaches it.
- A consumer that **cannot spill** takes free memory first come, first
served, and everything it holds is subtracted from the pool before the
spillable shares are computed.
Registering the join as unspillable therefore has three effects, all
wrong for an operator that spills:
1. It is not counted in `num_spillable`, so the sorts feeding it split
the pool as if the join needed nothing.
2. Every buffered batch it holds shrinks the sorts' shares, since it is
booked as unspillable memory.
3. It is never asked to spill until the whole pool is allocated.
So the operator that could spill cheaply squeezes the operators beside
it. As the join buffers a large key group, the sorts' shares fall toward
zero, they spill more than needed, and once a share drops below what a
sort must hold for its merge (`sort_spill_reservation_bytes`) the sort
fails outright. A query that would have finished by spilling the join
instead fails in the sort.
With the flag set, the join takes an even share, its buffered memory is
booked as spillable, and it spills when it exceeds that share, which is
the contract the pool's documentation describes for spillable operators.
Blast radius: `GreedyMemoryPool`, the default whenever a memory limit is
set, ignores the flag, so nothing changes there. `FairSpillPool` users
are affected, which includes `datafusion-cli --mem-pool-type fair`.
## What changes are included in this PR?
One line in `SortMergeJoinExec::execute`: `.with_can_spill(true)` on the
stream's `MemoryConsumer`, with a comment saying why. The bitwise stream
used by semi, anti and mark joins registers separately and is not
touched here; it may have the same omission and I have not checked.
## What is the testing strategy for this PR?
`stream_registers_as_a_spillable_consumer` in `sort_merge_join/tests.rs`
runs a small join against a recording pool that notes each consumer's
name and `can_spill` flag at registration, then asserts the stream's
entry is spillable. With the fix reverted, the test fails on "the
sort-merge join stream must register as able to spill".
The test deliberately does not run under `FairSpillPool`: what a stream
may hold there depends on which other consumers are alive at each
allocation, so an end-to-end assertion would depend on scheduling.
Reading the flag at registration is the deterministic form of the same
claim.
The existing sort-merge join suite passes unchanged (234 tests), and
clippy is clean with all targets and features.
## Are there any user-facing changes?
No API change. Under `FairSpillPool`, a sort-merge join now spills at
its fair share instead of consuming free memory first, so the sorts
feeding it keep their budget. No change under the default
`GreedyMemoryPool`.
---
.../src/joins/sort_merge_join/exec.rs | 4 +
.../src/joins/sort_merge_join/tests.rs | 214 ++++++++++++++++++++-
2 files changed, 212 insertions(+), 6 deletions(-)
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
index 911eca0a97..b0433250c0 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs
@@ -606,7 +606,11 @@ impl ExecutionPlan for SortMergeJoinExec {
let buffered = buffered.execute(partition, Arc::clone(&context))?;
let batch_size = context.session_config().batch_size();
+ // The stream spills its buffered batches when it cannot grow, so a
+ // pool that budgets spillable and unspillable consumers differently
+ // (`FairSpillPool`) has to know it can.
let reservation = MemoryConsumer::new(format!("SMJStream[{partition}]"))
+ .with_can_spill(true)
.register(context.memory_pool());
let spill_manager = SpillManager::new(
context.runtime_env(),
diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
index 4d817df738..ea9299abdd 100644
--- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
+++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs
@@ -34,6 +34,7 @@ use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn};
use crate::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec};
use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
use crate::projection::{ProjectionExec, ProjectionExpr};
+use crate::sorts::sort::SortExec;
use crate::spill::spill_manager::SpillManager;
use crate::test::TestMemoryExec;
use crate::test::exec::BarrierExec;
@@ -55,24 +56,28 @@ use bytes::Bytes;
use datafusion_common::JoinType::*;
use datafusion_common::instant::Instant;
use datafusion_common::{
- JoinSide, internal_err,
- test_util::{batches_to_sort_string, batches_to_string},
+ DataFusionError, JoinType, NullEquality, Result, ScalarValue, assert_batches_eq,
+ assert_contains,
};
use datafusion_common::{
- JoinType, NullEquality, Result, ScalarValue, assert_batches_eq, assert_contains,
+ JoinSide, internal_err,
+ test_util::{batches_to_sort_string, batches_to_string},
};
use datafusion_common_runtime::JoinSet;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::disk_manager::{
DiskManager, DiskManagerBuilder, DiskManagerMode,
};
-use datafusion_execution::memory_pool::MemoryConsumer;
+use datafusion_execution::memory_pool::{
+ FairSpillPool, MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool,
+};
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::Operator;
use datafusion_physical_expr::expressions::BinaryExpr;
use datafusion_physical_expr::expressions::Literal;
+use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
use datafusion_physical_expr_common::physical_expr::PhysicalExprRef;
use futures::{Stream, StreamExt};
use insta::assert_snapshot;
@@ -2465,6 +2470,203 @@ async fn overallocation_multi_batch_no_spill() -> Result<()> {
Ok(())
}
+/// The stream spills its buffered side when it cannot grow, so it has to be
+/// registered as a consumer that can spill. A `FairSpillPool` otherwise treats
+/// it as unspillable: it is left out of the fair share the spillable consumers
+/// split, and may take everything they have not yet claimed, starving the
+/// sorts the join usually runs on top of.
+#[tokio::test]
+async fn stream_registers_as_a_spillable_consumer() -> Result<()> {
+ /// Records how each consumer registered, and otherwise never limits anything.
+ #[derive(Debug, Default)]
+ struct RecordingPool {
+ inner: UnboundedMemoryPool,
+ registered: std::sync::Mutex<Vec<(String, bool)>>,
+ }
+ impl std::fmt::Display for RecordingPool {
+ fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+ write!(f, "RecordingPool")
+ }
+ }
+ impl MemoryPool for RecordingPool {
+ fn name(&self) -> &str {
+ "RecordingPool"
+ }
+ fn register(&self, consumer: &MemoryConsumer) {
+ self.registered
+ .lock()
+ .unwrap()
+ .push((consumer.name().to_string(), consumer.can_spill()));
+ }
+ fn grow(&self, reservation: &MemoryReservation, additional: usize) {
+ self.inner.grow(reservation, additional)
+ }
+ fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
+ self.inner.shrink(reservation, shrink)
+ }
+ fn try_grow(
+ &self,
+ reservation: &MemoryReservation,
+ additional: usize,
+ ) -> Result<()> {
+ self.inner.try_grow(reservation, additional)
+ }
+ fn reserved(&self) -> usize {
+ self.inner.reserved()
+ }
+ }
+
+ let pool = Arc::new(RecordingPool::default());
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::clone(&pool) as Arc<dyn MemoryPool>)
+ .build_arc()?;
+ let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime));
+
+ let left = build_table(
+ ("a1", &vec![1, 2]),
+ ("b1", &vec![1, 2]),
+ ("c1", &vec![7, 8]),
+ );
+ let right = build_table(
+ ("a2", &vec![1, 2]),
+ ("b2", &vec![1, 2]),
+ ("c2", &vec![9, 10]),
+ );
+ let on = vec![(
+ Arc::new(Column::new_with_schema("b1", &left.schema())?) as _,
+ Arc::new(Column::new_with_schema("b2", &right.schema())?) as _,
+ )];
+ let join = join(left, right, on, Inner)?;
+ common::collect(join.execute(0, task_ctx)?).await?;
+
+ let registered = pool.registered.lock().unwrap();
+ let (_, can_spill) = registered
+ .iter()
+ .find(|(name, _)| name == "SMJStream[0]")
+ .expect("the stream registers a reservation under its own name");
+ assert!(
+ can_spill,
+ "the sort-merge join stream must register as able to spill"
+ );
+ Ok(())
+}
+
+/// Runs the join under a `FairSpillPool` too small for its buffered side, on
+/// top of a sort that competes with it for that pool.
+///
+/// The buffered side is a single equal-key group larger than the pool, so the
+/// join spills part of it and holds the rest while it drains the streamed
+/// side. The streamed side is an in-memory sort of several batches whose merge
+/// asks the pool for every chunk it pulls in. Registered as unspillable, the
+/// join would take every byte the sort has not claimed before it starts
+/// spilling, and that merge would fail to allocate. As a spillable consumer
+/// it is held to its fair share and both finish.
+#[tokio::test]
+async fn fair_spill_pool_leaves_room_for_the_streamed_sort() -> Result<()> {
+ // Measured window: the join passes as a spillable consumer from 240 KB
+ // up, and still starves the sort as an unspillable one up to 704 KB.
+ const POOL_SIZE: usize = 384 * 1024;
+ const GROUP_KEY: i32 = i32::MAX;
+ const STREAMED_BATCHES: usize = 2;
+ const STREAMED_ROWS: usize = 512;
+ const BUFFERED_BATCHES: usize = 32;
+ const BUFFERED_ROWS: usize = 1024;
+
+ // Unsorted: every batch carries one row of the buffered group's key, which
+ // the sort moves to the end, so the join holds the whole buffered group
+ // while it pulls the streamed side through the sort's merge.
+ let streamed = build_table_from_batches(
+ (0..STREAMED_BATCHES)
+ .map(|batch| {
+ let base = (batch * STREAMED_ROWS) as i32;
+ let a: Vec<i32> = (0..STREAMED_ROWS as i32).map(|i| base + i).collect();
+ let b: Vec<i32> = (0..STREAMED_ROWS as i32)
+ .map(|i| {
+ if i == 0 {
+ GROUP_KEY
+ } else {
+ STREAMED_ROWS as i32 - i
+ }
+ })
+ .collect();
+ let c: Vec<i32> =
+ (0..STREAMED_ROWS as i32).map(|i| base + i + 1).collect();
+ build_table_i32(("a1", &a), ("b1", &b), ("c1", &c))
+ })
+ .collect(),
+ );
+ let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(
+ Column::new_with_schema("b1", &streamed.schema())?,
+ ))])
+ .unwrap();
+ let streamed = Arc::new(SortExec::new(ordering, streamed)) as Arc<dyn ExecutionPlan>;
+
+ // Already sorted: one group, larger than the pool.
+ let buffered = build_table_from_batches(
+ (0..BUFFERED_BATCHES)
+ .map(|batch| {
+ let base = (batch * BUFFERED_ROWS) as i32;
+ let a: Vec<i32> = (0..BUFFERED_ROWS as i32).map(|i| base + i).collect();
+ let b = vec![GROUP_KEY; BUFFERED_ROWS];
+ let c: Vec<i32> = (0..BUFFERED_ROWS as i32).map(|i| base - i).collect();
+ build_table_i32(("a2", &a), ("b2", &b), ("c2", &c))
+ })
+ .collect(),
+ );
+
+ let on = vec![(
+ Arc::new(Column::new_with_schema("b1", &streamed.schema())?) as _,
+ Arc::new(Column::new_with_schema("b2", &buffered.schema())?) as _,
+ )];
+ let session_config = SessionConfig::default()
+ .with_batch_size(128)
+ // Keep the sort in memory and merging chunk by chunk rather than
+ // sorting one concatenated batch in place.
+ .with_sort_in_place_threshold_bytes(0)
+ .with_sort_spill_reservation_bytes(0);
+
+ let run = |runtime| async {
+ let task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(session_config.clone())
+ .with_runtime(runtime),
+ );
+ let join = join(
+ Arc::clone(&streamed),
+ Arc::clone(&buffered),
+ on.clone(),
+ Inner,
+ )?;
+ let batches = common::collect(join.execute(0, task_ctx)?).await?;
+ let metrics = join.metrics().unwrap();
+ Ok::<_, DataFusionError>((batches, metrics.spill_count().unwrap()))
+ };
+
+ let fair = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::new(FairSpillPool::new(POOL_SIZE)))
+ .with_disk_manager_builder(
+ DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory),
+ )
+ .build_arc()?;
+ let (spilled, spill_count) = run(fair).await?;
+ assert!(
+ spill_count > 0,
+ "the join must have spilled its buffered side"
+ );
+ assert_eq!(
+ spilled.iter().map(|b| b.num_rows()).sum::<usize>(),
+ STREAMED_BATCHES * BUFFERED_BATCHES * BUFFERED_ROWS
+ );
+
+ let (unbounded, spill_count) = run(RuntimeEnvBuilder::new().build_arc()?).await?;
+ assert_eq!(spill_count, 0);
+ assert_eq!(
+ batches_to_sort_string(&spilled),
+ batches_to_sort_string(&unbounded)
+ );
+ Ok(())
+}
+
#[tokio::test]
async fn overallocation_single_batch_spill() -> Result<()> {
let left = build_table(
@@ -4477,7 +4679,7 @@ fn test_stream_resources(
inner_schema: SchemaRef,
metrics: &ExecutionPlanMetricsSet,
) -> (
- datafusion_execution::memory_pool::MemoryReservation,
+ MemoryReservation,
SpillManager,
Arc<datafusion_execution::runtime_env::RuntimeEnv>,
) {
@@ -6305,7 +6507,7 @@ impl SpillFile for PendingSpillFile {
tokio::fs::read(&path)
.await
.map(Bytes::from)
- .map_err(datafusion_common::DataFusionError::IoError)
+ .map_err(DataFusionError::IoError)
})
.flat_map(
|read_result| -> Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>> {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org