Coverage for src/qdrant_loader/core/conversion/postprocess.py: 95%

1"""Post-processing of the docling ``DoclingDocument`` for spreadsheet inputs. 

2 

3docling's Excel backend has two defects that hurt every downstream consumer 

4(markdown export *and* the HybridChunker, since both read the same structured 

5document). We repair them once, immediately after parsing, at the single choke 

6point where the :class:`~.outcome.ConvertedDocument` is built — so there is no 

7serialize/reparse round-trip and both views see the fix. 

8 

9Defect 1 — merged-cell duplication 

10 A merged cell is stored *once* in ``table.data.table_cells`` with correct 

11 span info, but ``TableData.grid`` (a computed property) — and every 

12 serializer that reads it — expands that one cell into *every* position it 

13 spans. A title merged across N columns becomes the same text N times. We 

14 keep the anchor cell (and its span metadata) intact and append explicit 

15 empty 1x1 filler cells for the non-anchor positions; because ``grid`` 

16 derives itself by overwriting positions in ``table_cells`` order, the 

17 later fillers blank the vacated positions while the anchor survives. 

18 

19 NOTE on object identity: ``grid`` is a ``@computed_field`` rebuilt on every 

20 access, and within it all spanned positions are the *same* ``TableCell`` 

21 object as the anchor. Mutating a grid position is therefore both futile 

22 (not persisted) and dangerous (would mutate the anchor). The repair has to 

23 live in ``table_cells``, which is what this does. 

24 

25Defect 2 — sheet names dropped 

26 Each sheet survives as a body group whose ``name`` is ``"sheet: <name>"``, 

27 but neither markdown export nor chunk contextualization surfaces it 

28 (markitdown emitted ``## Sheet: <name>`` headings). We insert a 

29 ``SECTION_HEADER`` text item as the first child of each sheet group, so the 

30 name becomes a markdown heading and rides into each chunk's heading path. 

31 

32Both repairs are strictly-better, default-on, and apply *only* to spreadsheet 

33inputs — prose formats (PDF/DOCX) are untouched. 

34""" 

35 

36from __future__ import annotations 

37 

38from typing import TYPE_CHECKING 

39 

40if TYPE_CHECKING: 

41 from docling_core.types.doc import DoclingDocument 

42 

43# docling's InputFormat.value for spreadsheets, and the prefix its Excel backend 

44# stamps onto each sheet group's name. 

45_XLSX_FORMAT = "xlsx" 

46_SHEET_PREFIX = "sheet:" 

47 

48 

49def postprocess_spreadsheet(document: DoclingDocument, source_format: str) -> None: 

50 """Repair the two XLSX defects in place. No-op for non-spreadsheet formats. 

51 

52 ``source_format`` is the ``ConvertedDocument.source_format`` (the docling 

53 ``InputFormat.value``); we gate on it so PDFs/DOCX are never touched. 

54 """ 

55 if source_format.lower() != _XLSX_FORMAT: 

56 return 

57 

58 _dedupe_merged_cells(document) 

59 _surface_sheet_headings(document) 

60 

61 

62def _dedupe_merged_cells(document: DoclingDocument) -> None: 

63 """Stop merged cells from being serialized once per spanned position. 

64 

65 For every multi-row/col cell, append empty 1x1 filler cells covering all 

66 non-anchor positions. ``TableData.grid`` rebuilds itself by writing each 

67 ``table_cells`` entry over its offset range in list order, so the fillers 

68 (added last) overwrite the duplicate positions while leaving the anchor — 

69 at ``(start_row_offset_idx, start_col_offset_idx)`` — and its span intact. 

70 """ 

71 from docling_core.types.doc.document import TableCell 

72 

73 for table in document.tables: 

74 data = table.data 

75 fillers: list[TableCell] = [] 

76 for cell in data.table_cells: 

77 if cell.col_span <= 1 and cell.row_span <= 1: 

78 continue 

79 anchor_row = cell.start_row_offset_idx 

80 anchor_col = cell.start_col_offset_idx 

81 for row in range(cell.start_row_offset_idx, cell.end_row_offset_idx): 

82 for col in range(cell.start_col_offset_idx, cell.end_col_offset_idx): 

83 if row == anchor_row and col == anchor_col: 

84 continue # keep the anchor's text exactly once 

85 fillers.append( 

86 TableCell( 

87 text="", 

88 row_span=1, 

89 col_span=1, 

90 start_row_offset_idx=row, 

91 end_row_offset_idx=row + 1, 

92 start_col_offset_idx=col, 

93 end_col_offset_idx=col + 1, 

94 ) 

95 ) 

96 if fillers: 

97 data.table_cells.extend(fillers) 

98 

99 

100def _surface_sheet_headings(document: DoclingDocument) -> None: 

101 """Insert a SECTION_HEADER as the first child of each sheet group. 

102 

103 docling naming changed across versions: older builds use 

104 ``"sheet: <name>"``, newer builds use the plain sheet name. We support both. 

105 

106 Inserting before the group's current first child (with ``after=False``) 

107 parents the header on the group and makes it lead the sheet. 

108 """ 

109 from docling_core.types.doc import DocItemLabel 

110 

111 for group in document.groups: 

112 name = group.name or "" 

113 if not name.strip() or not group.children: 

114 continue 

115 

116 lowered = name.lower() 

117 sheet_name = ( 

118 name.split(":", 1)[1].strip() 

119 if lowered.startswith(_SHEET_PREFIX) 

120 else name.strip() 

121 ) 

122 if not sheet_name: 

123 continue 

124 

125 first_child = group.children[0].resolve(document) 

126 document.insert_text( 

127 sibling=first_child, 

128 label=DocItemLabel.SECTION_HEADER, 

129 text=f"Sheet: {sheet_name}", 

130 after=False, 

131 )