Support bzip2 and xz compressed inputs in LLM file analysis

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

@@ -155,9 +155,10 @@ Supported Formats
 
 - Text-like: ``.log``, ``.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``, ``.json``, and ``.csv`` (``.log.gz``, ``.csv.bz2``, ``.json.xz``, ...).
+- Compression is not supported for ``.parquet``, ``.avro``, image, or PDF
+  inputs.
 
 Parquet and Avro readers require their corresponding optional extras:
 

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

@@ -18,11 +18,13 @@
 
 from __future__ import annotations
 
+import bz2
 import csv
 import gzip
 import io
 import json
 import logging
+import lzma
 from bisect import insort
 from dataclasses import dataclass
 from pathlib import PurePosixPath
@@ -38,7 +40,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
 
@@ -63,7 +65,12 @@
     "xz": "xz",
     "zst": "zstd",
 }
-_GZIP_SUPPORTED_FORMATS = frozenset({"csv", "json", "log"})
+_DECOMPRESSORS: dict[str, Callable[..., io.BufferedIOBase]] = {
+    "bzip2": bz2.open,
+    "gzip": gzip.open,
+    "xz": lzma.open,
+}
+_COMPRESSION_SUPPORTED_FORMATS = frozenset({"csv", "json", "log"})
 _TEXT_SAMPLE_HEAD_CHARS = 8_000
 _TEXT_SAMPLE_TAIL_CHARS = 2_000
 _MEDIA_TYPES = {
@@ -369,12 +376,12 @@ 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 _DECOMPRESSORS:
         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."
         )
@@ -535,10 +542,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

@@ -17,7 +17,9 @@
 from __future__ import annotations
 
 import builtins
+import bz2
 import gzip
+import lzma
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -333,6 +335,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),
@@ -354,19 +358,29 @@ 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", ["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"):
             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", "compress"),
+        [
+            ("gz", "gzip", gzip.compress),
+            ("bz2", "bzip2", bz2.compress),
+            ("xz", "xz", lzma.compress),
+        ],
+    )
+    def test_read_raw_bytes_decompresses(self, tmp_path, suffix, compression, compress):
+        path = tmp_path / f"events.log.{suffix}"
+        path.write_bytes(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"