Skip to content

Commit 8f3aee4

Browse files
kaxilkaxil
committed
Refactor DocumentLoaderOperator: streams, encoding, JSON shape, skip rules
Rebases onto main to recover the 0.3.0 release entries that were rolled back on the original branch, and applies the review feedback the user- side review surfaced. Operator - Replace the temp-file dance for PDF/DOCX bytes with in-memory streams. ``pypdf.PdfReader`` and ``docx.Document`` both accept binary streams, so ``source_bytes`` now goes through ``io.BytesIO`` directly. No more ``NamedTemporaryFile(delete=False)`` + ``os.unlink``. - Add ``encoding`` and ``encoding_errors`` parameters for non-UTF-8 input (Windows-1252 CSVs, files with a leading byte-order mark, ...). Failed decodes raise a ``ValueError`` that includes the offending file path so directory-mode runs are diagnosable. - Add ``json_text_field``: when set, the named key on each JSON item becomes the embedding text and every other key lands in ``metadata``. When unset, JSON dicts are flattened to ``"k: v, k: v"`` (matches the CSV parser) instead of being dumped back to JSON syntax tokens. - Directory-mode ``source_path`` now silently ignores files whose name starts with ``.`` (``.DS_Store``, editor swap files, ``.gitkeep``) and skips unknown-extension files with a warning rather than crashing on the first stray file. - ``glob.glob(source_path, recursive=True)`` so ``**`` patterns walk subdirectories (the docs already advertised this). - Auto-extracted metadata (``file_name``, ``file_path``, ``row_index``, ``item_index``, ``page_number``) now takes precedence over ``metadata_fields`` with the same key (via ``setdefault``). - Expanded ``template_fields`` to include ``file_type``, ``file_extensions``, ``parser`` so they can be driven from Jinja. - Hoisted ``AirflowOptionalProviderFeatureException`` import to the module top so the lazy ``pypdf`` / ``docx`` blocks are 2 lines each. Docs - Switched all inline ``code-block:: python`` snippets to ``exampleinclude::`` directives pointing at a new ``example_document_loader.py`` (basic, directory, bytes, ``json_text_field`` patterns), matching the convention every other operator in this provider uses. - New sections documenting encoding handling, metadata precedence, and the directory-mode skip rules (files whose name starts with a ``.`` / unknown-extension warn-and-skip). Tests - Dropped the tautological ``test_template_fields`` that just round- tripped the class attribute; replaced with a behavioural check confirming the templated fields are actually in the templated set. - New coverage for: dot-prefixed-name skip, unknown-extension warn + skip, ``encoding`` / ``encoding_errors``, ``json_text_field``, JSON dict flatten, CSV empty-cell skip, ``metadata_fields`` precedence (auto wins), recursive ``**`` glob. - PDF/DOCX bytes tests assert the library was called with a ``BytesIO``, locking in the no-temp-file behaviour.
1 parent ed80b29 commit 8f3aee4

6 files changed

Lines changed: 618 additions & 192 deletions

File tree

‎providers/common/ai/docs/operators/document_loader.rst‎

Lines changed: 136 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -22,96 +22,129 @@
2222

2323
Use :class:`~airflow.providers.common.ai.operators.document_loader.DocumentLoaderOperator`
2424
to parse files into ``list[dict(text, metadata)]`` for downstream embedding
25-
pipelines. The operator bridges Airflow's connectivity layer (hooks that
25+
pipelines. The operator bridges Airflow's connectivity layer (hooks that
2626
produce bytes or local files) and the AI embedding layer (operators that
2727
need structured text with metadata).
2828

29-
The operator is **framework-agnostic** — it has no dependency on LlamaIndex,
29+
The operator is **framework-agnostic** -- it has no dependency on LlamaIndex,
3030
LangChain, or any other AI framework.
3131

32-
Built-in formats
33-
----------------
32+
Basic usage
33+
-----------
3434

3535
``.txt``, ``.md``, ``.csv``, and ``.json`` are handled with zero extra
3636
dependencies:
3737

