Support bzip2 and xz compressed inputs in LLM file analysis

providers/common/ai/docs/operators/llm_file_analysis.rst

@@ -153,11 +153,20 @@ This operator also inherits ``LLMOperator``'s HITL review parameters --
 Supported Formats
 -----------------
 
-- Text-like: ``.log``, ``.json``, ``.csv``, ``.parquet``, ``.avro``
+- Text-like: ``.log``, ``.txt``, ``.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``.
-- Gzip is not supported for ``.parquet``, ``.avro``, image, or PDF inputs.
+- ``gzip``, ``bzip2``, and ``xz`` compressed text inputs are supported for
+  ``.log``, ``.txt``, ``.json``, and ``.csv`` (``.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:
 

providers/common/ai/src/airflow/providers/common/ai/utils/file_analysis.py

@@ -28,6 +28,17 @@
 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.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
 
@@ -64,7 +75,14 @@
     "xz": "xz",
     "zst": "zstd",
 }
-_GZIP_SUPPORTED_FORMATS = frozenset({"csv", "json", "log", "txt"})
+_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"})
 _TEXT_SAMPLE_HEAD_CHARS = 8_000
 _TEXT_SAMPLE_TAIL_CHARS = 2_000
 _MEDIA_TYPES = {
@@ -370,15 +388,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
 
 
@@ -536,10 +559,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:

providers/common/ai/tests/unit/common/ai/utils/test_file_analysis.py

@@ -30,6 +30,7 @@
     LLMFileAnalysisUnsupportedFormatError,
 )
 from airflow.providers.common.ai.utils.file_analysis import (
+    _DECOMPRESSORS,
     FileAnalysisRequest,
     _infer_partitions,
     _read_raw_bytes,
@@ -256,9 +257,14 @@ def test_total_size_limit_failure(self, tmp_path):
 
         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(
@@ -355,6 +361,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),
@@ -378,22 +386,77 @@ def test_detect_file_format_rejects_unsupported_compression(self, tmp_path):
         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