This is an automated email from the ASF dual-hosted git repository.
guan404ming 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 1e0f48ce6d0 Support bzip2 and xz compressed inputs in LLM file
analysis (#70302)
1e0f48ce6d0 is described below
commit 1e0f48ce6d08e537e4e4583445e7df2a9225bd88
Author: Guan-Ming Chiu <[email protected]>
AuthorDate: Thu Aug 27 20:14:24 2026 +0800
Support bzip2 and xz compressed inputs in LLM file analysis (#70302)
* Support bzip2 and xz compressed inputs in LLM file analysis
* Guard optional bz2/lzma imports in LLM file analysis
* Surface missing codec modules and pin multi-stream behavior
* Fix spelling check for xz in LLM file analysis docs
---
.../common/ai/docs/operators/llm_file_analysis.rst | 16 ++++-
.../providers/common/ai/utils/file_analysis.py | 39 +++++++---
.../unit/common/ai/utils/test_file_analysis.py | 83 +++++++++++++++++++---
3 files changed, 117 insertions(+), 21 deletions(-)
diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst
b/providers/common/ai/docs/operators/llm_file_analysis.rst
index bb019e2b6a0..ac688ebfaf4 100644
--- a/providers/common/ai/docs/operators/llm_file_analysis.rst
+++ b/providers/common/ai/docs/operators/llm_file_analysis.rst
@@ -167,9 +167,19 @@ Supported Formats
- Text-like: ``.log``, ``.txt``, ``.md``, ``.json``, ``.csv``, ``.parquet``,
``.avro``
- Multimodal: ``.png``, ``.jpg``, ``.jpeg``, ``.pdf`` when ``multi_modal=True``
-- Gzip-compressed text inputs are supported for ``.log.gz``, ``.json.gz``, and
- ``.csv.gz``, ``.txt.gz``, and ``.md.gz``.
-- Gzip is not supported for ``.parquet``, ``.avro``, image, or PDF inputs.
+- ``gzip``, ``bzip2``, and ``xz`` compressed text inputs are supported for
+ ``.log``, ``.json``, ``.csv``, ``.txt``, and ``.md`` (``.log.gz``,
``.csv.bz2``,
+ ``.json.xz``, ...).
+ ``bzip2`` and ``xz`` need a Python interpreter built with the ``bz2`` and
+ ``lzma`` modules; otherwise those inputs raise
+ ``AirflowOptionalProviderFeatureException``.
+- Compression is not supported for ``.parquet``, ``.avro``, image, or PDF
+ inputs.
+- For ``bzip2`` and ``xz``, concatenated streams are read in full, but any
+ data after a point that does not start a valid stream (for example ``xz``
+ Stream Padding between streams, or trailing garbage) is silently ignored,
+ so only the content up to that point is analyzed. ``gzip`` rejects such
+ files instead (``BadGzipFile``).
Parquet and Avro readers require their corresponding optional extras:
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py
b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py
index 6f6366b515e..f53a23fa63a 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py
@@ -28,6 +28,17 @@ from dataclasses import dataclass
from pathlib import PurePosixPath
from typing import TYPE_CHECKING, Any
+# bz2/lzma are optional CPython extensions and may be missing from some
interpreter builds
+try:
+ import bz2
+except ImportError:
+ bz2 = None # type: ignore[assignment]
+
+try:
+ import lzma
+except ImportError:
+ lzma = None # type: ignore[assignment]
+
from pydantic_ai.messages import BinaryContent
from airflow.providers.common.ai.exceptions import (
@@ -38,7 +49,7 @@ from airflow.providers.common.ai.exceptions import (
from airflow.providers.common.compat.sdk import
AirflowOptionalProviderFeatureException, ObjectStoragePath
if TYPE_CHECKING:
- from collections.abc import Sequence
+ from collections.abc import Callable, Sequence
from pydantic_ai.messages import UserContent
@@ -65,7 +76,14 @@ _COMPRESSION_SUFFIXES = {
"xz": "xz",
"zst": "zstd",
}
-_GZIP_SUPPORTED_FORMATS = frozenset({"csv", "json", "log", "txt", "md"})
+_CODEC_MODULES = {"bzip2": "bz2", "xz": "lzma"}
+_KNOWN_CODECS = frozenset({"gzip", *_CODEC_MODULES})
+_DECOMPRESSORS: dict[str, Callable[..., io.BufferedIOBase]] = {"gzip":
gzip.open}
+if bz2 is not None:
+ _DECOMPRESSORS["bzip2"] = bz2.open
+if lzma is not None:
+ _DECOMPRESSORS["xz"] = lzma.open
+_COMPRESSION_SUPPORTED_FORMATS = frozenset({"csv", "json", "log", "txt", "md"})
_TEXT_SAMPLE_HEAD_CHARS = 8_000
_TEXT_SAMPLE_TAIL_CHARS = 2_000
_MEDIA_TYPES = {
@@ -371,15 +389,20 @@ def detect_file_format(path: ObjectStoragePath) ->
tuple[str, str | None]:
raise LLMFileAnalysisUnsupportedFormatError(
f"Unsupported file format {detected!r} for {path}. Supported
formats: {', '.join(SUPPORTED_FILE_FORMATS)}."
)
- if compression and compression != "gzip":
+ if compression and compression not in _KNOWN_CODECS:
log.info("Rejecting file %s because compression=%s is not supported.",
path, compression)
raise LLMFileAnalysisUnsupportedFormatError(
f"Compression {compression!r} is not supported for file analysis."
)
- if compression == "gzip" and detected not in _GZIP_SUPPORTED_FORMATS:
+ if compression and detected not in _COMPRESSION_SUPPORTED_FORMATS:
raise LLMFileAnalysisUnsupportedFormatError(
f"Compression {compression!r} is not supported for {detected!r}
file analysis."
)
+ if compression and compression not in _DECOMPRESSORS:
+ raise AirflowOptionalProviderFeatureException(
+ f"Compression {compression!r} requires the
{_CODEC_MODULES[compression]!r} module, "
+ "which is missing from this Python build."
+ )
return detected, compression
@@ -537,10 +560,10 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows:
int, max_content_bytes
def _read_raw_bytes(path: ObjectStoragePath, *, compression: str | None,
max_bytes: int) -> bytes:
with path.open("rb") as handle:
- if compression == "gzip":
- with gzip.GzipFile(fileobj=handle) as gzip_handle:
- return _read_limited_bytes(gzip_handle, path=path,
max_bytes=max_bytes)
- return _read_limited_bytes(handle, path=path, max_bytes=max_bytes)
+ if compression is None:
+ return _read_limited_bytes(handle, path=path, max_bytes=max_bytes)
+ with _DECOMPRESSORS[compression](handle) as decompressed:
+ return _read_limited_bytes(decompressed, path=path,
max_bytes=max_bytes)
def _read_limited_bytes(handle: io.BufferedIOBase, *, path: ObjectStoragePath,
max_bytes: int) -> bytes:
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py
b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py
index bba7f56111a..b2a80186562 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py
@@ -30,6 +30,7 @@ from airflow.providers.common.ai.exceptions import (
LLMFileAnalysisUnsupportedFormatError,
)
from airflow.providers.common.ai.utils.file_analysis import (
+ _DECOMPRESSORS,
FileAnalysisRequest,
_infer_partitions,
_read_raw_bytes,
@@ -278,9 +279,14 @@ class TestBuildFileAnalysisRequest:
mock_prepare.assert_not_called()
- def test_gzip_expansion_respects_processed_content_limit(self, tmp_path):
- path = tmp_path / "big.log.gz"
- path.write_bytes(gzip.compress(b"A" * 20_000))
+ @pytest.mark.parametrize(
+ ("suffix", "module_name"),
+ [("gz", "gzip"), ("bz2", "bz2"), ("xz", "lzma")],
+ )
+ def test_compressed_expansion_respects_processed_content_limit(self,
tmp_path, suffix, module_name):
+ codec = pytest.importorskip(module_name)
+ path = tmp_path / f"big.log.{suffix}"
+ path.write_bytes(codec.compress(b"A" * 20_000))
with pytest.raises(LLMFileAnalysisLimitExceededError,
match="processed-content limit"):
build_file_analysis_request(
@@ -377,6 +383,8 @@ class TestFileAnalysisHelpers:
[
("events.csv", "csv", None),
("events.csv.gz", "csv", "gzip"),
+ ("events.csv.bz2", "csv", "bzip2"),
+ ("events.json.xz", "json", "xz"),
("dashboard.jpg", "jpg", None),
("report.pdf", "pdf", None),
("app", "log", None),
@@ -402,22 +410,77 @@ class TestFileAnalysisHelpers:
with pytest.raises(LLMFileAnalysisUnsupportedFormatError,
match="Compression"):
detect_file_format(ObjectStoragePath(str(path)))
- @pytest.mark.parametrize("filename", ["sample.parquet.gz",
"sample.avro.gz", "sample.png.gz"])
- def
test_detect_file_format_rejects_unsupported_gzip_format_combinations(self,
tmp_path, filename):
+ @pytest.mark.parametrize(
+ ("filename", "codec", "module_name"),
+ [("events.csv.bz2", "bzip2", "bz2"), ("events.json.xz", "xz", "lzma")],
+ )
+ def test_detect_file_format_without_codec_module(self, tmp_path, filename,
codec, module_name):
+ path = tmp_path / filename
+ path.write_bytes(b"content")
+ gz_path = tmp_path / "events.csv.gz"
+ gz_path.write_bytes(b"content")
+ plain_path = tmp_path / "events.csv"
+ plain_path.write_bytes(b"content")
+
+ with patch.dict(_DECOMPRESSORS, {"gzip": gzip.open}, clear=True):
+ with pytest.raises(
+ AirflowOptionalProviderFeatureException,
+ match=f"Compression '{codec}' requires the '{module_name}'
module",
+ ):
+ detect_file_format(ObjectStoragePath(str(path)))
+ assert detect_file_format(ObjectStoragePath(str(gz_path))) ==
("csv", "gzip")
+ assert detect_file_format(ObjectStoragePath(str(plain_path))) ==
("csv", None)
+
+ @pytest.mark.parametrize(
+ "filename", ["sample.parquet.gz", "sample.avro.bz2", "sample.png.xz",
"sample.pdf.gz"]
+ )
+ def
test_detect_file_format_rejects_unsupported_compression_format_combinations(self,
tmp_path, filename):
path = tmp_path / filename
path.write_bytes(b"content")
- with pytest.raises(LLMFileAnalysisUnsupportedFormatError, match="not
supported for"):
+ with pytest.raises(
+ LLMFileAnalysisUnsupportedFormatError, match=r"not supported for
'\w+' file analysis"
+ ):
detect_file_format(ObjectStoragePath(str(path)))
- def test_read_raw_bytes_decompresses_gzip(self, tmp_path):
- path = tmp_path / "events.log.gz"
- path.write_bytes(gzip.compress(b"line one\nline two\n"))
+ @pytest.mark.parametrize(
+ ("suffix", "compression", "module_name"),
+ [
+ ("gz", "gzip", "gzip"),
+ ("bz2", "bzip2", "bz2"),
+ ("xz", "xz", "lzma"),
+ ],
+ )
+ def test_read_raw_bytes_decompresses(self, tmp_path, suffix, compression,
module_name):
+ codec = pytest.importorskip(module_name)
+ path = tmp_path / f"events.log.{suffix}"
+ path.write_bytes(codec.compress(b"line one\nline two\n"))
- content = _read_raw_bytes(ObjectStoragePath(str(path)),
compression="gzip", max_bytes=1_024)
+ content = _read_raw_bytes(ObjectStoragePath(str(path)),
compression=compression, max_bytes=1_024)
assert content == b"line one\nline two\n"
+ @pytest.mark.parametrize(
+ ("suffix", "compression", "module_name", "separator", "expected"),
+ [
+ ("bz2", "bzip2", "bz2", b"", b"first\nsecond\n"),
+ ("bz2", "bzip2", "bz2", b"GARBAGE", b"first\n"),
+ ("xz", "xz", "lzma", b"", b"first\nsecond\n"),
+ ("xz", "xz", "lzma", b"\x00\x00\x00\x00", b"first\n"),
+ ],
+ ids=["bz2-concatenated", "bz2-trailing-garbage", "xz-concatenated",
"xz-stream-padding"],
+ )
+ def test_read_raw_bytes_multi_stream_behavior(
+ self, tmp_path, suffix, compression, module_name, separator, expected
+ ):
+ codec = pytest.importorskip(module_name)
+ path = tmp_path / f"events.log.{suffix}"
+ path.write_bytes(codec.compress(b"first\n") + separator +
codec.compress(b"second\n"))
+
+ content = _read_raw_bytes(ObjectStoragePath(str(path)),
compression=compression, max_bytes=1_024)
+
+ assert content == expected
+
def test_truncate_text_preserves_head_and_tail(self):
text = "A" * 9_000 + "B" * 3_000
