Skip to content

Commit a6b6154

Browse files
stevenhsdstevenhsd
authored
refactor: tweak hierarchy node design to allow iteration both directions on hierarchies (#155)
1 parent c2b8cf5 commit a6b6154

12 files changed

Lines changed: 1220 additions & 69 deletions

File tree

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

Lines changed: 12 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -387,36 +387,34 @@ def identify_and_remove_orphans(
387387

388388
def process_node(
389389
node: HierarchyNode,
390-
parent_entity_name: Optional[EntityName],
391390
orph_messages: Messages | None = None,
392391
):
393-
"""Recursive helper to process a node and its children."""
394-
current_entity_name = node.entity_name
392+
"""Identify orphans and remove in a given node"""
395393

396394
if orph_messages is None:
397395
orph_messages = []
398396

399-
if parent_entity_name is not None:
400-
self.logger.info(f"Identifying orphans in {current_entity_name}")
397+
if node.parent_entity is not None:
398+
self.logger.info(f"Identifying orphans in {node.entity_name}")
401399

402400
join_expr = " AND ".join(
403-
f"{parent_entity_name}.{k} = {current_entity_name}.{v}"
401+
f"{node.parent_entity}.{k} = {node.entity_name}.{v}"
404402
for k, v in node.join_fields.items()
405403
)
406404

407405
_, no_orphs = self.identify_orphans(
408406
entities=entities,
409407
config=OrphanIdentification(
410408
id=list(node.join_fields.values())[0],
411-
entity_name=current_entity_name,
412-
target_name=parent_entity_name,
409+
entity_name=node.entity_name,
410+
target_name=node.parent_entity,
413411
join_condition=join_expr,
414412
),
415413
)
416414

417415
if no_orphs > 0:
418416
self.logger.info(
419-
f"Removing records with missing parent from {current_entity_name}"
417+
f"Removing records with missing parent from {node.entity_name}"
420418
)
421419
location = list(node.join_fields.values())[0]
422420
with BackgroundMessageWriter(
@@ -428,7 +426,7 @@ def process_node(
428426
_orph_records = self.remove_orphans(
429427
entities=entities,
430428
config=OrphanRemoval(
431-
entity_name=current_entity_name,
429+
entity_name=node.entity_name,
432430
reporting=ReportingConfig(
433431
emit="record_failure",
434432
code=node.missing_parent_id_error_code,
@@ -441,7 +439,7 @@ def process_node(
441439
msg_writer.write_queue.put(
442440
[
443441
FeedbackMessage(
444-
entity=current_entity_name,
442+
entity=node.entity_name,
445443
record=record, # type: ignore
446444
error_location=location,
447445
error_message=node.missing_parent_id_error_message,
@@ -455,12 +453,9 @@ def process_node(
455453
]
456454
)
457455

458-
if node.children:
459-
for child_node in node.children:
460-
process_node(child_node, current_entity_name, orph_messages)
461-
462-
for root_node in entity_hierarchy.entity_trees.values():
463-
process_node(root_node, parent_entity_name=None)
456+
for tree in entity_hierarchy.entity_trees.values():
457+
for node in tree.iterate_root_down():
458+
process_node(node)
464459

465460
_orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME)
466461
if _orph_rel:

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

Lines changed: 25 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ class HierarchyNode(BaseModel):
1616
"""Stores entity hierarchy information"""
1717

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

29-
def get_descendents(self) -> list[str]:
30+
def get_descendents(self) -> list["HierarchyNode"]:
3031
"""Recursively list all descendents of the node"""
3132
descendents = []
3233
for node in self.children: # type: ignore
33-
descendents.append(node.entity_name)
34+
descendents.append(node)
3435
descendents.extend(node.get_descendents())
3536
return descendents
3637

38+
def get_descendent_names(self) -> list[str]:
39+
"""Recursively list all names of descendents of the node"""
40+
return [node.entity_name for node in self.get_descendents()]
41+
3742
def get_node(self, entity_name: str) -> Union["HierarchyNode", None]:
3843
"""Recursively search for node and return if found"""
3944
node = None
@@ -65,6 +70,20 @@ def as_dict(self) -> dict[str, dict[str, Any]]:
6570

6671
return {self.entity_name: ret_dict}
6772

73+
def _get_full_tree(self):
74+
"""Get all nodes in tree, including the root"""
75+
desc = self.get_descendents()
76+
desc.insert(0, self)
77+
return desc
78+
79+
def iterate_root_down(self):
80+
"""Iterate through nodes from root to lowest descendent"""
81+
yield from self._get_full_tree()
82+
83+
def iterate_lowest_descendent_up(self):
84+
"""Iterate through nodes from lowest descendent to root"""
85+
yield from self._get_full_tree()[::-1]
86+
6887

6988
class EntityHierarchy:
7089
"""Determines and stores entity hierarchy information from config"""
@@ -83,6 +102,7 @@ def determine_trees(
83102
top_level_parents: dict[EntityName, HierarchyNode] = {
84103
entity_name: HierarchyNode(
85104
entity_name=entity_name,
105+
parent_entity=None,
86106
**config.model_dump(
87107
exclude={
88108
"parent_entity",
@@ -102,6 +122,7 @@ def determine_trees(
102122
for entity_name in default_roots:
103123
top_level_parents[entity_name] = HierarchyNode(
104124
entity_name=entity_name,
125+
parent_entity=None,
105126
missing_parent_id_error_code=None,
106127
missing_parent_id_error_message=None,
107128
)
@@ -110,13 +131,11 @@ def determine_trees(
110131
for main_entity, parent_node in top_level_parents.items():
111132
if (
112133
linkage_detail.parent_entity == main_entity
113-
or linkage_detail.parent_entity in parent_node.get_descendents()
134+
or linkage_detail.parent_entity in parent_node.get_descendent_names()
114135
):
115136
parent_node.add_child_node(
116137
linkage_detail.parent_entity,
117-
HierarchyNode(
118-
entity_name=name, **linkage_detail.model_dump(exclude={"parent_entity"})
119-
),
138+
HierarchyNode(entity_name=name, **linkage_detail.model_dump()),
120139
)
121140
break
122141
else:

‎src/dve/pipeline/utils.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ def load_reader(
6060
if file_extension:
6161
err_msg = (
6262
f"The supplied file extension `{file_extension}`"
63-
+f" is not a supported file format for {model_name}."
63+
+ f" is not a supported file format for {model_name}."
6464
)
6565
else:
6666
err_msg = "No supplied file extension. Unable to parse file without a file extension."

‎tests/features/flights.feature‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ Feature: Pipeline tests using the flights dataset
1010
When I run the file transformation phase
1111
Then the country entity is stored as a parquet after the file_transformation phase
1212
And the airport entity is stored as a parquet after the file_transformation phase
13+
And the staff entity is stored as a parquet after the file_transformation phase
1314
And the flights entity is stored as a parquet after the file_transformation phase
1415
And the passengers entity is stored as a parquet after the file_transformation phase
1516
And the latest audit record for the submission is marked with processing status data_contract
@@ -36,19 +37,20 @@ Feature: Pipeline tests using the flights dataset
3637
When I run the file transformation phase
3738
Then the country entity is stored as a parquet after the file_transformation phase
3839
And the airport entity is stored as a parquet after the file_transformation phase
40+
And the staff entity is stored as a parquet after the file_transformation phase
3941
And the flights entity is stored as a parquet after the file_transformation phase
4042
And the passengers entity is stored as a parquet after the file_transformation phase
4143
And the latest audit record for the submission is marked with processing status data_contract
4244
When I run the data contract phase
4345
Then there are no file rejections from the data_contract phase
44-
And there are no record rejections from the data_contract phase
46+
And there is 1 record rejection from the data_contract phase
4547
When I run the business rules phase
4648
Then there are errors with the following details and associated error_count from the business_rules phase
4749
| ErrorType | ErrorCode | error_count |
48-
| record | C1 | 1 |
49-
| record | AG1 | 1 |
50-
| record | FG1 | 2 |
51-
| record | PG1 | 4 |
50+
| record | AG1 | 3 |
51+
| record | SG1 | 15 |
52+
| record | FG1 | 10 |
53+
| record | PG1 | 25 |
5254
When I run the error report phase
5355
Then An error report is produced
5456
# TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
@@ -66,6 +68,7 @@ Feature: Pipeline tests using the flights dataset
6668
When I run the file transformation phase
6769
Then the country entity is stored as a parquet after the file_transformation phase
6870
And the airport entity is stored as a parquet after the file_transformation phase
71+
And the staff entity is stored as a parquet after the file_transformation phase
6972
And the flights entity is stored as a parquet after the file_transformation phase
7073
And the passengers entity is stored as a parquet after the file_transformation phase
7174
And the latest audit record for the submission is marked with processing status data_contract
@@ -76,10 +79,10 @@ Feature: Pipeline tests using the flights dataset
7679
Then there are errors with the following details and associated error_count from the business_rules phase
7780
| ErrorType | ErrorCode | error_count |
7881
| record | F1 | 1 |
79-
| record | PG1 | 2 |
82+
| record | PG1 | 3 |
8083
When I run the error report phase
8184
Then An error report is produced
82-
# TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
85+
#TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
8386
# And The statistics entry for the submission shows the following information
8487
# | parameter | value |
8588
# | record_count | 1 |
@@ -95,6 +98,7 @@ Feature: Pipeline tests using the flights dataset
9598
Then the country entity is stored as a parquet after the file_transformation phase
9699
And the airport entity is stored as a parquet after the file_transformation phase
97100
And the flights entity is stored as a parquet after the file_transformation phase
101+
And the staff entity is stored as a parquet after the file_transformation phase
98102
And the passengers entity is stored as a parquet after the file_transformation phase
99103
And the latest audit record for the submission is marked with processing status data_contract
100104
When I run the data contract phase
@@ -104,13 +108,14 @@ Feature: Pipeline tests using the flights dataset
104108
Then there are errors with the following details and associated error_count from the business_rules phase
105109
| ErrorType | ErrorCode | error_count |
106110
| record | F1 | 1 |
107-
| record | PG1 | 2 |
111+
| record | PG1 | 3 |
108112
| record | P1 | 1 |
113+
| record | S1 | 7 |
109114
When I run the error report phase
110115
Then An error report is produced
111-
# TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
112-
# And The statistics entry for the submission shows the following information
113-
# | parameter | value |
114-
# | record_count | 1 |
115-
# | number_file_rejections | 0 |
116-
# | number_record_rejections | 1 |
116+
# TODO - fix the stats calculations as they're currently incorrect for hiearchical datasets
117+
# And The statistics entry for the submission shows the following information
118+
# | parameter | value |
119+
# | record_count | 1 |
120+
# | number_file_rejections | 0 |
121+
# | number_record_rejections | 1 |

‎tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -656,9 +656,11 @@ def test_identify_and_remove_orphans(self):
656656
children=[
657657
HierarchyNode(
658658
entity_name="passengers",
659+
parent_entity="flights",
659660
children=[
660661
HierarchyNode(
661662
entity_name="food",
663+
parent_entity="passengers",
662664
children=[],
663665
join_fields={"passenger_id": "passenger_id"},
664666
mandatory=False

‎tests/test_core_engine/test_hierarchy.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -260,8 +260,9 @@ def test_linkage_config_load():
260260
assert not children_001[0].children
261261
assert children_001[1].entity_name == "ds_101"
262262
assert dict_rep_001 == json.loads("""
263-
{
263+
{
264264
"ds_001": {
265+
"parent_entity": null,
265266
"join_fields": {},
266267
"mandatory": false,
267268
"no_valid_records_error_code": "NoValidRecords",
@@ -270,6 +271,7 @@ def test_linkage_config_load():
270271
"missing_parent_id_error_message": null,
271272
"children": {
272273
"ds_003": {
274+
"parent_entity": "ds_001",
273275
"join_fields": {
274276
"ds_001_id": "ds_001_id"
275277
},
@@ -281,6 +283,7 @@ def test_linkage_config_load():
281283
"children": {}
282284
},
283285
"ds_101": {
286+
"parent_entity": "ds_001",
284287
"join_fields": {
285288
"ds_001_id": "ds_001_id"
286289
},
@@ -291,6 +294,7 @@ def test_linkage_config_load():
291294
"missing_parent_id_error_message": "record removed as no parent",
292295
"children": {
293296
"ds_201": {
297+
"parent_entity": "ds_101",
294298
"join_fields": {
295299
"referral_id": "ds_101_id"
296300
},
@@ -301,6 +305,7 @@ def test_linkage_config_load():
301305
"missing_parent_id_error_message": "record removed as no parent",
302306
"children": {
303307
"ds_202": {
308+
"parent_entity": "ds_201",
304309
"join_fields": {
305310
"ds_201_id": "ds_201_id"
306311
},
@@ -327,7 +332,7 @@ def test_linkage_config_load():
327332
assert children_101[0].children[0].entity_name == "ds_202"
328333
assert not children_101[0].children[0].children
329334
assert dict_rep_101 == json.loads("""
330-
{
335+
{ "parent_entity": "ds_001",
331336
"join_fields": {
332337
"ds_001_id": "ds_001_id"
333338
},
@@ -338,6 +343,7 @@ def test_linkage_config_load():
338343
"missing_parent_id_error_message": "record removed as no parent",
339344
"children": {
340345
"ds_201": {
346+
"parent_entity": "ds_101",
341347
"join_fields": {
342348
"referral_id": "ds_101_id"
343349
},
@@ -348,6 +354,7 @@ def test_linkage_config_load():
348354
"missing_parent_id_error_message": "record removed as no parent",
349355
"children": {
350356
"ds_202": {
357+
"parent_entity": "ds_201",
351358
"join_fields": {
352359
"ds_201_id": "ds_201_id"
353360
},

0 commit comments

Comments
 (0)