feat(silo): add getDescendants function, add duplicate node error handling

anna-parker · anna-parker · commit fdb9fc52e243 · 2025-06-20T17:04:08.000+02:00
diff --git a/endToEndTests/test/queries/PhyloDescendentOf.json b/endToEndTests/test/queries/PhyloDescendentOf.json
@@ -0,0 +1,30 @@
+{
+  "testCaseName": "PhyloDescendentOf should return only descendent nodes of a given node",
+  "query": {
+    "action": {
+      "type": "Aggregated",
+      "groupByFields": ["gisaid_epi_isl"],
+      "orderByFields": [
+        {
+          "field": "gisaid_epi_isl",
+          "order": "ascending"
+        }
+      ]
+    },
+    "filterExpression": {
+      "type": "PhyloDescendentOf",
+      "column": "gisaid_epi_isl",
+      "searchExpression": "NODE_0000072"
+    }
+  },
+  "expectedQueryResult": [
+    {
+      "count": 1,
+      "gisaid_epi_isl": "EPI_ISL_1003849"
+    },
+    {
+      "count": 1,
+      "gisaid_epi_isl": "EPI_ISL_1260480"
+    }
+  ]
+}
diff --git a/src/silo/common/phylo_tree.cpp b/src/silo/common/phylo_tree.cpp
@@ -35,6 +35,11 @@ std::shared_ptr<TreeNode> parse_auspice_tree(
    }
 
    // Insert node into the map *after* children so it's fully constructed
+   if (node_map.find(node->node_id) != node_map.end()) {
+      throw silo::preprocessing::PreprocessingException(
+         fmt::format("Duplicate node ID found in Newick string: '{}'", node->node_id.string)
+      );
+   }
    node_map[node->node_id] = node;
    return node;
 }
@@ -150,6 +155,11 @@ std::shared_ptr<TreeNode> parseSubtree(
    node->node_id = parseLabel(sv);
    skipWhitespace(sv);
 
+   if (node_map.find(node->node_id) != node_map.end()) {
+      throw silo::preprocessing::PreprocessingException(
+         fmt::format("Duplicate node ID found in Newick string: '{}'", node->node_id.string)
+      );
+   }
    node_map[node->node_id] = node;
 
    return node;
@@ -179,7 +189,7 @@ PhyloTree PhyloTree::fromNewickString(const std::string& newick_string) {
       }
    } catch (const std::exception& e) {
       throw silo::preprocessing::PreprocessingException(
-         fmt::format("Error when parsing the Newick string: '{}'", newick_string)
+         fmt::format("Error when parsing the Newick string '{}': {}", newick_string, e.what())
       );
    }
 
@@ -207,7 +217,7 @@ PhyloTree PhyloTree::fromNewickFile(const std::filesystem::path& newick_path) {
       return fromNewickString(contents.str());
    } catch (const std::exception& e) {
       throw silo::preprocessing::PreprocessingException(
-         fmt::format("Error when parsing the Newick file: '{}'", newick_path.string())
+         fmt::format("Error when parsing the Newick file '{}': {}", newick_path.string(), e.what())
       );
    }
 }
