diff --git a/providers/common/ai/docs/operators/llm_file_analysis.rst b/providers/common/ai/docs/operators/llm_file_analysis.rst index 8b2733ab05acb..e0e9e4d80584e 100644 --- a/providers/common/ai/docs/operators/llm_file_analysis.rst +++ b/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: 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 f962119616e7a..632b2f9d2acd0 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 @@ -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: 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 32091ffa59b6d..3dd26d67e0977 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 @@ -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"