38-
.. code-block:: python
39-
40-
from airflow.providers.common.ai.operators.document_loader import DocumentLoaderOperator
38+
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_document_loader.py
39+
:language: python
40+
:start-after: [START howto_operator_document_loader_basic]
41+
:end-before: [END howto_operator_document_loader_basic]
4142

42-
load_docs = DocumentLoaderOperator(
43-
task_id="load_docs",
44-
source_path="/opt/airflow/data/articles/",
45-
)
46-
47-
CSV files produce one document per row. JSON files with a top-level array
48-
produce one document per element; a single JSON object produces one document.
43+
CSV files produce one document per row, with empty cells skipped. JSON files
44+
with a top-level array produce one document per element; a single JSON object
45+
produces one document. By default each dict is flattened into ``"key: value,
46+
key: value"`` text so the embedding sees content tokens rather than JSON
47+
syntax (see the ``json_text_field`` section below for the structured variant).
4948

5049
PDF parsing
5150
-----------
5251

53-
Install the ``pdf`` extra to parse PDF files via `pypdf <https://pypdf.readthedocs.io/>`__:
54-
55-
.. code-block:: bash
52+
Install the ``pdf`` extra to parse PDF files via
53+
`pypdf <https://pypdf.readthedocs.io/>`__::
5654

5755
pip install apache-airflow-providers-common-ai[pdf]
5856

59-
.. code-block:: python
60-
61-
load_pdfs = DocumentLoaderOperator(
62-
task_id="load_pdfs",
63-
source_path="/opt/airflow/data/reports/*.pdf",
64-
)
65-
66-
Each page with extractable text becomes a separate document. Empty pages are
67-
skipped. The ``page_number`` is included in the document metadata.
57+
Each page with extractable text becomes a separate document. Empty pages are
58+
skipped. ``page_number`` is included in the document metadata.
6859

6960
DOCX parsing
7061
------------
7162

7263
Install the ``docx`` extra to parse Word documents via
73-
`python-docx <https://python-docx.readthedocs.io/>`__:
74-
75-
.. code-block:: bash
64+
`python-docx <https://python-docx.readthedocs.io/>`__::
7665

7766
pip install apache-airflow-providers-common-ai[docx]
7867

79-
.. code-block:: python
80-
81-
load_word = DocumentLoaderOperator(
82-
task_id="load_word",
83-
source_path="/opt/airflow/data/specs/*.docx",
84-
)
85-
8668
All non-empty paragraphs are concatenated into a single document per file.
8769

8870
.. note::
8971

9072
DOCX extraction reads paragraph text only. Tables, headers, footers, and
91-
footnotes are not included. For comprehensive DOCX parsing, use a dedicated
92-
document extraction tool like Unstructured or Docling as a custom parser
73+
footnotes are not included. For richer DOCX parsing, use a dedicated
74+
extraction tool (``Unstructured``, ``docling``) as a custom parser
9375
backend.
9476

95-
Glob patterns and filtering
77+
Directory mode and filtering
9678
----------------------------
9779

