jayshrivastava commented on code in PR #26242:
URL: https://github.com/apache/datafusion/pull/26242#discussion_r4274694477
##########
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:
Can you comment why it's bigger? I guess it's because some `4096` size
buffers are re-used but not filled?
##########
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:
Adding comments around the math would make it easier to reason about
--
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]
