This is an automated email from the ASF dual-hosted git repository. kaxil pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
commit d4ba567a390b929391f9100fb9b042132a1848f4 Author: Kaxil Naik <[email protected]> AuthorDate: Wed Sep 30 07:04:26 2026 +0100 Tell the model how many rows a Parquet or Avro file holds (#73925) ObjectStorageToolset's read_file returned a Parquet or Avro file as its schema and first 20 rows, with no row count, and silently ignored offset and limit. Asked how many rows a file held, a model could only retry with an offset and get the same 20 rows back. The result now opens with the file's row count and says that offset and limit do not apply, ahead of the sample so that cutting a long result never drops it. Parquet's count comes from the file footer. An Avro file is read in one pass over its blocks: the sample comes from the first blocks and the count from every block header, so only the blocks the sample reaches have records decoded. A file whose schema is not a record now shows its values in the sample too. A corrupt deflate block anywhere in an Avro file, and corrupt .xz text, raise zlib.error and lzma.LZMAError, which are not OSError or ValueError. Both are now refused like any other unreadable file instead of failing the run. --- .../common/ai/docs/toolsets/object_storage.rst | 5 +- .../providers/common/ai/toolsets/object_storage.py | 24 ++++++-- .../providers/common/ai/utils/file_analysis.py | 69 ++++++++++++++++------ .../unit/common/ai/toolsets/test_object_storage.py | 65 ++++++++++++++++++-- 4 files changed, 136 insertions(+), 27 deletions(-) diff --git a/providers/common/ai/docs/toolsets/object_storage.rst b/providers/common/ai/docs/toolsets/object_storage.rst index 7dae611ecc2..fdc069bc41d 100644 --- a/providers/common/ai/docs/toolsets/object_storage.rst +++ b/providers/common/ai/docs/toolsets/object_storage.rst @@ -58,8 +58,9 @@ The three tools A text file a window of lines at a time. The model passes ``offset`` and ``limit`` as line numbers, and a result that stops early says which ``offset`` to continue from, the same shape as the sandbox's ``read_file``. A Parquet or Avro file comes - back as its schema and its first 20 rows. Text compressed as ``.gz``, ``.bz2`` or - ``.xz`` is decompressed first. + back as its row count, its schema and its first 20 rows; ``offset`` and ``limit`` do + not apply to it. To query one, such as summing a column, use :doc:`datafusion`. Text + compressed as ``.gz``, ``.bz2`` or ``.xz`` is decompressed first. Every path the model supplies is relative to the root. An absolute path, a path with a scheme such as ``s3://``, and a path that climbs out with ``..`` are refused, and on a diff --git a/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py index 1eeee652a28..bd688d14261 100644 --- a/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py +++ b/providers/common/ai/src/airflow/providers/common/ai/toolsets/object_storage.py @@ -18,7 +18,9 @@ from __future__ import annotations +import lzma import os +import zlib from datetime import datetime, timezone from pathlib import PurePosixPath from typing import TYPE_CHECKING, Any, Literal @@ -96,7 +98,8 @@ _DESCRIPTIONS = { GET_FILE_INFO: "Get the size and last-modified time of one file or directory.", READ_FILE: ( "Read a text file. Long files are returned a window at a time and the result tells you the " - "offset to continue from. For a Parquet or Avro file, returns its schema and first rows." + "offset to continue from. For a Parquet or Avro file, returns its row count, schema and first " + "rows; offset and limit do not apply to it." ), } @@ -123,7 +126,7 @@ class ObjectStorageToolset(AirflowToolset): is absolute, carries a scheme, or climbs out of the root with ``..`` is refused. ``read_file`` returns a text file a window of lines at a time, like the sandbox's own - ``read_file``, and a Parquet or Avro file as its schema and first rows. Compressed text + ``read_file``, and a Parquet or Avro file as its row count, schema and first rows. Compressed text (``.gz``, ``.bz2``, ``.xz``) is decompressed. Images, PDFs and other binary files are refused, as is any file larger than ``max_read_bytes``. So is a file that cannot be read, such as a corrupt one, or one the connection may not open: the model is told why, and the @@ -223,7 +226,14 @@ class ObjectStorageToolset(AirflowToolset): ) except ToolFailed: raise - except (OSError, EOFError, ValueError, AirflowOptionalProviderFeatureException) as e: + except ( + OSError, + EOFError, + ValueError, + zlib.error, + lzma.LZMAError, + AirflowOptionalProviderFeatureException, + ) as e: # Storage the connection may not read, a corrupt or mislabelled file, a codec this # Python build lacks: final for this path, so the model is told, not the task failed. raise ToolFailed(f"{relative or '/'!r} cannot be read: {type(e).__name__}: {e}") from None @@ -312,7 +322,13 @@ class ObjectStorageToolset(AirflowToolset): sample = sample_columnar_file( target, file_format=columnar, sample_rows=_SAMPLE_ROWS, max_bytes=self._max_read_bytes ) - return _cut(sample, self._max_output_bytes) + # First, so that cutting a long sample never drops it. + shown = min(_SAMPLE_ROWS, sample.total_rows) + header = ( + f"Rows: {sample.total_rows}. The schema and the first {shown} rows follow; " + f"offset and limit do not apply to {columnar.capitalize()} files.\n" + ) + return _cut(header + sample.text, self._max_output_bytes) data = read_bytes(target, compression=compression, max_bytes=self._max_read_bytes) except LLMFileAnalysisLimitExceededError: raise ToolFailed( 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 930a3fae0f6..eb402dfcd9f 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 @@ -21,6 +21,7 @@ from __future__ import annotations import csv import gzip import io +import itertools import json import logging from bisect import insort @@ -139,6 +140,16 @@ class _RenderResult: content_size_bytes: int +@dataclass +class ColumnarSample: + """A Parquet or Avro file described for a model.""" + + text: str + """Its schema and first rows.""" + total_rows: int + """How many rows the whole file holds.""" + + def build_file_analysis_request( *, file_path: str, @@ -484,14 +495,18 @@ def _render_csv( def sample_columnar_file( path: ObjectStoragePath, *, file_format: Literal["parquet", "avro"], sample_rows: int, max_bytes: int -) -> str: +) -> ColumnarSample: """ - Describe a Parquet or Avro file for a model: its schema and its first ``sample_rows`` rows. + Describe a Parquet or Avro file for a model: its schema, its first ``sample_rows`` rows and its row count. :raises LLMFileAnalysisLimitExceededError: if the file is larger than ``max_bytes``. """ - render = _render_parquet if file_format == "parquet" else _render_avro - return render(path, sample_rows=sample_rows, max_content_bytes=max_bytes).text + if file_format == "parquet": + result = _render_parquet(path, sample_rows=sample_rows, max_content_bytes=max_bytes) + else: + result = _render_avro(path, sample_rows=sample_rows, max_content_bytes=max_bytes, count_rows=True) + # Both count every row here: Parquet from its footer, Avro by reading every block header. + return ColumnarSample(text=result.text, total_rows=result.estimated_rows or 0) def _render_parquet(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes: int) -> _RenderResult: @@ -533,7 +548,9 @@ def _render_parquet(path: ObjectStoragePath, *, sample_rows: int, max_content_by ) -def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes: int) -> _RenderResult: +def _render_avro( + path: ObjectStoragePath, *, sample_rows: int, max_content_bytes: int, count_rows: bool = False +) -> _RenderResult: try: import fastavro except ImportError as exc: @@ -541,7 +558,7 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes "Avro analysis requires the `avro` extra for apache-airflow-providers-common-ai." ) from exc - sampled_rows: list[dict[str, Any]] = [] + sampled_rows: list[Any] = [] total_rows = 0 with path.open("rb") as handle: handle.seek(0, io.SEEK_END) @@ -551,18 +568,31 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes raise LLMFileAnalysisLimitExceededError( f"File {path} exceeds the configured processed-content limit: {content_size_bytes} bytes > {max_content_bytes} bytes." ) - reader = fastavro.reader(handle) - writer_schema = getattr(reader, "writer_schema", None) fully_read = False - if sample_rows > 0: - for record in reader: - total_rows += 1 - if isinstance(record, dict): - sampled_rows.append({str(key): value for key, value in record.items()}) - if total_rows >= sample_rows: - break - else: - fully_read = True + if count_rows: + # Each block header carries its record count, so only the blocks the sample reaches + # are decoded; the rest are counted. + blocks = fastavro.block_reader(handle) + writer_schema = blocks.writer_schema + for block in blocks: + total_rows += block.num_records + if len(sampled_rows) < sample_rows: + sampled_rows.extend( + _avro_sample_row(record) + for record in itertools.islice(block, sample_rows - len(sampled_rows)) + ) + fully_read = True + else: + reader = fastavro.reader(handle) + writer_schema = reader.writer_schema + if sample_rows > 0: + for record in reader: + total_rows += 1 + sampled_rows.append(_avro_sample_row(record)) + if total_rows >= sample_rows: + break + else: + fully_read = True payload = [ f"Schema: {dumps_masked(writer_schema, indent=2)}", "Sample rows:", @@ -575,6 +605,11 @@ def _render_avro(path: ObjectStoragePath, *, sample_rows: int, max_content_bytes ) +def _avro_sample_row(record: Any) -> Any: + # A file whose schema is not a record holds bare values, which are sampled as they are. + return {str(key): value for key, value in record.items()} if isinstance(record, dict) else record + + def read_bytes(path: ObjectStoragePath, *, compression: str | None, max_bytes: int) -> bytes: """ Read ``path``, decompressing it with ``compression``, and refuse more than ``max_bytes``. diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py index a02bb133b7a..989552db68f 100644 --- a/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py +++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_object_storage.py @@ -18,6 +18,7 @@ from __future__ import annotations import asyncio import gzip +import io import json import uuid from typing import Any @@ -235,11 +236,12 @@ class TestReadFile: assert result == "CREATE TABLE t (id INT);" - def test_a_corrupt_file_is_refused_instead_of_failing_the_run(self, storage): - (storage / "broken.txt.gz").write_bytes(b"not gzip at all") + @pytest.mark.parametrize("path", ["broken.txt.gz", "broken.txt.bz2", "broken.txt.xz"]) + def test_a_corrupt_file_is_refused_instead_of_failing_the_run(self, storage, path): + (storage / path).write_bytes(b"not compressed at all") with pytest.raises(ToolFailed, match="cannot be read"): - _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "broken.txt.gz"}) + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": path}) @pytest.mark.parametrize( ("path", "module", "error"), @@ -268,10 +270,65 @@ class TestReadFile: result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "orders.parquet"}) - assert result.startswith("Schema: id: int64, total: int64") + assert result.startswith( + "Rows: 100. The schema and the first 20 rows follow; " + "offset and limit do not apply to Parquet files.\n" + "Schema: id: int64, total: int64" + ) assert '"id": 19' in result assert '"id": 20' not in result + @pytest.mark.parametrize( + ("rows", "shown", "blocks"), + [ + pytest.param(50, 20, 4, id="longer-than-the-sample"), + pytest.param(5, 5, 1, id="shorter-than-the-sample"), + ], + ) + def test_reads_the_row_count_schema_and_first_rows_of_an_avro_file(self, storage, rows, shown, blocks): + fastavro = pytest.importorskip("fastavro") + schema = {"type": "record", "name": "order", "fields": [{"name": "id", "type": "int"}]} + with (storage / "orders.avro").open("wb") as out: + fastavro.writer(out, schema, [{"id": i} for i in range(rows)], sync_interval=16) + # Several blocks for the longer file, so the count has to add up more than one. + with (storage / "orders.avro").open("rb") as written: + assert len(list(fastavro.block_reader(written))) == blocks + + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "orders.avro"}) + + assert result.startswith( + f"Rows: {rows}. The schema and the first {shown} rows follow; " + "offset and limit do not apply to Avro files." + ) + assert f'"id": {shown - 1}' in result + assert f'"id": {shown}' not in result + + def test_samples_an_avro_file_whose_schema_is_not_a_record(self, storage): + fastavro = pytest.importorskip("fastavro") + with (storage / "totals.avro").open("wb") as out: + fastavro.writer(out, "long", [10, 20, 30]) + + result = _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "totals.avro"}) + + assert result.startswith("Rows: 3. The schema and the first 3 rows follow") + assert result.endswith("Sample rows:\n[\n 10,\n 20,\n 30\n]") + + def test_a_corrupt_avro_block_past_the_sample_is_refused(self, storage): + """Counting reaches every block, and a deflate block that does not inflate raises zlib.error.""" + fastavro = pytest.importorskip("fastavro") + schema = {"type": "record", "name": "order", "fields": [{"name": "note", "type": "string"}]} + buffer = io.BytesIO() + fastavro.writer( + buffer, schema, [{"note": f"row {i} " * 8} for i in range(50)], codec="deflate", sync_interval=256 + ) + data = bytearray(buffer.getvalue()) + # The last 16 bytes are the sync marker; the deflate data of the last block sits before it. + data[-40:-20] = b"\xff" * 20 + (storage / "broken.avro").write_bytes(bytes(data)) + + with pytest.raises(ToolFailed, match="cannot be read"): + _call(ObjectStorageToolset(f"file://{storage}"), "read_file", {"path": "broken.avro"}) + def test_reads_a_window_and_says_where_to_continue(self, storage): result = _call( ObjectStorageToolset(f"file://{storage}"),