98-
Pass a glob pattern to ``source_path`` to match multiple files. Use
99-
``file_extensions`` to limit which files are processed:
80+
Point ``source_path`` at a directory or pass a glob pattern (``**`` enables
81+
recursive matching). Combine with ``file_extensions`` to scope which files
82+
are processed:
83+
84+
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_document_loader.py
85+
:language: python
86+
:start-after: [START howto_operator_document_loader_directory]
87+
:end-before: [END howto_operator_document_loader_directory]
88+
89+
Directory-mode behavior when ``file_extensions`` is omitted:
90+
91+
- Files whose name starts with a ``.`` (``.DS_Store``, editor swap files,
92+
``.gitkeep``, ...) are silently ignored.
93+
- Files whose extension is not in the built-in dispatch map are skipped
94+
with a warning rather than crashing the operator. A glob pattern that
95+
matches an unknown extension is treated as intentional and parsed via
96+
the explicit ``parser`` argument.
97+
98+
Loading from bytes
99+
------------------
100+
101+
When upstream tasks produce file content as bytes (S3, GCS, HTTP, etc.),
102+
pass them via ``source_bytes`` and tell the operator how to interpret them
103+
with ``file_type``. ``source_bytes`` is not a template field because Jinja
104+
would render ``bytes`` as their ``repr`` text, which would break binary
105+
parsing:
106+
107+
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_document_loader.py
108+
:language: python
109+
:start-after: [START howto_operator_document_loader_bytes]
110+
:end-before: [END howto_operator_document_loader_bytes]
111+
112+
PDF and DOCX bytes are parsed via an in-memory stream -- no temporary files
113+
on disk.
114+
115+
Structured JSON ingestion
116+
-------------------------
117+
118+
For arrays of records where one field is the body and the rest are metadata
119+
(article ingestion, ticket exports, ...), set ``json_text_field`` to the key
120+
that holds the text. Every other key on the same item lands in ``metadata``:
121+
122+
.. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_document_loader.py
123+
:language: python
124+
:start-after: [START howto_operator_document_loader_json_field]
125+
:end-before: [END howto_operator_document_loader_json_field]
126+
127+
For **arbitrary API data** (Salesforce SOQL results, database query exports),
128+
a ``@task`` that maps fields to text and metadata is still appropriate when
129+
the field shape is more complex than what ``json_text_field`` covers:
100130

101131
.. code-block:: python
102132
103-
load_filtered = DocumentLoaderOperator(
104-
task_id="load_filtered",
105-
source_path="/opt/airflow/data/mixed/*",
106-
file_extensions=[".pdf", ".txt"],
107-
)
133+
@task
134+
def transform_cases(records: list[dict]) -> list[dict]:
135+
return [
136+
{
137+
"text": f"{r['Subject']}\n\n{r['Description']}",
138+
"metadata": {"case_id": r["Id"], "source": "salesforce"},
139+
}
140+
for r in records
141+
]
108142
109-
Composing with downstream operators
110-
------------------------------------
143+
Composing with downstream embedding operators
144+
---------------------------------------------
111145

112146
The output format (``list[dict(text, metadata)]``) is designed to feed
113-
directly into embedding operators. For example, with the LlamaIndex
114-
``EmbeddingOperator``:
147+
directly into embedding operators. With LlamaIndex's ``EmbeddingOperator``:
115148

116149
.. code-block:: python
117150
@@ -129,7 +162,7 @@ directly into embedding operators. For example, with the LlamaIndex
129162
load >> embed
130163
131164
Composing with Airflow providers
132-
---------------------------------
165+
--------------------------------
133166

134167
Use any Airflow provider to download files, then parse them with
135168
``DocumentLoaderOperator``:
@@ -152,57 +185,62 @@ Use any Airflow provider to download files, then parse them with
152185
153186
download >> load
154187
155-
For **structured API data** (Salesforce SOQL results, database query exports),
156-
a ``@task`` that maps fields to text and metadata is more appropriate than
157-
``DocumentLoaderOperator``, which is designed for binary file parsing:
158-
159-
.. code-block:: python
160-
161-
@task
162-
def transform_cases(records: list[dict]) -> list[dict]:
163-
return [
164-
{
165-
"text": f"{r['Subject']}\n\n{r['Description']}",
166-
"metadata": {"case_id": r["Id"], "source": "salesforce"},
167-
}
168-
for r in records
169-
]
170-
171-
Loading from bytes
172-
------------------
173-
174-
When upstream tasks produce file content as bytes, pass them directly via a
175-
``@task`` function. Note that ``source_bytes`` is not a template field because
176-
Jinja stringifies ``bytes`` to their ``repr``, which breaks binary parsing:
188+
Non-UTF-8 inputs
189+
----------------
177190