@@ -217,14 +227,60 @@ PhyloTree PhyloTree::fromFile(const std::filesystem::path& path) {
 
    std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
 
+   if (ext != ".nwk" && ext != ".json") {
+      throw silo::preprocessing::PreprocessingException(fmt::format(
+         "Error when parsing tree file: '{}'. Path must end with .nwk or .json", path.string()
+      ));
+   }
    if (ext == ".nwk") {
       return common::PhyloTree::fromNewickFile(path);
    } else if (ext == ".json") {
       return common::PhyloTree::fromAuspiceJSONFile(path);
    }
-   throw silo::preprocessing::PreprocessingException(fmt::format(
-      "Error when parsing tree file: '{}'. Path must end with .nwk or .json", path.string()
-   ));
+}
+
+void PhyloTree::validateNodeExists(const TreeNodeId& node_id) {
+   if (nodes.find(node_id) == nodes.end()) {
+      throw silo::preprocessing::PreprocessingException(
+         fmt::format("Node '{}' not found in the tree.", node_id.string)
+      );
+   }
+}
+
+void PhyloTree::validateNodeExists(const std::string& node_label) {
+   auto node_id = TreeNodeId{node_label};
+   validateNodeExists(node_id);
+}
+
+roaring::Roaring PhyloTree::getDescendants(const TreeNodeId& node_id) {
+   validateNodeExists(node_id);
+   auto child_it = nodes.find(node_id);
+   roaring::Roaring result_bitmap;
+   if (!child_it->second) {
+      throw silo::preprocessing::PreprocessingException("Node is null.");
+   }
+   std::function<void(const std::shared_ptr<TreeNode>&)> dfs =
+      [&](const std::shared_ptr<TreeNode>& current) {
+         if (!current)
+            return;
+         if (current->isLeaf()) {
+            if (current->row_index.has_value()) {
+               result_bitmap.add(current->row_index.value());
+            }
+         }
+         for (const auto& child : current->children) {
+            dfs(child);
+         }
+      };
+   if (child_it->second->isLeaf()) {
+      return result_bitmap;
+   }
+   dfs(child_it->second);
+   return result_bitmap;
+}
+
+roaring::Roaring PhyloTree::getDescendants(const std::string& node_label) {
+   return getDescendants(TreeNodeId{node_label});
 }
 
 }  // namespace silo::common
diff --git a/src/silo/common/phylo_tree.h b/src/silo/common/phylo_tree.h
@@ -6,22 +6,30 @@
 #include <nlohmann/json.hpp>
 
 #include "silo/common/tree_node_id.h"
+#include "silo/query_engine/batched_bitmap_reader.h"
 
 namespace silo::common {
 using silo::common::TreeNodeId;
 
 class TreeNode {
   public:
    TreeNodeId node_id;
+   std::optional<size_t> row_index;  // index of corresponding sequence in the database (will be
+                                     // empty for internal nodes)
    std::vector<std::shared_ptr<TreeNode>> children;
    std::optional<std::shared_ptr<TreeNode>> parent;
    int depth;
+
+   bool isLeaf() { return children.empty(); }
+   bool rowIndexExists() const { return row_index.has_value(); }
 };
 
 class PhyloTree {
   public:
    std::unordered_map<TreeNodeId, std::shared_ptr<TreeNode>> nodes;
 
+   // Functions for reading and parsing phylogenetic trees
+
    static PhyloTree fromAuspiceJSONFile(const std::filesystem::path& json_path);
 
    static PhyloTree fromAuspiceJSONString(const std::string& json_string);
@@ -31,6 +39,17 @@ class PhyloTree {
    static PhyloTree fromNewickString(const std::string& newick_string);
 
    static PhyloTree fromFile(const std::filesystem::path& path);
+
+   // Functions for querying the phylogenetic tree
+
+   void validateNodeExists(const TreeNodeId& node_id);
+
+   void validateNodeExists(const std::string& node_label);
+
+   // returns a bitmap of all descendants node{node_id} that are also in the database
+   roaring::Roaring getDescendants(const TreeNodeId& node_id);
+
+   roaring::Roaring getDescendants(const std::string& node_label);
 };
 
 }  // namespace silo::common
diff --git a/src/silo/common/phylo_tree.test.cpp b/src/silo/common/phylo_tree.test.cpp
@@ -50,6 +50,29 @@ TEST(PhyloTree, throwsOnInvalidJSON) {
    );
 }
 
