From e10b7972d3ba217efa476b987a7ea6173b48c43e Mon Sep 17 00:00:00 2001 From: georgeRobertson <50412379+georgeRobertson@users.noreply.github.com> Date: Wed, 30 Sep 2026 10:27:41 +0100 Subject: [PATCH] perf: remove orphan tracker as it's causing significant performance degradation --- src/dve/core_engine/backends/base/rules.py | 108 ++++++------------ .../implementations/duckdb/duckdb_helpers.py | 6 +- .../backends/implementations/duckdb/rules.py | 49 +++----- .../backends/implementations/spark/rules.py | 10 -- src/dve/core_engine/constants.py | 3 - tests/features/steps/steps_post_pipeline.py | 3 +- .../test_duckdb/test_rules.py | 7 +- 7 files changed, 57 insertions(+), 129 deletions(-) diff --git a/src/dve/core_engine/backends/base/rules.py b/src/dve/core_engine/backends/base/rules.py index 8508189..dbe2d9d 100644 --- a/src/dve/core_engine/backends/base/rules.py +++ b/src/dve/core_engine/backends/base/rules.py @@ -17,7 +17,6 @@ ) 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, @@ -36,7 +35,6 @@ Notification, OneToOneJoin, OrphanIdentification, - OrphanRemoval, ParentMetadata, RenameEntity, Rule, @@ -47,7 +45,6 @@ ) from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful from dve.core_engine.configuration.v1.hierarchy import 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 @@ -317,7 +314,7 @@ def join_header(self, entities: Entities, *, config: HeaderJoin) -> Messages: @abstractmethod def identify_orphans( self, entities: Entities, *, config: OrphanIdentification - ) -> tuple[Messages, int]: + ) -> Iterable: """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. @@ -330,18 +327,6 @@ def identify_orphans( """ 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 check_mandatory_group(self, entities: Entities, *, config: GroupIdentification) -> Iterator: """ @@ -395,72 +380,53 @@ def identify_and_remove_orphans( Processes recursively: removes orphans at each level, then processes children. """ - def process_node(node: HierarchyNode): + def process_node(node: HierarchyNode) -> bool: """Identify orphans and remove in a given node""" issues_found: bool = False if node.parent_entity is None: return issues_found - self.logger.info(f"Identifying orphans in {node.entity_name}") + self.logger.info(f"Checking for orphan records in {node.entity_name}") join_expr = " AND ".join( f"{node.parent_entity}.{k} = {node.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=node.entity_name, - target_name=node.parent_entity, - join_condition=join_expr, - ), - ) - - if no_orphs > 0: - self.logger.info(f"Removing records with missing parent from {node.entity_name}") - issues_found = True - 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=node.entity_name, - reporting=ReportingConfig( - emit="record_failure", - code=node.missing_parent_id_error_code, - message=node.missing_parent_id_error_message, - location=location, - ), + 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.identify_orphans( + entities=entities, + config=OrphanIdentification( + id=list(node.join_fields.values())[0], + entity_name=node.entity_name, + target_name=node.parent_entity, + join_condition=join_expr, + ), + ) + _messages = [ + FeedbackMessage( + entity=node.entity_name, + record=record, # type: ignore + error_location=location, + error_message=template_object( + node.missing_parent_id_error_message, record ), + failure_type="record", + error_type="record", + error_code=node.missing_parent_id_error_code, + reporting_field=location, + category="Parent Missing", ) - # moved to batch the write - risky if large number of - msg_writer.write_queue.put( - [ - FeedbackMessage( - entity=node.entity_name, - record=record, # type: ignore - error_location=location, - error_message=template_object( - node.missing_parent_id_error_message, record - ), - failure_type="record", - error_type="record", - error_code=node.missing_parent_id_error_code, - reporting_field=location, - category="Parent Missing", - ) - for record in _orph_records - ] - ) + for record in _orph_records + ] + msg_writer.write_queue.put(_messages) - return issues_found + return len(_messages) > 0 entity_issues_found: dict[EntityName, bool] = {} @@ -468,10 +434,6 @@ def process_node(node: HierarchyNode): for node in tree.iterate_root_down(): entity_issues_found[node.entity_name] = process_node(node) - _orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME) - if _orph_rel is not None: - del entities[ORPHANED_RECORD_ENTITY_NAME] - entities.update(entities) return [], entity_issues_found diff --git a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py index d9cd149..97ad3d3 100644 --- a/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py +++ b/src/dve/core_engine/backends/implementations/duckdb/duckdb_helpers.py @@ -3,6 +3,7 @@ """Helper objects for duckdb data contract implementation""" +import itertools from collections.abc import Generator, Iterator from dataclasses import is_dataclass from datetime import date, datetime, time @@ -99,8 +100,11 @@ def __call__(self): def table_exists(connection: DuckDBPyConnection, table_name: str) -> bool: """check if a table exists in a given DuckDBPyConnection""" - return table_name in map(lambda x: x[0], connection.sql("SHOW TABLES").fetchall()) + return table_name in get_all_existing_ddb_tables(connection) +def get_all_existing_ddb_tables(connection: DuckDBPyConnection) -> tuple[str]: + """Fetch all tables available ina given duckdb connection""" + return tuple(itertools.chain.from_iterable(connection.sql("SHOW TABLES").fetchall())) def relation_is_empty(relation: DuckDBPyRelation) -> bool: """Check if a duckdb relation is empty""" diff --git a/src/dve/core_engine/backends/implementations/duckdb/rules.py b/src/dve/core_engine/backends/implementations/duckdb/rules.py index 08b0d2d..edb7d71 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""" # pylint: disable=R0801 -from collections.abc import Callable, Iterator +from collections.abc import Callable, Iterable, Iterator from typing import get_type_hints from uuid import uuid4 @@ -53,11 +53,10 @@ Notification, OneToOneJoin, OrphanIdentification, - OrphanRemoval, SemiJoin, TableUnion, ) -from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME, RECORD_INDEX_COLUMN_NAME +from dve.core_engine.constants import 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 @@ -385,7 +384,7 @@ def identify_orphans( entities: DuckDBEntities, *, config: OrphanIdentification, - ) -> tuple[Messages, int]: + ) -> Iterable: """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. @@ -401,56 +400,34 @@ def identify_orphans( if relation_is_empty(source_rel): self.logger.info(f"{config.entity_name} is empty. Skipping orphan check.") - return [], 0 + return [] match_name = f"matched_{uuid4().hex}" target_rel = target_rel.select( StarExpression(exclude=[]), ConstantExpression(1).alias(match_name) ).set_alias(config.target_name) - pk, _fk = config.join_condition.split("=") - 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("*") + .select(RECORD_INDEX_COLUMN_NAME) + .set_alias("orphan") ) - _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] - 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 - else: - _no_orphans = 0 - self.logger.info(f"Found {_no_orphans} orphaned records in {config.entity_name}.") - return [], _no_orphans + if relation_is_empty(orphaned_rel): + self.logger.info( + f"Found 0 orphan records between {config.entity_name} and {config.target_name}" + ) + return [] - 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] - .filter(f"entity_name = '{config.entity_name}'") - .set_alias("orphan") - ) message_rel = ( entities[config.entity_name] .set_alias(config.entity_name) .join( - orphan_rel, + orphaned_rel, f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME} = orphan.{RECORD_INDEX_COLUMN_NAME}", # pylint: disable=C0301 "semi", ) @@ -459,7 +436,7 @@ def remove_orphans(self, entities: DuckDBEntities, *, config: OrphanRemoval) -> entities[config.entity_name] .set_alias(config.entity_name) .join( - orphan_rel, + orphaned_rel, f"{config.entity_name}.{RECORD_INDEX_COLUMN_NAME} = orphan.{RECORD_INDEX_COLUMN_NAME}", # pylint: disable=C0301 "anti", ) diff --git a/src/dve/core_engine/backends/implementations/spark/rules.py b/src/dve/core_engine/backends/implementations/spark/rules.py index 717151e..84df52b 100644 --- a/src/dve/core_engine/backends/implementations/spark/rules.py +++ b/src/dve/core_engine/backends/implementations/spark/rules.py @@ -43,7 +43,6 @@ Notification, OneToOneJoin, OrphanIdentification, - OrphanRemoval, SelectColumns, SemiJoin, TableUnion, @@ -378,15 +377,6 @@ def identify_orphans( entities[config.new_entity_name or config.entity_name] = result return [], 0 - def remove_orphans( - self, - entities: SparkEntities, - *, - config: OrphanRemoval, - ) -> Iterator: - # TODO - implement for spark - raise NotImplementedError - def check_mandatory_group( self, entities: SparkEntities, *, config: GroupIdentification ) -> Iterator: diff --git a/src/dve/core_engine/constants.py b/src/dve/core_engine/constants.py index 3581088..a2a4a65 100644 --- a/src/dve/core_engine/constants.py +++ b/src/dve/core_engine/constants.py @@ -6,6 +6,3 @@ 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_record_tracker" -"""Name of entity to keep track of records where there is a missing parent record""" diff --git a/tests/features/steps/steps_post_pipeline.py b/tests/features/steps/steps_post_pipeline.py index b3663a6..b284b3c 100644 --- a/tests/features/steps/steps_post_pipeline.py +++ b/tests/features/steps/steps_post_pipeline.py @@ -129,4 +129,5 @@ def check_entity_row_counts(context: Context): record = row.as_dict() entity_name = record["entity_name"] expected_count = int(record["row_count"]) - assert expected_count == read_output_parquet(processing_loc, entity_name, "business_rules").shape[0] \ No newline at end of file + output_df = read_output_parquet(processing_loc, entity_name, "business_rules") + assert expected_count == output_df.shape[0], output_df 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 10480a1..ecef834 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 @@ -40,7 +40,6 @@ SemiJoin, TableUnion, ) -from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME from dve.core_engine.configuration.v1.hierarchy import ( EntityHierarchy, HierarchyNode ) @@ -622,7 +621,7 @@ def test_identify_orphan_record_single_entity(self): ) rules = DuckDBStepImplementations(connection=cnn) - _msgs = rules.identify_orphans( + msgs = rules.identify_orphans( mod_entities.entities, config=OrphanIdentification( id="flight_id", @@ -631,9 +630,7 @@ def test_identify_orphan_record_single_entity(self): join_condition="passengers.flight_id = flights.flight_id" ) ) - result = mod_entities[ORPHANED_RECORD_ENTITY_NAME] - assert result.count("*").fetchone()[0] == 1 # type: ignore - assert result.select("entity_name").unique("*").count("*").fetchone()[0] == 1 # type: ignore + assert len(list(msgs)) == 1 def test_identify_and_remove_orphans(self): with duckdb.connect() as cnn: