geoffreyclaude commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4279865134
##########
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);
+
+ for consume_output in [false, true] {
+ let workspace = if consume_output {
+ // Also allow the merge's cursors, provided each run looks ahead.
+ 2 * loan + first_bytes
+ } else {
+ 2 * loan - first_bytes
+ };
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::new(GreedyMemoryPool::new(
+ 2 * input_bytes + workspace,
+ )))
+ .build_arc()?;
+ let mut sorter = ExternalSorter::new(
+ 0,
+ batch.schema(),
+ ordering.clone(),
+ batch_size,
+ workspace,
+ 0,
+ SpillCompression::Uncompressed,
+ &ExecutionPlanMetricsSet::new(),
+ Arc::clone(&runtime),
+ )?;
+ sorter.insert_batch(batch.clone()).await?;
+ sorter.insert_batch(batch.clone()).await?;
+ sorter.merge_reservation.free();
+ let stream = sorter.in_mem_sort_stream(false, false).await?;
+ // Buffered producers would leave two tasks waiting for the merge.
+ let alive = tokio::runtime::Handle::current()
+ .metrics()
+ .num_alive_tasks();
+ assert!(alive <= 1, "{alive} sort tasks remain before merging");
+ assert_eq!(
+ sorter.merge_pool.borrow(usize::MAX).size(),
+ workspace - 2 * (loan - first_bytes)
Review Comment:
Should hopefully be a tiny bit clearer now!
--
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]
