feat: Expose used MemoryPool in ResourcesExhausted error messages

Author: erenavsarogullari

datafusion-cli/tests/cli_integration.rs

@@ -228,11 +228,11 @@ fn test_cli_top_memory_consumers<'a>(
         "Consumer(can spill: bool) consumed XB, peak XB",
     );
     settings.add_filter(
-        r"Error: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total pool",
+        r"Error: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total '.*?' pool",
         "Error: Failed to allocate ",
     );
     settings.add_filter(
-        r"Resources exhausted: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total pool",
+        r"Resources exhausted: Failed to allocate additional .*? for .*? with .*? already allocated for this reservation - .*? remain available for the total '.*?' pool",
         "Resources exhausted: Failed to allocate",
     );
 

datafusion-cli/tests/snapshots/cli_top_memory_consumers@top2.snap

@@ -16,7 +16,7 @@ exit_code: 1
 [CLI_VERSION]
 Error: Not enough memory to continue external sort. Consider increasing the memory limit config: 'datafusion.runtime.memory_limit', or decreasing the config: 'datafusion.execution.sort_spill_reservation_bytes'.
 caused by
-Resources exhausted: Additional allocation failed for ExternalSorter[0] with top memory consumers (across reservations) as:
+Resources exhausted: Additional allocation failed for ExternalSorter[0] with top memory consumers (across reservations) using 'greedy' pool as:
   Consumer(can spill: bool) consumed XB, peak XB,
   Consumer(can spill: bool) consumed XB, peak XB.
 Error: Failed to allocate 

datafusion-cli/tests/snapshots/cli_top_memory_consumers@top3_default.snap

@@ -14,7 +14,7 @@ exit_code: 1
 [CLI_VERSION]
 Error: Not enough memory to continue external sort. Consider increasing the memory limit config: 'datafusion.runtime.memory_limit', or decreasing the config: 'datafusion.execution.sort_spill_reservation_bytes'.
 caused by
-Resources exhausted: Additional allocation failed for ExternalSorter[0] with top memory consumers (across reservations) as:
+Resources exhausted: Additional allocation failed for ExternalSorter[0] with top memory consumers (across reservations) using 'greedy' pool as:
   Consumer(can spill: bool) consumed XB, peak XB,
   Consumer(can spill: bool) consumed XB, peak XB,
   Consumer(can spill: bool) consumed XB, peak XB.

datafusion-examples/examples/execution_monitoring/memory_pool_tracking.rs

@@ -113,13 +113,14 @@ async fn automatic_usage_example() -> Result<()> {
                 Error: Not enough memory to continue external sort. Consider increasing the memory limit config: 'datafusion.runtime.memory_limit',
                 or decreasing the config: 'datafusion.execution.sort_spill_reservation_bytes'.
                 caused by
-                    Resources exhausted: Additional allocation failed with top memory consumers (across reservations) as:
+                    Resources exhausted: Additional allocation failed with top memory consumers (across reservations) using 'greedy' pool as:
                     ExternalSorterMerge[3]#112(can spill: false) consumed 10.0 MB, peak 10.0 MB,
                     ExternalSorterMerge[10]#147(can spill: false) consumed 10.0 MB, peak 10.0 MB,
                     ExternalSorter[1]#93(can spill: true) consumed 69.0 KB, peak 69.0 KB,
                     ExternalSorter[13]#155(can spill: true) consumed 67.6 KB, peak 67.6 KB,
                     ExternalSorter[8]#140(can spill: true) consumed 67.2 KB, peak 67.2 KB.
-                Error: Failed to allocate additional 10.0 MB for ExternalSorterMerge[0] with 0.0 B already allocated for this reservation - 7.1 MB remain available for the total pool
+                Error: Failed to allocate additional 10.0 MB for ExternalSorterMerge[0] with 0.0 B already allocated
+                for this reservation - 7.1 MB remain available for the total 'greedy' pool
              */
         }
     }

datafusion/core/tests/memory_limit/mod.rs

@@ -87,7 +87,7 @@ async fn group_by_none() {
         .with_query("select median(request_bytes) from t")
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:\n  AggregateStream",
+            "using GreedyMemoryPool with top memory consumers (across reservations) using 'greedy' pool as:\n  AggregateStream",
         ])
         .with_memory_limit(2_000)
         .run()