+TEST(PhyloTree, throwsOnInvalidAuspiceJSONDuplicateNodeId) {
+   EXPECT_THROW(
+      PhyloTree::fromAuspiceJSONString(R"({  
+  "version": "schema version",  
+  "meta": {},  
+  "tree": {  
+    "name": "ROOT",  
+    "children": [  
+      {  
+        "name": "CHILD",  
+        "children": [  
+          {  
+            "name": "CHILD"  
+          }  
+        ]  
+      }  
+    ]  
+  }  
+})"),
+      silo::preprocessing::PreprocessingException
+   );
+}
+
 TEST(PhyloTree, correctlyParsesFromNewick) {
    auto phylo_tree_file = PhyloTree::fromNewickString("((CHILD2)CHILD)ROOT;");
    ASSERT_EQ(phylo_tree_file.nodes.size(), 3);
@@ -95,3 +118,9 @@ TEST(PhyloTree, throwsOnInvalidNewickNoSemicolon) {
       silo::preprocessing::PreprocessingException
    );
 }
+
+TEST(PhyloTree, throwsOnInvalidNewickWithDuplicateNodeId) {
+   EXPECT_THROW(
+      PhyloTree::fromNewickString("((CHILD)CHILD)ROOT"), silo::preprocessing::PreprocessingException
+   );
+}
diff --git a/src/silo/query_engine/filter/expressions/phylo_tree_filter.cpp b/src/silo/query_engine/filter/expressions/phylo_tree_filter.cpp
@@ -0,0 +1,82 @@
+#include "silo/query_engine/filter/expressions/phylo_tree_filter.h"
+
+#include <optional>
+#include <utility>
+
+#include <fmt/format.h>
+#include <nlohmann/json.hpp>
+
+#include "silo/common/panic.h"
+#include "silo/common/string.h"
+#include "silo/database.h"
+#include "silo/query_engine/bad_request.h"
+#include "silo/query_engine/filter/expressions/expression.h"
+#include "silo/query_engine/filter/operators/bitmap_producer.h"
+#include "silo/query_engine/filter/operators/operator.h"
+#include "silo/storage/table_partition.h"
+
+namespace silo::query_engine::filter::expressions {
+
+PhyloChildFilter::PhyloChildFilter(std::string column_name, std::string internal_node)
+    : column_name(std::move(column_name)),
+      internal_node(std::move(internal_node)) {}
+
+std::string PhyloChildFilter::toString() const {
+   return fmt::format("column {} phylo_child_of {}", column_name, internal_node);
+};
+
+namespace {
+std::unique_ptr<silo::query_engine::filter::operators::Operator> createMatchingBitmap(
+   const storage::column::StringColumnPartition& string_column,
+   const std::string& internal_node,
+   size_t row_count
+) {
+   return std::make_unique<operators::BitmapProducer>(
+      [&, row_count]() {
+         roaring::Roaring result_bitmap = string_column.getDescendants(internal_node);
+         return CopyOnWriteBitmap(std::move(result_bitmap));
+      },
+      row_count
+   );
+}
+
+}  // namespace
+
+std::unique_ptr<silo::query_engine::filter::operators::Operator> PhyloChildFilter::compile(
+   const Database& /*database*/,
+   const storage::TablePartition& database_partition,
+   Expression::AmbiguityMode /*mode*/
+) const {
+   CHECK_SILO_QUERY(
+      database_partition.columns.string_columns.contains(column_name),
+      fmt::format("The database does not contain the string column '{}'", column_name)
+   );
+
+   SILO_ASSERT(database_partition.columns.string_columns.contains(column_name));
+   const auto& string_column = database_partition.columns.string_columns.at(column_name);
+   return createMatchingBitmap(string_column, internal_node, database_partition.sequence_count);
+}
+
+// NOLINTNEXTLINE(readability-identifier-naming)
+void from_json(const nlohmann::json& json, std::unique_ptr<PhyloChildFilter>& filter) {
+   CHECK_SILO_QUERY(
+      json.contains("column"), "The field 'column' is required in an PhyloChildFilter expression"
+   )
+   CHECK_SILO_QUERY(
+      json["column"].is_string(),
+      "The field 'column' in an PhyloChildFilter expression needs to be a string"
+   )
+   CHECK_SILO_QUERY(
+      json.contains("internal_node"),
+      "The field 'internal_node' is required in an PhyloChildFilter expression"
+   )
+   CHECK_SILO_QUERY(
+      json["internal_node"].is_string(),
+      "The field 'internal_node' in an PhyloChildFilter expression needs to be a string"
+   )
+   filter = std::make_unique<PhyloChildFilter>(
+      json["column"].get<std::string>(), json["internal_node"].get<std::string>()
+   );
+}
+
+}  // namespace silo::query_engine::filter::expressions
diff --git a/src/silo/query_engine/filter/expressions/phylo_tree_filter.h b/src/silo/query_engine/filter/expressions/phylo_tree_filter.h
@@ -0,0 +1,39 @@
+#pragma once
+
+#include <memory>
+#include <string>
+
+#include <nlohmann/json_fwd.hpp>
+
+#include "silo/database.h"
+#include "silo/query_engine/filter/expressions/expression.h"
+#include "silo/query_engine/filter/operators/operator.h"
+#include "silo/storage/table_partition.h"
+
+namespace silo::query_engine::filter::expressions {
+
+class PhyloChildFilter : public Expression {
+   std::string column_name;
+   std::string internal_node;
+
+  public:
+   explicit PhyloChildFilter(std::string column_name, std::string internal_node);
+
+   std::string toString() const override;
+
+   [[nodiscard]] std::unique_ptr<silo::query_engine::filter::operators::Operator> compile(
+      const Database& database,
+      const storage::TablePartition& database_partition,
+      AmbiguityMode mode
+   ) const override;
+
+  private:
+   std::optional<const roaring::Roaring*> getBitmapForValue(
+      const silo::storage::column::StringColumnPartition& phylo_tree_index_column
+   ) const;
+};
+
+// NOLINTNEXTLINE(readability-identifier-naming)
+void from_json(const nlohmann::json& json, std::unique_ptr<PhyloChildFilter>& filter);
+
+}  // namespace silo::query_engine::filter::expressions
diff --git a/src/silo/storage/column/string_column.cpp b/src/silo/storage/column/string_column.cpp
@@ -4,9 +4,12 @@
 
 #include "silo/common/bidirectional_map.h"
 #include "silo/common/string.h"