178-
.. code-block:: python
191+
The text parsers (``.txt`` / ``.md`` / ``.csv`` / ``.json``) and the bytes
192+
path default to UTF-8. To handle Windows-1252 CSVs, files with a leading
193+
``utf-8-sig`` byte-order mark, or any other encoding, set the ``encoding``
194+
parameter on the operator (and optionally ``encoding_errors="replace"`` to
195+
tolerate mixed-encoding sources at the cost of some character loss). A
196+
failed decode includes the offending file path in the error so
197+
directory-mode runs are easy to diagnose.
179198

180-
from airflow.decorators import task
181-
from airflow.providers.common.ai.operators.document_loader import DocumentLoaderOperator
199+
Metadata precedence
200+
-------------------
182201

183-
@task
184-
def parse_downloaded_pdf(raw_bytes: bytes) -> list[dict]:
185-
op = DocumentLoaderOperator(
186-
task_id="parse_pdf",
187-
source_bytes=raw_bytes,
188-
file_type=".pdf",
189-
)
190-
return op.execute(context=get_current_context())
202+
Auto-extracted metadata keys -- ``file_name``, ``file_path``, ``row_index``,
203+
``item_index``, ``page_number`` -- take precedence over keys with the same
204+
name in ``metadata_fields``. ``metadata_fields`` fills gaps; it never
205+
overwrites the auto-extracted shape.
191206

192207
Parameters
193208
----------
194209

195-
- ``source_path``: Local file path, directory, or glob pattern.
196-
Mutually exclusive with ``source_bytes``.
197-
- ``source_bytes``: Raw file bytes from XCom. Requires ``file_type``.
198-
Mutually exclusive with ``source_path``.
199-
- ``file_type``: File extension hint (e.g. ``".pdf"``). Required with
200-
``source_bytes``. Optional with ``source_path`` to override
201-
auto-detection.
202-
- ``parser``: Parsing backend. ``"auto"`` (default) selects from the file
203-
extension. Set explicitly to force a specific backend (e.g. ``"text"``
204-
to treat an unknown extension as plain text).
205-
- ``file_extensions``: Filter which files to process when ``source_path``
206-
matches multiple files.
207-
- ``metadata_fields``: Extra key-value pairs merged into every document's
208-
metadata dict.
210+
.. list-table::
211+
:header-rows: 1
212+
:widths: 25 75
213+
214+
* - Parameter
215+
- Description
216+
* - ``source_path``
217+
- Local file, directory, or glob pattern. ``**`` is recursive. Mutually
218+
exclusive with ``source_bytes``.
219+
* - ``source_bytes``
220+
- Raw file bytes from XCom. Requires ``file_type``. Mutually exclusive
221+
with ``source_path``. Not a template field (bytes don't survive Jinja).
222+
* - ``file_type``
223+
- File extension hint (e.g. ``".pdf"``). Required with ``source_bytes``;
224+
optional with ``source_path`` to override auto-detection.
225+
* - ``parser``
226+
- Parsing backend. ``"auto"`` (default) picks from the file extension.
227+
Set explicitly to force a backend (e.g. ``"text"`` to treat an
228+
unknown extension as plain text).
229+
* - ``file_extensions``
230+
- Filter for ``source_path`` directory or glob. When omitted in
231+
directory mode, files whose name starts with a ``.`` are ignored
232+
and unknown-extension files are skipped with a warning.
233+
* - ``metadata_fields``
234+
- Extra key-value pairs merged into every document's metadata. Does
235+
not override auto-extracted keys.
236+
* - ``encoding``
237+
- Text encoding for the bytes path and ``.txt`` / ``.md`` / ``.csv`` /
238+
``.json`` files. Defaults to ``"utf-8"``.
239+
* - ``encoding_errors``
240+
- How decode errors are handled (``"strict"`` / ``"replace"`` /
241+
``"ignore"``). Defaults to ``"strict"``.
242+
* - ``json_text_field``
243+
- When parsing JSON, treat this key as the embedding text; every other
244+
key on the same item lands in ``metadata``. When unset, dicts are
245+
flattened to ``"k: v, k: v"`` so the embedding sees content tokens
246+
rather than JSON syntax.

0 commit comments

Comments
 (0)