From bcc192c3a6cb8e07c315cfc4599614b249673450 Mon Sep 17 00:00:00 2001
From: Sahas Subramanian <sahas.subramanian@proton.me>
Date: Tue, 24 Jun 2025 11:36:16 +0200
Subject: [PATCH] Add a formula generator to coalesce metrics from multiple
 components

Signed-off-by: Sahas Subramanian <sahas.subramanian@proton.me>
---
 src/error.rs                                  |  3 +-
 src/graph/formulas.rs                         |  8 ++
 src/graph/formulas/generators.rs              |  1 +
 src/graph/formulas/generators/generic.rs      |  8 ++
 .../formulas/generators/generic/coalesce.rs   | 77 +++++++++++++++++++
 5 files changed, 96 insertions(+), 1 deletion(-)
 create mode 100644 src/graph/formulas/generators/generic.rs
 create mode 100644 src/graph/formulas/generators/generic/coalesce.rs

diff --git a/src/error.rs b/src/error.rs
index a8a2d6e..0197714 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -52,7 +52,8 @@ ErrorKind!(
     (Internal, internal),
     (InvalidComponent, invalid_component),
     (InvalidConnection, invalid_connection),
-    (InvalidGraph, invalid_graph)
+    (InvalidGraph, invalid_graph),
+    (MissingParameters, missing_parameters)
 );
 
 /// An error that can occur during the creation or traversal of a
diff --git a/src/graph/formulas.rs b/src/graph/formulas.rs
index 53e0b86..dedc9f1 100644
--- a/src/graph/formulas.rs
+++ b/src/graph/formulas.rs
@@ -51,6 +51,14 @@ where
         generators::pv::PVFormulaBuilder::try_new(self, pv_inverter_ids)?.build()
     }
 
+    /// Returns a string with the coalesce formula for the given component IDs.
+    ///
+    /// This formula uses the `COALESCE` function to return the first non-null
+    /// value from the components with the provided IDs.
+    pub fn coalesce(&self, component_ids: BTreeSet<u64>) -> Result<String, Error> {
+        generators::generic::CoalesceFormulaBuilder::try_new(self, component_ids)?.build()
+    }
+
     /// Returns a string representing the EV charger formula for the graph.
     pub fn ev_charger_formula(
         &self,
diff --git a/src/graph/formulas/generators.rs b/src/graph/formulas/generators.rs
index 54bfd58..81f3815 100644
--- a/src/graph/formulas/generators.rs
+++ b/src/graph/formulas/generators.rs
@@ -7,6 +7,7 @@ pub(super) mod battery;
 pub(super) mod chp;
 pub(super) mod consumer;
 pub(super) mod ev_charger;
+pub(super) mod generic;
 pub(super) mod grid;
 pub(super) mod producer;
 pub(super) mod pv;
diff --git a/src/graph/formulas/generators/generic.rs b/src/graph/formulas/generators/generic.rs
new file mode 100644
index 0000000..8c7934a
--- /dev/null
+++ b/src/graph/formulas/generators/generic.rs
@@ -0,0 +1,8 @@
+// License: MIT
+// Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+//! This module contains the methods for generating generic formulas.
+
+mod coalesce;
+
+pub(crate) use coalesce::CoalesceFormulaBuilder;
diff --git a/src/graph/formulas/generators/generic/coalesce.rs b/src/graph/formulas/generators/generic/coalesce.rs
new file mode 100644
index 0000000..9f04f5a
--- /dev/null
+++ b/src/graph/formulas/generators/generic/coalesce.rs
@@ -0,0 +1,77 @@
+// License: MIT
+// Copyright © 2025 Frequenz Energy-as-a-Service GmbH
+
+//! This module contains the methods for generating `COALESCE` formulas for
+//! measuring metrics from individual components, with fallback to other
+//! components.
+
+use std::collections::BTreeSet;
+
+use crate::{graph::formulas::expr::Expr, ComponentGraph, Edge, Error, Node};
+
+pub(crate) struct CoalesceFormulaBuilder {
+    component_ids: BTreeSet<u64>,
+}
+
+impl CoalesceFormulaBuilder {
+    pub fn try_new(
+        graph: &ComponentGraph<impl Node, impl Edge>,
+        component_ids: BTreeSet<u64>,
+    ) -> Result<Self, Error> {
+        if component_ids.is_empty() {
+            return Err(Error::missing_parameters("No component IDs specified."));
+        }
+        for component_id in &component_ids {
+            if graph.component(*component_id).is_err() {
+                return Err(Error::component_not_found(format!(
+                    "Component with ID {} not found in the graph.",
+                    component_id
+                )));
+            }
+        }
+        Ok(Self { component_ids })
+    }
+
+    /// Generates a formula that uses the `COALESCE` function to return the first
+    /// non-null value from the provided component IDs.
+    pub fn build(self) -> Result<String, Error> {
+        let expr = Expr::coalesce(
+            self.component_ids
+                .into_iter()
+                .map(|component_id| Expr::Component { component_id })
+                .collect(),
+        );
+        Ok(expr.to_string())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::graph::test_utils::ComponentGraphBuilder;
+
+    #[test]
+    fn test_coalesce_formula() -> Result<(), Error> {
+        let mut builder = ComponentGraphBuilder::new();
+        let grid = builder.grid();
+
+        // Add a grid meter and a battery chain behind it.
+        let grid_meter_1 = builder.meter();
+        builder.connect(grid, grid_meter_1);
+        let grid_meter_2 = builder.meter();
+        builder.connect(grid, grid_meter_2);
+
+        let graph = builder.build(None)?;
+        let formula = graph.coalesce(BTreeSet::from([1, 2]))?;
+        assert_eq!(formula, "COALESCE(#1, #2)");
+        let formula = graph.coalesce(BTreeSet::from([1]))?;
+        assert_eq!(formula, "COALESCE(#1)");
+        let formula = graph.coalesce(BTreeSet::from([]));
+        assert_eq!(
+            formula,
+            Err(Error::missing_parameters("No component IDs specified."))
+        );
+
+        Ok(())
+    }
+}