diff --git a/docs/changelog.rst b/docs/changelog.rst index 3955d9a97..34caa1f2d 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -329,8 +329,10 @@ v0.54.0 - SQL processing correctness and cleanup * Spanner adapter modules no longer expose module-level proxy lookup hooks. * Async migration squash now builds its internal migration runner with a real migration context, matching the synchronous command path. -* ObStore Arrow streaming no longer resolves cloud ``base_path`` twice for - async streams. +* Arrow batch streaming is now explicitly Parquet-only and reads one row group + at a time across local, fsspec, and obstore backends. Obstore streams through + its seekable reader without buffering the full object, resolves cloud + ``base_path`` only once, and closes readers deterministically. * ``sql.decode()`` now renders a trailing default argument as the ``ELSE`` clause documented for DECODE-style expressions. * Async drivers can use the statement-cache direct execution path when the diff --git a/docs/reference/storage.rst b/docs/reference/storage.rst index ae4f0bd2e..2a72074ff 100644 --- a/docs/reference/storage.rst +++ b/docs/reference/storage.rst @@ -6,6 +6,21 @@ Storage abstraction layer with multiple backend support (local filesystem, fsspec, obstore), configuration-based registration, and Arrow table import/export with CSV format support. +Parquet Batch Streaming +======================= + +The ``stream_arrow_sync()`` and ``stream_arrow_async()`` backend methods stream +Parquet files in file and row-group order. They accept a keyword-only +``batch_size`` (default ``65_536``) which controls the maximum rows in each +record batch. Each read is restricted to one Parquet row group, so the I/O bound +is one row group rather than one record batch. Choose the Parquet row-group size +when writing files according to the memory bound required while reading them. + +These methods intentionally support only ``file_format="parquet"``. Use the +regular Arrow read APIs for CSV, Arrow IPC, JSON, and JSONL payloads. Closing a +sync generator or calling ``aclose()`` on its async iterator closes the active +storage reader. + Pipelines ========= diff --git a/sqlspec/protocols.py b/sqlspec/protocols.py index e87000178..18c9fe0dd 100644 --- a/sqlspec/protocols.py +++ b/sqlspec/protocols.py @@ -4,7 +4,7 @@ and runtime isinstance() checks. """ -from typing import TYPE_CHECKING, Any, ClassVar, Protocol, overload, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Protocol, overload, runtime_checkable from typing_extensions import Self @@ -543,7 +543,9 @@ def write_arrow_sync(self, path: "str | Path", table: "ArrowTable", **kwargs: An msg = "Arrow writing not implemented" raise NotImplementedError(msg) - def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> "Iterator[ArrowRecordBatch]": + def stream_arrow_sync( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> "Iterator[ArrowRecordBatch]": """Stream Arrow record batches from matching objects synchronously.""" msg = "Arrow streaming not implemented" raise NotImplementedError(msg) @@ -621,7 +623,9 @@ async def write_arrow_async(self, path: "str | Path", table: "ArrowTable", **kwa raise NotImplementedError(msg) # NOTE: Returns AsyncIterator directly; this is intentionally not async def. - def stream_arrow_async(self, pattern: str, **kwargs: Any) -> "AsyncIterator[ArrowRecordBatch]": + def stream_arrow_async( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> "AsyncIterator[ArrowRecordBatch]": """Stream Arrow record batches from matching objects.""" msg = "Async arrow streaming not implemented" raise NotImplementedError(msg) diff --git a/sqlspec/storage/_arrow_stream.py b/sqlspec/storage/_arrow_stream.py new file mode 100644 index 000000000..99951e3d9 --- /dev/null +++ b/sqlspec/storage/_arrow_stream.py @@ -0,0 +1,34 @@ +"""Shared helpers for bounded Parquet batch streaming.""" + +from pathlib import PurePath +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Iterator + + from sqlspec.typing import ArrowRecordBatch + +__all__ = ("iter_parquet_row_groups", "validate_parquet_stream_options") + +_NON_PARQUET_SUFFIXES = frozenset({".arrow", ".csv", ".feather", ".ipc", ".json", ".jsonl", ".ndjson"}) + + +def validate_parquet_stream_options(pattern: str, file_format: str, batch_size: int) -> None: + """Validate a Parquet streaming request before storage is accessed.""" + if file_format != "parquet": + msg = f"Arrow batch streaming supports only Parquet files; received file_format={file_format!r}" + raise ValueError(msg) + if batch_size <= 0: + msg = f"batch_size must be greater than zero; received {batch_size}" + raise ValueError(msg) + + suffix = PurePath(pattern).suffix.lower() + if suffix in _NON_PARQUET_SUFFIXES: + msg = f"Arrow batch streaming supports only Parquet files; pattern {pattern!r} selects {suffix} files" + raise ValueError(msg) + + +def iter_parquet_row_groups(parquet_file: Any, *, batch_size: int, **kwargs: Any) -> "Iterator[ArrowRecordBatch]": + """Yield batches while limiting each PyArrow read to one row group.""" + for row_group in range(parquet_file.num_row_groups): + yield from parquet_file.iter_batches(batch_size=batch_size, row_groups=[row_group], **kwargs) diff --git a/sqlspec/storage/backends/base.py b/sqlspec/storage/backends/base.py index 55b1eb9cd..2281fde4f 100644 --- a/sqlspec/storage/backends/base.py +++ b/sqlspec/storage/backends/base.py @@ -6,7 +6,7 @@ import contextlib from abc import abstractmethod from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast from mypy_extensions import mypyc_attr from typing_extensions import Self @@ -58,17 +58,38 @@ def _read_chunk_or_sentinel(file_obj: Any, chunk_size: int) -> Any: class AsyncArrowBatchIterator: """Async iterator wrapper for sync Arrow batch iterators.""" - __slots__ = ("_sync_iter",) + __slots__ = ("_closed", "_sync_iter") def __init__(self, sync_iterator: "Iterator[ArrowRecordBatch]") -> None: self._sync_iter = sync_iterator + self._closed = False def __aiter__(self) -> "AsyncArrowBatchIterator": return self + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, exc_type: "type[BaseException] | None", exc_val: "BaseException | None", exc_tb: "TracebackType | None" + ) -> None: + await self.aclose() + + async def aclose(self) -> None: + """Close the underlying generator and its active storage reader.""" + if self._closed: + return + self._closed = True + close = getattr(self._sync_iter, "close", None) + if close is not None: + await asyncio.get_running_loop().run_in_executor(None, close) + def _sync_next(self) -> "ArrowRecordBatch": + if self._closed: + raise _StopAsync() result = _next_or_sentinel(self._sync_iter) if result is _EXHAUSTED: + self._closed = True raise _StopAsync() return cast("ArrowRecordBatch", result) @@ -227,7 +248,9 @@ def write_arrow_sync(self, path: str, table: "ArrowTable", **kwargs: Any) -> Non raise NotImplementedError @abstractmethod - def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> "Iterator[ArrowRecordBatch]": + def stream_arrow_sync( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> "Iterator[ArrowRecordBatch]": """Stream Arrow record batches from storage synchronously.""" raise NotImplementedError @@ -300,6 +323,8 @@ async def write_arrow_async(self, path: str, table: "ArrowTable", **kwargs: Any) # NOTE: Returns AsyncIterator directly; keep in sync with ObjectStoreProtocol. @abstractmethod - def stream_arrow_async(self, pattern: str, **kwargs: Any) -> "AsyncIterator[ArrowRecordBatch]": + def stream_arrow_async( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> "AsyncIterator[ArrowRecordBatch]": """Stream Arrow record batches from storage asynchronously.""" raise NotImplementedError diff --git a/sqlspec/storage/backends/fsspec.py b/sqlspec/storage/backends/fsspec.py index 3d9f74daf..02b388e00 100644 --- a/sqlspec/storage/backends/fsspec.py +++ b/sqlspec/storage/backends/fsspec.py @@ -3,11 +3,12 @@ from collections.abc import AsyncIterator, Iterator from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload from urllib.parse import urlparse from mypy_extensions import mypyc_attr +from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options from sqlspec.storage._paths import resolve_storage_path from sqlspec.storage._utils import _log_storage_event, import_pyarrow_parquet from sqlspec.storage.backends.base import AsyncArrowBatchIterator, AsyncThreadedBytesIterator @@ -396,18 +397,23 @@ def stream_read_sync(self, path: "str | Path", chunk_size: "int | None" = None, break yield cast("bytes", chunk) - def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> Iterator["ArrowRecordBatch"]: + def stream_arrow_sync( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> Iterator["ArrowRecordBatch"]: """Stream Arrow record batches from storage synchronously. Args: pattern: The glob pattern to match. - **kwargs: Additional arguments to pass to the glob method. + file_format: Storage format. Only Parquet supports bounded batch streaming. + batch_size: Maximum number of rows in each yielded record batch. + **kwargs: Additional arguments passed to PyArrow batch iteration. Yields: Arrow record batches from matching files. """ + validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() - for obj_path in self.glob_sync(pattern, **kwargs): + for obj_path in self.glob_sync(pattern): file_handle = execute_sync_storage_operation( partial(self.fs.open, obj_path, mode="rb"), backend=self.backend_type, @@ -421,7 +427,7 @@ def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> Iterator["ArrowRecor operation="stream_arrow", path=str(obj_path), ) - yield from parquet_file.iter_batches() # pyright: ignore[reportUnknownMemberType] + yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) async def read_bytes_async(self, path: "str | Path", **kwargs: Any) -> bytes: """Read bytes from storage asynchronously.""" @@ -456,17 +462,23 @@ async def stream_read_async( return AsyncThreadedBytesIterator(file_obj, chunk_size) - def stream_arrow_async(self, pattern: str, **kwargs: Any) -> AsyncIterator["ArrowRecordBatch"]: + def stream_arrow_async( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> AsyncIterator["ArrowRecordBatch"]: """Stream Arrow record batches from storage asynchronously. Args: pattern: The glob pattern to match. - **kwargs: Additional arguments to pass to the glob method. + file_format: Storage format. Only Parquet supports bounded batch streaming. + batch_size: Maximum number of rows in each yielded record batch. + **kwargs: Additional arguments passed to PyArrow batch iteration. Returns: AsyncIterator yielding Arrow record batches. """ - return AsyncArrowBatchIterator(self.stream_arrow_sync(pattern, **kwargs)) + return AsyncArrowBatchIterator( + self.stream_arrow_sync(pattern, file_format=file_format, batch_size=batch_size, **kwargs) + ) async def read_text_async(self, path: "str | Path", encoding: str = "utf-8", **kwargs: Any) -> str: """Read text from storage asynchronously.""" diff --git a/sqlspec/storage/backends/local.py b/sqlspec/storage/backends/local.py index bcb36b11a..a9a6c4ec4 100644 --- a/sqlspec/storage/backends/local.py +++ b/sqlspec/storage/backends/local.py @@ -8,12 +8,13 @@ from collections.abc import AsyncIterator, Iterator from functools import partial from pathlib import Path -from typing import TYPE_CHECKING, Any, ClassVar, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload from urllib.parse import unquote, urlparse from mypy_extensions import mypyc_attr from sqlspec.exceptions import FileNotFoundInStorageError +from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options from sqlspec.storage._paths import strip_windows_drive_prefix from sqlspec.storage._utils import import_pyarrow_parquet from sqlspec.storage.backends.base import AsyncArrowBatchIterator, AsyncThreadedBytesIterator @@ -284,12 +285,15 @@ def write_arrow_sync(self, path: "str | Path", table: "ArrowTable", **kwargs: An path=str(resolved), ) - def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> Iterator["ArrowRecordBatch"]: + def stream_arrow_sync( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> Iterator["ArrowRecordBatch"]: """Stream Arrow record batches from files matching pattern synchronously. Yields: Arrow record batches from matching files. """ + validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() files = self.glob_sync(pattern) for file_path in files: @@ -301,7 +305,7 @@ def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> Iterator["ArrowRecor operation="stream_arrow", path=resolved_str, ) - yield from parquet_file.iter_batches() # pyright: ignore[reportUnknownMemberType] + yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) @property def supports_signing(self) -> bool: @@ -409,17 +413,23 @@ async def write_arrow_async(self, path: "str | Path", table: "ArrowTable", **kwa """ await async_(self.write_arrow_sync)(path, table, **kwargs) - def stream_arrow_async(self, pattern: str, **kwargs: Any) -> AsyncIterator["ArrowRecordBatch"]: + def stream_arrow_async( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> AsyncIterator["ArrowRecordBatch"]: """Stream Arrow record batches asynchronously. Args: pattern: Glob pattern to match files. + file_format: Storage format. Only Parquet supports bounded batch streaming. + batch_size: Maximum number of rows in each yielded record batch. **kwargs: Additional arguments passed to stream_arrow_sync(). Returns: AsyncIterator yielding Arrow record batches. """ - return AsyncArrowBatchIterator(self.stream_arrow_sync(pattern, **kwargs)) + return AsyncArrowBatchIterator( + self.stream_arrow_sync(pattern, file_format=file_format, batch_size=batch_size, **kwargs) + ) @overload async def sign_async(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ... diff --git a/sqlspec/storage/backends/obstore.py b/sqlspec/storage/backends/obstore.py index cef3e7449..13790dc70 100644 --- a/sqlspec/storage/backends/obstore.py +++ b/sqlspec/storage/backends/obstore.py @@ -11,12 +11,14 @@ from datetime import timedelta from functools import partial from pathlib import Path, PurePosixPath -from typing import TYPE_CHECKING, Any, ClassVar, Final, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, overload from urllib.parse import urlparse from mypy_extensions import mypyc_attr +from typing_extensions import Self from sqlspec.exceptions import StorageOperationFailedError +from sqlspec.storage._arrow_stream import iter_parquet_row_groups, validate_parquet_stream_options from sqlspec.storage._paths import is_file_destination, resolve_storage_path from sqlspec.storage._utils import _log_storage_event, import_pyarrow, import_pyarrow_parquet from sqlspec.storage.backends.base import AsyncArrowBatchIterator, AsyncObStoreStreamIterator @@ -37,6 +39,56 @@ __all__ = ("ObStoreBackend",) +class _ObStoreFileProxy: + """Complete obstore's seekable reader interface for PyArrow.""" + + __slots__ = ("_closed", "_reader") + + def __init__(self, reader: Any) -> None: + self._reader = reader + self._closed = False + + @property + def closed(self) -> bool: + return self._closed + + def readable(self) -> bool: + return not self._closed + + def seekable(self) -> bool: + return not self._closed and bool(self._reader.seekable()) + + def writable(self) -> bool: + return False + + def read(self, size: int = -1) -> bytes: + if size < 0: + return cast("bytes", self._reader.readall()) + return cast("bytes", self._reader.read(size)) + + def readinto(self, buffer: Any) -> int: + data = self.read(len(buffer)) + buffer[: len(data)] = data + return len(data) + + def seek(self, offset: int, whence: int = 0) -> int: + return cast("int", self._reader.seek(offset, whence)) + + def tell(self) -> int: + return cast("int", self._reader.tell()) + + def close(self) -> None: + if not self._closed: + self._closed = True + self._reader.close() + + def __enter__(self) -> Self: + return self + + def __exit__(self, *_: Any) -> None: + self.close() + + @mypyc_attr(allow_interpreted_subclasses=True) class ObStoreBackend: """Object storage backend using obstore. @@ -54,8 +106,8 @@ class ObStoreBackend: whereas cloud stores use base_path as a prefix. - Native Streaming: Uses obstore's native streaming yielding Buffer objects, which are converted to bytes. - - Seekable Streams: PyArrow's ParquetFile requires a seekable file, so we wrap - the buffered stream accordingly (e.g. using io.BytesIO). + - Seekable Streams: PyArrow's ParquetFile reads through obstore's seekable + ``open_reader`` interface without draining the object into memory. - Thread Offloading: Uses async_() with a storage limiter to offload blocking PyArrow serialization/parsing to a thread pool, preventing event loop blocking. """ @@ -329,7 +381,11 @@ def glob_sync(self, pattern: str, **kwargs: Any) -> "list[str]": Lists all objects and filters them client-side using the pattern. """ - resolved_pattern = resolve_storage_path(pattern, self.base_path, self.protocol, strip_file_scheme=True) + resolved_pattern = ( + pattern + if self._is_local_store + else resolve_storage_path(pattern, self.base_path, self.protocol, strip_file_scheme=True) + ) all_objects = self.list_objects_sync(recursive=True, **kwargs) if "**" in pattern: @@ -486,32 +542,32 @@ def stream_read_sync(self, path: "str | Path", chunk_size: "int | None" = None, for chunk in result.stream(min_chunk_size=chunk_size): yield bytes(chunk) - def stream_arrow_sync(self, pattern: str, **kwargs: Any) -> "Iterator[ArrowRecordBatch]": + def stream_arrow_sync( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> "Iterator[ArrowRecordBatch]": """Stream Arrow record batches using obstore's native streaming synchronously. - For each matching file, streams data through a buffered wrapper - that PyArrow can read directly without loading the entire file. + For each matching file, PyArrow reads through obstore's seekable reader. Yields: - Chunks of bytes from the file, with size determined by chunk_size (default: 65536 bytes). + Arrow record batches in file and row-group order. """ + from obstore import open_reader + + validate_parquet_stream_options(pattern, file_format, batch_size) pq = import_pyarrow_parquet() - for obj_path in self.glob_sync(pattern, **kwargs): - resolved_path = resolve_storage_path(obj_path, self.base_path, self.protocol, strip_file_scheme=True) - result = execute_sync_storage_operation( - partial(self.store.get, resolved_path), + for obj_path in self.glob_sync(pattern): + reader = execute_sync_storage_operation( + partial(open_reader, self.store, obj_path), backend=self.backend_type, - operation="stream_arrow", - path=resolved_path, + operation="stream_open", + path=obj_path, ) - - buffer = io.BytesIO() - for chunk in result.stream(): - buffer.write(chunk) - buffer.seek(0) - - parquet_file = pq.ParquetFile(buffer) - yield from parquet_file.iter_batches() + with _ObStoreFileProxy(reader) as stream: + parquet_file = execute_sync_storage_operation( + partial(pq.ParquetFile, stream), backend=self.backend_type, operation="stream_arrow", path=obj_path + ) + yield from iter_parquet_row_groups(parquet_file, batch_size=batch_size, **kwargs) @property def supports_signing(self) -> bool: @@ -837,17 +893,23 @@ def _serialize() -> bytes: path=resolved_path, ) - def stream_arrow_async(self, pattern: str, **kwargs: Any) -> AsyncIterator["ArrowRecordBatch"]: + def stream_arrow_async( + self, pattern: str, *, file_format: Literal["parquet"] = "parquet", batch_size: int = 65_536, **kwargs: Any + ) -> AsyncIterator["ArrowRecordBatch"]: """Stream Arrow record batches from storage asynchronously. Args: pattern: Glob pattern to match files. + file_format: Storage format. Only Parquet supports bounded batch streaming. + batch_size: Maximum number of rows in each yielded record batch. **kwargs: Additional arguments passed to stream_arrow_sync(). Returns: AsyncIterator yielding Arrow record batches. """ - return AsyncArrowBatchIterator(self.stream_arrow_sync(pattern, **kwargs)) + return AsyncArrowBatchIterator( + self.stream_arrow_sync(pattern, file_format=file_format, batch_size=batch_size, **kwargs) + ) @overload async def sign_async(self, paths: str, expires_in: int = 3600, for_upload: bool = False) -> str: ... diff --git a/tests/unit/storage/test_arrow_streaming.py b/tests/unit/storage/test_arrow_streaming.py new file mode 100644 index 000000000..3e1cb1c45 --- /dev/null +++ b/tests/unit/storage/test_arrow_streaming.py @@ -0,0 +1,444 @@ +"""Cross-backend tests for bounded Parquet streaming.""" + +import io +from collections.abc import Iterator +from pathlib import Path +from typing import Any, cast + +import pytest + +from sqlspec.storage._arrow_stream import iter_parquet_row_groups +from sqlspec.storage.backends.local import LocalStore +from sqlspec.typing import FSSPEC_INSTALLED, OBSTORE_INSTALLED, PYARROW_INSTALLED + + +class _TrackedParquetFile: + num_row_groups = 3 + + def __init__(self) -> None: + self.calls: list[tuple[int, list[int]]] = [] + + def iter_batches(self, *, batch_size: int, row_groups: list[int], **kwargs: Any) -> Iterator[Any]: + _ = kwargs + self.calls.append((batch_size, row_groups)) + yield row_groups[0] + + +def test_row_group_iterator_does_not_touch_later_groups_before_first_batch() -> None: + parquet_file = _TrackedParquetFile() + batches = iter_parquet_row_groups(parquet_file, batch_size=17) + + assert cast("Any", next(batches)) == 0 + assert parquet_file.calls == [(17, [0])] + + assert cast("Any", list(batches)) == [1, 2] + assert parquet_file.calls == [(17, [0]), (17, [1]), (17, [2])] + + +@pytest.mark.parametrize("batch_size", [0, -1]) +def test_local_stream_rejects_nonpositive_batch_size_before_glob( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, batch_size: int +) -> None: + store = LocalStore(str(tmp_path)) + monkeypatch.setattr(LocalStore, "glob_sync", lambda *_args, **_kwargs: pytest.fail("storage accessed")) + + with pytest.raises(ValueError, match="batch_size must be greater than zero"): + list(store.stream_arrow_sync("*.parquet", batch_size=batch_size)) + + +@pytest.mark.parametrize("pattern", ["*.csv", "data.jsonl", "data.arrow", "data.ipc"]) +def test_local_stream_rejects_recognized_non_parquet_suffix_before_glob( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, pattern: str +) -> None: + store = LocalStore(str(tmp_path)) + monkeypatch.setattr(LocalStore, "glob_sync", lambda *_args, **_kwargs: pytest.fail("storage accessed")) + + with pytest.raises(ValueError, match="supports only Parquet"): + list(store.stream_arrow_sync(pattern)) + + +def test_local_stream_rejects_non_parquet_format_before_glob(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + store = LocalStore(str(tmp_path)) + monkeypatch.setattr(LocalStore, "glob_sync", lambda *_args, **_kwargs: pytest.fail("storage accessed")) + + with pytest.raises(ValueError, match="file_format='csv'"): + list(store.stream_arrow_sync("*", file_format=cast("Any", "csv"))) + + +@pytest.mark.skipif(not PYARROW_INSTALLED, reason="PyArrow missing") +def test_local_stream_preserves_multi_file_multi_row_group_order(tmp_path: Path) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + pq.write_table(pa.table({"value": [0, 1, 2]}), tmp_path / "a.parquet", row_group_size=2) + pq.write_table(pa.table({"value": [3, 4, 5]}), tmp_path / "b.parquet", row_group_size=2) + store = LocalStore(str(tmp_path)) + + batches = list(store.stream_arrow_sync("*.parquet", batch_size=1)) + + assert [value for batch in batches for value in batch.column(0).to_pylist()] == list(range(6)) + assert [batch.num_rows for batch in batches] == [1, 1, 1, 1, 1, 1] + + +@pytest.mark.skipif(not PYARROW_INSTALLED, reason="PyArrow missing") +async def test_local_async_stream_matches_sync(tmp_path: Path) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + pq.write_table(pa.table({"value": [0, 1, 2, 3]}), tmp_path / "data.parquet", row_group_size=2) + store = LocalStore(str(tmp_path)) + + sync_values = [ + value for batch in store.stream_arrow_sync("*.parquet", batch_size=1) for value in batch[0].to_pylist() + ] + async_values = [ + value async for batch in store.stream_arrow_async("*.parquet", batch_size=1) for value in batch[0].to_pylist() + ] + + assert async_values == sync_values == [0, 1, 2, 3] + + +@pytest.mark.skipif(not FSSPEC_INSTALLED or not PYARROW_INSTALLED, reason="fsspec or PyArrow missing") +def test_fsspec_stream_preserves_base_path_and_order(tmp_path: Path) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + from sqlspec.storage.backends.fsspec import FSSpecBackend + + data_path = tmp_path / "nested" + data_path.mkdir() + pq.write_table(pa.table({"value": [1, 2, 3]}), data_path / "data.parquet", row_group_size=2) + + store = FSSpecBackend(f"file://{tmp_path}", base_path="nested") + values = [value for batch in store.stream_arrow_sync("*.parquet", batch_size=1) for value in batch[0].to_pylist()] + + assert values == [1, 2, 3] + + +@pytest.mark.skipif(not OBSTORE_INSTALLED or not PYARROW_INSTALLED, reason="obstore or PyArrow missing") +def test_obstore_stream_preserves_base_path_without_full_object_drain(tmp_path: Path) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + from sqlspec.storage.backends.obstore import ObStoreBackend + + data_path = tmp_path / "nested" + data_path.mkdir() + pq.write_table(pa.table({"value": [1, 2, 3]}), data_path / "data.parquet", row_group_size=2) + + store = ObStoreBackend(f"file://{tmp_path}", base_path="nested") + values = [value for batch in store.stream_arrow_sync("*.parquet", batch_size=1) for value in batch[0].to_pylist()] + + assert values == [1, 2, 3] + + +@pytest.mark.skipif(not OBSTORE_INSTALLED, reason="obstore missing") +def test_obstore_stream_closes_reader_on_early_generator_close(monkeypatch: pytest.MonkeyPatch) -> None: + import obstore + + from sqlspec.storage.backends import obstore as backend_module + from sqlspec.storage.backends.obstore import ObStoreBackend + + class Reader: + closed = False + + def close(self) -> None: + self.closed = True + + def seekable(self) -> bool: + return True + + class ParquetFile: + num_row_groups = 2 + + def __init__(self, _stream: Any) -> None: + pass + + def iter_batches(self, **kwargs: Any) -> Iterator[Any]: + yield kwargs["row_groups"][0] + + reader = Reader() + opened_paths: list[str] = [] + + def open_reader(_store: Any, path: str) -> Reader: + opened_paths.append(path) + return reader + + monkeypatch.setattr(obstore, "open_reader", open_reader) + monkeypatch.setattr(backend_module, "import_pyarrow_parquet", lambda: type("PQ", (), {"ParquetFile": ParquetFile})) + monkeypatch.setattr(ObStoreBackend, "glob_sync", lambda *_args, **_kwargs: ["mybase/data.parquet"]) + store = ObStoreBackend("memory://", base_path="mybase") + + batches = store.stream_arrow_sync("*.parquet") + assert cast("Any", next(batches)) == 0 + cast("Any", batches).close() + + assert reader.closed + assert opened_paths == ["mybase/data.parquet"] + + +@pytest.mark.skipif(not OBSTORE_INSTALLED, reason="obstore missing") +def test_obstore_stream_closes_reader_when_batch_iteration_fails(monkeypatch: pytest.MonkeyPatch) -> None: + import obstore + + from sqlspec.storage.backends import obstore as backend_module + from sqlspec.storage.backends.obstore import ObStoreBackend + + class Reader: + closed = False + + def close(self) -> None: + self.closed = True + + def seekable(self) -> bool: + return True + + class ParquetFile: + num_row_groups = 1 + + def __init__(self, _stream: Any) -> None: + pass + + def iter_batches(self, **kwargs: Any) -> Iterator[Any]: + _ = kwargs + msg = "reader failed" + raise OSError(msg) + yield + + reader = Reader() + monkeypatch.setattr(obstore, "open_reader", lambda *_args, **_kwargs: reader) + monkeypatch.setattr(backend_module, "import_pyarrow_parquet", lambda: type("PQ", (), {"ParquetFile": ParquetFile})) + monkeypatch.setattr(ObStoreBackend, "glob_sync", lambda *_args, **_kwargs: ["data.parquet"]) + store = ObStoreBackend("memory://") + + with pytest.raises(OSError, match="reader failed"): + list(store.stream_arrow_sync("*.parquet")) + + assert reader.closed + + +@pytest.mark.skipif(not OBSTORE_INSTALLED or not PYARROW_INSTALLED, reason="obstore or PyArrow missing") +def test_obstore_first_batch_does_not_read_later_row_group_data( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import obstore + import pyarrow as pa + import pyarrow.parquet as pq + + from sqlspec.storage.backends.obstore import ObStoreBackend + + path = tmp_path / "tracked.parquet" + pq.write_table( + pa.table({"value": [f"row-{index}-" + ("x" * 65_536) for index in range(12)]}), + path, + row_group_size=4, + compression="none", + ) + payload = path.read_bytes() + metadata = pq.ParquetFile(path).metadata + row_group_ranges: list[tuple[int, int]] = [] + for index in range(metadata.num_row_groups): + column = metadata.row_group(index).column(0) + dictionary_offset = column.dictionary_page_offset + assert dictionary_offset is not None + start = min(dictionary_offset, column.data_page_offset) + row_group_ranges.append((start, start + column.total_compressed_size)) + + class TrackedReader(io.BytesIO): + def __init__(self, data: bytes) -> None: + super().__init__(data) + self.read_ranges: list[tuple[int, int]] = [] + + def read(self, size: int | None = -1) -> bytes: + start = self.tell() + data = super().read(size) + self.read_ranges.append((start, start + len(data))) + return data + + def readall(self) -> bytes: + return self.read() + + reader = TrackedReader(payload) + monkeypatch.setattr(obstore, "open_reader", lambda *_args, **_kwargs: reader) + monkeypatch.setattr(ObStoreBackend, "glob_sync", lambda *_args, **_kwargs: ["tracked.parquet"]) + store = ObStoreBackend("memory://") + + batches = store.stream_arrow_sync("*.parquet", batch_size=2) + assert next(batches)[0].to_pylist() == ["row-0-" + ("x" * 65_536), "row-1-" + ("x" * 65_536)] + + # Parquet discovery requires a suffix read for footer metadata. Exclude only + # that request, then prove the first batch fetched row-group 0 data without + # fetching a byte range belonging to either later row group. + data_reads = [(start, end) for start, end in reader.read_ranges if end != len(payload)] + first_start, first_end = row_group_ranges[0] + assert any(read_start < first_end and first_start < read_end for read_start, read_end in data_reads) + assert not any( + read_start < group_end and group_start < read_end + for read_start, read_end in data_reads + for group_start, group_end in row_group_ranges[1:] + ) + cast("Any", batches).close() + + +@pytest.mark.parametrize("backend_name", ["local", "fsspec", "obstore"]) +@pytest.mark.skipif(not PYARROW_INSTALLED, reason="PyArrow missing") +async def test_backend_matrix_multi_file_row_group_sync_async_parity(tmp_path: Path, backend_name: str) -> None: + import pyarrow as pa + import pyarrow.parquet as pq + + nested = tmp_path / "nested" + nested.mkdir() + pq.write_table(pa.table({"value": [0, 1, 2]}), nested / "a.parquet", row_group_size=2) + pq.write_table(pa.table({"value": [3, 4, 5]}), nested / "b.parquet", row_group_size=2) + + if backend_name == "local": + store: Any = LocalStore(str(tmp_path), base_path="nested") + elif backend_name == "fsspec": + if not FSSPEC_INSTALLED: + pytest.skip("fsspec missing") + from sqlspec.storage.backends.fsspec import FSSpecBackend + + store = FSSpecBackend(f"file://{tmp_path}", base_path="nested") + else: + if not OBSTORE_INSTALLED: + pytest.skip("obstore missing") + from sqlspec.storage.backends.obstore import ObStoreBackend + + store = ObStoreBackend(f"file://{tmp_path}", base_path="nested") + + sync_batches = list(store.stream_arrow_sync("*.parquet", batch_size=1)) + async_batches = [batch async for batch in store.stream_arrow_async("*.parquet", batch_size=1)] + + assert [value for batch in sync_batches for value in batch[0].to_pylist()] == list(range(6)) + assert [value for batch in async_batches for value in batch[0].to_pylist()] == list(range(6)) + assert [batch.num_rows for batch in sync_batches] == [1] * 6 + + +@pytest.mark.parametrize("backend_name", ["local", "fsspec", "obstore"]) +def test_backend_matrix_validation_precedes_object_open( + tmp_path: Path, backend_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + if backend_name == "local": + backend_type: Any = LocalStore + store: Any = LocalStore(str(tmp_path)) + elif backend_name == "fsspec": + if not FSSPEC_INSTALLED: + pytest.skip("fsspec missing") + from sqlspec.storage.backends.fsspec import FSSpecBackend + + backend_type = FSSpecBackend + store = FSSpecBackend("memory") + else: + if not OBSTORE_INSTALLED: + pytest.skip("obstore missing") + from sqlspec.storage.backends.obstore import ObStoreBackend + + backend_type = ObStoreBackend + store = ObStoreBackend("memory://") + + monkeypatch.setattr(backend_type, "glob_sync", lambda *_args, **_kwargs: pytest.fail("storage accessed")) + + with pytest.raises(ValueError, match="supports only Parquet"): + list(store.stream_arrow_sync("*.csv")) + with pytest.raises(ValueError, match="batch_size must be greater than zero"): + list(store.stream_arrow_sync("*.parquet", batch_size=0)) + + +@pytest.mark.skipif(not OBSTORE_INSTALLED, reason="obstore missing") +async def test_obstore_async_close_closes_active_reader(monkeypatch: pytest.MonkeyPatch) -> None: + import obstore + + from sqlspec.storage.backends import obstore as backend_module + from sqlspec.storage.backends.obstore import ObStoreBackend + + class Reader: + closed = False + + def close(self) -> None: + self.closed = True + + def seekable(self) -> bool: + return True + + class ParquetFile: + num_row_groups = 2 + + def __init__(self, _stream: Any) -> None: + pass + + def iter_batches(self, **kwargs: Any) -> Iterator[Any]: + yield kwargs["row_groups"][0] + + reader = Reader() + monkeypatch.setattr(obstore, "open_reader", lambda *_args, **_kwargs: reader) + monkeypatch.setattr(backend_module, "import_pyarrow_parquet", lambda: type("PQ", (), {"ParquetFile": ParquetFile})) + monkeypatch.setattr(ObStoreBackend, "glob_sync", lambda *_args, **_kwargs: ["data.parquet"]) + iterator = ObStoreBackend("memory://").stream_arrow_async("*.parquet") + + assert cast("Any", await anext(iterator)) == 0 + await cast("Any", iterator).aclose() + + assert reader.closed + + +@pytest.mark.skipif(not OBSTORE_INSTALLED, reason="obstore missing") +async def test_obstore_async_reader_exception_closes_reader(monkeypatch: pytest.MonkeyPatch) -> None: + import obstore + + from sqlspec.storage.backends import obstore as backend_module + from sqlspec.storage.backends.obstore import ObStoreBackend + + class Reader: + closed = False + + def close(self) -> None: + self.closed = True + + def seekable(self) -> bool: + return True + + class ParquetFile: + num_row_groups = 1 + + def __init__(self, _stream: Any) -> None: + pass + + def iter_batches(self, **kwargs: Any) -> Iterator[Any]: + _ = kwargs + raise OSError("async reader failed") + yield + + reader = Reader() + monkeypatch.setattr(obstore, "open_reader", lambda *_args, **_kwargs: reader) + monkeypatch.setattr(backend_module, "import_pyarrow_parquet", lambda: type("PQ", (), {"ParquetFile": ParquetFile})) + monkeypatch.setattr(ObStoreBackend, "glob_sync", lambda *_args, **_kwargs: ["data.parquet"]) + iterator = ObStoreBackend("memory://").stream_arrow_async("*.parquet") + + with pytest.raises(OSError, match="async reader failed"): + await anext(iterator) + + assert reader.closed + + +def test_async_arrow_iterator_close_closes_active_sync_generator() -> None: + from sqlspec.storage.backends.base import AsyncArrowBatchIterator + + closed = False + + def batches() -> Iterator[Any]: + nonlocal closed + try: + yield object() + yield object() + finally: + closed = True + + async def exercise() -> None: + iterator = AsyncArrowBatchIterator(batches()) + await anext(iterator) + await iterator.aclose() + + import asyncio + + asyncio.run(exercise()) + assert closed