kaxil commented on code in PR #73926:
URL: https://github.com/apache/airflow/pull/73926#discussion_r4320234365
##########
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:
`row.cells` resolves each `w:vMerge` continuation by walking up to the row
where the merge starts and yielding that cell once per column of its own span,
so three things get past the new check. A continuation with `gridSpan` 1 under
a cell spanning two columns still yields two cells, so the sum check passes
while the row grows: staircased over 100 rows of a 20000-column grid, a 39 KB
file produced 302M characters in 300 s with 1.4 GB peak. Valid merges are
quadratic: 300 rows with 10 columns merged down took 177 s in breeze on 1.2.0,
and at about 1000 rows the table is dropped on `RecursionError`. And the merged
cell's text, nested tables included, is re-extracted on every row it spans, so
8-row merges nested 6 deep took 185 s from a 37 KB file.
Walking `row._tr.tc_lst` once and carrying `{grid_offset: (text, span)}`
from the row above fixes all three: a continuation reuses that text when the
offset and span match, and raises `ValueError` otherwise. I tried it and those
files take under 0.2 s each, a 2000-row merge is read instead of skipped, and
the existing tests pass apart from the two that assert python-docx's `no tr
above topmost tr` message. The nested case still repeats the text 8^depth times
(262144 copies of one value at depth 6), so it also needs a cap on each table's
text.
##########
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)
Review Comment:
I suggested capping at `len(table.columns)`, but I got that wrong.
`len(table.columns)` comes from `w:tblGrid`. Apache POI's `XWPFTable` never
writes that element (the code is commented out in `createEmptyTable`), and
python-docx raises `InvalidXmlError` without it. So every table in a
POI-generated file is now skipped.
The count also comes from the file, so it doesn't bound anything: 2000
one-cell rows with `gridAfter` under a 5000-column grid still produce 30M
characters from a 38 KB file.
Could we fall back to the widest row when the grid is missing, and cap the
width at a fixed limit?
--
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]
