Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 12 additions & 17 deletions src/dve/core_engine/backends/base/rules.py
Original file line number Diff line number Diff line change
Expand Up @@ -387,36 +387,34 @@ def identify_and_remove_orphans(

def process_node(
node: HierarchyNode,
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
"""Identify orphans and remove in a given node"""

if orph_messages is None:
orph_messages = []

if parent_entity_name is not None:
self.logger.info(f"Identifying orphans in {current_entity_name}")
if node.parent_entity is not None:
self.logger.info(f"Identifying orphans in {node.entity_name}")

join_expr = " AND ".join(
f"{parent_entity_name}.{k} = {current_entity_name}.{v}"
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=current_entity_name,
target_name=parent_entity_name,
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 {current_entity_name}"
f"Removing records with missing parent from {node.entity_name}"
)
location = list(node.join_fields.values())[0]
with BackgroundMessageWriter(
Expand All @@ -428,7 +426,7 @@ def process_node(
_orph_records = self.remove_orphans(
entities=entities,
config=OrphanRemoval(
entity_name=current_entity_name,
entity_name=node.entity_name,
reporting=ReportingConfig(
emit="record_failure",
code=node.missing_parent_id_error_code,
Expand All @@ -441,7 +439,7 @@ def process_node(
msg_writer.write_queue.put(
[
FeedbackMessage(
entity=current_entity_name,
entity=node.entity_name,
record=record, # type: ignore
error_location=location,
error_message=node.missing_parent_id_error_message,
Expand All @@ -455,12 +453,9 @@ def process_node(
]
)

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)
for tree in entity_hierarchy.entity_trees.values():
for node in tree.iterate_root_down():
process_node(node)

_orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME)
if _orph_rel:
Expand Down
31 changes: 25 additions & 6 deletions src/dve/core_engine/configuration/v1/hierarchy.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ class HierarchyNode(BaseModel):
"""Stores entity hierarchy information"""

entity_name: str
parent_entity: Optional[str] = None
children: list["HierarchyNode"] = Field(default_factory=list)
mandatory: bool = False
join_fields: dict[str, str] = Field(default_factory=dict)
Expand All @@ -26,14 +27,18 @@ class HierarchyNode(BaseModel):
"Records removed due to no valid parent record"
)

def get_descendents(self) -> list[str]:
def get_descendents(self) -> list["HierarchyNode"]:
"""Recursively list all descendents of the node"""
descendents = []
for node in self.children: # type: ignore
descendents.append(node.entity_name)
descendents.append(node)
descendents.extend(node.get_descendents())
return descendents

def get_descendent_names(self) -> list[str]:
"""Recursively list all names of descendents of the node"""
return [node.entity_name for node in self.get_descendents()]

def get_node(self, entity_name: str) -> Union["HierarchyNode", None]:
"""Recursively search for node and return if found"""
node = None
Expand Down Expand Up @@ -65,6 +70,20 @@ def as_dict(self) -> dict[str, dict[str, Any]]:

return {self.entity_name: ret_dict}

def _get_full_tree(self):
"""Get all nodes in tree, including the root"""
desc = self.get_descendents()
desc.insert(0, self)
return desc

def iterate_root_down(self):
"""Iterate through nodes from root to lowest descendent"""
yield from self._get_full_tree()

def iterate_lowest_descendent_up(self):
"""Iterate through nodes from lowest descendent to root"""
yield from self._get_full_tree()[::-1]


class EntityHierarchy:
"""Determines and stores entity hierarchy information from config"""
Expand All @@ -83,6 +102,7 @@ def determine_trees(
top_level_parents: dict[EntityName, HierarchyNode] = {
entity_name: HierarchyNode(
entity_name=entity_name,
parent_entity=None,
**config.model_dump(
exclude={
"parent_entity",
Expand All @@ -102,6 +122,7 @@ def determine_trees(
for entity_name in default_roots:
top_level_parents[entity_name] = HierarchyNode(
entity_name=entity_name,
parent_entity=None,
missing_parent_id_error_code=None,
missing_parent_id_error_message=None,
)
Expand All @@ -110,13 +131,11 @@ def determine_trees(
for main_entity, parent_node in top_level_parents.items():
if (
linkage_detail.parent_entity == main_entity
or linkage_detail.parent_entity in parent_node.get_descendents()
or linkage_detail.parent_entity in parent_node.get_descendent_names()
):
parent_node.add_child_node(
linkage_detail.parent_entity,
HierarchyNode(
entity_name=name, **linkage_detail.model_dump(exclude={"parent_entity"})
),
HierarchyNode(entity_name=name, **linkage_detail.model_dump()),
)
break
else:
Expand Down
2 changes: 1 addition & 1 deletion src/dve/pipeline/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ def load_reader(
if file_extension:
err_msg = (
f"The supplied file extension `{file_extension}`"
+f" is not a supported file format for {model_name}."
+ f" is not a supported file format for {model_name}."
)
else:
err_msg = "No supplied file extension. Unable to parse file without a file extension."
Expand Down
33 changes: 19 additions & 14 deletions tests/features/flights.feature
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ Feature: Pipeline tests using the flights dataset
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 staff 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
Expand All @@ -36,19 +37,20 @@ Feature: Pipeline tests using the flights dataset
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 staff 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
And there is 1 record rejection 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 |
| record | AG1 | 3 |
| record | SG1 | 15 |
| record | FG1 | 10 |
| record | PG1 | 25 |
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
Expand All @@ -66,6 +68,7 @@ Feature: Pipeline tests using the flights dataset
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 staff 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
Expand All @@ -76,10 +79,10 @@ Feature: Pipeline tests using the flights dataset
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 | PG1 | 3 |
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
#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 |
Expand All @@ -95,6 +98,7 @@ Feature: Pipeline tests using the flights dataset
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 staff 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
Expand All @@ -104,13 +108,14 @@ Feature: Pipeline tests using the flights dataset
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 | PG1 | 3 |
| record | P1 | 1 |
| record | S1 | 7 |
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 |
# 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 |
Original file line number Diff line number Diff line change
Expand Up @@ -656,9 +656,11 @@ def test_identify_and_remove_orphans(self):
children=[
HierarchyNode(
entity_name="passengers",
parent_entity="flights",
children=[
HierarchyNode(
entity_name="food",
parent_entity="passengers",
children=[],
join_fields={"passenger_id": "passenger_id"},
mandatory=False
Expand Down
11 changes: 9 additions & 2 deletions tests/test_core_engine/test_hierarchy.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,8 +260,9 @@ def test_linkage_config_load():
assert not children_001[0].children
assert children_001[1].entity_name == "ds_101"
assert dict_rep_001 == json.loads("""
{
{
"ds_001": {
"parent_entity": null,
"join_fields": {},
"mandatory": false,
"no_valid_records_error_code": "NoValidRecords",
Expand All @@ -270,6 +271,7 @@ def test_linkage_config_load():
"missing_parent_id_error_message": null,
"children": {
"ds_003": {
"parent_entity": "ds_001",
"join_fields": {
"ds_001_id": "ds_001_id"
},
Expand All @@ -281,6 +283,7 @@ def test_linkage_config_load():
"children": {}
},
"ds_101": {
"parent_entity": "ds_001",
"join_fields": {
"ds_001_id": "ds_001_id"
},
Expand All @@ -291,6 +294,7 @@ def test_linkage_config_load():
"missing_parent_id_error_message": "record removed as no parent",
"children": {
"ds_201": {
"parent_entity": "ds_101",
"join_fields": {
"referral_id": "ds_101_id"
},
Expand All @@ -301,6 +305,7 @@ def test_linkage_config_load():
"missing_parent_id_error_message": "record removed as no parent",
"children": {
"ds_202": {
"parent_entity": "ds_201",
"join_fields": {
"ds_201_id": "ds_201_id"
},
Expand All @@ -327,7 +332,7 @@ def test_linkage_config_load():
assert children_101[0].children[0].entity_name == "ds_202"
assert not children_101[0].children[0].children
assert dict_rep_101 == json.loads("""
{
{ "parent_entity": "ds_001",
"join_fields": {
"ds_001_id": "ds_001_id"
},
Expand All @@ -338,6 +343,7 @@ def test_linkage_config_load():
"missing_parent_id_error_message": "record removed as no parent",
"children": {
"ds_201": {
"parent_entity": "ds_101",
"join_fields": {
"referral_id": "ds_101_id"
},
Expand All @@ -348,6 +354,7 @@ def test_linkage_config_load():
"missing_parent_id_error_message": "record removed as no parent",
"children": {
"ds_202": {
"parent_entity": "ds_201",
"join_fields": {
"ds_201_id": "ds_201_id"
},
Expand Down
Loading
Loading