From d910010db1bd012546ceed1d3046b5cbd5bea435 Mon Sep 17 00:00:00 2001 From: Libin Liang Date: Thu, 9 Jul 2026 15:05:01 -0400 Subject: [PATCH 1/4] - Add table cell hierarchy tree construction. --- kensho_kenverters/CHANGELOG.md | 6 + kensho_kenverters/constants.py | 2 + kensho_kenverters/output_to_tables.py | 181 +++++++++ .../tests/test_output_to_tables.py | 349 ++++++++++++++++++ 4 files changed, 538 insertions(+) diff --git a/kensho_kenverters/CHANGELOG.md b/kensho_kenverters/CHANGELOG.md index 276eaf8..29394a5 100644 --- a/kensho_kenverters/CHANGELOG.md +++ b/kensho_kenverters/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## v3.1.2 + +### Added + +* Add codes to construct table cell hierarchy tree. + ## v3.1.1 ### Added diff --git a/kensho_kenverters/constants.py b/kensho_kenverters/constants.py index 76edf0e..e546316 100644 --- a/kensho_kenverters/constants.py +++ b/kensho_kenverters/constants.py @@ -98,3 +98,5 @@ class ContentCategory(Enum): RELATIONS_BETWEEN_ITEMS = {"support"} EMPTY_STRING = "" + +ROW_KEY_PARENT_RELATION = "row_key_parent" diff --git a/kensho_kenverters/output_to_tables.py b/kensho_kenverters/output_to_tables.py index 22baa00..d56f8b7 100644 --- a/kensho_kenverters/output_to_tables.py +++ b/kensho_kenverters/output_to_tables.py @@ -9,6 +9,7 @@ from .constants import ( EMPTY_STRING, + ROW_KEY_PARENT_RELATION, TABLE_CONTENT_CATEGORIES, AnnotationType, ContentCategory, @@ -19,8 +20,10 @@ ContentModel, LocationModel, LocationType, + RelationAnnotationModel, Table, TableCategoryType, + TableCellHierarchyTreeModel, TableGridAndStructure, TableStructureAnnotationModel, ) @@ -253,6 +256,184 @@ def convert_uid_grid_to_content_grid( return content_grid +# --------- Table Hierarchy Tree --------- + + +def _build_table_cell_hierarchy_tree_node( + cell_uid: str, + cell_uid_to_annotation: dict[str, TableStructureAnnotationModel], + parent_to_children: dict[str, list[str]], + all_projected_row_header_uids: set[str], + row_index_to_annotations: dict[int, list[TableStructureAnnotationModel]], +) -> TableCellHierarchyTreeModel: + """Recursively build a hierarchy tree node for a projected row header cell. + + Args: + cell_uid: the uid of the projected row header cell to build a node for. + cell_uid_to_annotation: mapping from cell uid to its table structure annotation. + parent_to_children: mapping from parent uid to its children uids + (from row_key_parent relations). + all_projected_row_header_uids: set of all projected row header uids across all tables. + row_index_to_annotations: mapping from row index to all cell annotations in that row. + + Returns: + a TableCellHierarchyTreeModel node for the given cell uid. + """ + children_uids = parent_to_children.get(cell_uid, []) + + # Separate children into projected row headers (tree children) and + # data row cells (contents) + child_nodes: list[TableCellHierarchyTreeModel] = [] + content_annotations: list[TableStructureAnnotationModel] = [] + for child_uid in children_uids: + # If the child is a projected header, make a new child node + if child_uid in all_projected_row_header_uids: + child_nodes.append( + _build_table_cell_hierarchy_tree_node( + child_uid, + cell_uid_to_annotation, + parent_to_children, + all_projected_row_header_uids, + row_index_to_annotations, + ) + ) + else: + # If the child is a regular row header, assign all annotations in the row + # to the contents + child_annotation = cell_uid_to_annotation.get(child_uid) + if child_annotation: + row_index = child_annotation.data.index[0] + content_annotations.extend(row_index_to_annotations.get(row_index, [])) + + return TableCellHierarchyTreeModel( + node_uid=cell_uid, + node_type=ContentCategory.TABLE_CELL.value, + children=child_nodes, + contents=content_annotations, + ) + + +def _get_table_uid_to_table_cell_hierarchy_tree( + table_uid_to_cells_mapping: dict[str, list[ContentModel]], + table_cell_annotations: list[TableStructureAnnotationModel], + relation_annotations: list[RelationAnnotationModel], +) -> dict[str, TableCellHierarchyTreeModel]: + """Build a TableCellHierarchyTreeModel for each table uid. + + The root node represents the table itself. Its children are the top-level projected row + header cells. Each projected row header node's children are its sub-level projected row + headers (from row_key_parent relations), and its contents are the cell annotations for the + leftmost cells (regular row headers) that belong to that projected row header and the cell + annotations in the same row. + + Args: + table_uid_to_cells_mapping: mapping of table uid to cells (ContentModel) in that table. + table_cell_annotations: list of table structure annotations. + relation_annotations: list of relation annotations (row_key_parent relations define + the hierarchy). + + Returns: + a mapping of table uid to the TableCellHierarchyTreeModel representing the table hierarchy. + """ + # Build mapping from cell uid to its table uid + cell_uid_to_table_uid: dict[str, str] = {} + for table_uid, cells in table_uid_to_cells_mapping.items(): + for cell in cells: + cell_uid_to_table_uid[cell.uid] = table_uid + + # Build mapping from cell uid to its annotation + cell_uid_to_annotation: dict[str, TableStructureAnnotationModel] = {} + for annotation in table_cell_annotations: + for uid in annotation.content_uids: + cell_uid_to_annotation[uid] = annotation + + # Extract row_key_parent relations and group by table uid + # In row_key_parent: source is parent, target is child + parent_to_children: dict[str, list[str]] = defaultdict(list) + child_uids: set[str] = set() + for relation in relation_annotations: + if relation.data.relation_type == ROW_KEY_PARENT_RELATION: + parent_uid = relation.data.source_content_uid + child_uid = relation.data.target_content_uid + parent_to_children[parent_uid].append(child_uid) + child_uids.add(child_uid) + + # Identify projected row header uids per table + table_uid_to_projected_row_header_uids: dict[str, list[str]] = defaultdict(list) + for table_uid, cells in table_uid_to_cells_mapping.items(): + for cell in cells: + cell_ann = cell_uid_to_annotation.get(cell.uid) + if cell_ann and cell_ann.data.is_projected_row_header: + table_uid_to_projected_row_header_uids[table_uid].append(cell.uid) + + # Set of all projected row header uids for quick lookup + all_projected_row_header_uids: set[str] = set() + for uids in table_uid_to_projected_row_header_uids.values(): + all_projected_row_header_uids.update(uids) + + # Build the tree for each table + result: dict[str, TableCellHierarchyTreeModel] = {} + for table_uid in table_uid_to_cells_mapping: + # Build row_index_to_annotations for this table + row_index_to_annotations: dict[int, list[TableStructureAnnotationModel]] = ( + defaultdict(list) + ) + for cell in table_uid_to_cells_mapping[table_uid]: + cell_ann = cell_uid_to_annotation.get(cell.uid) + if cell_ann: + row_index_to_annotations[cell_ann.data.index[0]].append(cell_ann) + + projected_uids = table_uid_to_projected_row_header_uids.get(table_uid, []) + # Top-level projected row headers are those that are not children of any other + top_level_uids = [uid for uid in projected_uids if uid not in child_uids] + + # Build child nodes for top-level projected row headers + top_level_nodes = [ + _build_table_cell_hierarchy_tree_node( + uid, + cell_uid_to_annotation, + parent_to_children, + all_projected_row_header_uids, + row_index_to_annotations, + ) + for uid in top_level_uids + ] + + # Collect row indices already represented in the hierarchy tree + # (either as children of a projected row header, or as projected row headers + # themselves). These rows will be excluded from the table root's contents, + # since they already appear as nodes or contents within the tree. + excluded_row_indices: set[int] = set() + for child_uid in child_uids: + child_ann = cell_uid_to_annotation.get(child_uid) + if child_ann and cell_uid_to_table_uid.get(child_uid) == table_uid: + excluded_row_indices.add(child_ann.data.index[0]) + for uid in projected_uids: + proj_ann = cell_uid_to_annotation.get(uid) + if proj_ann: + excluded_row_indices.add(proj_ann.data.index[0]) + + # Remaining rows (not excluded, not column headers) belong to the table + table_contents: list[TableStructureAnnotationModel] = [] + for row_index, annotations in row_index_to_annotations.items(): + if row_index in excluded_row_indices: + continue + # Skip rows that contain column headers + if any(ann.data.is_column_header for ann in annotations): + continue + table_contents.extend(annotations) + + # The root node represents the table itself + result[table_uid] = TableCellHierarchyTreeModel( + node_uid=table_uid, + node_type=ContentCategory.TABLE.value, + children=top_level_nodes, + contents=table_contents, + ) + + return result + + # --------- Main API --------- diff --git a/kensho_kenverters/tests/test_output_to_tables.py b/kensho_kenverters/tests/test_output_to_tables.py index f3b4e97..1ccb296 100644 --- a/kensho_kenverters/tests/test_output_to_tables.py +++ b/kensho_kenverters/tests/test_output_to_tables.py @@ -1,20 +1,29 @@ import json import os +from collections import defaultdict from typing import Any, ClassVar from unittest import TestCase +from ..constants import ContentCategory from ..extract_output_models import ( AnnotationDataModel, Cell, + ContentModel, LocationModel, + RelationAnnotationModel, + TableCellHierarchyTreeModel, TableGridAndStructure, TableStructureAnnotationModel, ) from ..output_to_tables import ( + _build_table_cell_hierarchy_tree_node, + _get_table_uid_to_table_cell_hierarchy_tree, build_table_grids, extract_pd_dfs_from_output, extract_pd_dfs_with_locs_and_table_structure_from_output, + get_table_uid_to_cells_mapping, ) +from ..utils import load_output_to_pydantic OUTPUT_FILE_PATH = os.path.join( os.path.dirname(__file__), "data", "extract_output.json" @@ -4566,3 +4575,343 @@ def test_build_table_grids_figure_extracted_table_structure(self) -> None: {"content_tree": content, "annotations": annotations}, True ) self.assertEqual(expected_tables_grid_and_structure, tables_grid_and_structure) + + +def _make_table_structure_annotation( + content_uid: str, + row: int, + col: int, + row_span: int = 1, + col_span: int = 1, + is_column_header: bool = False, + is_projected_row_header: bool = False, +) -> TableStructureAnnotationModel: + """Create a TableStructureAnnotationModel for testing.""" + return TableStructureAnnotationModel( + content_uids=[content_uid], + data=AnnotationDataModel( + index=(row, col), + span=(row_span, col_span), + is_column_header=is_column_header, + is_projected_row_header=is_projected_row_header, + ), + type="table_structure", + locations=None, + ) + + +def _make_content_model_dict( + uid: str, content: str, node_type: str = "TABLE_CELL" +) -> dict[str, Any]: + """Create a ContentModel-compatible dict for testing.""" + return { + "uid": uid, + "type": node_type, + "content": content, + "children": [], + "locations": None, + } + + +def _build_simple_table_document() -> dict[str, Any]: + """Build a simple table document for testing hierarchy. + + Table structure: + Row 0: Column headers ["Name", "Value"] + Row 1: Projected row header "ASSETS" (spans 2 cols) + Row 2: Projected row header "Current Assets" (spans 2 cols) + Row 3: Data row ["Cash", "100"] (child of "Current Assets") + Row 4: Data row ["Securities", "200"] (child of "Current Assets") + Row 5: Data row ["Equipment", "500"] (child of "ASSETS", not via "Current Assets") + + Hierarchy: + TABLE + ASSETS + Current Assets + Cash | 100 + Securities | 200 + Equipment | 500 + """ + return { + "annotations": [ + # Column headers (row 0) + { + "content_uids": ["c1"], + "data": { + "index": [0, 0], + "span": [1, 1], + "is_column_header": True, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + { + "content_uids": ["c2"], + "data": { + "index": [0, 1], + "span": [1, 1], + "is_column_header": True, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + # Projected row header "ASSETS" (row 1, spans 2 cols) + { + "content_uids": ["c3"], + "data": { + "index": [1, 0], + "span": [1, 2], + "is_column_header": False, + "is_projected_row_header": True, + }, + "type": "table_structure", + }, + # Projected row header "Current Assets" (row 2, spans 2 cols) + { + "content_uids": ["c4"], + "data": { + "index": [2, 0], + "span": [1, 2], + "is_column_header": False, + "is_projected_row_header": True, + }, + "type": "table_structure", + }, + # Data row "Cash" (row 3) + { + "content_uids": ["c5"], + "data": { + "index": [3, 0], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + { + "content_uids": ["c6"], + "data": { + "index": [3, 1], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + # Data row "Securities" (row 4) + { + "content_uids": ["c7"], + "data": { + "index": [4, 0], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + { + "content_uids": ["c8"], + "data": { + "index": [4, 1], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + # Data row "Equipment" (row 5) + { + "content_uids": ["c9"], + "data": { + "index": [5, 0], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + { + "content_uids": ["c10"], + "data": { + "index": [5, 1], + "span": [1, 1], + "is_column_header": False, + "is_projected_row_header": False, + }, + "type": "table_structure", + }, + # Relations + { + "data": { + "relation_type": "row_key_parent", + "source_content_uid": "c3", + "target_content_uid": "c4", + }, + "type": "relation", + }, + { + "data": { + "relation_type": "row_key_parent", + "source_content_uid": "c4", + "target_content_uid": "c5", + }, + "type": "relation", + }, + { + "data": { + "relation_type": "row_key_parent", + "source_content_uid": "c4", + "target_content_uid": "c7", + }, + "type": "relation", + }, + { + "data": { + "relation_type": "row_key_parent", + "source_content_uid": "c3", + "target_content_uid": "c9", + }, + "type": "relation", + }, + ], + "content_tree": { + "uid": "0", + "type": "DOCUMENT", + "content": None, + "children": [ + { + "uid": "t1", + "type": "TABLE", + "content": None, + "children": [ + _make_content_model_dict("c1", "Name"), + _make_content_model_dict("c2", "Value"), + _make_content_model_dict("c3", "ASSETS"), + _make_content_model_dict("c4", "Current Assets"), + _make_content_model_dict("c5", "Cash"), + _make_content_model_dict("c6", "100"), + _make_content_model_dict("c7", "Securities"), + _make_content_model_dict("c8", "200"), + _make_content_model_dict("c9", "Equipment"), + _make_content_model_dict("c10", "500"), + ], + } + ], + }, + } + + +class TestBuildTableCellHierarchyTreeNode(TestCase): + """Tests for _build_table_cell_hierarchy_tree_node.""" + + def test_leaf_node_no_children(self) -> None: + """A projected row header with only data row children produces contents, no child nodes.""" + cell_uid_to_annotation = { + "c4": _make_table_structure_annotation( + "c4", 2, 0, is_projected_row_header=True + ), + "c5": _make_table_structure_annotation("c5", 3, 0), + "c6": _make_table_structure_annotation("c6", 3, 1), + "c7": _make_table_structure_annotation("c7", 4, 0), + "c8": _make_table_structure_annotation("c8", 4, 1), + } + parent_to_children = {"c4": ["c5", "c7"]} + all_projected_row_header_uids = {"c3", "c4"} + row_index_to_annotations = defaultdict(list) + row_index_to_annotations[3] = [ + cell_uid_to_annotation["c5"], + cell_uid_to_annotation["c6"], + ] + row_index_to_annotations[4] = [ + cell_uid_to_annotation["c7"], + cell_uid_to_annotation["c8"], + ] + + node = _build_table_cell_hierarchy_tree_node( + "c4", + cell_uid_to_annotation, + parent_to_children, + all_projected_row_header_uids, + row_index_to_annotations, + ) + + self.assertEqual(node.node_uid, "c4") + self.assertEqual(node.node_type, ContentCategory.TABLE_CELL.value) + self.assertEqual(len(node.children), 0) + self.assertEqual(len(node.contents), 4) + + def test_node_with_projected_row_header_children(self) -> None: + """A projected row header with sub-headers produces child nodes.""" + cell_uid_to_annotation = { + "c3": _make_table_structure_annotation( + "c3", 1, 0, is_projected_row_header=True + ), + "c4": _make_table_structure_annotation( + "c4", 2, 0, is_projected_row_header=True + ), + "c9": _make_table_structure_annotation("c9", 5, 0), + "c10": _make_table_structure_annotation("c10", 5, 1), + } + parent_to_children = {"c3": ["c4", "c9"]} + all_projected_row_header_uids = {"c3", "c4"} + row_index_to_annotations = defaultdict(list) + row_index_to_annotations[5] = [ + cell_uid_to_annotation["c9"], + cell_uid_to_annotation["c10"], + ] + + node = _build_table_cell_hierarchy_tree_node( + "c3", + cell_uid_to_annotation, + dict(parent_to_children), + all_projected_row_header_uids, + row_index_to_annotations, + ) + + self.assertEqual(node.node_uid, "c3") + self.assertEqual(len(node.children), 1) + self.assertEqual(node.children[0].node_uid, "c4") + self.assertEqual(len(node.contents), 2) + + +class TestGetTableUidToTableCellHierarchyTree(TestCase): + """Tests for _get_table_uid_to_table_cell_hierarchy_tree.""" + + def test_builds_tree_for_simple_table(self) -> None: + """Test building hierarchy tree from annotations and relations.""" + doc = _build_simple_table_document() + parsed = load_output_to_pydantic(doc) + + table_uid_to_cells_mapping = get_table_uid_to_cells_mapping(parsed.content_tree) + table_cell_annotations = [ + ann + for ann in parsed.annotations + if isinstance(ann, TableStructureAnnotationModel) + ] + relation_annotations = [ + ann + for ann in parsed.annotations + if not isinstance(ann, TableStructureAnnotationModel) + ] + + result = _get_table_uid_to_table_cell_hierarchy_tree( + table_uid_to_cells_mapping, + table_cell_annotations, + relation_annotations, + ) + + self.assertIn("t1", result) + tree = result["t1"] + self.assertEqual(tree.node_uid, "t1") + self.assertEqual(tree.node_type, ContentCategory.TABLE.value) + self.assertEqual(len(tree.children), 1) + + assets_node = tree.children[0] + self.assertEqual(assets_node.node_uid, "c3") + self.assertEqual(len(assets_node.children), 1) + self.assertEqual(len(assets_node.contents), 2) + + current_assets_node = assets_node.children[0] + self.assertEqual(current_assets_node.node_uid, "c4") + self.assertEqual(len(current_assets_node.children), 0) + self.assertEqual(len(current_assets_node.contents), 4) From 062ca32535a13525d1e1913e5fd56f1d62a2c79b Mon Sep 17 00:00:00 2001 From: Libin Liang Date: Thu, 9 Jul 2026 15:08:55 -0400 Subject: [PATCH 2/4] - Update pyproject.toml. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index ad60d34..b7daf34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "kensho_kenverters" -version = "3.1.1" +version = "3.1.2" description = "Extract Output Translator Tools" readme = "README.md" authors = ["Valerie Faucon-Morin "] From f1ed6b08dc49b2a1efe9ae7c0228e665ac7a7c1c Mon Sep 17 00:00:00 2001 From: Libin Liang Date: Thu, 9 Jul 2026 15:15:18 -0400 Subject: [PATCH 3/4] - fix linting issue. --- kensho_kenverters/tests/test_output_to_tables.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/kensho_kenverters/tests/test_output_to_tables.py b/kensho_kenverters/tests/test_output_to_tables.py index 1ccb296..10d641b 100644 --- a/kensho_kenverters/tests/test_output_to_tables.py +++ b/kensho_kenverters/tests/test_output_to_tables.py @@ -8,10 +8,7 @@ from ..extract_output_models import ( AnnotationDataModel, Cell, - ContentModel, LocationModel, - RelationAnnotationModel, - TableCellHierarchyTreeModel, TableGridAndStructure, TableStructureAnnotationModel, ) From e084999a13a9da4d9138e53a4a5c88b1b6376a19 Mon Sep 17 00:00:00 2001 From: Libin Liang Date: Thu, 9 Jul 2026 21:45:46 -0400 Subject: [PATCH 4/4] - adjust the changelog. --- kensho_kenverters/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kensho_kenverters/CHANGELOG.md b/kensho_kenverters/CHANGELOG.md index 29394a5..dce130f 100644 --- a/kensho_kenverters/CHANGELOG.md +++ b/kensho_kenverters/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -* Add codes to construct table cell hierarchy tree. +* Add code to construct table cell hierarchy tree. ## v3.1.1