[GitHub] [arrow-datafusion] yjshen opened a new pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen opened a new pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526


   # Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Closes #587 .
   
    # Rationale for this change
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   When DataFusion processes a single partition, it will keep allocating memory until the OS or the container system kills it. To make it worse, concurrently executing partitions or even simultaneously running plans will compete for available memory until all memory is exhausted. It is more challenging to meet the memory requirements for all operators of each partition when it is running. None of the partitions or plans would run to finish. 
   
   Therefore, the ability to control the total memory usage of the process as a whole, and at the same time, allocate the available memory to each execution partition is extremely important. Under this guarantee: when the memory is sufficient, the operator can acquire as much of the memory to do the computation; when the memory is tight, the operator can be downgraded to use the disk to store some intermediate results (spilling to disk) and use a limited memory for execution.
   
   
   
   # What changes are included in this PR?
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   The proposed memory management architecture is the following:
   
   1. User designates max execution memory by setting `RuntimeConfig.max_memory` and `RuntimeConfig.memory_fraction` (float64 between 0..1). The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
   2. The entities that take up memory during its execution are called `Memory Consumers`. Operators or others are encouraged to register themselves to the memory manager and report its usage through `mem_used()`.
   3. There are two kinds of consumers:
      - `Controlling` consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
      - `Tracking` consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
   4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
      `(pool_size - all_tracking_used) / active_num_controlling_consumers`.
   
   ```
               Memory Space for the DataFusion Lib / Process of `pool_size`
      ┌──────────────────────────────────────────────z─────────────────────────────┐
      │                                              z                             │
      │                                              z                             │
      │               Controlling                    z          Tracking           │
      │            Memory Consumers                  z       Memory Consumers      │
      │                                              z                             │
      │                                              z                             │
      └──────────────────────────────────────────────z─────────────────────────────┘
   
   ```
   
   # Are there any user-facing changes?
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   Users could limit the max memory used for DataFusion through `RuntimeConfig::max_memory` and `RuntimeConfig::memory_fraction`
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->
   
   # Note
   
   In addition to the proposed memory manager as well as the runtime that plumbing the execute API, an `ExternalSortExec` is implemented to illustrate the API usage.
   
   TODOs:
   - [ ] Tests for `external sorter`


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780779376



