You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "Weijun-H (via GitHub)" <gi...@apache.org> on 2026/09/21 11:42:19 UTC
[PR] Limit sort merge output batches by byte target [datafusion]
Weijun-H opened a new pull request, #25565:
URL: https://github.com/apache/datafusion/pull/25565
## What
Adds a sorter-local byte target for spill merge output batches. Before `BatchBuilder` materializes a merge output batch, it estimates the largest supported output prefix that fits the target, reserves that output-construction budget, and interleaves only that prefix.
The target is currently derived from the selected spill runs' recorded maximum batch memory. This keeps a merge pass from rebuilding a multi-row output batch that is substantially wider than the spilled input batches it is merging.
## Policy / Compatibility
This does not try to make sort memory a hard resident-memory cap. It is a pre-construction accounting and output-shaping guard for avoidably large multi-row sort output batches.
A multi-row reservation failure retries with a smaller prefix. A single row may exceed the byte target, but it must still obtain its construction reservation; if that final reservation fails, the query returns the existing memory error. That compatibility point is related to the policy discussion in #25183.
Unsupported schemas retain the current path. In this first cut, the byte target is applied only when every output column is a supported flat type; any unsupported column makes the whole batch fall back to current behavior. View arrays and containers such as Struct/List/Map are intentionally not estimated here. #23390 is relevant background for retained-buffer concerns around view arrays.
## Tests
The focused tests cover:
- selecting the largest supported prefix under the byte target;
- output-construction reservation against the real pool for spill merge;
- bitmap-capacity accounting for nullable and Boolean output;
- unsupported mixed schemas falling back as a whole batch;
- a wide Binary spill merge that must split output batches by byte target while preserving complete sorted output.
No performance improvement is claimed here.
Related to #25183.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Limit spill merge output batches by byte target [datafusion]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#discussion_r4093032619
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -272,13 +334,238 @@ impl BatchBuilder {
return Ok(None);
}
+ let Some((target_batch_bytes, mut estimated_bytes)) =
+ self.target_batch_bytes.and_then(|target| {
+ self.estimated_prefix_bytes(self.indices.len())
+ .map(|estimated| (target, estimated))
+ })
+ else {
+ let (rows_to_emit, columns) = retry_interleave(
+ self.indices.len(),
+ self.indices.len(),
+ |rows_to_emit| self.try_interleave_columns(&self.indices[..rows_to_emit]),
+ )?;
+
+ return Ok(Some(self.finish_record_batch(rows_to_emit, columns)?));
+ };
+
+ let initial_rows_to_emit = self.indices.len();
+ let mut rows_to_emit =
+ if initial_rows_to_emit <= 1 || estimated_bytes <= target_batch_bytes {
+ initial_rows_to_emit
+ } else {
+ self.largest_prefix_under_target(target_batch_bytes)
+ .unwrap_or(1)
+ };
+
+ if rows_to_emit != initial_rows_to_emit {
+ estimated_bytes = self
+ .estimated_prefix_bytes(rows_to_emit)
+ .expect("a smaller prefix has the same supported arrays");
+ }
+
+ loop {
+ match try_grow_reservation_to_at_least(
+ &mut self.output_construction_reservation,
+ estimated_bytes,
+ ) {
+ Ok(()) => break,
+ Err(_) if rows_to_emit > 1 => {
+ let failed_bytes = estimated_bytes;
+ rows_to_emit /= 2;
+ estimated_bytes = self
+ .estimated_prefix_bytes(rows_to_emit)
+ .expect("a smaller prefix has the same supported arrays");
+ warn!(
+ "Could not reserve {failed_bytes} bytes for sort output, retrying with {rows_to_emit} rows requiring {estimated_bytes} bytes"
+ );
+ }
+ Err(e) => return Err(e),
+ }
+ }
+
let (rows_to_emit, columns) =
- retry_interleave(self.indices.len(), self.indices.len(), |rows_to_emit| {
+ match retry_interleave(rows_to_emit, rows_to_emit, |rows_to_emit| {
self.try_interleave_columns(&self.indices[..rows_to_emit])
- })?;
+ }) {
+ Ok(value) => value,
+ Err(e) => {
+ self.output_construction_reservation.free();
+ return Err(e);
+ }
+ };
Ok(Some(self.finish_record_batch(rows_to_emit, columns)?))
}
+
+ fn largest_prefix_under_target(&self, target_batch_bytes: usize) -> Option<usize> {
+ if self.estimated_prefix_bytes(1)? > target_batch_bytes {
+ return Some(1);
+ }
+
+ let mut low = 1;
+ let mut high = self.indices.len();
+ while low < high {
+ let mid = low + (high - low).div_ceil(2);
+ if self.estimated_prefix_bytes(mid)? <= target_batch_bytes {
+ low = mid;
+ } else {
+ high = mid - 1;
+ }
+ }
+ Some(low)
+ }
+
+ fn estimated_prefix_bytes(&self, rows_to_emit: usize) -> Option<usize> {
+ let mut total = 0usize;
+ for column_idx in 0..self.schema.fields.len() {
+ total = total.checked_add(
+ self.column_prefix_memory_upper_bound(column_idx, rows_to_emit)?,
+ )?;
+ }
+ Some(total)
+ }
+
+ fn column_prefix_memory_upper_bound(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ let data_type = self.schema.field(column_idx).data_type();
+
+ match data_type {
+ DataType::Null => Some(0),
+ DataType::Boolean => bitmap_buffer_bytes(rows_to_emit)?
+ .checked_add(self.validity_buffer_bytes(column_idx, rows_to_emit)?),
+ DataType::Binary => self.byte_array_prefix_memory_upper_bound::<BinaryType>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::LargeBinary => self
+ .byte_array_prefix_memory_upper_bound::<LargeBinaryType>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::Utf8 => self.byte_array_prefix_memory_upper_bound::<Utf8Type>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::LargeUtf8 => self
+ .byte_array_prefix_memory_upper_bound::<LargeUtf8Type>(
+ column_idx,
+ rows_to_emit,
+ ),
+ DataType::FixedSizeBinary(width) => usize::try_from(*width)
+ .ok()
+ .and_then(|width| {
+ rows_to_emit
+ .checked_mul(width)
+ .and_then(aligned_buffer_bytes)
+ })
+ .and_then(|values| {
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add(values)
+ }),
+ _ => fixed_width(data_type).and_then(|width| {
+ rows_to_emit.checked_mul(width).and_then(|values| {
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add(values)
+ })
+ }),
+ }
+ }
+
+ /// Arrow allocates an output validity buffer when any input array for this
+ /// column contains nulls, even if the selected rows are all valid.
+ fn validity_buffer_bytes(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ if self
+ .batches
+ .iter()
+ .any(|(_, batch)| batch.column(column_idx).null_count() > 0)
+ {
+ bitmap_buffer_bytes(rows_to_emit)
+ } else {
+ Some(0)
+ }
+ }
+
+ fn byte_array_prefix_memory_upper_bound<T: ByteArrayType>(
+ &self,
+ column_idx: usize,
+ rows_to_emit: usize,
+ ) -> Option<usize> {
+ let mut values_len = 0usize;
+ for (batch_idx, row_idx) in &self.indices[..rows_to_emit] {
+ let array = self.batches[*batch_idx].1.column(column_idx);
+ let array = array.as_any().downcast_ref::<GenericByteArray<T>>()?;
+ values_len =
+ values_len.checked_add(array.value_length(*row_idx).as_usize())?;
+ }
+
+ self.validity_buffer_bytes(column_idx, rows_to_emit)?
+ .checked_add((rows_to_emit + 1).checked_mul(size_of::<T::Offset>())?)?
+ .checked_add(values_len)
Review Comment:
`values_len` is not an upper bound on the allocation made for an interleaved byte-array values buffer: Arrow's mutable byte buffer can retain alignment padding, just as the `FixedSizeBinary` branch accounts for above. This can reserve the unaligned estimate, allocate the batch, and only then fail (or emit a multi-row batch above the byte target) when `finish_record_batch` observes the larger capacity. Round the values-buffer estimate before using it for prefix selection and pre-allocation admission.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Limit spill merge output batches by byte target [datafusion]
Posted by "flarion-weijun (via GitHub)" <gi...@apache.org>.
flarion-weijun commented on PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#issuecomment-5863263682
Thanks for the detailed review! @andygrove Addressed all five points:
1. Restored fixed-width batch sizes, with matching reservations and a smaller target when memory is tight.
2. Matched spill targets to Arrow interleave allocation estimates.
3. Restored exact batch and metric assertions.
4. Replaced repeated variable-width prefix scans with incremental accumulation, preserving the fixed-width fast path.
5. Removed post-allocation reservation growth; estimation failures now fail before allocation.
Also addressed the nits and documented the stricter memory-accounting behavior. Ready for another look.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Limit spill merge output batches by byte target [datafusion]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#issuecomment-5812254300
## [Codecov](https://app.codecov.io/gh/apache/datafusion/pull/25565?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `90.20888%` with `75 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 82.50%. Comparing base ([`6574a8c`](https://app.codecov.io/gh/apache/datafusion/commit/6574a8c3b2c93ef9689a677c1cc1c0a4a844617f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`f9c4788`](https://app.codecov.io/gh/apache/datafusion/commit/f9c4788e72a1576719910f1235ae927f30eeab16?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 46 commits behind head on main.
| [Files with missing lines](https://app.codecov.io/gh/apache/datafusion/pull/25565?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [datafusion/physical-plan/src/sorts/builder.rs](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fsorts%2Fbuilder.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9zb3J0cy9idWlsZGVyLnJz) | 88.54% | [31 Missing and 14 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...usion/physical-plan/src/sorts/multi\_level\_merge.rs](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fsorts%2Fmulti_level_merge.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9zb3J0cy9tdWx0aV9sZXZlbF9tZXJnZS5ycw==) | 91.54% | [12 Missing and 11 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [datafusion/physical-plan/src/sorts/sort.rs](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fsorts%2Fsort.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9zb3J0cy9zb3J0LnJz) | 90.90% | [0 Missing and 5 partials :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...afusion/physical-plan/src/sorts/streaming\_merge.rs](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&filepath=datafusion%2Fphysical-plan%2Fsrc%2Fsorts%2Fstreaming_merge.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-ZGF0YWZ1c2lvbi9waHlzaWNhbC1wbGFuL3NyYy9zb3J0cy9zdHJlYW1pbmdfbWVyZ2UucnM=) | 93.33% | [2 Missing :warning: ](https://app.codecov.io/gh/apache/datafusion/pull/25565?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## main #25565 +/- ##
==========================================
+ Coverage 82.42% 82.50% +0.08%
==========================================
Files 1138 1140 +2
Lines 435501 438385 +2884
Branches 435501 438385 +2884
==========================================
+ Hits 358955 361691 +2736
- Misses 54841 54875 +34
- Partials 21705 21819 +114
```
</details>
[:umbrella: View full report in Codecov by Harness](https://app.codecov.io/gh/apache/datafusion/pull/25565?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache).
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Limit spill merge output batches by byte target [datafusion]
Posted by "Weijun-H (via GitHub)" <gi...@apache.org>.
Weijun-H commented on PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#issuecomment-6407683400
Thanks @andygrove ! Merged main and resolved the conflicts in c466a1780f88, tightened both assertions, and updated the accounting tests and snapshots for the separate output reservation. Memory limits, results, and rewrite-count assertions are unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Limit spill merge output batches by byte target [datafusion]
Posted by "Weijun-H (via GitHub)" <gi...@apache.org>.
Weijun-H commented on PR #25565:
URL: https://github.com/apache/datafusion/pull/25565#issuecomment-5863273880
Thanks for the detailed review! @andygrove Addressed all five points:
1. Restored fixed-width batch sizes, with matching reservations and a smaller target when memory is tight.
2. Matched spill targets to Arrow interleave allocation estimates.
3. Restored exact batch and metric assertions.
4. Replaced repeated variable-width prefix scans with incremental accumulation, preserving the fixed-width fast path.
5. Removed post-allocation reservation growth; estimation failures now fail before allocation.
Also addressed the nits and documented the stricter memory-accounting behavior. Ready for another look.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org