This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f0e82be3163 Add json line support for DocumentLoaderOperator (#72246)
f0e82be3163 is described below
commit f0e82be31634fd784abce417f4b89f31ae1af4ae
Author: Jeff(Wei-Hao) Lu <[email protected]>
AuthorDate: Wed Sep 23 07:07:45 2026 +0800
Add json line support for DocumentLoaderOperator (#72246)
* add jsonl support
* fix mypy
* raise error with correct line
* Split JSON Lines on newline rather than splitlines()
str.splitlines() also breaks on U+2028, U+2029 and U+0085, which are legal
unescaped characters inside a JSON string (JSON only forbids raw
U+0000-U+001F
there). A valid record containing one was torn in half, reported as invalid,
and every later line number was shifted -- which matters because scraped web
text is a common input for this operator. A trailing \r from a CRLF file is
JSON whitespace, so json.loads still tolerates it.
Generated-by: Claude Opus 5
* add source hint to make error msg contain file path
---------
Co-authored-by: jeff3071 <[email protected]>
---
.../common/ai/docs/operators/document_loader.rst | 35 +++++------
.../common/ai/operators/document_loader.py | 52 +++++++++++++----
.../common/ai/operators/test_document_loader.py | 67 ++++++++++++++++++++++
3 files changed, 126 insertions(+), 28 deletions(-)
diff --git a/providers/common/ai/docs/operators/document_loader.rst
b/providers/common/ai/docs/operators/document_loader.rst
index a699304ecc3..600831f36ec 100644
--- a/providers/common/ai/docs/operators/document_loader.rst
+++ b/providers/common/ai/docs/operators/document_loader.rst
@@ -32,8 +32,8 @@ LangChain, or any other AI framework.
Basic usage
-----------
-``.txt``, ``.md``, ``.csv``, and ``.json`` are handled with zero extra
-dependencies:
+``.txt``, ``.md``, ``.csv``, ``.json``, and ``.jsonl`` are handled with zero
+extra dependencies:
.. exampleinclude::
/../../ai/src/airflow/providers/common/ai/example_dags/example_document_loader.py
:language: python
@@ -45,6 +45,9 @@ with a top-level array produce one document per element; a
single JSON object
produces one document. By default each dict is flattened into ``"key: value,
key: value"`` text so the embedding sees content tokens rather than JSON
syntax (see the ``json_text_field`` section below for the structured variant).
+JSON Lines files produce one document per non-empty line. A JSON array on a
+single line remains one document rather than being expanded, and
+``json_text_field`` applies to each line just as it does to JSON records.
PDF parsing
-----------
@@ -153,8 +156,8 @@ Format coverage roadmap
-----------------------
The current built-in dispatch covers ``.txt``, ``.md``, ``.csv``, ``.json``,
-``.pdf``, ``.docx``. Additional formats are deferred to follow-ups, each
-gated behind its own extra so users only install what they need:
+``.jsonl``, ``.pdf``, ``.docx``. Additional formats are deferred to follow-ups,
+each gated behind its own extra so users only install what they need:
- ``.pptx`` via ``python-pptx``
- ``.epub`` via ``ebooklib``
@@ -233,13 +236,13 @@ download-then-parse pattern still works:
Non-UTF-8 inputs
----------------
-The text parsers (``.txt`` / ``.md`` / ``.csv`` / ``.json``) and the bytes
-path default to UTF-8. To handle Windows-1252 CSVs, files with a leading
-``utf-8-sig`` byte-order mark, or any other encoding, set the ``encoding``
-parameter on the operator (and optionally ``encoding_errors="replace"`` to
-tolerate mixed-encoding sources at the cost of some character loss). A
-failed decode includes the offending file path in the error so
-directory-mode runs are easy to diagnose.
+The text parsers (``.txt`` / ``.md`` / ``.csv`` / ``.json`` / ``.jsonl``) and
+the bytes path default to UTF-8. To handle Windows-1252 CSVs, files with a
+leading ``utf-8-sig`` byte-order mark, or any other encoding, set the
+``encoding`` parameter on the operator (and optionally
+``encoding_errors="replace"`` to tolerate mixed-encoding sources at the cost
+of some character loss). A failed decode includes the offending file path in
+the error so directory-mode runs are easy to diagnose.
Metadata precedence
-------------------
@@ -287,12 +290,12 @@ Parameters
not override auto-extracted keys.
* - ``encoding``
- Text encoding for the bytes path and ``.txt`` / ``.md`` / ``.csv`` /
- ``.json`` files. Defaults to ``"utf-8"``.
+ ``.json`` / ``.jsonl`` files. Defaults to ``"utf-8"``.
* - ``encoding_errors``
- How decode errors are handled (``"strict"`` / ``"replace"`` /
``"ignore"``). Defaults to ``"strict"``.
* - ``json_text_field``
- - When parsing JSON, treat this key as the embedding text; every other
- key on the same item lands in ``metadata``. When unset, dicts are
- flattened to ``"k: v, k: v"`` so the embedding sees content tokens
- rather than JSON syntax.
+ - When parsing JSON or JSON Lines, treat this key as the embedding text;
+ every other key on the same item lands in ``metadata``. When unset,
+ dicts are flattened to ``"k: v, k: v"`` so the embedding sees content
+ tokens rather than JSON syntax.
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
index c426f104758..42f630ffcf5 100644
---
a/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
+++
b/providers/common/ai/src/airflow/providers/common/ai/operators/document_loader.py
@@ -50,9 +50,9 @@ class DocumentLoaderOperator(BaseOperator):
with metadata). Framework-agnostic: no LlamaIndex, LangChain, or other
AI framework dependency.
- Built-in parsers handle ``.txt``, ``.md``, ``.csv``, and ``.json`` with
- zero extra dependencies. PDF and DOCX support require optional packages
- installable via extras::
+ Built-in parsers handle ``.txt``, ``.md``, ``.csv``, ``.json``, and
+ ``.jsonl`` with zero extra dependencies. PDF and DOCX support require
+ optional packages installable via extras::
pip install apache-airflow-providers-common-ai[pdf] # pypdf
pip install apache-airflow-providers-common-ai[docx] # python-docx
@@ -92,17 +92,18 @@ class DocumentLoaderOperator(BaseOperator):
document's ``metadata`` dict. Auto-extracted fields such as
``file_name``, ``file_path``, ``row_index``, ``item_index``, and
``page_number`` take precedence over keys with the same name.
- :param encoding: Text encoding used for ``.txt``/``.md``/``.csv``/``.json``
- and for the bytes path. Defaults to ``"utf-8"``.
+ :param encoding: Text encoding used for
+ ``.txt``/``.md``/``.csv``/``.json``/``.jsonl`` and for the bytes path.
+ Defaults to ``"utf-8"``.
:param encoding_errors: How decode errors are handled. Defaults to
``"strict"``; set to ``"replace"`` or ``"ignore"`` to tolerate
mixed-encoding inputs at the cost of some character loss.
- :param json_text_field: When parsing JSON, treat this key as the
- embedding text and put every other key into ``metadata``. Applies
- to each item when the top-level JSON is a list, or to the object
- when it is a single dict. When ``None`` (default), the operator
- flattens dicts into ``"k: v, k: v"`` text (same shape as the CSV
- parser).
+ :param json_text_field: When parsing JSON or JSON Lines, treat this key
+ as the embedding text and put every other key into ``metadata``.
+ Applies to each item when the top-level JSON is a list, to the object
+ when it is a single dict, or to each JSON Lines record. When ``None``
+ (default), the operator flattens dicts into ``"k: v, k: v"`` text
+ (same shape as the CSV parser).
"""
template_fields: Sequence[str] = (
@@ -119,6 +120,7 @@ class DocumentLoaderOperator(BaseOperator):
".md": "text",
".csv": "csv",
".json": "json",
+ ".jsonl": "jsonl",
".pdf": "pypdf",
".docx": "python-docx",
}
@@ -296,11 +298,14 @@ class DocumentLoaderOperator(BaseOperator):
if backend == "python-docx":
return self._parse_docx_stream(io.BytesIO(raw))
- text = self._decode(raw, source_hint=f"<bytes:{ext}>")
+ source_hint = f"<bytes:{ext}>"
+ text = self._decode(raw, source_hint=source_hint)
if backend == "csv":
return self._parse_csv_text(text)
if backend == "json":
return self._parse_json_text(text)
+ if backend == "jsonl":
+ return self._parse_json_lines_text(text, source_hint=source_hint)
return [{"text": text, "metadata": {}}]
def _parse_file(self, file_path: Path, ext: str) -> list[dict[str, Any]]:
@@ -312,6 +317,8 @@ class DocumentLoaderOperator(BaseOperator):
return self._parse_csv(file_path)
if backend == "json":
return self._parse_json(file_path)
+ if backend == "jsonl":
+ return self._parse_json_lines(file_path)
if backend == "pypdf":
with file_path.open("rb") as fh:
return self._parse_pdf_stream(fh)
@@ -373,6 +380,27 @@ class DocumentLoaderOperator(BaseOperator):
return [self._json_item_to_doc(item, item_index=idx) for idx, item
in enumerate(data)]
return [self._json_item_to_doc(data, item_index=None)]
+ def _parse_json_lines(self, file_path: Path) -> list[dict[str, Any]]:
+ return self._parse_json_lines_text(self._read_text(file_path),
source_hint=str(file_path))
+
+ def _parse_json_lines_text(self, text: str, *, source_hint: str) ->
list[dict[str, Any]]:
+ documents: list[dict[str, Any]] = []
+ # split("\n") rather than splitlines(): JSON Lines is defined with \n,
while splitlines()
+ # also breaks on U+2028/U+2029/U+0085, which are legal unescaped
characters inside a JSON
+ # string and would tear a valid record in half. A trailing \r from
CRLF is JSON whitespace.
+ for line_number, line in enumerate(text.split("\n"), start=1):
+ if not line.strip():
+ continue
+ try:
+ item = json.loads(line)
+ except json.JSONDecodeError as e:
+ raise ValueError(
+ f"Failed to parse {source_hint}: "
+ f"invalid JSON on line {line_number}, column {e.colno}:
{e.msg}"
+ ) from e
+ documents.append(self._json_item_to_doc(item,
item_index=len(documents)))
+ return documents
+
def _json_item_to_doc(self, item: Any, *, item_index: int | None) ->
dict[str, Any]:
metadata: dict[str, Any] = {}
if item_index is not None:
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
b/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
index 19628d5c058..60031e3edfb 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_document_loader.py
@@ -215,6 +215,73 @@ class TestJsonParser:
assert len(result) == 2
+class TestJsonLinesParser:
+ def test_one_document_per_non_empty_line(self, tmp_path):
+ f = tmp_path / "items.jsonl"
+ f.write_text('{"title": "First"}\n\n[1, 2]\n', encoding="utf-8")
+
+ op = DocumentLoaderOperator(task_id="test", source_path=str(f))
+ result = op.execute(context=MagicMock())
+
+ assert len(result) == 2
+ assert result[0]["text"] == "title: First"
+ assert result[1]["text"] == "[1, 2]"
+ assert [doc["metadata"]["item_index"] for doc in result] == [0, 1]
+
+ def test_source_bytes_uses_json_text_field(self):
+ raw = b'{"body": "First", "source": "a"}\n{"body": "Second", "source":
"b"}\n'
+ op = DocumentLoaderOperator(
+ task_id="test",
+ source_bytes=raw,
+ file_type=".jsonl",
+ json_text_field="body",
+ )
+ result = op.execute(context=MagicMock())
+
+ assert [doc["text"] for doc in result] == ["First", "Second"]
+ assert [doc["metadata"]["source"] for doc in result] == ["a", "b"]
+
+ def test_invalid_json_line(self):
+ raw = b'{"valid": true}\n\n{"invalid": }\n'
+ op = DocumentLoaderOperator(task_id="test", source_bytes=raw,
file_type=".jsonl")
+
+ with pytest.raises(
+ ValueError, match=r"Failed to parse <bytes:\.jsonl>: invalid JSON
on line 3, column 13"
+ ):
+ op.execute(context=MagicMock())
+
+ def test_invalid_json_line_in_file_names_source(self, tmp_path):
+ f = tmp_path / "invalid.jsonl"
+ f.write_text('{"valid": true}\n{"invalid": }\n', encoding="utf-8")
+ op = DocumentLoaderOperator(task_id="test", source_path=str(f))
+
+ with pytest.raises(ValueError, match=r"invalid JSON on line 2, column
13") as exc_info:
+ op.execute(context=MagicMock())
+
+ assert str(f) in str(exc_info.value)
+
+ def test_line_separator_inside_string_is_not_a_record_break(self):
+ # U+2028 is a legal unescaped character inside a JSON string, but
str.splitlines()
+ # breaks on it, which would tear this record in two and shift every
later line number.
+ raw = '{"a": "x\u2028y"}\n{"b": 1}\n'.encode()
+ op = DocumentLoaderOperator(task_id="test", source_bytes=raw,
file_type=".jsonl")
+
+ result = op.execute(context=MagicMock())
+
+ assert len(result) == 2
+ assert result[0]["text"] == "a: x\u2028y"
+ assert result[1]["text"] == "b: 1"
+
+ def test_directory_recognizes_jsonl_extension(self, tmp_path):
+ (tmp_path / "items.jsonl").write_text('{"name": "kept"}\n',
encoding="utf-8")
+
+ op = DocumentLoaderOperator(task_id="test", source_path=str(tmp_path))
+ result = op.execute(context=MagicMock())
+
+ assert len(result) == 1
+ assert result[0]["text"] == "name: kept"
+
+
def _make_mock_pypdf_module(mock_reader):
"""Create a fake pypdf module with a PdfReader that returns mock_reader."""
mock_module = MagicMock()
