Feature/129 predicate memory (#338)

* Add ascii_tree and byte_size as dependencies

* Add option to show memory footprint of chase tables

* Fix docstring

---------

Co-authored-by: Alex Ivliev <[email protected]>

nemo-cli/src/cli.rs

@@ -125,6 +125,9 @@ pub struct CliApp {
     /// Display detailed timing information
     #[arg(long = "detailed-timing", default_value = "false")]
     pub detailed_timing: bool,
+    /// Display detailed memory information
+    #[arg(long = "detailed-memory", default_value = "false")]
+    pub detailed_memory: bool,
     /// Specify directory for input files.
     #[arg(short = 'I', long = "input-dir")]
     pub input_directory: Option<PathBuf>,

nemo-cli/src/main.rs

@@ -182,6 +182,10 @@ fn run(mut cli: CliApp) -> Result<(), Error> {
         );
     }
 
+    if cli.detailed_memory {
+        println!("\n{}", engine.memory_usage());
+    }
+
     Ok(())
 }
 

nemo-physical/src/management/database.rs

@@ -243,6 +243,18 @@ enum TableStatus {
     Reference(TableId, Permutation),
 }
 
+impl ByteSized for TableStatus {
+    fn size_bytes(&self) -> ByteSize {
+        match self {
+            TableStatus::Present(map) => map
+                .iter()
+                .map(|(_, s)| s.size_bytes())
+                .fold(ByteSize(0), |acc, x| acc + x),
+            TableStatus::Reference(_, _) => ByteSize(0),
+        }
+    }
+}
+
 /// Manages tables under different orders.
 /// Also has the capability of representing a table as a reordered version of another.
 #[derive(Debug, Default)]
@@ -1236,6 +1248,21 @@ impl DatabaseInstance {
             }
         };
     }
+
+    /// Return the amount of memory cosumed by the table under the given [`TableId`].
+    /// This also includes additional index structures but excludes tables that are currently stored on disk.
+    ///
+    /// # Panics
+    /// Panics if the given id does not exist.
+    pub fn memory_consumption(&self, id: TableId) -> ByteSize {
+        let status = self
+            .storage_handler
+            .map
+            .get(&id)
+            .expect("Function assumes that there is a table with the given id.");
+
+        status.size_bytes()
+    }
 }
 
 impl ByteSized for DatabaseInstance {

nemo/Cargo.toml

@@ -39,6 +39,8 @@ oxiri = "0.2.2"
 tokio = { version = "1.29.1", features = [ "rt" ] }
 reqwest = { version = "0.11.18" }
 num = "0.4.0"
+bytesize = "1.2"
+ascii_tree = "0.1.1"
 
 [dev-dependencies]
 env_logger = "*"

nemo/src/execution/execution_engine.rs

@@ -16,7 +16,7 @@ use crate::{
         Identifier, Program, TermOperation,
     },
     program_analysis::analysis::ProgramAnalysis,
-    table_manager::TableManager,
+    table_manager::{MemoryUsage, TableManager},
 };
 
 use super::{rule_execution::RuleExecution, selection_strategy::strategy::RuleSelectionStrategy};
@@ -374,4 +374,9 @@ impl<Strategy: RuleSelectionStrategy> ExecutionEngine<Strategy> {
 
         result
     }
+
+    /// Return the amount of consumed memory for the tables used by the chase.
+    pub fn memory_usage(&self) -> MemoryUsage {
+        self.table_manager.memory_usage()
+    }
 }

nemo/src/table_manager.rs

@@ -1,6 +1,8 @@
 //! Managing of tables
 
 use super::model::{Identifier, PrimitiveType};
+
+use bytesize::ByteSize;
 use nemo_physical::{
     datatypes::data_value::DataValueIteratorT,
     management::{
@@ -266,6 +268,69 @@ impl SubtableExecutionPlan {
     }
 }
 
+/// Stores information about memory usage of predicates
+#[derive(Debug)]
+pub struct MemoryUsage {
+    name: String,
+    memory: ByteSize,
+
+    sub_blocks: Vec<MemoryUsage>,
+}
+
+impl MemoryUsage {
+    /// Create a new [`MemoryUsage`].
+    pub fn new(name: &str, memory: ByteSize) -> Self {
+        Self {
+            name: name.to_string(),
+            memory,
+            sub_blocks: Vec::new(),
+        }
+    }
+
+    /// Create a new [`MemoryUsage`] block.
+    pub fn new_block(name: &str) -> Self {
+        Self::new(name, ByteSize(0))
+    }
+
+    /// Add a sub-block.
+    pub fn add_sub_block(&mut self, sub_block: MemoryUsage) {
+        self.memory += sub_block.memory;
+        self.sub_blocks.push(sub_block)
+    }
+
+    fn format_node(&self) -> String {
+        format!("{} ({})", self.name, self.memory)
+    }
+
+    fn ascii_tree_recursive(usage: &Self) -> ascii_tree::Tree {
+        if usage.sub_blocks.is_empty() {
+            ascii_tree::Tree::Leaf(vec![usage.format_node()])
+        } else {
+            let mut sorted_sub_blocks: Vec<&MemoryUsage> = usage.sub_blocks.iter().collect();
+            sorted_sub_blocks.sort_by(|a, b| b.memory.partial_cmp(&a.memory).unwrap());
+
+            ascii_tree::Tree::Node(
+                usage.format_node(),
+                sorted_sub_blocks
+                    .into_iter()
+                    .map(Self::ascii_tree_recursive)
+                    .collect(),
+            )
+        }
+    }
+
+    /// Return an [`ascii_tree::Tree`] representation.
+    pub fn ascii_tree(&self) -> ascii_tree::Tree {
+        Self::ascii_tree_recursive(self)
+    }
+}
+
+impl std::fmt::Display for MemoryUsage {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        ascii_tree::write_tree(f, &self.ascii_tree())
+    }
+}
+
 /// Manager object for handling tables that are the result
 /// of a seminaive existential rules evaluation process.
 #[derive(Debug)]
@@ -538,6 +603,31 @@ impl TableManager {
     pub fn get_dict(&self) -> Ref<'_, Dict> {
         self.database.get_dict_constants()
     }
+
+    /// Return the current [`MemoryUsage`].
+    pub fn memory_usage(&self) -> MemoryUsage {
+        let mut result = MemoryUsage::new_block("Chase");
+
+        for (identifier, subtable_handler) in &self.predicate_subtables {
+            let mut predicate_usage = MemoryUsage::new_block(&identifier.to_string());
+
+            for (step, id) in &subtable_handler.single {
+                let memory = self.database.memory_consumption(*id);
+                predicate_usage.add_sub_block(MemoryUsage::new(&format!("Step {}", step), memory));
+            }
+            for (steps, id) in &subtable_handler.combined {
+                let memory = self.database.memory_consumption(*id);
+                predicate_usage.add_sub_block(MemoryUsage::new(
+                    &format!("Steps {}-{}", steps.start, steps.start + steps.len),
+                    memory,
+                ));
+            }
+
+            result.add_sub_block(predicate_usage)
+        }
+
+        result
+    }
 }
 
 #[cfg(test)]