feat: add orphan record identification and removal

Author: georgeRobertson

src/dve/core_engine/backends/base/rules.py

@@ -3,7 +3,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 from typing import Any, ClassVar, Generic, NoReturn, Optional, TypeVar
 from uuid import uuid4
 
@@ -17,6 +17,7 @@
 )
 from dve.core_engine.backends.base.core import get_entity_type
 from dve.core_engine.backends.exceptions import render_error
+from dve.core_engine.backends.metadata.reporting import ReportingConfig
 from dve.core_engine.backends.metadata.rules import (
     AbstractStep,
     Aggregation,
@@ -34,6 +35,7 @@
     Notification,
     OneToOneJoin,
     OrphanIdentification,
+    OrphanRemoval,
     ParentMetadata,
     RenameEntity,
     Rule,
@@ -43,8 +45,15 @@
     TableUnion,
 )
 from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful
+from dve.core_engine.configuration.v1.hierarchy import (
+    ChildHierarchyNode,
+    EntityHierarchy,
+    HierarchyNode,
+)
+from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME
 from dve.core_engine.exceptions import CriticalProcessingError
 from dve.core_engine.loggers import get_logger
+from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.type_hints import URI, DVEStageName, EntityName, Messages, TemplateVariables
 
 T_contra = TypeVar("T_contra", bound=AbstractStep, contravariant=True)
@@ -307,7 +316,10 @@ def join_header(self, entities: Entities, *, config: HeaderJoin) -> Messages:
         """
         raise NotImplementedError
 
-    def identify_orphans(self, entities: Entities, *, config: OrphanIdentification) -> Messages:
+    @abstractmethod
+    def identify_orphans(
+        self, entities: Entities, *, config: OrphanIdentification
+    ) -> tuple[Messages, int]:
         """Identify records in an entity which don't have at least one corresponding
         match in the target. A new boolean column will be added to `entity` ('IsOrphaned')
         indicating whether the condition matched.
@@ -320,6 +332,18 @@ def identify_orphans(self, entities: Entities, *, config: OrphanIdentification)
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def remove_orphans(self, entities: Entities, *, config: OrphanRemoval) -> Iterator:
+        """
+        Remove orphaned records from an entity based on the orphans found in
+        identify_orphans method. Returns a generator objects with the records removed
+        for generating feedback messages from.
+
+        This may not be implemented by some backends.
+
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def union(self, entities: Entities, *, config: TableUnion) -> Messages:
         """Union two entities together, taking the columns from each by name.
@@ -352,6 +376,100 @@ def notify(self, entities: Entities, *, config: Notification) -> Messages:
 
         """
 
+    def identify_and_remove_orphans(
+        self,
+        working_directory: URI,
+        entities: Entities,
+        entity_hierarchy: EntityHierarchy,
+        key_fields: Optional[dict[str, list[str]]] = None,
+    ) -> Messages:
+        """
+        Identifies and removes orphan records by traversing the EntityHierarchy object.
+        An orphan is a child record whose parent FK does not exist in the parent entity.
+        Processes recursively: removes orphans at each level, then processes children.
+        """
+
+        def process_node(
+            node: HierarchyNode | ChildHierarchyNode,
+            parent_entity_name: Optional[EntityName],
+            orph_messages: Messages | None = None,
+        ):
+            """Recursive helper to process a node and its children."""
+            current_entity_name = node.entity_name
+
+            if orph_messages is None:
+                orph_messages = []
+
+            if isinstance(node, ChildHierarchyNode) and parent_entity_name is not None:
+                self.logger.info(f"Identifying orphans in {current_entity_name}")
+
+                join_expr = " AND ".join(
+                    f"{parent_entity_name}.{k} = {current_entity_name}.{v}"
+                    for k, v in node.join_fields.items()
+                )
+
+                _, no_orphs = self.identify_orphans(
+                    entities=entities,
+                    config=OrphanIdentification(
+                        id=list(node.join_fields.values())[0],
+                        entity_name=current_entity_name,
+                        target_name=parent_entity_name,
+                        join_condition=join_expr,
+                    ),
+                )
+
+                if no_orphs > 0:
+                    self.logger.info(f"Removing orphan records from {current_entity_name}")
+                    location = list(node.join_fields.values())[0]
+                    with BackgroundMessageWriter(
+                        working_directory=working_directory,
+                        dve_stage=self.__stage_name__,
+                        key_fields=key_fields,
+                        logger=self.logger,
+                    ) as msg_writer:
+                        _orph_records = self.remove_orphans(
+                            entities=entities,
+                            config=OrphanRemoval(
+                                entity_name=current_entity_name,
+                                reporting=ReportingConfig(
+                                    emit="record_failure",
+                                    code=node.orphaned_records_error_code,
+                                    message=node.orphaned_records_error_message,
+                                    location=location,
+                                )
+                            )
+                        )
+                        for record in _orph_records:
+                            msg_writer.write_queue.put([
+                                FeedbackMessage(
+                                    entity=current_entity_name,
+                                    record=record,  # type: ignore
+                                    error_location=location,
+                                    error_message=node.orphaned_records_error_message,
+                                    failure_type="record",
+                                    error_type="record",
+                                    error_code=node.orphaned_records_error_code,
+                                    reporting_field=location,
+                                    category="Parent Missing",
+                                )
+                            ])
+
+            if node.children:
+                for child_node in node.children:
+                    process_node(child_node, current_entity_name, orph_messages)
+
+
+        for root_node in entity_hierarchy.entity_trees.values():
+            process_node(root_node, parent_entity_name=None)
+
+        _orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME)
+        if _orph_rel:
+            del entities[ORPHANED_RECORD_ENTITY_NAME]
+
+        entities.update(entities)
+
+        return []
+
     # pylint: disable=R0912,R0914
     def apply_sync_filters(
         self,

src/dve/core_engine/backends/implementations/duckdb/rules.py

@@ -1,6 +1,6 @@
 """Business rule definitions for duckdb backend"""
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import get_type_hints
 from uuid import uuid4
 
@@ -50,9 +50,11 @@
     Notification,
     OneToOneJoin,
     OrphanIdentification,
+    OrphanRemoval,
     SemiJoin,
     TableUnion,
 )
+from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME, RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.functions import implementations as functions
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.templating import template_object
@@ -375,8 +377,11 @@ def join_header(self, entities: DuckDBEntities, *, config: HeaderJoin) -> Messag
         return []
 
     def identify_orphans(
-        self, entities: DuckDBEntities, *, config: OrphanIdentification
-    ) -> Messages:
+        self,
+        entities: DuckDBEntities,
+        *,
+        config: OrphanIdentification,
+    ) -> tuple[Messages, int]:
         """Identify records in an entity which don't have at least one corresponding
         match in the target. A new boolean column will be added to `entity` ('IsOrphaned')
         indicating whether the condition matched.
