This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new cfc4cbb97 chore: Refactor Memory Pools (#1662)
cfc4cbb97 is described below
commit cfc4cbb97ee40298a818e390fad33e9bded7ded0
Author: Emily Matheys <[email protected]>
AuthorDate: Tue Apr 22 21:44:41 2025 +0300
chore: Refactor Memory Pools (#1662)
* chore: Refactor Memory Pools
* merge mod lines
* add licenses to new rust files
* fix typo
---
native/core/src/execution/jni_api.rs | 210 ++-------------------
native/core/src/execution/memory_pools/config.rs | 105 +++++++++++
.../fair_pool.rs} | 0
native/core/src/execution/memory_pools/mod.rs | 111 +++++++++++
.../core/src/execution/memory_pools/task_shared.rs | 60 ++++++
.../unified_pool.rs} | 0
native/core/src/execution/mod.rs | 5 +-
7 files changed, 289 insertions(+), 202 deletions(-)
diff --git a/native/core/src/execution/jni_api.rs
b/native/core/src/execution/jni_api.rs
index e339d42f7..2dff70ee2 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -17,12 +17,10 @@
//! Define JNI APIs which can be called from Java/Scala.
-use super::{serde, utils::SparkArrowConvert, CometMemoryPool};
+use super::{serde, utils::SparkArrowConvert};
use arrow::array::RecordBatch;
use arrow::datatypes::DataType as ArrowDataType;
-use datafusion::execution::memory_pool::{
- FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
UnboundedMemoryPool,
-};
+use datafusion::execution::memory_pool::MemoryPool;
use datafusion::{
execution::{disk_manager::DiskManagerConfig, runtime_env::RuntimeEnv},
physical_plan::{display::DisplayableExecutionPlan,
SendableRecordBatchStream},
@@ -40,7 +38,7 @@ use jni::{
};
use std::path::PathBuf;
use std::time::{Duration, Instant};
-use std::{collections::HashMap, sync::Arc, task::Poll};
+use std::{sync::Arc, task::Poll};
use crate::{
errors::{try_unwrap_or_throw, CometError, CometResult},
@@ -60,16 +58,16 @@ use jni::{
objects::GlobalRef,
sys::{jboolean, jdouble, jintArray, jobjectArray, jstring},
};
-use std::num::NonZeroUsize;
-use std::sync::Mutex;
use tokio::runtime::Runtime;
-use crate::execution::fair_memory_pool::CometFairMemoryPool;
+use crate::execution::memory_pools::{
+ create_memory_pool, handle_task_shared_pool_release,
parse_memory_pool_config, MemoryPoolConfig,
+};
use crate::execution::operators::ScanExec;
use crate::execution::shuffle::{read_ipc_compressed, CompressionCodec};
use crate::execution::spark_plan::SparkPlan;
use log::info;
-use once_cell::sync::{Lazy, OnceCell};
+use once_cell::sync::Lazy;
static TOKIO_RUNTIME: Lazy<Runtime> = Lazy::new(|| {
let mut builder = tokio::runtime::Builder::new_multi_thread();
@@ -130,51 +128,6 @@ struct ExecutionContext {
pub memory_pool_config: MemoryPoolConfig,
}
-#[derive(PartialEq, Eq)]
-enum MemoryPoolType {
- Unified,
- FairUnified,
- Greedy,
- FairSpill,
- GreedyTaskShared,
- FairSpillTaskShared,
- GreedyGlobal,
- FairSpillGlobal,
- Unbounded,
-}
-
-struct MemoryPoolConfig {
- pool_type: MemoryPoolType,
- pool_size: usize,
-}
-
-impl MemoryPoolConfig {
- fn new(pool_type: MemoryPoolType, pool_size: usize) -> Self {
- Self {
- pool_type,
- pool_size,
- }
- }
-}
-
-/// The per-task memory pools keyed by task attempt id.
-static TASK_SHARED_MEMORY_POOLS: Lazy<Mutex<HashMap<i64, PerTaskMemoryPool>>> =
- Lazy::new(|| Mutex::new(HashMap::new()));
-
-struct PerTaskMemoryPool {
- memory_pool: Arc<dyn MemoryPool>,
- num_plans: usize,
-}
-
-impl PerTaskMemoryPool {
- fn new(memory_pool: Arc<dyn MemoryPool>) -> Self {
- Self {
- memory_pool,
- num_plans: 0,
- }
- }
-}
-
/// Accept serialized query plan and return the address of the native query
plan.
/// # Safety
/// This function is inherently unsafe since it deals with raw pointers passed
from JNI.
@@ -321,134 +274,6 @@ fn prepare_datafusion_session_context(
Ok(session_ctx)
}
-fn parse_memory_pool_config(
- off_heap_mode: bool,
- memory_pool_type: String,
- memory_limit: i64,
- memory_limit_per_task: i64,
-) -> CometResult<MemoryPoolConfig> {
- let pool_size = memory_limit as usize;
- let memory_pool_config = if off_heap_mode {
- match memory_pool_type.as_str() {
- "fair_unified" =>
MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size),
- "default" | "unified" => {
- // the `unified` memory pool interacts with Spark's memory
pool to allocate
- // memory therefore does not need a size to be explicitly set.
The pool size
- // shared with Spark is set by `spark.memory.offHeap.size`.
- MemoryPoolConfig::new(MemoryPoolType::Unified, 0)
- }
- _ => {
- return Err(CometError::Config(format!(
- "Unsupported memory pool type for off-heap mode: {}",
- memory_pool_type
- )))
- }
- }
- } else {
- // Use the memory pool from DF
- let pool_size_per_task = memory_limit_per_task as usize;
- match memory_pool_type.as_str() {
- "fair_spill_task_shared" => {
- MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared,
pool_size_per_task)
- }
- "default" | "greedy_task_shared" => {
- MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared,
pool_size_per_task)
- }
- "fair_spill_global" => {
- MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal,
pool_size)
- }
- "greedy_global" =>
MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size),
- "fair_spill" => MemoryPoolConfig::new(MemoryPoolType::FairSpill,
pool_size_per_task),
- "greedy" => MemoryPoolConfig::new(MemoryPoolType::Greedy,
pool_size_per_task),
- "unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0),
- _ => {
- return Err(CometError::Config(format!(
- "Unsupported memory pool type for on-heap mode: {}",
- memory_pool_type
- )))
- }
- }
- };
- Ok(memory_pool_config)
-}
-
-fn create_memory_pool(
- memory_pool_config: &MemoryPoolConfig,
- comet_task_memory_manager: Arc<GlobalRef>,
- task_attempt_id: i64,
-) -> Arc<dyn MemoryPool> {
- const NUM_TRACKED_CONSUMERS: usize = 10;
- match memory_pool_config.pool_type {
- MemoryPoolType::Unified => {
- // Set Comet memory pool for native
- let memory_pool = CometMemoryPool::new(comet_task_memory_manager);
- Arc::new(TrackConsumersPool::new(
- memory_pool,
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- }
- MemoryPoolType::FairUnified => {
- // Set Comet fair memory pool for native
- let memory_pool =
- CometFairMemoryPool::new(comet_task_memory_manager,
memory_pool_config.pool_size);
- Arc::new(TrackConsumersPool::new(
- memory_pool,
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- }
- MemoryPoolType::Greedy => Arc::new(TrackConsumersPool::new(
- GreedyMemoryPool::new(memory_pool_config.pool_size),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- )),
- MemoryPoolType::FairSpill => Arc::new(TrackConsumersPool::new(
- FairSpillPool::new(memory_pool_config.pool_size),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- )),
- MemoryPoolType::GreedyGlobal => {
- static GLOBAL_MEMORY_POOL_GREEDY: OnceCell<Arc<dyn MemoryPool>> =
OnceCell::new();
- let memory_pool = GLOBAL_MEMORY_POOL_GREEDY.get_or_init(|| {
- Arc::new(TrackConsumersPool::new(
- GreedyMemoryPool::new(memory_pool_config.pool_size),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- });
- Arc::clone(memory_pool)
- }
- MemoryPoolType::FairSpillGlobal => {
- static GLOBAL_MEMORY_POOL_FAIR: OnceCell<Arc<dyn MemoryPool>> =
OnceCell::new();
- let memory_pool = GLOBAL_MEMORY_POOL_FAIR.get_or_init(|| {
- Arc::new(TrackConsumersPool::new(
- FairSpillPool::new(memory_pool_config.pool_size),
- NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- });
- Arc::clone(memory_pool)
- }
- MemoryPoolType::GreedyTaskShared | MemoryPoolType::FairSpillTaskShared
=> {
- let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
- let per_task_memory_pool =
- memory_pool_map.entry(task_attempt_id).or_insert_with(|| {
- let pool: Arc<dyn MemoryPool> =
- if memory_pool_config.pool_type ==
MemoryPoolType::GreedyTaskShared {
- Arc::new(TrackConsumersPool::new(
-
GreedyMemoryPool::new(memory_pool_config.pool_size),
-
NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- } else {
- Arc::new(TrackConsumersPool::new(
-
FairSpillPool::new(memory_pool_config.pool_size),
-
NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
- ))
- };
- PerTaskMemoryPool::new(pool)
- });
- per_task_memory_pool.num_plans += 1;
- Arc::clone(&per_task_memory_pool.memory_pool)
- }
- MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()),
- }
-}
-
/// Prepares arrow arrays for output.
fn prepare_output(
env: &mut JNIEnv,
@@ -643,22 +468,11 @@ pub extern "system" fn
Java_org_apache_comet_Native_releasePlan(
// Update metrics
update_metrics(&mut env, execution_context)?;
- if execution_context.memory_pool_config.pool_type ==
MemoryPoolType::FairSpillTaskShared
- || execution_context.memory_pool_config.pool_type ==
MemoryPoolType::GreedyTaskShared
- {
- // Decrement the number of native plans using the per-task shared
memory pool, and
- // remove the memory pool if the released native plan is the last
native plan using it.
- let task_attempt_id = execution_context.task_attempt_id;
- let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
- if let Some(per_task_memory_pool) =
memory_pool_map.get_mut(&task_attempt_id) {
- per_task_memory_pool.num_plans -= 1;
- if per_task_memory_pool.num_plans == 0 {
- // Drop the memory pool from the per-task memory pool map
if there are no
- // more native plans using it.
- memory_pool_map.remove(&task_attempt_id);
- }
- }
- }
+ handle_task_shared_pool_release(
+ execution_context.memory_pool_config.pool_type,
+ execution_context.task_attempt_id,
+ );
+
let _: Box<ExecutionContext> = Box::from_raw(execution_context);
Ok(())
})
diff --git a/native/core/src/execution/memory_pools/config.rs
b/native/core/src/execution/memory_pools/config.rs
new file mode 100644
index 000000000..cd93ea046
--- /dev/null
+++ b/native/core/src/execution/memory_pools/config.rs
@@ -0,0 +1,105 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::errors::{CometError, CometResult};
+
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub(crate) enum MemoryPoolType {
+ Unified,
+ FairUnified,
+ Greedy,
+ FairSpill,
+ GreedyTaskShared,
+ FairSpillTaskShared,
+ GreedyGlobal,
+ FairSpillGlobal,
+ Unbounded,
+}
+
+impl MemoryPoolType {
+ pub(crate) fn is_task_shared(&self) -> bool {
+ matches!(
+ self,
+ MemoryPoolType::GreedyTaskShared |
MemoryPoolType::FairSpillTaskShared
+ )
+ }
+}
+
+pub(crate) struct MemoryPoolConfig {
+ pub(crate) pool_type: MemoryPoolType,
+ pub(crate) pool_size: usize,
+}
+
+impl MemoryPoolConfig {
+ pub(crate) fn new(pool_type: MemoryPoolType, pool_size: usize) -> Self {
+ Self {
+ pool_type,
+ pool_size,
+ }
+ }
+}
+
+pub(crate) fn parse_memory_pool_config(
+ off_heap_mode: bool,
+ memory_pool_type: String,
+ memory_limit: i64,
+ memory_limit_per_task: i64,
+) -> CometResult<MemoryPoolConfig> {
+ let pool_size = memory_limit as usize;
+ let memory_pool_config = if off_heap_mode {
+ match memory_pool_type.as_str() {
+ "fair_unified" =>
MemoryPoolConfig::new(MemoryPoolType::FairUnified, pool_size),
+ "default" | "unified" => {
+ // the `unified` memory pool interacts with Spark's memory
pool to allocate
+ // memory therefore does not need a size to be explicitly set.
The pool size
+ // shared with Spark is set by `spark.memory.offHeap.size`.
+ MemoryPoolConfig::new(MemoryPoolType::Unified, 0)
+ }
+ _ => {
+ return Err(CometError::Config(format!(
+ "Unsupported memory pool type for off-heap mode: {}",
+ memory_pool_type
+ )))
+ }
+ }
+ } else {
+ // Use the memory pool from DF
+ let pool_size_per_task = memory_limit_per_task as usize;
+ match memory_pool_type.as_str() {
+ "fair_spill_task_shared" => {
+ MemoryPoolConfig::new(MemoryPoolType::FairSpillTaskShared,
pool_size_per_task)
+ }
+ "default" | "greedy_task_shared" => {
+ MemoryPoolConfig::new(MemoryPoolType::GreedyTaskShared,
pool_size_per_task)
+ }
+ "fair_spill_global" => {
+ MemoryPoolConfig::new(MemoryPoolType::FairSpillGlobal,
pool_size)
+ }
+ "greedy_global" =>
MemoryPoolConfig::new(MemoryPoolType::GreedyGlobal, pool_size),
+ "fair_spill" => MemoryPoolConfig::new(MemoryPoolType::FairSpill,
pool_size_per_task),
+ "greedy" => MemoryPoolConfig::new(MemoryPoolType::Greedy,
pool_size_per_task),
+ "unbounded" => MemoryPoolConfig::new(MemoryPoolType::Unbounded, 0),
+ _ => {
+ return Err(CometError::Config(format!(
+ "Unsupported memory pool type for on-heap mode: {}",
+ memory_pool_type
+ )))
+ }
+ }
+ };
+ Ok(memory_pool_config)
+}
diff --git a/native/core/src/execution/fair_memory_pool.rs
b/native/core/src/execution/memory_pools/fair_pool.rs
similarity index 100%
rename from native/core/src/execution/fair_memory_pool.rs
rename to native/core/src/execution/memory_pools/fair_pool.rs
diff --git a/native/core/src/execution/memory_pools/mod.rs
b/native/core/src/execution/memory_pools/mod.rs
new file mode 100644
index 000000000..39b5a0ca3
--- /dev/null
+++ b/native/core/src/execution/memory_pools/mod.rs
@@ -0,0 +1,111 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+mod config;
+mod fair_pool;
+mod task_shared;
+mod unified_pool;
+
+use datafusion::execution::memory_pool::{
+ FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
UnboundedMemoryPool,
+};
+use fair_pool::CometFairMemoryPool;
+use jni::objects::GlobalRef;
+use once_cell::sync::OnceCell;
+use std::num::NonZeroUsize;
+use std::sync::Arc;
+use unified_pool::CometMemoryPool;
+
+pub(crate) use config::*;
+pub(crate) use task_shared::*;
+
+pub(crate) fn create_memory_pool(
+ memory_pool_config: &MemoryPoolConfig,
+ comet_task_memory_manager: Arc<GlobalRef>,
+ task_attempt_id: i64,
+) -> Arc<dyn MemoryPool> {
+ const NUM_TRACKED_CONSUMERS: usize = 10;
+ match memory_pool_config.pool_type {
+ MemoryPoolType::Unified => {
+ // Set Comet memory pool for native
+ let memory_pool = CometMemoryPool::new(comet_task_memory_manager);
+ Arc::new(TrackConsumersPool::new(
+ memory_pool,
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ }
+ MemoryPoolType::FairUnified => {
+ // Set Comet fair memory pool for native
+ let memory_pool =
+ CometFairMemoryPool::new(comet_task_memory_manager,
memory_pool_config.pool_size);
+ Arc::new(TrackConsumersPool::new(
+ memory_pool,
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ }
+ MemoryPoolType::Greedy => Arc::new(TrackConsumersPool::new(
+ GreedyMemoryPool::new(memory_pool_config.pool_size),
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ )),
+ MemoryPoolType::FairSpill => Arc::new(TrackConsumersPool::new(
+ FairSpillPool::new(memory_pool_config.pool_size),
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ )),
+ MemoryPoolType::GreedyGlobal => {
+ static GLOBAL_MEMORY_POOL_GREEDY: OnceCell<Arc<dyn MemoryPool>> =
OnceCell::new();
+ let memory_pool = GLOBAL_MEMORY_POOL_GREEDY.get_or_init(|| {
+ Arc::new(TrackConsumersPool::new(
+ GreedyMemoryPool::new(memory_pool_config.pool_size),
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ });
+ Arc::clone(memory_pool)
+ }
+ MemoryPoolType::FairSpillGlobal => {
+ static GLOBAL_MEMORY_POOL_FAIR: OnceCell<Arc<dyn MemoryPool>> =
OnceCell::new();
+ let memory_pool = GLOBAL_MEMORY_POOL_FAIR.get_or_init(|| {
+ Arc::new(TrackConsumersPool::new(
+ FairSpillPool::new(memory_pool_config.pool_size),
+ NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ });
+ Arc::clone(memory_pool)
+ }
+ MemoryPoolType::GreedyTaskShared | MemoryPoolType::FairSpillTaskShared
=> {
+ let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
+ let per_task_memory_pool =
+ memory_pool_map.entry(task_attempt_id).or_insert_with(|| {
+ let pool: Arc<dyn MemoryPool> =
+ if memory_pool_config.pool_type ==
MemoryPoolType::GreedyTaskShared {
+ Arc::new(TrackConsumersPool::new(
+
GreedyMemoryPool::new(memory_pool_config.pool_size),
+
NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ } else {
+ Arc::new(TrackConsumersPool::new(
+
FairSpillPool::new(memory_pool_config.pool_size),
+
NonZeroUsize::new(NUM_TRACKED_CONSUMERS).unwrap(),
+ ))
+ };
+ PerTaskMemoryPool::new(pool)
+ });
+ per_task_memory_pool.num_plans += 1;
+ Arc::clone(&per_task_memory_pool.memory_pool)
+ }
+ MemoryPoolType::Unbounded => Arc::new(UnboundedMemoryPool::default()),
+ }
+}
diff --git a/native/core/src/execution/memory_pools/task_shared.rs
b/native/core/src/execution/memory_pools/task_shared.rs
new file mode 100644
index 000000000..01bbce2bb
--- /dev/null
+++ b/native/core/src/execution/memory_pools/task_shared.rs
@@ -0,0 +1,60 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::execution::memory_pools::MemoryPoolType;
+use datafusion::execution::memory_pool::MemoryPool;
+use once_cell::sync::Lazy;
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+/// The per-task memory pools keyed by task attempt id.
+pub(crate) static TASK_SHARED_MEMORY_POOLS: Lazy<Mutex<HashMap<i64,
PerTaskMemoryPool>>> =
+ Lazy::new(|| Mutex::new(HashMap::new()));
+
+pub(crate) struct PerTaskMemoryPool {
+ pub(crate) memory_pool: Arc<dyn MemoryPool>,
+ pub(crate) num_plans: usize,
+}
+
+impl PerTaskMemoryPool {
+ pub(crate) fn new(memory_pool: Arc<dyn MemoryPool>) -> Self {
+ Self {
+ memory_pool,
+ num_plans: 0,
+ }
+ }
+}
+
+// This function reduces the refcount of a per-task memory pool when a native
plan is released.
+// If the refcount reaches zero, the memory pool is removed from the map and
dropped.
+pub(crate) fn handle_task_shared_pool_release(pool_type: MemoryPoolType,
task_attempt_id: i64) {
+ if !pool_type.is_task_shared() {
+ return;
+ }
+
+ // Decrement the number of native plans using the per-task shared memory
pool, and
+ // remove the memory pool if the released native plan is the last native
plan using it.
+ let mut memory_pool_map = TASK_SHARED_MEMORY_POOLS.lock().unwrap();
+ if let Some(per_task_memory_pool) =
memory_pool_map.get_mut(&task_attempt_id) {
+ per_task_memory_pool.num_plans -= 1;
+ if per_task_memory_pool.num_plans == 0 {
+ // Drop the memory pool from the per-task memory pool map if there
are no
+ // more native plans using it.
+ memory_pool_map.remove(&task_attempt_id);
+ }
+ }
+}
diff --git a/native/core/src/execution/memory_pool.rs
b/native/core/src/execution/memory_pools/unified_pool.rs
similarity index 100%
rename from native/core/src/execution/memory_pool.rs
rename to native/core/src/execution/memory_pools/unified_pool.rs
diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs
index 23b16f5f9..9c4956939 100644
--- a/native/core/src/execution/mod.rs
+++ b/native/core/src/execution/mod.rs
@@ -27,12 +27,9 @@ pub(crate) mod sort;
pub(crate) mod spark_plan;
pub(crate) mod util;
pub use datafusion_comet_spark_expr::timezone;
+mod memory_pools;
pub(crate) mod utils;
-mod fair_memory_pool;
-mod memory_pool;
-pub use memory_pool::*;
-
#[cfg(test)]
mod tests {
#[test]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
