From d55a63f75587989006ab2fdede33b41444cf1e29 Mon Sep 17 00:00:00 2001 From: georgeRobertson <50412379+georgeRobertson@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:40:45 +0100 Subject: [PATCH] feat: add orphan record identification and removal --- src/dve/core_engine/backends/base/rules.py | 122 +++++++++++- .../backends/implementations/duckdb/rules.py | 80 +++++--- .../backends/implementations/spark/rules.py | 17 +- .../core_engine/backends/metadata/rules.py | 8 +- .../core_engine/configuration/v1/hierarchy.py | 10 +- src/dve/core_engine/constants.py | 3 + src/dve/core_engine/type_hints.py | 2 +- src/dve/pipeline/pipeline.py | 15 +- tests/features/flights.feature | 116 +++++++++++ .../test_duckdb/test_rules.py | 187 ++++++++++-------- .../test_spark/test_rules.py | 3 + tests/testdata/flights/flights.dischema.json | 149 ++++++++++++++ tests/testdata/flights/missing_country_id.xml | 48 +++++ tests/testdata/flights/missing_flight_id.xml | 48 +++++ .../mixture_of_group_rej_and_bi_rej.xml | 53 +++++ tests/testdata/flights/perfect_flights.xml | 49 +++++ 16 files changed, 785 insertions(+), 125 deletions(-) create mode 100644 tests/features/flights.feature create mode 100644 tests/testdata/flights/flights.dischema.json create mode 100644 tests/testdata/flights/missing_country_id.xml create mode 100644 tests/testdata/flights/missing_flight_id.xml create mode 100644 tests/testdata/flights/mixture_of_group_rej_and_bi_rej.xml create mode 100644 tests/testdata/flights/perfect_flights.xml 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