xlsongc commented on code in PR #73926:
URL: https://github.com/apache/airflow/pull/73926#discussion_r4424552476
##########
providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py:
##########
@@ -450,21 +454,123 @@ def _parse_pdf_stream(self, stream: BinaryIO) ->
list[dict[str, Any]]:
documents.append({"text": text, "metadata": {"page_number":
page_num + 1}})
return documents
- def _parse_docx_stream(self, stream: BinaryIO) -> list[dict[str, Any]]:
+ def _parse_docx_stream(self, stream: BinaryIO, *, source_hint: str) ->
list[dict[str, Any]]:
"""
- Parse a DOCX stream into documents.
+ Parse a DOCX stream into a single document.
- Extracts paragraph text only. Tables, headers, footers, and footnotes
- are not included. For richer DOCX parsing, plug in a dedicated
- extraction tool (``Unstructured``, ``docling``) as a custom parser
- backend.
+ Paragraphs and tables in the document body are extracted in document
+ order. Each table row becomes one "| cell | cell |" line, and a nested
+ table is flattened into its cell. Headers, footers, footnotes, content
+ controls, text boxes, and pending tracked insertions are not included.
"""
try:
from docx import Document
+ from docx.oxml.exceptions import InvalidXmlError
+ from docx.table import Table
except ImportError as e:
raise AirflowOptionalProviderFeatureException(e)
doc = Document(stream)
- paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
- text = "\n\n".join(paragraphs)
+ blocks = []
+ for block in doc.iter_inner_content():
+ if isinstance(block, Table):
+ rows = self._get_docx_table_rows_or_none(
+ block, Table, InvalidXmlError, source_hint=source_hint
+ )
+ if rows is None:
+ continue
+ text = "\n".join(f"| {' | '.join(cells)} |" for cells in rows)
+ else:
+ text = block.text
+ if text.strip():
+ blocks.append(text)
+ text = "\n\n".join(blocks)
return [{"text": text, "metadata": {}}]
+
+ def _get_docx_table_rows_or_none(
+ self,
+ table: Table,
+ table_cls: type[Table],
+ invalid_xml_error_cls: type[InvalidXmlError],
+ *,
+ source_hint: str,
+ ) -> list[list[str]] | None:
+ try:
+ return self._get_docx_table_rows(table, table_cls,
invalid_xml_error_cls, source_hint=source_hint)
+ except (ValueError, invalid_xml_error_cls, RecursionError) as e:
+ # Malformed table XML raises ValueError (a vMerge with no cell
above, or a gridSpan
+ # that does not fit the table) or InvalidXmlError (a missing
required attribute).
+ # python-docx walks vertical merges up recursively, so one
spanning about 1000 rows
+ # raises RecursionError.
+ self.log.warning("Skipping a table in %s that python-docx could
not read: %r", source_hint, e)
+ return None
+
+ def _get_docx_table_rows(
+ self,
+ table: Table,
+ table_cls: type[Table],
+ invalid_xml_error_cls: type[InvalidXmlError],
+ *,
+ source_hint: str,
+ ) -> list[list[str]]:
+ # w:gridBefore and w:gridAfter come straight from the XML and can be
any integer, so the
+ # empty cells added for them only fill the row up to the column count.
+ column_count = len(table.columns)
+
+ rows = []
+ for row in table.rows:
+ # row.cells yields a cell once for every grid column its
w:gridSpan covers, all at
+ # once: spans that add up to more than the table width are built
in full, and a span
+ # below 1 drops the cell and its text. Both are rejected before
row.cells is read.
+ spanned_columns = 0
+ for index, tc in enumerate(row._tr.tc_lst, start=1):
+ spanned_columns += tc.grid_span
+ if tc.grid_span < 1 or spanned_columns > column_count:
+ raise ValueError(
+ f"w:gridSpan {tc.grid_span} of cell {index} does not
fit the {column_count} table columns"
+ )
+ cells: list[str] = []
+ previous_cell = None
+ for cell in row.cells:
Review Comment:
Done, as you suggested: each row is read from `row._tr.tc_lst` once,
carrying `{grid_offset: (text, span)}` down from the row above. A continuation
reuses that text when the offset and span match and raises `ValueError`
otherwise, so the staircased file is skipped with a warning. The 1100-row merge
in `test_docx_long_vertical_merge_read_in_full` is now read in full in about
0.2 s, `test_docx_vertical_merge_text_extracted_once` checks that a merged
cell's text is extracted once, and
`test_docx_vmerge_continuation_not_lining_up_skipped_with_warning` covers
continuations narrower and wider than the cell above. The two `no tr above
topmost tr` expectations now assert the new message, since a first-row
continuation fails the same check.
Each table's text is now capped at 1,000,000 characters
(`_DOCX_TABLE_MAX_CHARS`), counted as rendered: cell text plus 3 characters per
cell, so padded empty cells and merged text repeated on each row count too. A
longer table is skipped with a warning. 1M is about 6,000 rows of 10 columns.
Your 8-row merges nested 6 deep (262,144 copies, 786K characters) stay under it
and read in under 0.01 s; 7 deep (6.3M) and 8 deep (50M) are skipped. Every
nesting level is checked, so the deepest table over the cap is dropped and the
tables around it keep their other cells, the same rule as for unreadable nested
tables. `test_docx_table_text_over_limit_skipped_with_warning`,
`test_docx_padding_and_repeated_merge_text_count_toward_text_limit` and
`test_docx_nested_table_text_counts_toward_parent_table_limit` fail without the
cap.
The cap is per table, so output still grows with the number of tables: 100
copies of the 6-deep table in one 40 KB file give 79M characters. That is
linear in the XML now (about 130x) rather than 8^depth. I can add a
document-wide budget here if you'd prefer. A hard bound would also need a
decompressed-size check, since python-docx reads each zip part whole before
this code runs; I'd open a separate issue for that.
Two related changes. The warning now reads `Skipping a table in <file>:
<error>`, since the width and text limits are not python-docx failures. And
`RecursionError` is no longer caught: the walk doesn't recurse over merges, and
nested tables can't get deep enough, because python-docx parses without
`huge_tree`, so libxml2 rejects a file nested more than 256 elements deep (83
tables, about 260 frames). The mocked `RecursionError` test is removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
