Merge pull request #18 from kensho-technologies/val/remove-table-validations

Remove validations for rows/cols for tables

Author: valerie-fauconmorin-kensho

kensho_kenverters/CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## v1.2.3
+
+Remove table validation for rows and columns to not fail downstream of Extract model failures
+
 ## v1.2.2
 
 Use underscores instead of hyphens for setuptools 78

kensho_kenverters/tables_utils.py

@@ -9,13 +9,27 @@
 from kensho_kenverters.extract_output_models import AnnotationDataModel, AnnotationModel
 
 
-def _check_complete_set(integer_set: set[int]) -> bool:
-    """Check that the set of integers contains all integers between 0 and its max."""
-    return integer_set == set(range(max(integer_set, default=-1) + 1))
+def _create_empty_annotation(row: int, col: int) -> AnnotationModel:
+    """Create an empty annotation."""
+    return AnnotationModel(
+        type=AnnotationType.TABLE_STRUCTURE.value,
+        content_uids=[],
+        data=AnnotationDataModel(
+            span=(1, 1),
+            index=(row, col),
+        ),
+        locations=None,
+    )
+
+
+def _validate_annotations(
+    duplicated_annotations: list[AnnotationModel], max_row: int, max_col: int
+) -> list[AnnotationModel]:
+    """Validate duplicated annotations.
 
+    Fill with empty annotations if rows or columns are missing.
+    """
 
-def _validate_annotations(duplicated_annotations: Sequence[AnnotationModel]) -> None:
-    """Validate duplicated annotations."""
     # Check all spans are 1 (annotations are duplicated)
     all_spans = [annotation.data.span for annotation in duplicated_annotations]
     if any(span != (1, 1) for span in all_spans):
@@ -26,13 +40,13 @@ def _validate_annotations(duplicated_annotations: Sequence[AnnotationModel]) ->
     if len(set(all_indices)) != len(all_indices):
         raise ValueError("Overlapping indices in table.")
 
-    # Check no empty rows or columns
-    all_rows = set(index[0] for index in all_indices)
-    all_columns = set(index[1] for index in all_indices)
-    if not _check_complete_set(all_rows):
-        raise ValueError("Empty row in table.")
-    if not _check_complete_set(all_columns):
-        raise ValueError("Empty column in table.")
+    # Add any missing cells
+    for row in range(max_row + 1):
+        for col in range(max_col + 1):
+            if (row, col) not in all_indices:
+                duplicated_annotations.append(_create_empty_annotation(row, col))
+
+    return duplicated_annotations
 
 
 def duplicate_spanning_annotations(
@@ -52,6 +66,8 @@ def duplicate_spanning_annotations(
         duplicated annotations. Duplicated annotations must all have span (1, 1).
     """
     duplicated_annotations = []
+    max_row = 0
+    max_col = 0
     for annotation in annotations:
         data = annotation.data
         row_span, col_span = data.span
@@ -76,9 +92,10 @@ def duplicate_spanning_annotations(
                     locations=annotation.locations,
                 )
                 duplicated_annotations.append(new_annotation)
+                max_row = max(max_row, row_index + row_span_index)
+                max_col = max(max_col, col_index + col_span_index)
 
-    _validate_annotations(duplicated_annotations)
-    return duplicated_annotations
+    return _validate_annotations(duplicated_annotations, max_row, max_col)
 
 
 def get_table_shape(

kensho_kenverters/tests/test_tables_utils.py

@@ -57,7 +57,7 @@ def test_convert_tpdle_to_pd_df(self) -> None:
         )
         self.assertEqual(expected_df.to_csv(), converted_df.to_csv())
 
-    def test_duplicate_spanning_annotations(self) -> None:
+    def test_duplicate_spanning_annotations_1_spans(self) -> None:
         # Test no duplication when all spans are 1
         duplicated = duplicate_spanning_annotations(
             self.parsed_serialized_document.annotations
@@ -417,6 +417,7 @@ def test_duplicate_spanning_annotations(self) -> None:
 
         self.assertEqual(duplicated, expected_duplicated)
 
+    def test_duplicate_spanning_annotations_greater_1_spans(self) -> None:
         # Test properly duplicated when cells have a span > 1
         annotations = self.parsed_serialized_document.annotations
         annotations.pop(-1)
@@ -775,3 +776,53 @@ def test_duplicate_spanning_annotations(self) -> None:
         ]
         duplicated = duplicate_spanning_annotations(annotations)
         self.assertEqual(expected_duplicated, duplicated)
+
+    def test_duplicate_spanning_annotations_missing_cells(self) -> None:
+        # Test missing cells
+        annotations = [
+            AnnotationModel(
+                content_uids=["7"],
+                data=AnnotationDataModel(index=(0, 0), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.16008,
+                        y=0.40464,
+                        page_number=0,
+                    )
+                ],
+            ),
+            AnnotationModel(
+                content_uids=["8"],
+                data=AnnotationDataModel(index=(0, 1), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.16008,
+                        y=0.80464,
+                        page_number=0,
+                    )
+                ],
+            ),
+            AnnotationModel(
+                content_uids=["9"],
+                data=AnnotationDataModel(index=(3, 0), span=(1, 1)),
+                type="table_structure",
+                locations=[
+                    LocationModel(
+                        height=0.01188,
+                        width=0.22128,
+                        x=0.7008,
+                        y=0.40464,
+                        page_number=0,
+                    )
+                ],
+            ),
+        ]
+
+        duplicated = duplicate_spanning_annotations(annotations)
+        self.assertEqual(len(duplicated), 8)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "kensho_kenverters"
-version = "1.2.2"
+version = "1.2.3"
 description = "Extract Output Translator Tools"
 readme = "README.md"
 authors = ["Valerie Faucon-Morin <valerie.fauconmorin@kensho.com>"]