@@ -99,7 +99,7 @@ async fn group_by_row_hash() {
     TestCase::new()
         .with_query("select count(*) from t GROUP BY response_bytes")
         .with_expected_errors(vec![
-            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n  GroupedHashAggregateStream"
+            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  GroupedHashAggregateStream"
         ])
         .with_memory_limit(2_000)
         .run()
@@ -112,7 +112,7 @@ async fn group_by_hash() {
         // group by dict column
         .with_query("select count(*) from t GROUP BY service, host, pod, container")
         .with_expected_errors(vec![
-            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n  GroupedHashAggregateStream"
+            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  GroupedHashAggregateStream"
         ])
         .with_memory_limit(1_000)
         .run()
@@ -126,7 +126,7 @@ async fn join_by_key_multiple_partitions() {
         .with_query("select t1.* from t t1 JOIN t t2 ON t1.service = t2.service")
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:\n  HashJoinInput",
+            "with top memory consumers (across reservations) using 'greedy' pool as:\n  HashJoinInput",
         ])
         .with_memory_limit(1_000)
         .with_config(config)
@@ -141,7 +141,7 @@ async fn join_by_key_single_partition() {
         .with_query("select t1.* from t t1 JOIN t t2 ON t1.service = t2.service")
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:\n  HashJoinInput",
+            "with top memory consumers (across reservations) using 'greedy' pool as:\n  HashJoinInput",
         ])
         .with_memory_limit(1_000)
         .with_config(config)
@@ -154,7 +154,7 @@ async fn join_by_expression() {
     TestCase::new()
         .with_query("select t1.* from t t1 JOIN t t2 ON t1.service != t2.service")
         .with_expected_errors(vec![
-           "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n  NestedLoopJoinLoad[0]",
+           "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  NestedLoopJoinLoad[0]",
         ])
         .with_memory_limit(1_000)
         .run()
@@ -167,7 +167,7 @@ async fn cross_join() {
         .with_query("select t1.*, t2.* from t t1 CROSS JOIN t t2")
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:\n  CrossJoinExec",
+            "with top memory consumers (across reservations) using 'greedy' pool as:\n  CrossJoinExec",
         ])
         .with_memory_limit(1_000)
         .run()
@@ -223,7 +223,7 @@ async fn symmetric_hash_join() {
             "select t1.* from t t1 JOIN t t2 ON t1.pod = t2.pod AND t1.time = t2.time",
         )
         .with_expected_errors(vec![
-            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n  SymmetricHashJoinStream",
+            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  SymmetricHashJoinStream",
         ])
         .with_memory_limit(1_000)
         .with_scenario(Scenario::AccessLogStreaming)
@@ -241,7 +241,7 @@ async fn sort_preserving_merge() {
     // so only a merge is needed
         .with_query("select * from t ORDER BY a ASC NULLS LAST, b ASC NULLS LAST LIMIT 10")
         .with_expected_errors(vec![
-            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n  SortPreservingMergeExec",
+            "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  SortPreservingMergeExec",
         ])
         // provide insufficient memory to merge
         .with_memory_limit(partition_size / 2)
@@ -321,7 +321,7 @@ async fn sort_spill_reservation() {
     test.clone()
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:",
+            "with top memory consumers (across reservations) using 'greedy' pool as:",
             "B for ExternalSorterMerge",
         ])
         .with_config(config)
@@ -352,7 +352,7 @@ async fn oom_recursive_cte() {
         )
         .with_expected_errors(vec![
             "Resources exhausted: Additional allocation failed",
-            "with top memory consumers (across reservations) as:\n  RecursiveQuery",
+            "with top memory consumers (across reservations) using 'greedy' pool as:\n  RecursiveQuery",
         ])
         .with_memory_limit(2_000)
         .run()
@@ -404,7 +404,7 @@ async fn oom_with_tracked_consumer_pool() {
         .with_expected_errors(vec![
             "Failed to allocate additional",
             "for ParquetSink(ArrowColumnWriter)",
-            "Additional allocation failed", "with top memory consumers (across reservations) as:\n  ParquetSink(ArrowColumnWriter)"
+            "Additional allocation failed", "with top memory consumers (across reservations) using 'greedy' pool as:\n  ParquetSink(ArrowColumnWriter)"
         ])
         .with_memory_pool(Arc::new(
             TrackConsumersPool::new(

datafusion/execution/src/memory_pool/mod.rs

@@ -178,6 +178,9 @@ pub use pool::*;
 /// * [`TrackConsumersPool`]: Wraps another [`MemoryPool`] and tracks consumers,
 ///   providing better error messages on the largest memory users.
 pub trait MemoryPool: Send + Sync + std::fmt::Debug {
+    /// Return pool name
+    fn name(&self) -> &str;
+
     /// Registers a new [`MemoryConsumer`]
     ///
     /// Note: Subsequent calls to [`Self::grow`] must be made to reserve memory