##########
File path: ballista/rust/executor/src/collect.rs
##########
@@ -75,11 +76,12 @@ impl ExecutionPlan for CollectExec {
     async fn execute(
         &self,
         partition: usize,
-    ) -> Result<Pin<Box<dyn RecordBatchStream + Send + Sync>>> {
+        runtime: Arc<RuntimeEnv>,

Review comment:
       I think threading a parameter like this  all the way through makes sense 👍 

##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,157 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+    /// If runtime env has initialized
+    initialized: Arc<Mutex<bool>>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    #[allow(clippy::mutex_atomic)]
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+            initialized: Arc::new(Mutex::new(false)),
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: Arc<dyn MemoryConsumer>) {
+        {
+            let mut initialized = self.initialized.lock().unwrap();
+            if !*initialized {
+                self.memory_manager.initialize();
+                *initialized = true;
+            }
+        }
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.

Review comment:
       ```suggestion
       /// Max execution memory allowed for DataFusion. Defaults
       /// to `usize::MAX`
   ```

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       I think the idea of periodically polling tracking consumers is reasonable.
   
   I am a little worried about a task that polls based on some clock interval, however -- it is likely that the frequency will be too fast or two slow.
   
   What about updating tracking consumers every call to `try_grow?` or query to the memory manager for total memory used?

##########
File path: datafusion/src/execution/disk_manager.rs
##########
@@ -0,0 +1,108 @@
+// 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.
+
+//! Manages files generated during query execution, files are
+//! hashed among the directories listed in RuntimeConfig::local_dirs.
+
+use crate::error::{DataFusionError, Result};
+use log::info;
+use std::collections::hash_map::DefaultHasher;
+use std::fs;
+use std::fs::File;
+use std::hash::{Hash, Hasher};
+use std::path::{Path, PathBuf};
+use uuid::Uuid;

Review comment:
       Given DataFusion already uses `rand` I think we could avoid the `uuid` dependency by using that directly
   
   Something like:
   
   
   ```
   /// Return 10 character random string
   pub fn dir_name() -> String {
       thread_rng()
           .sample_iter(&Alphanumeric)
           .take(10)
           .map(char::from)
           .collect()
   }
   ```
   
   From https://github.com/influxdata/influxdb_iox/blob/3cda6b6c0f5c72b710fa73aec581ec079700fd11/influxdb_iox/tests/end_to_end_cases/scenario.rs#L286-L293
   

##########
File path: datafusion/src/execution/disk_manager.rs
##########
@@ -0,0 +1,108 @@
+// 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.
+
+//! Manages files generated during query execution, files are
+//! hashed among the directories listed in RuntimeConfig::local_dirs.
+
+use crate::error::{DataFusionError, Result};
+use log::info;
+use std::collections::hash_map::DefaultHasher;
+use std::fs;
+use std::fs::File;
+use std::hash::{Hash, Hasher};
+use std::path::{Path, PathBuf};
+use uuid::Uuid;
+
+/// Manages files generated during query execution, e.g. spill files generated
+/// while processing dataset larger than available memory.
+pub struct DiskManager {
+    local_dirs: Vec<String>,
+}
+
+impl DiskManager {
+    /// Create local dirs inside user provided dirs through conf
+    pub fn new(conf_dirs: &[String]) -> Result<Self> {

Review comment:
       Since this PR adds the `tempdir` dependency to datafusion, I think you could probably use `TempDir::new` and `TempDir::new_in` here rather than managing `String`s directly
   
   https://docs.rs/tempdir/latest/tempdir/struct.TempDir.html#method.new_in
   
   That would avoid the need for directory creation retry as well as it would handle cleaning up the files after execution
   
   I think we could switch to using `TempDirs` in a following PR however, as the API here is solid

##########
File path: datafusion/src/error.rs
##########
@@ -129,6 +132,9 @@ impl Display for DataFusionError {
             DataFusionError::Execution(ref desc) => {
                 write!(f, "Execution error: {}", desc)
             }
+            DataFusionError::ResourcesExhausted(ref desc) => {

Review comment:
       👍 

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {

Review comment:
       I am a little confused about why  `try_grow` and `spill` are part of this trait
   
   The main use of this  trait seems to be for the `MemoryManager` to interact with `ExecutionPlans` (to ask them for their memory usage, etc), but `try_grow` and `spill` call back to the memory manager
   
   By calling back to the `MemoryManager`, I think this means there is a Ref count cycle between `ExecutionPlan` and `MemoryManager` (they both have `Arc`s pointing at each other, so the ref counts will never go to zero)
   
   Looking at how `try_grow` and `spill` are used in `ExternalSort`, those methods aren't called from the `MemoryManager` at all.
   
   Thus, I suggest removing `try_grow`, `memory_manager()` and `spill` method from this trait (and put them directly into `ExternalSort`)

##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,157 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+    /// If runtime env has initialized
+    initialized: Arc<Mutex<bool>>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    #[allow(clippy::mutex_atomic)]
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+            initialized: Arc::new(Mutex::new(false)),
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: Arc<dyn MemoryConsumer>) {
+        {
+            let mut initialized = self.initialized.lock().unwrap();
+            if !*initialized {
+                self.memory_manager.initialize();
+                *initialized = true;
+            }
+        }
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.
+    pub max_memory: usize,
+    /// The fraction of total memory used for execution.
+    /// The purpose of this config is to set aside memory for untracked data structures,
+    /// and imprecise size estimation during memory acquisition.
+    pub memory_fraction: f64,
+    /// Local dirs to store temporary files during execution.
+    pub local_dirs: Vec<String>,
+}
+
+impl RuntimeConfig {
+    /// New with default values
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Customize batch size
+    pub fn with_batch_size(mut self, n: usize) -> Self {
+        // batch size must be greater than zero
+        assert!(n > 0);
+        self.batch_size = n;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_max_execution_memory(mut self, max_memory: usize) -> Self {
+        assert!(max_memory > 0);
+        self.max_memory = max_memory;
+        self
+    }
+
+    /// Customize exec memory fraction
+    pub fn with_memory_fraction(mut self, fraction: f64) -> Self {
+        assert!(fraction > 0f64 && fraction < 1f64);
+        self.memory_fraction = fraction;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_local_dirs(mut self, local_dirs: Vec<String>) -> Self {
+        assert!(!local_dirs.is_empty());
+        self.local_dirs = local_dirs;
+        self
+    }
+}
+
+impl Default for RuntimeConfig {
+    fn default() -> Self {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let path = tmp_dir.path().to_str().unwrap().to_string();
+        std::mem::forget(tmp_dir);
+
+        Self {
+            batch_size: 8192,
+            max_memory: usize::MAX,

Review comment:
       I agree with this default 👍 
   
   ```suggestion
               // Effectively "no limit"
               max_memory: usize::MAX,
   ```

##########
File path: datafusion/src/execution/disk_manager.rs
##########
@@ -0,0 +1,108 @@
+// 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.
+
+//! Manages files generated during query execution, files are
+//! hashed among the directories listed in RuntimeConfig::local_dirs.
+
+use crate::error::{DataFusionError, Result};
+use log::info;
+use std::collections::hash_map::DefaultHasher;
+use std::fs;
+use std::fs::File;
+use std::hash::{Hash, Hasher};
+use std::path::{Path, PathBuf};
+use uuid::Uuid;
+
+/// Manages files generated during query execution, e.g. spill files generated
+/// while processing dataset larger than available memory.
+pub struct DiskManager {
+    local_dirs: Vec<String>,
+}
+
+impl DiskManager {
+    /// Create local dirs inside user provided dirs through conf
+    pub fn new(conf_dirs: &[String]) -> Result<Self> {
+        let local_dirs = create_local_dirs(conf_dirs)?;
+        info!(
+            "Created local dirs {:?} as DataFusion working directory",
+            local_dirs
+        );
+        Ok(Self { local_dirs })
+    }
+
+    /// Create a file in conf dirs in randomized manner and return the file path
+    pub fn create_tmp_file(&self) -> Result<String> {
+        create_tmp_file(&self.local_dirs)
+    }
+
+    #[allow(dead_code)]
+    fn cleanup_resource(&mut self) -> Result<()> {

Review comment:
       This seems to be disconnected

##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,157 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::{Arc, Mutex};
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+    /// If runtime env has initialized
+    initialized: Arc<Mutex<bool>>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    #[allow(clippy::mutex_atomic)]
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+            initialized: Arc::new(Mutex::new(false)),
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: Arc<dyn MemoryConsumer>) {
+        {
+            let mut initialized = self.initialized.lock().unwrap();
+            if !*initialized {
+                self.memory_manager.initialize();
+                *initialized = true;
+            }
+        }
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.
+    pub max_memory: usize,
+    /// The fraction of total memory used for execution.
+    /// The purpose of this config is to set aside memory for untracked data structures,
+    /// and imprecise size estimation during memory acquisition.

Review comment:
       ```suggestion
       /// and imprecise size estimation during memory acquisition.
       /// Defaults to 0.7
   ```

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();
+                consumers.insert(id, consumer);
+            }
+            ConsumerType::Tracking => {
+                let mut trackers = self.trackers.lock().unwrap();
+                trackers.insert(id, consumer);
+            }
+        }
+    }
+
+    /// Grow memory attempt from a consumer, return if we could grant that much to it
+    async fn try_grow(
+        self: &Arc<Self>,
+        required: usize,
+        current: usize,
+        consumer_id: &MemoryConsumerId,
+    ) -> bool {
+        let max_per_op = {
+            let total_available =
+                self.pool_size - self.trackers_total_usage.load(Ordering::SeqCst);
+            let ops = self.consumers.lock().unwrap().len();
+            (total_available / ops) as usize
+        };
+        let granted = required + current < max_per_op;
+        info!(
+            "trying to acquire {} whiling holding {} from {}, got: {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            consumer_id,
+            granted,
+        );
+        granted
+    }
+
+    /// Drop a memory consumer from memory usage tracking
+    pub(crate) fn drop_consumer(self: &Arc<Self>, id: &MemoryConsumerId) {
+        // find in consumers first
+        {
+            let mut consumers = self.consumers.lock().unwrap();
+            if consumers.contains_key(id) {
+                consumers.remove(id);
+                return;
+            }

Review comment:
       ```suggestion
               if let Some(_) = consumers.remove(id) {
                   return;
               }
   ```

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {

Review comment:
       I didn't see any code that registered any Tracking consumers yet. 
   
   In terms of plumbing, what do you think about:
   1. making all `ExecutionPlans` `MemoryConsumers` and providing default implementations (that reported 0 usage)
   2. Registering all ExecutionPlans somehow as MemoryConsumers as part of physical plan creation?
   
   That way all implementations of ExecutionPlan could report their usage without having to explicitly register themselves with the memory manager. Also the manager could report on how many operators were not providing any statistics, etc

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();
+                consumers.insert(id, consumer);
+            }
+            ConsumerType::Tracking => {
+                let mut trackers = self.trackers.lock().unwrap();
+                trackers.insert(id, consumer);
+            }
+        }
+    }
+
+    /// Grow memory attempt from a consumer, return if we could grant that much to it
+    async fn try_grow(
+        self: &Arc<Self>,
+        required: usize,
+        current: usize,
+        consumer_id: &MemoryConsumerId,
+    ) -> bool {
+        let max_per_op = {
+            let total_available =
+                self.pool_size - self.trackers_total_usage.load(Ordering::SeqCst);
+            let ops = self.consumers.lock().unwrap().len();
+            (total_available / ops) as usize
+        };
+        let granted = required + current < max_per_op;
+        info!(
+            "trying to acquire {} whiling holding {} from {}, got: {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            consumer_id,
+            granted,
+        );
+        granted
+    }
+
+    /// Drop a memory consumer from memory usage tracking
+    pub(crate) fn drop_consumer(self: &Arc<Self>, id: &MemoryConsumerId) {
+        // find in consumers first
+        {
+            let mut consumers = self.consumers.lock().unwrap();
+            if consumers.contains_key(id) {
+                consumers.remove(id);
+                return;
+            }
+        }
+        {
+            let mut trackers = self.trackers.lock().unwrap();
+            if trackers.contains_key(id) {
+                let removed = trackers.remove(id);
+                match removed {
+                    None => {}
+                    Some(tracker) => {
+                        let usage = tracker.mem_used();
+                        self.trackers_total_usage.fetch_sub(usage, Ordering::SeqCst);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Shutdown
+    pub(crate) fn shutdown(&mut self) {
+        let maybe_handle = self.join_handle.lock().unwrap().take();
+        match maybe_handle {
+            None => {}
+            Some(handle) => handle.abort(),
+        }
+    }
+
+    #[allow(dead_code)]

Review comment:
       Recommend making this `impl Display for MemoryManager` instead

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       I think the two categories `controlling` and `tracking` categories make sense to me
   
   It seems like the key difference is that `controlling` consumers work with the resource manager to limit their resource usage (aka ask for more memory, or spill to disk, etc) where as `tracking` consumers just report their usage to the resource manager. 
   
   Maybe renaming `Controlling` -->  `Requesting` would make it clearer?
   
   But I think `Controlling` is also fine with some documentation clarification

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       I left some comments above, but I also think `Controlling` and `Tracking` is fine 
   
   Maybe using `Requesting` instead of `Controlling` would be clearer

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,

Review comment:
       As mentioned above, I think there is a refcount cycle here
   
   Specifically, `ExternalSort` has a `RuntimeEnv` which has an `Arc<MemoryManager>` but the MemoryManager has a `Arc` back to the `ExternalSort`.
   
   Thus I suggest changing this and the trackers below to be
   
   ```suggestion
       consumers: Arc<Mutex<HashMap<MemoryConsumerId, Weak<dyn MemoryConsumer>>>>,
   ```
   
   to break cycle 

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();
+                consumers.insert(id, consumer);
+            }
+            ConsumerType::Tracking => {
+                let mut trackers = self.trackers.lock().unwrap();
+                trackers.insert(id, consumer);
+            }
+        }
+    }
+
+    /// Grow memory attempt from a consumer, return if we could grant that much to it
+    async fn try_grow(
+        self: &Arc<Self>,
+        required: usize,
+        current: usize,
+        consumer_id: &MemoryConsumerId,
+    ) -> bool {
+        let max_per_op = {
+            let total_available =
+                self.pool_size - self.trackers_total_usage.load(Ordering::SeqCst);
+            let ops = self.consumers.lock().unwrap().len();
+            (total_available / ops) as usize
+        };
+        let granted = required + current < max_per_op;
+        info!(
+            "trying to acquire {} whiling holding {} from {}, got: {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            consumer_id,
+            granted,
+        );
+        granted
+    }
+
+    /// Drop a memory consumer from memory usage tracking
+    pub(crate) fn drop_consumer(self: &Arc<Self>, id: &MemoryConsumerId) {
+        // find in consumers first
+        {
+            let mut consumers = self.consumers.lock().unwrap();
+            if consumers.contains_key(id) {
+                consumers.remove(id);
+                return;
+            }
+        }
+        {
+            let mut trackers = self.trackers.lock().unwrap();
+            if trackers.contains_key(id) {
+                let removed = trackers.remove(id);
+                match removed {
+                    None => {}
+                    Some(tracker) => {
+                        let usage = tracker.mem_used();
+                        self.trackers_total_usage.fetch_sub(usage, Ordering::SeqCst);
+                    }
+                }
+            }

Review comment:
       ```suggestion
               if let Some(tracker) = trackers.remove(id) {
                           let usage = tracker.mem_used();
                           self.trackers_total_usage.fetch_sub(usage, Ordering::SeqCst);
               }
   ```




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r784032226



##########
File path: datafusion/src/execution/mod.rs
##########
@@ -19,4 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
+pub mod disk_manager;

Review comment:
       I agree that it makes sense to keep it private (or at least make it clearer in docs that the API is likely to change in the short term)




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783720514



##########
File path: datafusion/src/execution/mod.rs
##########
@@ -19,4 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
+pub mod disk_manager;

Review comment:
       I may need to expose df generated dirs to outer modules, for instance, write shuffle data in df in stage-1 and read shuffle data in reducers run in JVM. I think it's no harm to expose the module out?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780702085



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       Controlling and Tracking seems fine to me. Perhaps we can wait for feed backs from other reviewers.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

xudong963 commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783622538



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.
+    /// Defaults to `usize::MAX`
+    pub max_memory: usize,
+    /// The fraction of total memory used for execution.
+    /// The purpose of this config is to set aside memory for untracked data structures,
+    /// and imprecise size estimation during memory acquisition.
+    /// Defaults to 0.7
+    pub memory_fraction: f64,
+    /// Local dirs to store temporary files during execution.
+    pub local_dirs: Vec<String>,
+}
+
+impl RuntimeConfig {
+    /// New with default values
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Customize batch size
+    pub fn with_batch_size(mut self, n: usize) -> Self {
+        // batch size must be greater than zero
+        assert!(n > 0);
+        self.batch_size = n;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_max_execution_memory(mut self, max_memory: usize) -> Self {
+        assert!(max_memory > 0);
+        self.max_memory = max_memory;
+        self
+    }
+
+    /// Customize exec memory fraction
+    pub fn with_memory_fraction(mut self, fraction: f64) -> Self {
+        assert!(fraction > 0f64 && fraction <= 1f64);
+        self.memory_fraction = fraction;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_local_dirs(mut self, local_dirs: Vec<String>) -> Self {
+        assert!(!local_dirs.is_empty());
+        self.local_dirs = local_dirs;
+        self
+    }
+}
+
+impl Default for RuntimeConfig {
+    fn default() -> Self {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let path = tmp_dir.path().to_str().unwrap().to_string();
+        std::mem::forget(tmp_dir);
+
+        Self {

Review comment:
       Make sense to me. We can deal with it uniformly in the following PR




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb merged pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb merged pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526


   


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

xudong963 commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783186162



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.
+    /// Defaults to `usize::MAX`
+    pub max_memory: usize,
+    /// The fraction of total memory used for execution.
+    /// The purpose of this config is to set aside memory for untracked data structures,
+    /// and imprecise size estimation during memory acquisition.
+    /// Defaults to 0.7
+    pub memory_fraction: f64,
+    /// Local dirs to store temporary files during execution.
+    pub local_dirs: Vec<String>,
+}
+
+impl RuntimeConfig {
+    /// New with default values
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Customize batch size
+    pub fn with_batch_size(mut self, n: usize) -> Self {
+        // batch size must be greater than zero
+        assert!(n > 0);
+        self.batch_size = n;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_max_execution_memory(mut self, max_memory: usize) -> Self {
+        assert!(max_memory > 0);
+        self.max_memory = max_memory;
+        self
+    }
+
+    /// Customize exec memory fraction
+    pub fn with_memory_fraction(mut self, fraction: f64) -> Self {
+        assert!(fraction > 0f64 && fraction <= 1f64);
+        self.memory_fraction = fraction;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_local_dirs(mut self, local_dirs: Vec<String>) -> Self {
+        assert!(!local_dirs.is_empty());
+        self.local_dirs = local_dirs;
+        self
+    }
+}
+
+impl Default for RuntimeConfig {
+    fn default() -> Self {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let path = tmp_dir.path().to_str().unwrap().to_string();
+        std::mem::forget(tmp_dir);
+
+        Self {

Review comment:
       It is better to define constants for default values of `batch_size` and `memory_fraction` than bare numbers

##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

Review comment:
       Why use anonymous lifetime? Maybe it's unnecessary.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1009924630


   > I think there is a gap between ExecPlan and MemoryConsumer. Since an execute method would be called multiple times with different partition, it's always the SendableRecordBatchStream such as SortPreservingMergeStream, CrossJoinStream that takes up memory. Should I make it like:
   
   This is a good point (that the memory management is done on a per-partition basis rather than a per `ExecutionPlan` basis. I need to think 🤔  about it some more. 
   
   I would recommend we don't change `SendableRecordBatchStream` which is complicated enough as is. 
   
   I will make time today to review this PR again thoroughly -- thank you @yjshen  I think we are close


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1009955495


   @tustvold Thanks for bringing it up. I find the stream a single place to have all runtime entities be auto-registered to the memory manager at once. Maybe a wrapper over the stream could achieve the goal? 


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r782802400



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,490 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Condvar, Mutex, Weak};
+
+static CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    CONSUMER_ID.fetch_add(1, Ordering::SeqCst)
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow its memory usage by requesting more from the memory manager or
+    /// shrinks its memory usage when we can no more assign available memory to it.
+    /// Examples are spillable sorter, spillable hashmap, etc.
+    Requesting,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Requesting`).
+pub trait MemoryConsumer: Send + Sync {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current = self.mem_used();
+        info!(
+            "trying to acquire {} whiling holding {} from consumer {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            self.id(),
+        );
+
+        let can_grow_directly = self
+            .memory_manager()
+            .can_grow_directly(required, current)
+            .await;
+        if !can_grow_directly {
+            info!(
+                "Failed to grow memory of {} directly from consumer {}, spilling first ...",
+                human_readable_size(required),
+                self.id()
+            );
+            let freed = self.spill().await?;
+            self.memory_manager().record_free(freed);
+        }
+        self.memory_manager().record_acquire(required);

Review comment:
       Fixed, mind taking another look?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r784025939



##########
File path: datafusion/src/execution/mod.rs
##########
@@ -19,4 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
+pub mod disk_manager;

Review comment:
       If we public more API, we should take more effort to maintain the api back compatibility feature.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780632583



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       I'm a little short of naming 😂 . What do you think are proper words that describe the two kinds of consumers?
   The current `Controlling / controller` and `Tracking / tracker` seems less informative.
    
   




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] tustvold commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

tustvold commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1009925733


   > Should I make SendableRecordBatchStream pin arc instead of pin box and register each stream arc to runtime at each execute() last line?
   
   Not fully caught up, but how would you consume from such a thing? You need a mutable reference to poll a stream? Streams, like iterators, are not meant to be shared.
   
   _As an aside the Sync constraint on SendableRecordBatchStream is potentially extraneous for this reason, you can't do much with a shared stream anyway so requiring share-ability between threads imposes unnecessary implementation constraints_


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780618595



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       Am i correct that this background refresh process is needed because tracking consumer memory updates are managed internally instead of through MemoryManager?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       nitpick, the wording might be a little bit misleading. it looks like tracking consumers can also shrink and grow its memory usage as well through heap allocation and deallocation? or do we expect Tracking consumers to only grow used memory at runtime?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();
+                consumers.insert(id, consumer);
+            }
+            ConsumerType::Tracking => {
+                let mut trackers = self.trackers.lock().unwrap();
+                trackers.insert(id, consumer);
+            }
+        }
+    }
+
+    /// Grow memory attempt from a consumer, return if we could grant that much to it
+    async fn try_grow(
+        self: &Arc<Self>,
+        required: usize,
+        current: usize,
+        consumer_id: &MemoryConsumerId,
+    ) -> bool {
+        let max_per_op = {
+            let total_available =
+                self.pool_size - self.trackers_total_usage.load(Ordering::SeqCst);

Review comment:
       due to background refresh delay, this is best memory usage counting right? if so, it's worth documenting.

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       nitpick, since tracker is also a type of consumer, perhaps this field should be named controllers instead?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {

Review comment:
       nitpick, perhaps better to implement this through the `std::fmt::Display` trait.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen edited a comment on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen edited a comment on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1008690139


   @houqp @alamb  Thanks for your detailed and insightful review!
   
   ### Resolved:
   
   - The maintained total trackers' memory and the background threads that update it are removed. Instead, total tracker memory is collected each time the memory manager runs its `can_grow` method.
   - Renamed controlling consumer to requesting consumer, requester in short.
   - Use Weak<dyn MemoryConsumer> in MemoryManager now.
   - Use TempDir instead of manually retries of creating scratch dirs. Use rand instead of `uuid` crate.
   
   ### To discuss:
   > I didn't see any code that registered any Tracking consumers yet.
   
   There is one in `SortMergingStream` while merging multiple partial order results from spill files and the last piece of batches that are still in memory. The last piece is created as in-memory batches backed `StreamWrapper`, and just reporting its usage as in-memory batches total size.
   > 
   > In terms of plumbing, what do you think about:
   > 
   > 1. making all `ExecutionPlans` `MemoryConsumers` and providing default implementations (that reported 0 usage)
   > 2. Registering all ExecutionPlans somehow as MemoryConsumers as part of physical plan creation?
   > 
   > That way all implementations of ExecutionPlan could report their usage without having to explicitly register themselves with the memory manager. Also the manager could report on how many operators were not providing any statistics, etc
   
   I think there is a gap between ExecPlan and MemoryConsumer.  Since an `execute` method would be called multiple times with different `partition`, it's always the `SendableRecordBatchStream` such as `SortPreservingMergeStream`, `CrossJoinStream` that takes up memory. Should I make it like:
   
   ```rust
   
   /// Trait for types that stream [arrow::record_batch::RecordBatch]
   pub trait RecordBatchStream: Stream<Item = ArrowResult<RecordBatch>> + MemoryConsumer {
       /// Returns the schema of this `RecordBatchStream`.
       ///
       /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
       /// stream should have the same schema as returned from this method.
       fn schema(&self) -> SchemaRef;
   }
   
   /// Trait for a stream of record batches.
   pub type SendableRecordBatchStream = Pin<Arc<dyn RecordBatchStream + Send + Sync>>;
   
   ```
   Should I make SendableRecordBatchStream pin arc instead of pin box and register each stream arc to runtime at each `execute()` last line? Also register consumers through:
   
   ```rs
   pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
   ```
   may sometimes be awkward:
   
   ```rs
   runtime.register_consumer(&(streams.clone() as Arc<dyn MemoryConsumer>));
   ``` 
   
   Any thoughts?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1011084878


   FWIW I also plan to run the TPCH benchmarks on this PR and will post the results (I don't expect any changes)


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 removed a comment on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 removed a comment on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1011780122


   > FWIW I also plan to run the TPCH benchmarks on this PR and will post the results (I don't expect any changes)
   
   You will post the result in this review?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012347334


   > Will you post the result in this review?
   
   Yes, I will do so. 


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012677522


   > 
   I talked with @yjshen yesterday, maybe the external sort has been implemented as a draft in his branch. @alamb 
   You can fill an issue to track the follow-up tasks, it's clear for all contributors.


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012442236


   Here are the results of my comparison: the `simple_mm` branch appears to be about 10% faster for reasons I don't understand
   
   Setup:
   1. 10G TPCH data (Scale Factor 10)
   2. 16 core / 64G mem machine in google cloud ("Cascade Lake" architecture)
   3. Ran q1  which is a basic select / predicate / orderby (query below)
   
   Benchmark command:
   ```shell
   cd benchmarks
   cargo run --release  --bin tpch -- benchmark datafusion --partitions 16 -m  --iterations 10 --path /data/tpch_data_10G/ --format tbl --query 1
   ```
   
   ## `master`
   Compared master at 14176ffb50307b1d550729c8334658293e057f87 (arrow-datafusion) (merge base of `simple_mm`)
   
   ```
   Query 1 iteration 0 took 550.8 ms
   Query 1 iteration 1 took 542.1 ms
   Query 1 iteration 2 took 533.0 ms
   Query 1 iteration 3 took 539.4 ms
   Query 1 iteration 4 took 543.0 ms
   Query 1 iteration 5 took 538.5 ms
   Query 1 iteration 6 took 537.9 ms
   Query 1 iteration 7 took 536.6 ms
   Query 1 iteration 8 took 537.5 ms
   Query 1 iteration 9 took 539.8 ms
   Query 1 avg time: 539.86 ms
   ```
   
   ## `yjshen/simple_mm`
   yjshen/simple_mm at 04dca98e8d2cdf433d1959a015d0569fd6e3d0c3
   
   ```
   Query 1 iteration 0 took 500.2 ms
   Query 1 iteration 1 took 492.2 ms
   Query 1 iteration 2 took 489.1 ms
   Query 1 iteration 3 took 488.4 ms
   Query 1 iteration 4 took 489.2 ms
   Query 1 iteration 5 took 485.6 ms
   Query 1 iteration 6 took 488.2 ms
   Query 1 iteration 7 took 489.5 ms
   Query 1 iteration 8 took 491.4 ms
   Query 1 iteration 9 took 489.7 ms
   Query 1 avg time: 490.36 ms
   ```
   
   ## Query
   
   Query 1
   ```sql
   select
       l_returnflag,
       l_linestatus,
       sum(l_quantity) as sum_qty,
       sum(l_extendedprice) as sum_base_price,
       sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
       sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge,
       avg(l_quantity) as avg_qty,
       avg(l_extendedprice) as avg_price,
       avg(l_discount) as avg_disc,
       count(*) as count_order
   from
       lineitem
   where
           l_shipdate <= date '1998-09-02'
   group by
       l_returnflag,
       l_linestatus
   order by
       l_returnflag,
       l_linestatus;
   ```
   


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1007756450


   Thanks @yjshen  -- I'll try and give this a good look over the weekend


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1007756450


   Thanks @yjshen  -- I'll try and give this a good look over the weekend


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] tustvold edited a comment on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

tustvold edited a comment on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1009925733


   > Should I make SendableRecordBatchStream pin arc instead of pin box and register each stream arc to runtime at each execute() last line?
   
   Not fully caught up, but how would you consume from such a thing? You need a mutable reference to poll a stream? Streams, like iterators, are not meant to be shared.
   
   _As an aside the Sync constraint on SendableRecordBatchStream is potentially extraneous for this reason, you can't do much with a shared stream anyway, so requiring share-ability between threads imposes unnecessary implementation constraints_


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783701976



##########
File path: datafusion/src/execution/mod.rs
##########
@@ -19,4 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
+pub mod disk_manager;

Review comment:
       Do we need to public this api for the lib?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r782692980



##########
File path: datafusion/src/physical_plan/sorts/in_mem_sort.rs
##########
@@ -0,0 +1,241 @@
+// 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 std::collections::BinaryHeap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::{
+    array::{make_array as make_arrow_array, MutableArrayData},
+    compute::SortOptions,
+    datatypes::SchemaRef,
+    error::{ArrowError, Result as ArrowResult},
+    record_batch::RecordBatch,
+};
+use futures::Stream;
+
+use crate::error::Result;
+use crate::physical_plan::metrics::BaselineMetrics;
+use crate::physical_plan::sorts::{RowIndex, SortKeyCursor};
+use crate::physical_plan::{
+    expressions::PhysicalSortExpr, PhysicalExpr, RecordBatchStream,
+};
+
+/// Merge buffered, self-sorted record batches to get an order.
+///
+/// Internally, it uses MinHeap to reduce extra memory consumption
+/// by not concatenating all batches into one and sorting it as done by `SortExec`.

Review comment:
       The main difference between `InMemSortStream` and `SortPreservingMergeStream` lies in assuming different numbers of "entity" (batches for IMSS and streams for SPMS) merged. 
   
   Since `InMemSort` means to merge much more partially ordered "entities", the sorter should reduce the num of comparison for each item pop-up. Hence a MinHeap was introduced.
   
   On the other hand, `InMemSort` is more specialized to have each "entity" only one record batch, therefore simplified logic compared to consider stream continuation in `SortPreservingMergeStream`.
   
   Currently, the common parts `SortKeyCursor` and `RowIndex` for both sorts are extracted to `sorts/mod.rs` reduce duplication.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783354735



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,490 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Condvar, Mutex, Weak};
+
+static CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    CONSUMER_ID.fetch_add(1, Ordering::SeqCst)
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow its memory usage by requesting more from the memory manager or
+    /// shrinks its memory usage when we can no more assign available memory to it.
+    /// Examples are spillable sorter, spillable hashmap, etc.
+    Requesting,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Requesting`).
+pub trait MemoryConsumer: Send + Sync {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current = self.mem_used();
+        info!(
+            "trying to acquire {} whiling holding {} from consumer {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            self.id(),
+        );
+
+        let can_grow_directly = self
+            .memory_manager()
+            .can_grow_directly(required, current)
+            .await;
+        if !can_grow_directly {
+            info!(
+                "Failed to grow memory of {} directly from consumer {}, spilling first ...",
+                human_readable_size(required),
+                self.id()
+            );
+            let freed = self.spill().await?;
+            self.memory_manager().record_free(freed);
+        }
+        self.memory_manager().record_acquire(required);

Review comment:
       fix looks good to me :+1: 




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780631979



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       Yes, I'm creating a background task that runs periodically to update `tracking consumers'` total memory usage, to avoid `controlling consumers` to ask for available memory frequently.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780634649



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       Again, do you think the two`controlling` and `tracking` categories and their usage pattern make sense to you? Any suggestions on the naming?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780618595



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       Am i correct that this background refresh process is needed because tracking consumer memory updates are managed internally instead of through MemoryManager?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       nitpick, the wording might be a little bit misleading. it looks like tracking consumers can also shrink and grow its memory usage as well through heap allocation and deallocation? or do we expect Tracking consumers to only grow used memory at runtime?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();
+                consumers.insert(id, consumer);
+            }
+            ConsumerType::Tracking => {
+                let mut trackers = self.trackers.lock().unwrap();
+                trackers.insert(id, consumer);
+            }
+        }
+    }
+
+    /// Grow memory attempt from a consumer, return if we could grant that much to it
+    async fn try_grow(
+        self: &Arc<Self>,
+        required: usize,
+        current: usize,
+        consumer_id: &MemoryConsumerId,
+    ) -> bool {
+        let max_per_op = {
+            let total_available =
+                self.pool_size - self.trackers_total_usage.load(Ordering::SeqCst);

Review comment:
       due to background refresh delay, this is best memory usage counting right? if so, it's worth documenting.

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       nitpick, since tracker is also a type of consumer, perhaps this field should be named controllers instead?

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {

Review comment:
       nitpick, perhaps better to implement this through the `std::fmt::Display` trait.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012347733


   Ok, I am going to fire up my benchmark machine, get some numbers, and assuming they look good merge this PR


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] xudong963 commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

xudong963 commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783608865



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

Review comment:
       The examples in https://doc.rust-lang.org/core/fmt/struct.Formatter.html don't add anonymous lifetime. Add and unadd both can be compiled. I'm just curious why you added it here. I don't have a strong preference between the two. 




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r784058382



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

Review comment:
       Thanks @xudong963 I've removed the anonymous lifetime in the lastest commit




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1011780122


   > FWIW I also plan to run the TPCH benchmarks on this PR and will post the results (I don't expect any changes)
   
   You will post the result in this review?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

liukun4515 commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1011780340


   > FWIW I also plan to run the TPCH benchmarks on this PR and will post the results (I don't expect any changes)
   
   Will you post the result in this review?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012716591


   Thank you all again for helping me with the initial document proposal as well as insightful reviews in this PR ❤️
   
   > Shall I file follow on tickets for the next step?
   
   Yes, that would be great! please open issues that you have in mind. I have a bunch of ideas as follow-ups as well. I think we already have the foundation for many exciting features to come. How do you like to open an umbrella issue as well as sub-task issues? I could file sub-tasks issues under it as well.
   
   > I am particularly interested in ensuring we consolidate the Sort code.  (so there is only a single sort operator that does in memory sorting if it has enough memory budget but then spills to disk if needed). I would enjoy helping make this happen (perhaps by writing some tests?)
   
   Thanks, I can do the initial consolidation, please join at any time or just take it over. depends on your time schedule.
   


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

alamb commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r782585936



##########
File path: datafusion/src/physical_plan/sorts/in_mem_sort.rs
##########
@@ -0,0 +1,241 @@
+// 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 std::collections::BinaryHeap;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use arrow::{
+    array::{make_array as make_arrow_array, MutableArrayData},
+    compute::SortOptions,
+    datatypes::SchemaRef,
+    error::{ArrowError, Result as ArrowResult},
+    record_batch::RecordBatch,
+};
+use futures::Stream;
+
+use crate::error::Result;
+use crate::physical_plan::metrics::BaselineMetrics;
+use crate::physical_plan::sorts::{RowIndex, SortKeyCursor};
+use crate::physical_plan::{
+    expressions::PhysicalSortExpr, PhysicalExpr, RecordBatchStream,
+};
+
+/// Merge buffered, self-sorted record batches to get an order.
+///
+/// Internally, it uses MinHeap to reduce extra memory consumption
+/// by not concatenating all batches into one and sorting it as done by `SortExec`.

Review comment:
       I need to study the connection between `SortExec`, `SortPreservingMergeStream` and `InMemSortStream` some more to fully get understand this. I can't help by think that `InMemSortStream` is doing the same thing as `SortPreservingMergeStream` -- and I wonder if we can reuse that same code

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,490 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Condvar, Mutex, Weak};
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}

Review comment:
       I don't think we need mut / unsafe here. This works for me locally; 
   ```suggestion
   static CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
   
   fn next_id() -> usize {
       CONSUMER_ID.fetch_add(1, Ordering::SeqCst)
   }
   ```




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1008690139


   @houqp @alamb  Thanks for your detailed and insightful review!
   
   ### Resolved:
   
   - The maintained total trackers' memory and the background threads that update it are removed. Instead, total tracker memory is collected each time the memory manager runs its `can_grow` method.
   - Renamed controlling consumer to requesting consumer, requester in short.
   - Use Weak<dyn MemoryConsumer> in MemoryManager now.
   - Use TempDir instead of manually retries of creating scratch dirs. Use rand instead of `uuid` crate.
   
   ### To discuss:
   > I didn't see any code that registered any Tracking consumers yet.
   
   There is one in `SortMergingStream` while merging multiple partial order results from spill files and the last piece of batches that are still in memory. The last piece is created as in-memory batches backed `StreamWrapper`, and just reporting its usage as in-memory batches total size.
   > 
   > In terms of plumbing, what do you think about:
   > 
   > 1. making all `ExecutionPlans` `MemoryConsumers` and providing default implementations (that reported 0 usage)
   > 2. Registering all ExecutionPlans somehow as MemoryConsumers as part of physical plan creation?
   > 
   > That way all implementations of ExecutionPlan could report their usage without having to explicitly register themselves with the memory manager. Also the manager could report on how many operators were not providing any statistics, etc
   
   I think there is a gap between ExecPlan and MemoryConsumer.  Since an `execute` method would be called multiple times with different `partition`, it's always the `SendableRecordBatchStream` such as `SortPreservingMergeStream`, `CrossJoinStream` that takes up memory. Should I make it like:
   
   ```rust
   
   /// Trait for types that stream [arrow::record_batch::RecordBatch]
   pub trait RecordBatchStream: Stream<Item = ArrowResult<RecordBatch>> + MemoryConsumer {
       /// Returns the schema of this `RecordBatchStream`.
       ///
       /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
       /// stream should have the same schema as returned from this method.
       fn schema(&self) -> SchemaRef;
   }
   
   /// Trait for a stream of record batches.
   pub type SendableRecordBatchStream = Pin<Arc<dyn RecordBatchStream + Send + Sync>>;
   
   ```
   Should I make SendableRecordBatchStream pin arc instead of pin box and register each stream arc to runtime at each `execute()` last line? Also register consumers through:
   
   ```
       pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
   ``
   may sometimes be awkward:
   
   ```
           runtime.register_consumer(&(streams.clone() as Arc<dyn MemoryConsumer>));
   ``` 
   
   Any thoughts?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780702404



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       I haven't put much thought into this yet, but I am curious what are your thoughts on having tracking consumers also report memory usage update directly to the memory manager? Basically similar to what we have with the controlling consumers, but without the capability to force them to spill.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780701856



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       Ha, ok, this explanation makes a lot more sense to me now. When I initially read the comment, I was thinking `consumers themselves can grow or shrink`. So perhaps we can just reword it to `consumers that can be controlled by memory manager to grow and shrink`.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780631979



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       Yes, I'm creating a background task that runs periodically to update `tracking consumers'` total memory usage, to avoid `controlling consumers` to ask for available memory frequently.

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {
+            let mut interval_timer = tokio::time::interval(Duration::from_secs(60));
+            loop {
+                interval_timer.tick().await;
+                manager.update_tracker_total();
+            }
+        });
+        let _ = self.join_handle.lock().unwrap().insert(handle);
+    }
+
+    /// Register a new memory consumer for memory usage tracking
+    pub(crate) fn register_consumer(self: &Arc<Self>, consumer: Arc<dyn MemoryConsumer>) {
+        let id = consumer.id().clone();
+        match consumer.type_() {
+            ConsumerType::Controlling => {
+                let mut consumers = self.consumers.lock().unwrap();

Review comment:
       I'm a little short of naming 😂 . What do you think are proper words that describe the two kinds of consumers?
   The current `Controlling / controller` and `Tracking / tracker` seems less informative.
    
   

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       I use controlling consumers to describe consumers like external-sorter, hash-aggregator/joiner, or such pipeline breaker operations' consumers that would use a huge amount of memory, and we could control their memory usage pattern by forcing them to spill sometime when no more memory is available. 
   
   Tracking consumers now describe a category of entities that may use some memory during execution, but we do not control them, just check their usage updates periodically. They are free to grow or shrink memory usage, the only demand for the implementors is to be accurate as possible for the `mem_used()` report.
   
   Tracking consumers would occupy tiny memory space, such as a constant internal buffer for each batch calculation, or much larger memory if they are coming from the previous `controlling consumers`.  For example, the last part of in-memory batches for the `external sort` or the already built, partial memory-residential hashtables for hash join probing are tracking consumers we would like to only remember its memory usage, but not trigger its spill to reduce system complexity.

##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       Again, do you think the two`controlling` and `tracking` categories and their usage pattern make sense to you? Any suggestions on the naming?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] liukun4515 commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

liukun4515 commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1007988928


   @yjshen  thanks, this is milestone pr for memory controller in datafusion.


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780991701



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       There are no more maintained `tracker_total` and no more background maintaining tasks. Memory manager decides total tracker memory each time its `can_grow` is called now.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780988323



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {

Review comment:
       Maybe it's `try_grow` in both `MemoryManager` and `MemoryConsumer` misleading. 
   
   In MemoryConsumer, The above three `try_grow()` `spill()` `memory_manager()`  act as follows:
   `try_grow()` asks for more memory from `memory_manger()` , if not granted, call `spill()` on itself.  
   
   In MemoryManger, I renamed the function `can_grow()`:
   `can_grow` first gathers all tracker consumers' total memory, then the active number of memory requesters, at last, it tells if the requesters could be fulfilled.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] tustvold edited a comment on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

tustvold edited a comment on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1010171638


   > Maybe a wrapper over the stream could achieve the goal
   
   My instinct would be to suggest having the shared ref internal to the stream implementation, instead of a wrapper. Otherwise I suspect you will run into borrow checker, pinning, and async pain. This would also avoid needing to make breaking changes to SendableRecordBatchStream?
   
   Another thing to potentially think about is that many of the operators aren't actually streams, rather they spawn a tokio task and then return an mpsc queue. There will need to be some accounting of both data buffered in the queue, and data in the operators "task". My gut feeling is this is going to require adding some sort of RAII tracking field to `RecordBatch` or possibly `Buffer` but I'm not really sure...


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] tustvold commented on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

tustvold commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1010171638


   > Maybe a wrapper over the stream could achieve the goal
   
   My instinct would be to suggest having the shared ref internal to the stream implementation, instead of a wrapper. Otherwise I suspect you will run into borrow checker, and async pain. This would also avoid needing to make breaking changes to SendableRecordBatchStream?
   
   Another thing to potentially think about is that many of the operators aren't actually streams, rather they spawn a tokio task and then return an mpsc queue. There will need to be some accounting of both data buffered in the queue, and data in the operators "task". My gut feeling is this is going to require adding some sort of RAII tracking field to `RecordBatch` or possibly `Buffer` but I'm not really sure...


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen edited a comment on pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen edited a comment on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1008690139


   @houqp @alamb  Thanks for your detailed and insightful review!
   
   ### Resolved:
   
   - The maintained total trackers' memory and the background threads that update it are removed. Instead, total tracker memory is collected each time the memory manager runs its `can_grow` method.
   - Renamed controlling consumer to requesting consumer, requester in short.
   - Use Weak<dyn MemoryConsumer> in MemoryManager now.
   - Use TempDir instead of manually retries of creating scratch dirs. Use rand instead of `uuid` crate.
   
   ### To discuss:
   > I didn't see any code that registered any Tracking consumers yet.
   
   There is one in `SortMergingStream` while merging multiple partial order results from spill files and the last piece of batches that are still in memory. The last piece is created as in-memory batches backed `StreamWrapper`, and just reporting its usage as in-memory batches total size.
   > 
   > In terms of plumbing, what do you think about:
   > 
   > 1. making all `ExecutionPlans` `MemoryConsumers` and providing default implementations (that reported 0 usage)
   > 2. Registering all ExecutionPlans somehow as MemoryConsumers as part of physical plan creation?
   > 
   > That way all implementations of ExecutionPlan could report their usage without having to explicitly register themselves with the memory manager. Also the manager could report on how many operators were not providing any statistics, etc
   
   I think there is a gap between ExecPlan and MemoryConsumer.  Since an `execute` method would be called multiple times with different `partition`, it's always the `SendableRecordBatchStream` such as `SortPreservingMergeStream`, `CrossJoinStream` that takes up memory. Should I make it like:
   
   ```rust
   
   /// Trait for types that stream [arrow::record_batch::RecordBatch]
   pub trait RecordBatchStream: Stream<Item = ArrowResult<RecordBatch>> + MemoryConsumer {
       /// Returns the schema of this `RecordBatchStream`.
       ///
       /// Implementation of this trait should guarantee that all `RecordBatch`'s returned by this
       /// stream should have the same schema as returned from this method.
       fn schema(&self) -> SchemaRef;
   }
   
   /// Trait for a stream of record batches.
   pub type SendableRecordBatchStream = Pin<Arc<dyn RecordBatchStream + Send + Sync>>;
   
   ```
   Should I make SendableRecordBatchStream pin arc instead of pin box and register each stream arc to runtime at each `execute()` last line? Also register consumers through:
   
   ```rs
       pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
   ```
   may sometimes be awkward:
   
   ```rs
       runtime.register_consumer(&(streams.clone() as Arc<dyn MemoryConsumer>));
   ``` 
   
   Any thoughts?


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780634037



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution

Review comment:
       I use controlling consumers to describe consumers like external-sorter, hash-aggregator/joiner, or such pipeline breaker operations' consumers that would use a huge amount of memory, and we could control their memory usage pattern by forcing them to spill sometime when no more memory is available. 
   
   Tracking consumers now describe a category of entities that may use some memory during execution, but we do not control them, just check their usage updates periodically. They are free to grow or shrink memory usage, the only demand for the implementors is to be accurate as possible for the `mem_used()` report.
   
   Tracking consumers would occupy tiny memory space, such as a constant internal buffer for each batch calculation, or much larger memory if they are coming from the previous `controlling consumers`.  For example, the last part of in-memory batches for the `external sort` or the already built, partial memory-residential hashtables for hash join probing are tracking consumers we would like to only remember its memory usage, but not trigger its spill to reduce system complexity.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r780728793



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,320 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+use tokio::task;
+use tokio::task::JoinHandle;
+
+static mut CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    unsafe { CONSUMER_ID.fetch_add(1, Ordering::SeqCst) }
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow or shrink its memory usage during execution
+    /// such as spillable sorter, spillable hashmap, etc.
+    Controlling,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Controlling`).
+pub trait MemoryConsumer: Send + Sync + Debug {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current_usage = self.mem_used();
+        let can_grow = self
+            .memory_manager()
+            .try_grow(required, current_usage, self.id())
+            .await;
+        if !can_grow {
+            info!(
+                "Failed to grow memory of {} from {}, spilling...",
+                human_readable_size(required),
+                self.id()
+            );
+            self.spill().await?;
+        }
+        Ok(())
+    }
+
+    /// Spill in-memory buffers to disk, free memory
+    async fn spill(&self) -> Result<()>;
+
+    /// Current memory used by this consumer
+    fn mem_used(&self) -> usize;
+
+    /// Current status of the consumer
+    fn str_repr(&self) -> String {
+        let mem = self.mem_used();
+        format!(
+            "{}[{}]: {}",
+            self.name(),
+            self.id(),
+            human_readable_size(mem)
+        )
+    }
+}
+
+/*
+The memory management architecture is the following:
+
+1. User designates max execution memory by setting RuntimeConfig.max_memory and RuntimeConfig.memory_fraction (float64 between 0..1).
+   The actual max memory DataFusion could use `pool_size =  max_memory * memory_fraction`.
+2. The entities that take up memory during its execution are called 'Memory Consumers'. Operators or others are encouraged to
+   register themselves to the memory manager and report its usage through `mem_used()`.
+3. There are two kinds of consumers:
+   - 'Controlling' consumers that would acquire memory during its execution and release memory through `spill` if no more memory is available.
+   - 'Tracking' consumers that exist for reporting purposes to provide a more accurate memory usage estimation for memory consumers.
+4. Controlling and tracking consumers share the pool. Each controlling consumer could acquire a maximum of
+   (pool_size - all_tracking_used) / active_num_controlling_consumers.
+
+            Memory Space for the DataFusion Lib / Process of `pool_size`
+   ┌──────────────────────────────────────────────z─────────────────────────────┐
+   │                                              z                             │
+   │                                              z                             │
+   │               Controlling                    z          Tracking           │
+   │            Memory Consumers                  z       Memory Consumers      │
+   │                                              z                             │
+   │                                              z                             │
+   └──────────────────────────────────────────────z─────────────────────────────┘
+*/
+
+/// Manage memory usage during physical plan execution
+pub struct MemoryManager {
+    consumers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers: Arc<Mutex<HashMap<MemoryConsumerId, Arc<dyn MemoryConsumer>>>>,
+    trackers_total_usage: AtomicUsize,
+    pool_size: usize,
+    join_handle: Arc<Mutex<Option<JoinHandle<()>>>>,
+}
+
+impl MemoryManager {
+    /// Create new memory manager based on max available pool_size
+    pub fn new(pool_size: usize) -> Self {
+        info!(
+            "Creating memory manager with initial size {}",
+            human_readable_size(pool_size)
+        );
+        Self {
+            consumers: Arc::new(Mutex::new(HashMap::new())),
+            trackers: Arc::new(Mutex::new(HashMap::new())),
+            trackers_total_usage: AtomicUsize::new(0),
+            pool_size,
+            join_handle: Arc::new(Mutex::new(None)),
+        }
+    }
+
+    fn update_tracker_total(self: &Arc<Self>) {
+        let trackers = self.trackers.lock().unwrap();
+        if trackers.len() > 0 {
+            let sum = trackers.values().fold(0usize, |acc, y| acc + y.mem_used());
+            drop(trackers);
+            self.trackers_total_usage.store(sum, Ordering::SeqCst);
+        }
+    }
+
+    /// Initialize
+    pub(crate) fn initialize(self: &Arc<Self>) {
+        let manager = self.clone();
+        let handle = task::spawn(async move {

Review comment:
       The main reason is to reduce interaction with the memory manager during one's execution, to reduce complexity as well as eliminate synchronization needs.
   
   - For tracking consumers that were converted from controlling consumers. for example, the hashtable size / partial sort in-mem size is known when tracking consumer is created or transformed to, then no more need for them to acquire memory or interact with memory manager. 
   - For other tracking consumers with internal computational buffers. One can report its usage by simply updating its internal state `mem_used`, no extra function calls or interaction with the memory manager during execution.
   




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] houqp commented on a change in pull request #1526: A simplified memory manager for query execution

[GitBox <gi...@apache.org>]

houqp commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r782743336



##########
File path: datafusion/src/execution/memory_manager.rs
##########
@@ -0,0 +1,490 @@
+// 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.
+
+//! Manages all available memory during query execution
+
+use crate::error::Result;
+use async_trait::async_trait;
+use hashbrown::HashMap;
+use log::info;
+use std::fmt;
+use std::fmt::{Debug, Display, Formatter};
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Condvar, Mutex, Weak};
+
+static CONSUMER_ID: AtomicUsize = AtomicUsize::new(0);
+
+fn next_id() -> usize {
+    CONSUMER_ID.fetch_add(1, Ordering::SeqCst)
+}
+
+/// Type of the memory consumer
+pub enum ConsumerType {
+    /// consumers that can grow its memory usage by requesting more from the memory manager or
+    /// shrinks its memory usage when we can no more assign available memory to it.
+    /// Examples are spillable sorter, spillable hashmap, etc.
+    Requesting,
+    /// consumers that are not spillable, counting in for only tracking purpose.
+    Tracking,
+}
+
+#[derive(Clone, Debug, Hash, Eq, PartialEq)]
+/// Id that uniquely identifies a Memory Consumer
+pub struct MemoryConsumerId {
+    /// partition the consumer belongs to
+    pub partition_id: usize,
+    /// unique id
+    pub id: usize,
+}
+
+impl MemoryConsumerId {
+    /// Auto incremented new Id
+    pub fn new(partition_id: usize) -> Self {
+        let id = next_id();
+        Self { partition_id, id }
+    }
+}
+
+impl Display for MemoryConsumerId {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "{}:{}", self.partition_id, self.id)
+    }
+}
+
+#[async_trait]
+/// A memory consumer that either takes up memory (of type `ConsumerType::Tracking`)
+/// or grows/shrinks memory usage based on available memory (of type `ConsumerType::Requesting`).
+pub trait MemoryConsumer: Send + Sync {
+    /// Display name of the consumer
+    fn name(&self) -> String;
+
+    /// Unique id of the consumer
+    fn id(&self) -> &MemoryConsumerId;
+
+    /// Ptr to MemoryManager
+    fn memory_manager(&self) -> Arc<MemoryManager>;
+
+    /// Partition that the consumer belongs to
+    fn partition_id(&self) -> usize {
+        self.id().partition_id
+    }
+
+    /// Type of the consumer
+    fn type_(&self) -> &ConsumerType;
+
+    /// Grow memory by `required` to buffer more data in memory,
+    /// this may trigger spill before grow when the memory threshold is
+    /// reached for this consumer.
+    async fn try_grow(&self, required: usize) -> Result<()> {
+        let current = self.mem_used();
+        info!(
+            "trying to acquire {} whiling holding {} from consumer {}",
+            human_readable_size(required),
+            human_readable_size(current),
+            self.id(),
+        );
+
+        let can_grow_directly = self
+            .memory_manager()
+            .can_grow_directly(required, current)
+            .await;
+        if !can_grow_directly {
+            info!(
+                "Failed to grow memory of {} directly from consumer {}, spilling first ...",
+                human_readable_size(required),
+                self.id()
+            );
+            let freed = self.spill().await?;
+            self.memory_manager().record_free(freed);
+        }
+        self.memory_manager().record_acquire(required);

Review comment:
       `request_total` could change after we called `can_grow_directly` and before we call `record_acquire` here. If we are doing best effort accounting, then this is not the end of the world, but just wanted to point this out in case it's not expected.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783583158



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {

Review comment:
       Ah, my IDE generates it. I think it aligns with the method signature and the doc. Is it recommended to remove?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r783611819



##########
File path: datafusion/src/execution/runtime_env.rs
##########
@@ -0,0 +1,149 @@
+// 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.
+
+//! Execution runtime environment that tracks memory, disk and various configurations
+//! that are used during physical plan execution.
+
+use crate::error::Result;
+use crate::execution::disk_manager::DiskManager;
+use crate::execution::memory_manager::{MemoryConsumer, MemoryConsumerId, MemoryManager};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+
+#[derive(Clone)]
+/// Execution runtime environment
+pub struct RuntimeEnv {
+    /// Runtime configuration
+    pub config: RuntimeConfig,
+    /// Runtime memory management
+    pub memory_manager: Arc<MemoryManager>,
+    /// Manage temporary files during query execution
+    pub disk_manager: Arc<DiskManager>,
+}
+
+impl Debug for RuntimeEnv {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "RuntimeEnv")
+    }
+}
+
+impl RuntimeEnv {
+    /// Create env based on configuration
+    pub fn new(config: RuntimeConfig) -> Result<Self> {
+        let memory_manager = Arc::new(MemoryManager::new(
+            (config.max_memory as f64 * config.memory_fraction) as usize,
+        ));
+        let disk_manager = Arc::new(DiskManager::new(&config.local_dirs)?);
+        Ok(Self {
+            config,
+            memory_manager,
+            disk_manager,
+        })
+    }
+
+    /// Get execution batch size based on config
+    pub fn batch_size(&self) -> usize {
+        self.config.batch_size
+    }
+
+    /// Register the consumer to get it tracked
+    pub fn register_consumer(&self, memory_consumer: &Arc<dyn MemoryConsumer>) {
+        self.memory_manager.register_consumer(memory_consumer);
+    }
+
+    /// Drop the consumer from get tracked
+    pub fn drop_consumer(&self, id: &MemoryConsumerId) {
+        self.memory_manager.drop_consumer(id)
+    }
+}
+
+impl Default for RuntimeEnv {
+    fn default() -> Self {
+        RuntimeEnv::new(RuntimeConfig::new()).unwrap()
+    }
+}
+
+#[derive(Clone)]
+/// Execution runtime configuration
+pub struct RuntimeConfig {
+    /// Default batch size while creating new batches, it's especially useful
+    /// for buffer-in-memory batches since creating tiny batches would results
+    /// in too much metadata memory consumption.
+    pub batch_size: usize,
+    /// Max execution memory allowed for DataFusion.
+    /// Defaults to `usize::MAX`
+    pub max_memory: usize,
+    /// The fraction of total memory used for execution.
+    /// The purpose of this config is to set aside memory for untracked data structures,
+    /// and imprecise size estimation during memory acquisition.
+    /// Defaults to 0.7
+    pub memory_fraction: f64,
+    /// Local dirs to store temporary files during execution.
+    pub local_dirs: Vec<String>,
+}
+
+impl RuntimeConfig {
+    /// New with default values
+    pub fn new() -> Self {
+        Default::default()
+    }
+
+    /// Customize batch size
+    pub fn with_batch_size(mut self, n: usize) -> Self {
+        // batch size must be greater than zero
+        assert!(n > 0);
+        self.batch_size = n;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_max_execution_memory(mut self, max_memory: usize) -> Self {
+        assert!(max_memory > 0);
+        self.max_memory = max_memory;
+        self
+    }
+
+    /// Customize exec memory fraction
+    pub fn with_memory_fraction(mut self, fraction: f64) -> Self {
+        assert!(fraction > 0f64 && fraction <= 1f64);
+        self.memory_fraction = fraction;
+        self
+    }
+
+    /// Customize exec size
+    pub fn with_local_dirs(mut self, local_dirs: Vec<String>) -> Self {
+        assert!(!local_dirs.is_empty());
+        self.local_dirs = local_dirs;
+        self
+    }
+}
+
+impl Default for RuntimeConfig {
+    fn default() -> Self {
+        let tmp_dir = tempfile::tempdir().unwrap();
+        let path = tmp_dir.path().to_str().unwrap().to_string();
+        std::mem::forget(tmp_dir);
+
+        Self {

Review comment:
       I think this is already the default of config and meant to be overwritten if needed. And a single place for these defaults? 
   
   BTW, I think it is worth reconsidering and restructuring the multiple configs and their usages. `ExecutionConfig`, `PhysicalPlanConfig` and `RuntimeConfig`. 
   
   At least we should not pass `target_batch_size` during query planning since we already have runtimeEnv plumbing through the execute() API now, will create follow-up PR once we've merged this one.




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] yjshen commented on a change in pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

yjshen commented on a change in pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#discussion_r784057787



##########
File path: datafusion/src/execution/mod.rs
##########
@@ -19,4 +19,7 @@
 
 pub mod context;
 pub mod dataframe_impl;
+pub mod disk_manager;

Review comment:
       Thanks, I've changed these two into `pub(crate)`




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow-datafusion] alamb commented on pull request #1526: Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation

[GitBox <gi...@apache.org>]

alamb commented on pull request #1526:
URL: https://github.com/apache/arrow-datafusion/pull/1526#issuecomment-1012543096


   I think while not perfect this PR is a step in the right direction  towards being able to handle queries that need to spill. 🚀  thank you @yjshen 
   
   Shall I file follow on tickets for the next step?
   
   I am particularly interested in ensuring we consolidate the Sort code (so there is only a single sort operator that does in memory sorting if it has enough memory budget but then spills to disk if needed). I would enjoy helping make this happen (perhaps by writing some tests?)
   
   


-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org