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 sorted `StringView` chunk retains the shared string buffers, and
summing the chunks' memory sizes counts those buffers repeatedly. I’ve added a
comment explaining this.
Related: #25820 fixes shared-buffer counting for sort inputs, but explicitly
leaves `StringView` accounting for follow-up work (I ran into a conflict with
that PR when rebasing over the latest `main`.)
--
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]