@@ -390,41 +395,60 @@ def identify_orphans(
         target_rel: DuckDBPyRelation = entities[config.target_name]
         target_rel = target_rel.set_alias(config.target_name)
 
-        key_name = f"key_{uuid4().hex}"
-        source_rel = source_rel.select(f"*, row_number() over () as {key_name}").set_alias(
-            config.entity_name
-        )
         match_name = f"matched_{uuid4().hex}"
         target_rel = target_rel.select(
             StarExpression(exclude=[]), ConstantExpression(1).alias(match_name)
         ).set_alias(config.target_name)
 
-        joined_rel: DuckDBPyRelation = source_rel.join(
-            target_rel, condition=config.join_condition, how="left"
-        ).aggregate(f"{key_name}, coalesce(count({match_name})==0, TRUE) AS IsOrphaned")
+        pk, _fk = config.join_condition.split("=")
 
-        if "IsOrphaned" not in source_rel.columns:
-            result: DuckDBPyRelation = source_rel.join(
-                joined_rel, condition=key_name, how="left"
-            ).select(StarExpression(exclude=[key_name]))
-        else:
-            result = source_rel.set_alias("source").join(
-                joined_rel.set_alias("joined"),
-                condition=f"source.{key_name} = joined.{key_name}",
-                how="left",
+        orphaned_rel: DuckDBPyRelation = (
+            source_rel.join(target_rel, condition=config.join_condition, how="left")
+            .aggregate(
+                f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME}, {config.entity_name}.{config.id}, coalesce(count({match_name}), 0)==0 AS IsOrphaned"  # pylint: disable=C0301
             )
+            .filter("IsOrphaned")
+            .select(
+                RECORD_INDEX_COLUMN_NAME,
+                ConstantExpression(config.entity_name).alias("entity_name"),
+                ConstantExpression(pk.strip().rsplit(".")[1]).alias("pk"),
+                ColumnExpression(config.id).alias("pk_value"),  # type: ignore
+            )
+            .unique("*")
+        )
+        _orph_records: tuple[int] = orphaned_rel.count(RECORD_INDEX_COLUMN_NAME).fetchone()  # type: ignore # pylint: disable=C0301
+        if _orph_records:
+            _no_orphans = _orph_records[0]
+        else:
+            _no_orphans = 0
+        self.logger.info(f"Found {_no_orphans} orphaned records in {config.entity_name}.")
 
-            columns = {name: f"source.{name}" for name in source_rel.columns}
-            if "IsOrphaned" in source_rel.columns:
-                columns["IsOrphaned"] = ColumnExpression("source.IsOrphaned") | ColumnExpression("joined.IsOrphaned")  # type: ignore # pylint: disable=line-too-long
-            columns.pop(key_name, None)
-
-            result = result.select(
-                ",".join([f"{column} as {name}" for name, column in columns.items()])
+        if entities.get(ORPHANED_RECORD_ENTITY_NAME) is not None:
+            entities[ORPHANED_RECORD_ENTITY_NAME] = entities[ORPHANED_RECORD_ENTITY_NAME].union(
+                orphaned_rel
             )
+        else:
+            entities[ORPHANED_RECORD_ENTITY_NAME] = orphaned_rel
+        return [], _no_orphans
+
+    def remove_orphans(self, entities: DuckDBEntities, *, config: OrphanRemoval) -> Iterator:
+        """Method to remove identified orphans in the orphan tracker entity."""
+        orphan_rel = entities[ORPHANED_RECORD_ENTITY_NAME].set_alias("orphan")
+        filtered_rel = (
+            entities[config.entity_name]
+            .set_alias(config.entity_name)
+            .join(
+                orphan_rel,
+                f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME} = orphan.{RECORD_INDEX_COLUMN_NAME}",  # pylint: disable=C0301
+                "anti",
+            )
+        )
 
-        entities[config.new_entity_name or config.entity_name] = result
-        return []
+        entities[config.entity_name] = filtered_rel
+
+        return duckdb_rel_to_dictionaries(
+            orphan_rel.filter(f"entity_name = '{config.entity_name}'")
+        )
 
     def union(self, entities: DuckDBEntities, *, config: TableUnion) -> Messages:
         """Union two entities together, taking the columns from each by name.

src/dve/core_engine/backends/implementations/spark/rules.py

@@ -1,6 +1,6 @@
 """Step implementations in Spark."""
 
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from typing import Optional
 from uuid import uuid4
 
@@ -41,6 +41,7 @@
     Notification,
     OneToOneJoin,
     OrphanIdentification,
+    OrphanRemoval,
     SelectColumns,
     SemiJoin,
     TableUnion,
@@ -338,7 +339,8 @@ def union(self, entities: SparkEntities, *, config: TableUnion) -> Messages:
 
     def identify_orphans(
         self, entities: SparkEntities, *, config: OrphanIdentification
-    ) -> Messages:
+    ) -> tuple[Messages, int]:
+        # TODO - adjust this to new setup of identify and remove orphans
         source_df: DataFrame = entities[config.entity_name]
         source_df = source_df.alias(config.entity_name)
         target_df: DataFrame = entities[config.target_name]
@@ -371,7 +373,16 @@ def identify_orphans(
             result = result.select(*[column.alias(name) for name, column in columns.items()])
 
         entities[config.new_entity_name or config.entity_name] = result
-        return []
+        return [], 0
+
+    def remove_orphans(
+        self,
+        entities: SparkEntities,
+        *,
+        config: OrphanRemoval,
+    ) -> Iterator:
+        # TODO - implement for spark
+        raise NotImplementedError
 
     def filter(self, entities: SparkEntities, *, config: ImmediateFilter) -> Messages:
         """Filter an entity immediately, and do not emit any messages.

src/dve/core_engine/backends/metadata/rules.py

@@ -553,11 +553,17 @@ class OrphanIdentification(AbstractConditionalJoin):
 
     """
 
-
 Step = Union[AbstractStep, Literal["sync"]]
 """A step within a rule. This is either a rule config or the literal string 'sync'."""
 
 
+class OrphanRemoval(BaseStep):
+    """Remove an orphan record from the `entity`."""
+
+    reporting: ReportingConfig
+    """The reporting information for the row removal."""
+
+
 class Rule(BaseModel):
     """A rule, made up of multiple steps."""
 

src/dve/core_engine/configuration/v1/hierarchy.py

@@ -16,12 +16,12 @@ class HierarchyNode(BaseModel):
     """Stores entity hierarchy information"""
 
     entity_name: str
-    children: list["HierarchyNode"] = Field(default_factory=list)
+    children: Optional[list["ChildHierarchyNode"]] = Field(default_factory=list)
 
     def get_descendents(self) -> list[str]:
         """Recursively list all descendents of the node"""
         descendents = []
-        for node in self.children:
+        for node in self.children:  # type: ignore
             descendents.append(node.entity_name)
             descendents.extend(node.get_descendents())
         return descendents
@@ -31,7 +31,7 @@ def get_node(self, entity_name: str) -> Union["HierarchyNode", None]:
         node = None
         if self.entity_name == entity_name:
             return self
-        for child in self.children:
+        for child in self.children:  # type: ignore
             node = child.get_node(entity_name)
             if node:
                 break
@@ -48,8 +48,8 @@ def add_child_node(self, parent_entity: str, child_info: "HierarchyNode") -> Non
 
     def as_dict(self) -> dict[str, dict[str, Any]]:
         """Get dictionary representation of entity hierarchy"""
-        child_dict = {}
-        for node in self.children:
+        child_dict: dict[str, dict[str, Any]] = {}
+        for node in self.children:  # type: ignore
             child_dict.update(node.as_dict())
 
         ret_dict = self.model_dump(exclude={"entity_name", "children"})

src/dve/core_engine/constants.py

@@ -6,3 +6,6 @@
 CONTRACT_ERROR_VALUE_FIELD_NAME: str = "__error_value"
 """The name of the field that can be used to extract the field value that caused
    a pydantic validation error"""
+
+ORPHANED_RECORD_ENTITY_NAME: str = "orphaned_records_tracker"
+"""Name to keep track of identified orphaned records"""