DerGut commented on code in PR #15692:
URL: https://github.com/apache/datafusion/pull/15692#discussion_r2041251481
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1552,6 +1593,62 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_batch_reservation_error() -> Result<()> {
+ // Pick a memory limit and sort_spill_reservation that make the first
batch reservation fail.
+ // These values assume that the ExternalSorter will reserve 800 bytes
for the first batch.
+ let expected_batch_reservation = 800;
+ let merge_reservation: usize = 0; // Set to 0 for simplicity
+ let memory_limit: usize = expected_batch_reservation +
merge_reservation - 1; // Just short of what we need
+
+ let session_config =
+
SessionConfig::new().with_sort_spill_reservation_bytes(merge_reservation);
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_limit(memory_limit, 1.0)
+ .build_arc()?;
+ let task_ctx = Arc::new(
+ TaskContext::default()
+ .with_session_config(session_config)
+ .with_runtime(runtime),
+ );
+
+ let plan = test::scan_partitioned(1);
+
+ // Read the first record batch to assert that our memory limit and
sort_spill_reservation
+ // settings trigger the test scenario.
+ {
+ let mut stream = plan.execute(0, Arc::clone(&task_ctx))?;
+ let first_batch = stream.next().await.unwrap()?;
+ let batch_reservation =
get_reserved_byte_for_record_batch(&first_batch);
+
+ assert_eq!(batch_reservation, expected_batch_reservation);
+ assert!(memory_limit < (merge_reservation + batch_reservation));
+ }
+
+ let sort_exec = Arc::new(SortExec::new(
+ LexOrdering::new(vec![PhysicalSortExpr {
+ expr: col("i", &plan.schema())?,
+ options: SortOptions::default(),
+ }]),
+ plan,
+ ));
+
+ let result = collect(
+ Arc::clone(&sort_exec) as Arc<dyn ExecutionPlan>,
+ Arc::clone(&task_ctx),
+ )
+ .await;
+
+ let err = result.unwrap_err();
+ assert!(
+ matches!(err, DataFusionError::ResourcesExhausted(_)),
Review Comment:
Note, this is now failing because it gets a `DataFusionError::Context`
instead. I'm still pondering whether to unwrap the error here, or add something
like an `DataFusionError::is(&self, target: &DataFusionError) -> bool` to match
the nested error.
I don't think matching against `Context` instead would be interesting here.
I'd rather test the cause of the error, less so whether it has added context or
not.
--
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-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
