Traversal optimizations (#7244)

Co-authored-by: Adrian Galvan <galvana@uci.edu>

Author: galvana

changelog/7244.yaml

@@ -0,0 +1,4 @@
+type: Changed
+description: Optimized graph traversal algorithms for improved performance with large dataset graphs
+pr: 7244
+labels: []

src/fides/api/graph/node_filters.py

@@ -109,10 +109,19 @@ def _compute_reachable_nodes(
 
         try:
             # Create a traversal object with just this identity
+            # Skip verification in the constructor and explicitly call traverse()
+            # This avoids the overhead of _verify_traversal() when we're going
+            # to traverse anyway to detect unreachable nodes.
             with suppress_logging():
                 # Suppress the logs since we don't want to flood the logs
                 # with traversal info for each identity we want to evaluate
-                BaseTraversal(self.graph, {identity_key: "dummy_value"})
+                traversal = BaseTraversal(
+                    self.graph,
+                    {identity_key: "dummy_value"},
+                    skip_verification=True,
+                )
+                # Now explicitly traverse to detect unreachable nodes
+                traversal.traverse({}, lambda n, m: None)
 
             # If successful, all nodes are reachable
             self.reachable_by_identity[identity_key] = all_addresses

src/fides/api/graph/traversal.py

@@ -190,12 +190,46 @@ def build_consent_networkx_digraph(
     return networkx_graph
 
 
+def compute_all_descendants(
+    graph: networkx.DiGraph,
+) -> Dict[CollectionAddress, Set[CollectionAddress]]:
+    """
+    Compute descendants for ALL nodes in O(N+E) using reverse topological order.
+
+    This is much more efficient than calling networkx.descendants() for each node,
+    which would be O(N * (N+E)) = O(N²) for a graph with N nodes.
+
+    By processing in reverse topological order (leaves first), we can compute
+    each node's descendants as the union of its children's descendants plus
+    its direct children.
+
+    Returns a Dict mapping each node (CollectionAddress) in the graph to the
+    set of all nodes that are reachable from it (i.e. its transitive successors).
+    This is used to populate the ``all_descendant_tasks`` field on each
+    RequestTask so that any node can quickly determine every downstream task
+    that must complete before the overall request is finished.
+    """
+    all_descendants: Dict[CollectionAddress, Set[CollectionAddress]] = {
+        node: set() for node in graph.nodes
+    }
+
+    # Process nodes in reverse topological order (leaves first)
+    for node in reversed(list(networkx.topological_sort(graph))):
+        # This node's descendants = union of (each child + child's descendants)
+        for child in graph.successors(node):
+            all_descendants[node].add(child)
+            all_descendants[node].update(all_descendants[child])
+
+    return all_descendants
+
+
 def base_task_data(
     graph: networkx.DiGraph,
     dataset_graph: DatasetGraph,
     privacy_request: PrivacyRequest,
     node: CollectionAddress,
     traversal_nodes: Dict[CollectionAddress, TraversalNode],
+    all_descendants: Dict[CollectionAddress, Set[CollectionAddress]],
 ) -> Dict:
     """Build a dictionary of common RequestTask attributes that are shared for building
     access, consent, and erasure tasks"""
@@ -236,7 +270,7 @@ def base_task_data(
             [downstream.value for downstream in graph.successors(node)]
         ),
         "all_descendant_tasks": sorted(
-            [descend.value for descend in list(networkx.descendants(graph, node))]
+            [descend.value for descend in all_descendants.get(node, set())]
         ),
         "collection_address": node.value,
         "dataset_name": node.dataset,
@@ -270,6 +304,9 @@ def persist_new_access_request_tasks(
         traversal_nodes, end_nodes, traversal
     )
 
+    # Pre-compute all descendants in O(N+E) instead of O(N²)
+    all_descendants = compute_all_descendants(graph)
+
     for node in list(networkx.topological_sort(graph)):
         if privacy_request.get_existing_request_task(
             session, action_type=ActionType.access, collection_address=node
@@ -280,7 +317,12 @@ def persist_new_access_request_tasks(
             session,
             data={
                 **base_task_data(
-                    graph, dataset_graph, privacy_request, node, traversal_nodes
+                    graph,
+                    dataset_graph,
+                    privacy_request,
+                    node,
+                    traversal_nodes,
+                    all_descendants,
                 ),
                 "access_data": (
                     [traversal.seed_data] if node == ROOT_COLLECTION_ADDRESS else []
@@ -314,6 +356,9 @@ def persist_initial_erasure_request_tasks(
     )
     graph: networkx.DiGraph = build_erasure_networkx_digraph(traversal_nodes, end_nodes)
 
+    # Pre-compute all descendants in O(N+E) instead of O(N²)
+    all_descendants = compute_all_descendants(graph)
+
     for node in list(networkx.topological_sort(graph)):
         if privacy_request.get_existing_request_task(
             session, action_type=ActionType.erasure, collection_address=node
@@ -324,7 +369,12 @@ def persist_initial_erasure_request_tasks(
             session,
             data={
                 **base_task_data(
-                    graph, dataset_graph, privacy_request, node, traversal_nodes
+                    graph,
+                    dataset_graph,
+                    privacy_request,
+                    node,
+                    traversal_nodes,
+                    all_descendants,
                 ),
                 "action_type": ActionType.erasure,
             },
@@ -406,6 +456,9 @@ def persist_new_consent_request_tasks(
     """
     graph: networkx.DiGraph = build_consent_networkx_digraph(traversal_nodes)
 
+    # Pre-compute all descendants in O(N+E) instead of O(N²)
+    all_descendants = compute_all_descendants(graph)
+
     for node in list(networkx.topological_sort(graph)):
         if privacy_request.get_existing_request_task(
             session, action_type=ActionType.consent, collection_address=node
@@ -415,7 +468,12 @@ def persist_new_consent_request_tasks(
             session,
             data={
                 **base_task_data(
-                    graph, dataset_graph, privacy_request, node, traversal_nodes
+                    graph,
+                    dataset_graph,
+                    privacy_request,
+                    node,
+                    traversal_nodes,
+                    all_descendants,
                 ),
                 # Consent nodes take in identity data from their upstream root node
                 "access_data": ([identity] if node == ROOT_COLLECTION_ADDRESS else []),

src/fides/api/task/create_request_tasks.py

@@ -89,4 +89,8 @@ class ExecutionSettings(FidesSettings):
         default=False,
         description="Whether the memory watchdog is enabled to monitor and gracefully terminate tasks that approach memory limits.",
     )
+    use_legacy_traversal: bool = Field(
+        default=False,
+        description="When enabled, falls back to the legacy traversal algorithm. Intended as a temporary safety net in case of regressions with the optimized traversal.",
+    )
     model_config = SettingsConfigDict(env_prefix=ENV_PREFIX)