geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4279860629
##########
datafusion/physical-plan/src/sorts/sort/spill_tests.rs:
##########
@@ -522,6 +522,86 @@ async fn
test_chunked_string_view_spill_can_use_workspace() -> Result<()> {
check_chunked_string_view_workspace(true).await
}
+/// Two sort jobs share a tight memory budget. The first must return borrowed
+/// workspace when it hands off its first batch, so the second can prepare
+/// before the merge starts consuming output. During merging, looking ahead to
+/// the next batch must also return workspace for the merge's cursors.
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_sort_preparation_returns_first_batch_workspace() -> Result<()> {
+ let batch_size = 1024;
+ let values: ArrayRef = Arc::new(StringViewArray::from_iter_values(
+ (0..4096)
+ .rev()
+ .map(|i| format!("row-{i:08}-{}", "x".repeat(87))),
+ ));
+ let batch = RecordBatch::try_from_iter([("key", values)])?;
+ let ordering: LexOrdering =
[PhysicalSortExpr::new_default(Arc::new(Column::new(
+ "key", 0,
+ )))]
+ .into();
+ let input_bytes = get_reserved_bytes_for_record_batch(&batch)?;
+ let sorted_batches = sort_batch_chunked(&batch, &ordering, batch_size)?;
+ let loan = sorted_batches
+ .iter()
+ .map(get_record_batch_memory_size)
+ .sum::<usize>()
+ - input_bytes;
+ let first_bytes = get_record_batch_memory_size(&sorted_batches[0]);
+ assert!(loan > first_bytes);
Review Comment:
Each output chunk retains the shared string buffers, so summing the
per-batch estimates counts those buffers repeatedly. Added a comment explaining
this.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
