fix partition id

yjshen · Jan 4, 2022 · 0843c77 · 0843c77
1 parent 0f5e377
commit 0843c77
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/datafusion/src/execution/memory_management/allocation_strategist.rs b/datafusion/src/execution/memory_management/allocation_strategist.rs
@@ -17,7 +17,6 @@
 
 //! Execution Memory Pool that guarantees a memory allocation strategy
 
-use crate::execution::memory_management::MemoryConsumerId;
 use async_trait::async_trait;
 use hashbrown::HashMap;
 use log::{info, warn};
@@ -189,7 +188,7 @@ impl MemoryAllocationStrategist for FairStrategist {
             if to_grant < required && current_mem + to_grant < min_memory_per_partition {
                 info!(
                     "{:?} waiting for at least 1/2N of pool to be free",
-                    consumer
+                    partition_id
                 );
                 let _ = self.notify.notified().await;
             } else {
@@ -214,12 +213,12 @@ impl MemoryAllocationStrategist for FairStrategist {
         } else {
             let mut partition_usage = self.memory_usage.write().await;
             if granted_size > real_size {
-                *partition_usage.entry(consumer.partition_id).or_insert(0) -=
+                *partition_usage.entry(partition_id).or_insert(0) -=
                     granted_size - real_size;
             } else {
                 // TODO: this would have caused OOM already if size estimation ahead is much smaller than
                 // that of actual allocation
-                *partition_usage.entry(consumer.partition_id).or_insert(0) +=
+                *partition_usage.entry(partition_id).or_insert(0) +=
                     real_size - granted_size;
             }
         }

diff --git a/datafusion/src/execution/memory_management/mod.rs b/datafusion/src/execution/memory_management/mod.rs
@@ -116,7 +116,7 @@ impl MemoryManager {
         consumer: &MemoryConsumerId,
     ) {
         self.strategist
-            .update_usage(granted_size, real_size, consumer)
+            .update_usage(granted_size, real_size, consumer.partition_id)
             .await
     }