+#include "silo/common/tree_node_id.h"
+#include "silo/initialize/initialize_exception.h"
 
 using silo::common::String;
 using silo::common::STRING_SIZE;
+using silo::common::TreeNodeId;
 
 namespace silo::storage::column {
 
@@ -39,6 +42,18 @@ StringColumnPartition::StringColumnPartition(StringColumnMetadata* metadata)
 void StringColumnPartition::insert(const std::string& value) {
    const String<STRING_SIZE> tmp(value, metadata->dictionary);
    values.push_back(tmp);
+   if (metadata->phylo_tree.has_value()) {
+      auto child_it = (metadata->phylo_tree->nodes).find(TreeNodeId{value});
+      if (child_it == metadata->phylo_tree->nodes.end()) {
+         return;
+      }
+      if (child_it->second->rowIndexExists()) {
+         throw silo::initialize::InitializeException(
+            fmt::format("Node '{}' already exists in the phylogenetic tree.", value)
+         );
+      }
+      child_it->second->row_index = values.size() - 1;
+   }
 }
 
 void StringColumnPartition::insertNull() {
diff --git a/src/silo/storage/column/string_column.h b/src/silo/storage/column/string_column.h
diff --git a/src/silo/storage/column/string_column.test.cpp b/src/silo/storage/column/string_column.test.cpp
diff --git a/testBaseData/unitTestDummyDataset/phylogenetic_tree.nwk b/testBaseData/unitTestDummyDataset/phylogenetic_tree.nwk
