diff --git a/src/dve/core_engine/backends/base/rules.py b/src/dve/core_engine/backends/base/rules.py
index 9b6b4fe..1340e32 100644
--- a/src/dve/core_engine/backends/base/rules.py
+++ b/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,
diff --git a/src/dve/core_engine/backends/implementations/duckdb/rules.py b/src/dve/core_engine/backends/implementations/duckdb/rules.py
index c4277d9..4479846 100644
--- a/src/dve/core_engine/backends/implementations/duckdb/rules.py
+++ b/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.
diff --git a/src/dve/core_engine/backends/implementations/spark/rules.py b/src/dve/core_engine/backends/implementations/spark/rules.py
index 825ee15..ff15b52 100644
--- a/src/dve/core_engine/backends/implementations/spark/rules.py
+++ b/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.
diff --git a/src/dve/core_engine/backends/metadata/rules.py b/src/dve/core_engine/backends/metadata/rules.py
index f3a6305..9b96a14 100644
--- a/src/dve/core_engine/backends/metadata/rules.py
+++ b/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."""
diff --git a/src/dve/core_engine/configuration/v1/hierarchy.py b/src/dve/core_engine/configuration/v1/hierarchy.py
index 5964270..15dedcf 100644
--- a/src/dve/core_engine/configuration/v1/hierarchy.py
+++ b/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"})
diff --git a/src/dve/core_engine/constants.py b/src/dve/core_engine/constants.py
index a2a4a65..7afcae7 100644
--- a/src/dve/core_engine/constants.py
+++ b/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"""
diff --git a/src/dve/core_engine/type_hints.py b/src/dve/core_engine/type_hints.py
index 154ada6..e369ff4 100644
--- a/src/dve/core_engine/type_hints.py
+++ b/src/dve/core_engine/type_hints.py
@@ -133,7 +133,7 @@
"""A string indicating the field that the error pertains to."""
FieldValue = Optional[Any]
"""The value that caused the error."""
-ErrorCategory = Literal["Blank", "Wrong format", "Bad value", "Bad file"]
+ErrorCategory = Literal["Blank", "Wrong format", "Bad value", "Bad file", "Parent Missing"]
"""A string indicating the category of the error."""
RecordIndex = Optional[int]
"""The record index that the error relates to (if applicable)"""
diff --git a/src/dve/pipeline/pipeline.py b/src/dve/pipeline/pipeline.py
index a9be3ff..eaf7661 100644
--- a/src/dve/pipeline/pipeline.py
+++ b/src/dve/pipeline/pipeline.py
@@ -34,6 +34,7 @@
from dve.core_engine.backends.readers.utilities import get_all_model_fields
from dve.core_engine.backends.types import EntityType
from dve.core_engine.backends.utilities import stringify_model
+from dve.core_engine.configuration.v1.hierarchy import EntityHierarchy
from dve.core_engine.exceptions import CriticalProcessingError
from dve.core_engine.loggers import get_logger
from dve.core_engine.message import FeedbackMessage
@@ -596,8 +597,13 @@ def apply_business_rules( # pylint: disable=R0914
key_fields = {model: conf.reporting_fields for model, conf in model_config.items()}
+ entity_hierarchy = EntityHierarchy.from_engine_config(config)
+
_errors_uri, rules_success = self.step_implementations.apply_rules( # type: ignore
- working_directory, entity_manager, rules, key_fields
+ working_directory,
+ entity_manager,
+ rules,
+ key_fields,
)
rule_messages = load_feedback_messages(
@@ -636,6 +642,13 @@ def apply_business_rules( # pylint: disable=R0914
projected
)
+ self.step_implementations.identify_and_remove_orphans( # type: ignore
+ working_directory,
+ entity_manager.entities,
+ entity_hierarchy,
+ key_fields,
+ )
+
submission_status.number_of_records = self.get_entity_count(
entity=entity_manager.entities[f"""Original{rules.global_variables.get(
'entity',
diff --git a/tests/features/flights.feature b/tests/features/flights.feature
new file mode 100644
index 0000000..ea710e9
--- /dev/null
+++ b/tests/features/flights.feature
@@ -0,0 +1,116 @@
+Feature: Pipeline tests using the flights dataset
+ Test hierarchical record rejection and ensuring that records are removed correctly including
+ any "orphan" records generated from record removal in parent entities.
+
+ Scenario: A perfect flights file
+ Given I submit the flights file perfect_flights.xml for processing
+ And A duckdb pipeline is configured with schema file 'flights.dischema.json'
+ And I add initial audit entries for the submission
+ Then the latest audit record for the submission is marked with processing status file_transformation
+ When I run the file transformation phase
+ Then the country entity is stored as a parquet after the file_transformation phase
+ And the airport entity is stored as a parquet after the file_transformation phase
+ And the flights entity is stored as a parquet after the file_transformation phase
+ And the passengers entity is stored as a parquet after the file_transformation phase
+ And the latest audit record for the submission is marked with processing status data_contract
+ When I run the data contract phase
+ Then there are no file rejections from the data_contract phase
+ And there are no record rejections from the data_contract phase
+ When I run the business rules phase
+ Then there are no file rejections from the business_rules phase
+ And there are no record rejections from the business_rules phase
+ When I run the error report phase
+ Then An error report is produced
+ # TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
+ # And The statistics entry for the submission shows the following information
+ # | parameter | value |
+ # | record_count | 1 |
+ # | number_file_rejections | 0 |
+ # | number_record_rejections | 0 |
+
+ Scenario: A flights submission where the root record is rejected
+ Given I submit the flights file missing_country_id.xml for processing
+ And A duckdb pipeline is configured with schema file 'flights.dischema.json'
+ And I add initial audit entries for the submission
+ Then the latest audit record for the submission is marked with processing status file_transformation
+ When I run the file transformation phase
+ Then the country entity is stored as a parquet after the file_transformation phase
+ And the airport entity is stored as a parquet after the file_transformation phase
+ And the flights entity is stored as a parquet after the file_transformation phase
+ And the passengers entity is stored as a parquet after the file_transformation phase
+ And the latest audit record for the submission is marked with processing status data_contract
+ When I run the data contract phase
+ Then there are no file rejections from the data_contract phase
+ And there are no record rejections from the data_contract phase
+ When I run the business rules phase
+ Then there are errors with the following details and associated error_count from the business_rules phase
+ | ErrorType | ErrorCode | error_count |
+ | record | C1 | 1 |
+ | record | AG1 | 1 |
+ | record | FG1 | 2 |
+ | record | PG1 | 4 |
+ When I run the error report phase
+ Then An error report is produced
+ # TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
+ # And The statistics entry for the submission shows the following information
+ # | parameter | value |
+ # | record_count | 1 |
+ # | number_file_rejections | 0 |
+ # | number_record_rejections | 1 |
+
+ Scenario: A flights submission where a child primary key is rejected
+ Given I submit the flights file missing_flight_id.xml for processing
+ And A duckdb pipeline is configured with schema file 'flights.dischema.json'
+ And I add initial audit entries for the submission
+ Then the latest audit record for the submission is marked with processing status file_transformation
+ When I run the file transformation phase
+ Then the country entity is stored as a parquet after the file_transformation phase
+ And the airport entity is stored as a parquet after the file_transformation phase
+ And the flights entity is stored as a parquet after the file_transformation phase
+ And the passengers entity is stored as a parquet after the file_transformation phase
+ And the latest audit record for the submission is marked with processing status data_contract
+ When I run the data contract phase
+ Then there are no file rejections from the data_contract phase
+ And there are no record rejections from the data_contract phase
+ When I run the business rules phase
+ Then there are errors with the following details and associated error_count from the business_rules phase
+ | ErrorType | ErrorCode | error_count |
+ | record | F1 | 1 |
+ | record | PG1 | 2 |
+ When I run the error report phase
+ Then An error report is produced
+ # TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
+ # And The statistics entry for the submission shows the following information
+ # | parameter | value |
+ # | record_count | 1 |
+ # | number_file_rejections | 0 |
+ # | number_record_rejections | 1 |
+
+ Scenario: A flights submission with a mixture of group and record rejections
+ Given I submit the flights file mixture_of_group_rej_and_bi_rej.xml for processing
+ And A duckdb pipeline is configured with schema file 'flights.dischema.json'
+ And I add initial audit entries for the submission
+ Then the latest audit record for the submission is marked with processing status file_transformation
+ When I run the file transformation phase
+ Then the country entity is stored as a parquet after the file_transformation phase
+ And the airport entity is stored as a parquet after the file_transformation phase
+ And the flights entity is stored as a parquet after the file_transformation phase
+ And the passengers entity is stored as a parquet after the file_transformation phase
+ And the latest audit record for the submission is marked with processing status data_contract
+ When I run the data contract phase
+ Then there are no file rejections from the data_contract phase
+ And there are no record rejections from the data_contract phase
+ When I run the business rules phase
+ Then there are errors with the following details and associated error_count from the business_rules phase
+ | ErrorType | ErrorCode | error_count |
+ | record | F1 | 1 |
+ | record | PG1 | 2 |
+ | record | P1 | 1 |
+ When I run the error report phase
+ Then An error report is produced
+ # TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
+ # And The statistics entry for the submission shows the following information
+ # | parameter | value |
+ # | record_count | 1 |
+ # | number_file_rejections | 0 |
+ # | number_record_rejections | 1 |
diff --git a/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py b/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py
index 35007a9..e55228d 100644
--- a/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py
+++ b/tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py
@@ -1,6 +1,7 @@
"""Test DuckDB backend steps."""
# pylint: disable=redefined-outer-name,unused-import,line-too-long
+import tempfile
from pathlib import Path
from typing import Iterator, List, Optional, Set, Tuple, Type
@@ -39,6 +40,9 @@
SemiJoin,
TableUnion,
)
+from dve.core_engine.configuration.v1.hierarchy import (
+ ChildHierarchyNode, EntityHierarchy, HierarchyNode
+)
from dve.core_engine.type_hints import MultipleExpressions
from tests.test_core_engine.test_backends.fixtures import (
duckdb_connection,
@@ -581,91 +585,106 @@ def test_header_multi_rows_raises(
DUCKDB_STEP_BACKEND.join_header(entities, config=header_join)
-def test_orphans_planets_satellites(
- planets_rel: DuckDBPyRelation, largest_satellites_rel: DuckDBPyRelation
-):
- """Test a basic orphan idenfitication from satellites to planets."""
- # Each satellite _must_ have a planet.
- join = OrphanIdentification(
- entity_name="satellites",
- target_name="planets",
- join_condition="satellites.planet == planets.planet",
- )
- entities = EntityManager(
- {
- "planets": planets_rel.filter(ColumnExpression("Planet") != ConstantExpression("Mars")),
- "satellites": largest_satellites_rel,
- }
- )
-
- DUCKDB_STEP_BACKEND.evaluate(entities, config=join)
- actual_rel = (
- entities["satellites"]
- .filter(ColumnExpression("IsOrphaned"))
- .select(ColumnExpression("name"))
- )
- actual_rows = sorted(actual_rel.df().to_dict(orient="records"), key=lambda row: row["name"])
-
- expected_rel = largest_satellites_rel.filter(
- ColumnExpression("Planet") == ConstantExpression("Mars")
- ).select(ColumnExpression("name"))
- expected_rows = sorted(expected_rel.df().to_dict(orient="records"), key=lambda row: row["name"])
-
- assert actual_rows == expected_rows
-
-
-def test_chained_orphans_planets_satellites(
- planets_rel: DuckDBPyRelation, largest_satellites_rel: DuckDBPyRelation
-):
- """Test a basic chained orphan idenfitication from satellites to planets."""
- join = OrphanIdentification(
- entity_name="satellites",
- target_name="planets",
- join_condition="satellites.planet == planets.planet",
- )
- entities = EntityManager(
- {
- "planets": planets_rel.filter(ColumnExpression("planet") != ConstantExpression("Mars")),
- "satellites": largest_satellites_rel,
- }
- )
- DUCKDB_STEP_BACKEND.evaluate(entities, config=join)
- entities["planets"] = planets_rel.filter(
- ColumnExpression("planet") != ConstantExpression("Earth")
- )
- DUCKDB_STEP_BACKEND.evaluate(entities, config=join)
-
- actual_rel = (
- entities["satellites"]
- .filter(ColumnExpression("IsOrphaned"))
- .select(ColumnExpression("name"))
- )
- actual_rows = sorted(actual_rel.df().to_dict(orient="records"), key=lambda row: row["name"])
-
- expected_rel = largest_satellites_rel.filter(
- ColumnExpression("Planet").isin(ConstantExpression("Mars"), ConstantExpression("Earth"))
- ).select(ColumnExpression("name"))
- expected_rows = sorted(expected_rel.df().to_dict(orient="records"), key=lambda row: row["name"])
-
- assert actual_rows == expected_rows
-
-
-def test_orphans_missing_entities_raises(
- planets_rel: DuckDBPyRelation, satellites_rel: DuckDBPyRelation
-):
- """Test that trying to join orphans from missing entities raises correctly."""
- join = OrphanIdentification(
- entity_name="planets",
- target_name="satellites",
- join_condition="planets.planet == satellites.planet",
- )
+class TestOrphanRecords:
+ """
+ Check that Orphan records identification and removal is working as expected.
- entities = EntityManager({"planets": planets_rel})
- with pytest.raises(MissingEntity):
- DUCKDB_STEP_BACKEND.identify_orphans(entities, config=join)
- entities = EntityManager({"satellites": satellites_rel})
- with pytest.raises(MissingEntity):
- DUCKDB_STEP_BACKEND.identify_orphans(entities, config=join)
+ Current scenarios are:
+ Flight ID 1 = Perfect Record - no orphans
+ Flight ID 2 = Record rejected at the flights entity, therefore, two expected orphans in passengers and food entities.
+ """
+ mod_flights_df = pl.DataFrame([
+ {'flight_id': 1, '__record_index__': 1},
+ ])
+ mod_passengers_df = pl.DataFrame([
+ {'flight_id': 1, 'passenger_id': 1, '__record_index__': 1},
+ {'flight_id': 1, 'passenger_id': 2, '__record_index__': 2},
+ {'flight_id': 2, 'passenger_id': 3, '__record_index__': 3},
+ ])
+ mod_food_df = pl.DataFrame([
+ {'passenger_id': 1, 'food_id': 1, '__record_index__': 1},
+ {'passenger_id': 1, 'food_id': 2, '__record_index__': 2},
+ {'passenger_id': 3, 'food_id': 3, '__record_index__': 3},
+ ])
+
+ def test_identify_orphan_record_single_entity(self):
+ """Ensure that a single one-to-one check works to identify orphan records."""
+ with duckdb.connect() as cnn:
+ cnn.register("mod_flights", self.mod_flights_df)
+ cnn.register("mod_passengers", self.mod_passengers_df)
+
+ mod_entities = EntityManager(
+ entities={
+ "flights": cnn.sql("SELECT * FROM mod_flights"),
+ "passengers": cnn.sql("SELECT * FROM mod_passengers"),
+ }
+ )
+
+ rules = DuckDBStepImplementations(connection=cnn)
+ _msgs = rules.identify_orphans(
+ mod_entities.entities,
+ config=OrphanIdentification(
+ id="flight_id",
+ entity_name="passengers",
+ target_name="flights",
+ join_condition="passengers.flight_id = flights.flight_id"
+ )
+ )
+ result = mod_entities["orphaned_records_tracker"]
+ assert result.count("*").fetchone()[0] == 1 # type: ignore
+ assert result.select("entity_name").unique("*").count("*").fetchone()[0] == 1 # type: ignore
+
+ def test_identify_and_remove_orphans(self):
+ with duckdb.connect() as cnn:
+ cnn.register("mod_flights", self.mod_flights_df)
+ cnn.register("mod_passengers", self.mod_passengers_df)
+ cnn.register("mod_food", self.mod_food_df)
+
+ mod_entities = EntityManager(
+ entities={
+ "flights": cnn.sql("SELECT * FROM mod_flights"),
+ "passengers": cnn.sql("SELECT * FROM mod_passengers"),
+ "food": cnn.sql("SELECT * FROM mod_food"),
+ }
+ )
+
+ rules = DuckDBStepImplementations(connection=cnn)
+ hierarchy = EntityHierarchy({
+ "flights": HierarchyNode(
+ entity_name="flights",
+ children=[
+ ChildHierarchyNode(
+ entity_name="passengers",
+ children=[
+ ChildHierarchyNode(
+ entity_name="food",
+ children=[],
+ join_fields={"passenger_id": "passenger_id"},
+ mandatory=False
+ )
+ ],
+ join_fields={"flight_id": "flight_id"},
+ mandatory=True
+ )
+ ]
+ )
+ })
+
+ with tempfile.TemporaryDirectory() as wd:
+ rules.identify_and_remove_orphans(
+ wd,
+ mod_entities.entities,
+ hierarchy
+ )
+
+ flights_rel = mod_entities["flights"]
+ assert flights_rel.select("__record_index__").unique("*").count("*").fetchone()[0] == 1 # type: ignore
+
+ passenger_rel = mod_entities["passengers"]
+ assert passenger_rel.select("__record_index__").unique("*").count("*").fetchone()[0] == 2 # type: ignore
+
+ food_rel = mod_entities["food"]
+ assert food_rel.select("__record_index__").unique("*").count("*").fetchone()[0] == 2 # type: ignore
def test_has_match_planets_satellites(
diff --git a/tests/test_core_engine/test_backends/test_implementations/test_spark/test_rules.py b/tests/test_core_engine/test_backends/test_implementations/test_spark/test_rules.py
index 673e611..654dc08 100644
--- a/tests/test_core_engine/test_backends/test_implementations/test_spark/test_rules.py
+++ b/tests/test_core_engine/test_backends/test_implementations/test_spark/test_rules.py
@@ -568,6 +568,7 @@ def test_header_multi_rows_raises(planets_df: DataFrame, value_literal_1_header:
SPARK_STEP_BACKEND.join_header(entities, config=header_join)
+@pytest.mark.skip(reason="Logic is no longer valid")
def test_orphans_planets_satellites(planets_df: DataFrame, largest_satellites_df: DataFrame):
"""Test a basic orphan idenfitication from satellites to planets."""
# Each satellite _must_ have a planet.
@@ -593,6 +594,7 @@ def test_orphans_planets_satellites(planets_df: DataFrame, largest_satellites_df
assert actual_rows == expected_rows
+@pytest.mark.skip(reason="Logic is no longer valid")
def test_chained_orphans_planets_satellites(
planets_df: DataFrame, largest_satellites_df: DataFrame
):
@@ -623,6 +625,7 @@ def test_chained_orphans_planets_satellites(
assert actual_rows == expected_rows
+@pytest.mark.skip(reason="Logic is no longer valid")
def test_orphans_missing_entities_raises(planets_df: DataFrame, satellites_df: DataFrame):
"""Test that trying to join orphans from missing entities raises correctly."""
join = OrphanIdentification(
diff --git a/tests/testdata/flights/flights.dischema.json b/tests/testdata/flights/flights.dischema.json
new file mode 100644
index 0000000..faac68f
--- /dev/null
+++ b/tests/testdata/flights/flights.dischema.json
@@ -0,0 +1,149 @@
+{
+ "contract": {
+ "schemas": {
+ "passengers": {
+ "fields": {
+ "flight_id": "int",
+ "passenger_id": "int",
+ "passenger_name": "str"
+ }
+ }
+ },
+ "datasets": {
+ "country": {
+ "fields": {
+ "country_id": "int",
+ "country_name": "str"
+ },
+ "reader_config": {
+ ".xml": {
+ "reader": "DuckDBXMLStreamReader",
+ "kwargs": {
+ "record_tag": "country",
+ "root_tag": "country"
+ }
+ }
+ },
+ "key_field": "country_id"
+ },
+ "airport": {
+ "fields": {
+ "country_id": "int",
+ "airport_id": "int",
+ "airport_name": "str",
+ "postcode": "str"
+ },
+ "reader_config": {
+ ".xml": {
+ "reader": "DuckDBXMLStreamReader",
+ "kwargs": {
+ "record_tag": "airport",
+ "root_tag": "country"
+ }
+ }
+ },
+ "key_field": "airport_id"
+ },
+ "flights": {
+ "fields": {
+ "airport_id": "int",
+ "flight_id": "int",
+ "destination": "str"
+ },
+ "reader_config": {
+ ".xml": {
+ "reader": "DuckDBXMLStreamReader",
+ "kwargs": {
+ "record_tag": "flight",
+ "root_tag": "country"
+ }
+ }
+ },
+ "key_field": "flight_id"
+ },
+ "passengers": {
+ "fields": {
+ "flight_id": "int",
+ "passenger_id": "int",
+ "passenger_name": "str"
+ },
+ "reader_config": {
+ ".xml": {
+ "reader": "DuckDBXMLStreamReader",
+ "kwargs": {
+ "record_tag": "passenger",
+ "root_tag": "country"
+ }
+ }
+ },
+ "key_field": "passenger_id"
+ }
+ }
+ },
+ "transformations": {
+ "filters": [
+ {
+ "entity": "country",
+ "name": "country_id_missing",
+ "expression": "country_id IS NOT NULL",
+ "failure_type": "record",
+ "failure_message": "Record Rejected - Country is missing an id",
+ "reporting_field": "country_id",
+ "reporting_entity": "country",
+ "category": "Blank",
+ "error_code": "C1"
+ },
+ {
+ "entity": "flights",
+ "name": "flight_missing_id",
+ "expression": "flight_id IS NOT NULL",
+ "failure_type": "record",
+ "failure_message": "Record Rejected - Flight is missing an id",
+ "reporting_field": "flight_id",
+ "reporting_entity": "flights",
+ "category": "Blank",
+ "error_code": "F1"
+ },
+ {
+ "entity": "passengers",
+ "name": "passenger_name_is_null",
+ "expression": "passenger_name IS NOT NULL",
+ "failure_type": "record",
+ "failure_message": "Record Rejected - Passenger Name is missing",
+ "reporting_field": "passenger_name",
+ "reporting_entity": "passengers",
+ "category": "Blank",
+ "error_code": "P1"
+ }
+ ]
+ },
+ "entity_relationships": {
+ "airport": {
+ "parent_entity": "country",
+ "join_fields": {
+ "country_id": "country_id"
+ },
+ "mandatory": true,
+ "orphaned_records_error_code": "AG1",
+ "orphaned_records_error_message": "Group rejected - No valid country group found country"
+ },
+ "flights": {
+ "parent_entity": "airport",
+ "join_fields": {
+ "airport_id": "airport_id"
+ },
+ "mandatory": false,
+ "orphaned_records_error_code": "FG1",
+ "orphaned_records_error_message": "Group rejected - No valid airport group found for airport"
+ },
+ "passengers": {
+ "parent_entity": "flights",
+ "join_fields": {
+ "flight_id": "flight_id"
+ },
+ "mandatory": false,
+ "orphaned_records_error_code": "PG1",
+ "orphaned_records_error_message": "Group rejected - No valid flight group found for passenger"
+ }
+ }
+}
\ No newline at end of file
diff --git a/tests/testdata/flights/missing_country_id.xml b/tests/testdata/flights/missing_country_id.xml
new file mode 100644
index 0000000..5c73640
--- /dev/null
+++ b/tests/testdata/flights/missing_country_id.xml
@@ -0,0 +1,48 @@
+
+
+ England
+
+
+ 1
+ 1
+ Heathrow
+ TW6 1EW
+
+
+ 1
+ 1
+ Paris
+
+
+ 1
+ 1
+ John
+
+
+ 1
+ 2
+ Jane
+
+
+
+
+ 1
+ 2
+ Madrid
+
+
+ 2
+ 3
+ Homer
+
+
+ 2
+ 4
+ Marge
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/tests/testdata/flights/missing_flight_id.xml b/tests/testdata/flights/missing_flight_id.xml
new file mode 100644
index 0000000..f9969f1
--- /dev/null
+++ b/tests/testdata/flights/missing_flight_id.xml
@@ -0,0 +1,48 @@
+
+
+ 1
+ England
+
+
+ 1
+ 1
+ Heathrow
+ TW6 1EW
+
+
+ 1
+ Paris
+
+
+ 1
+ 1
+ John
+
+
+ 1
+ 2
+ Jane
+
+
+
+
+ 1
+ 2
+ Madrid
+
+
+ 2
+ 3
+ Homer
+
+
+ 2
+ 4
+ Marge
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/tests/testdata/flights/mixture_of_group_rej_and_bi_rej.xml b/tests/testdata/flights/mixture_of_group_rej_and_bi_rej.xml
new file mode 100644
index 0000000..411b2da
--- /dev/null
+++ b/tests/testdata/flights/mixture_of_group_rej_and_bi_rej.xml
@@ -0,0 +1,53 @@
+
+
+ 1
+ England
+
+
+ 1
+ 1
+ Heathrow
+ TW6 1EW
+
+
+ 1
+ Paris
+
+
+ 1
+ 1
+ John
+
+
+ 1
+ 2
+ Jane
+
+
+
+
+ 1
+ 2
+ Madrid
+
+
+ 2
+ 3
+ Homer
+
+
+ 2
+ 4
+ Marge
+
+
+ 2
+ 5
+
+
+
+
+
+
+
+
\ No newline at end of file
diff --git a/tests/testdata/flights/perfect_flights.xml b/tests/testdata/flights/perfect_flights.xml
new file mode 100644
index 0000000..a581a6f
--- /dev/null
+++ b/tests/testdata/flights/perfect_flights.xml
@@ -0,0 +1,49 @@
+
+
+ 1
+ England
+
+
+ 1
+ 1
+ Heathrow
+ TW6 1EW
+
+
+ 1
+ 1
+ Paris
+
+
+ 1
+ 1
+ John
+
+
+ 1
+ 2
+ Jane
+
+
+
+
+ 1
+ 2
+ Madrid
+
+
+ 2
+ 3
+ Homer
+
+
+ 2
+ 4
+ Marge
+
+
+
+
+
+
+
\ No newline at end of file