Sort-tpch Q5 fails when memory-limit is intermediate while succeeds with smaller memory

[ding-young]

Describe the bug

Related to #14692

Thanks to #15700, some queries are now able to succeed under tighter memory limits. However, there are still failing cases. One strange situation is that some queries succeed under stricter memory limits but fail when given more generous limits.

For example, sort-tpch Q5 fails with 100M limit while succeeds with 30M limit.

Query 5 failed: Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
  ExternalSorterMerge[0]#1(can spill: false) consumed 99.9 MB,
  ExternalSorter[0]#0(can spill: true) consumed 0.0 B.
Error: Failed to allocate additional 160.0 KB for ExternalSorterMerge[0] with 39.1 MB already allocated for this reservation - 109.7 KB remain available for the total pool
Failed Queries: 5 

To Reproduce

• 100M, partitions 1 (fails)

cargo run --profile release-nonlto --bin dfbench -- sort-tpch --path benchmarks/data/tpch_sf1 --partitions 1 --memory-limit 100M --query 5

• 30M, partitions 1 (succeeds with spill)

cargo run --profile release-nonlto --bin dfbench -- sort-tpch --path benchmarks/data/tpch_sf1 --partitions 1 --memory-limit 30M --query 5

Expected behavior

sort-tpch Q5 succeeds with both 30M and 100M memory-limit

Additional context

sort-tpch Q6, which requires more memory than Q5, succeeds with the same configuration. Therefore, Q5 isn't a query that’s inherently too memory-intensive.

[ding-young]

take

[rluvaton]

Do we know where it fails?

is it in the MultiLevelMergeBuilder? we should also add some context in the error below when we can't reserve memory even for the minimum required spill files to merge:

datafusion/datafusion/physical-plan/src/sorts/multi_level_merge.rs

Line 381 in 0183244

return Err(err);

we know it's not in as we would see in the context:

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Lines 767 to 777 in 0183244

	/// Wraps the error with a context message suggesting settings to tweak.
	/// This is meant to be used with DataFusionError::ResourcesExhausted only.
	fn err_with_oom_context(e: DataFusionError) -> DataFusionError {
	match e {
	DataFusionError::ResourcesExhausted(_) => e.context(
	"Not enough memory to continue external sort. \
	Consider increasing the memory limit, or decreasing sort_spill_reservation_bytes"
	),
	// This is not an OOM error, so just return it as is.
	_ => e,
	}

but it can also be in:

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Lines 684 to 692 in 0183244

	StreamingMergeBuilder::new()
	.with_streams(streams)
	.with_schema(Arc::clone(&self.schema))
	.with_expressions(&self.expr.clone())
	.with_metrics(metrics)
	.with_batch_size(self.batch_size)
	.with_fetch(None)
	.with_reservation(self.merge_reservation.new_empty())
	.build()

which the SortPreservingMergeStream also allocate memory

[rluvaton]

Also, did this query failed before my change?

[rluvaton]

Also, note that my fuzz tests all set with_sort_spill_reservation_bytes to 1:

datafusion/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs

Line 59 in 0183244

.with_sort_spill_reservation_bytes(1),

and I think that we should not reserve memory for merge in:

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Lines 569 to 570 in 0183244

	// Reserve headroom for next sort/merge
	self.reserve_memory_for_merge()?;

when called here:

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Lines 341 to 346 in 0183244

	// Sort `in_mem_batches` and spill it first. If there are many
	// `in_mem_batches` and the memory limit is almost reached, merging
	// them with the spilled files at the same time might cause OOM.
	if !self.in_mem_batches.is_empty() {
	self.sort_and_spill_in_mem_batches().await?;
	}

[ding-young]

Also, did this query failed before my change?

Yes, I think the failure is not related to multi level merge.

It fails at

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Line 538 in 0183244

let batch = batch?;

I'm taking a closer look at polling SortPreservingMergeStream

[rluvaton]

what type of memory pool is it?

datafusion/datafusion/physical-plan/src/sorts/multi_level_merge.rs

Lines 355 to 356 in 0183244

	// For memory pools that are not shared this is good, for other this is not
	// and there should be some upper limit to memory reservation so we won't starve the system

[rluvaton]

Also, did this query failed before my change?

Yes, I think the failure is not related to multi level merge.

Great,

It fails at

datafusion/datafusion/physical-plan/src/sorts/sort.rs

Line 538 in 0183244

let batch = batch?;

I'm taking a closer look at polling SortPreservingMergeStream

can it be because of sliced data?

[rluvaton]

I would be happy to review the fix PR

[ding-young]

I would be happy to review the fix PR

I'll ping you again whenever I bring up with the fix 😄