fix: bound Parquet storage streaming

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
       - id: mixed-line-ending
       - id: trailing-whitespace
   - repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: "v0.16.0"
+    rev: "v0.16.1"
     hooks:
       - id: ruff
         args: ["--fix"]

docs/changelog.rst

@@ -9,7 +9,7 @@ important operational fixes.
 Recent Updates
 ==============
 
-Unreleased
+v0.58.0 - Configuration and storage correctness
 ------------------------------------------------------------------------------
 
 **Changed:**
@@ -19,9 +19,36 @@ Unreleased
   valid key. A misspelling such as ``version_table`` instead of
   ``version_table_name`` was previously accepted and silently ignored, leaving
   the setting at its default. Remove or correct unrecognized keys to upgrade.
+* Storage writes reject formats that cannot carry the payload being written,
+  raising :class:`~sqlspec.exceptions.StorageCapabilityError` before any
+  encoding or storage I/O. Row writes accept only ``json`` and ``jsonl``; Arrow
+  table writes accept only ``parquet``, ``arrow-ipc``, and ``csv``. A mismatched
+  format previously wrote one payload type under another format label. Read APIs
+  still accept all five formats.
+* ``stream_arrow_sync()`` and ``stream_arrow_async()`` accept only
+  ``file_format="parquet"`` and raise
+  :class:`~sqlspec.exceptions.StorageCapabilityError` for other formats. Use the
+  regular Arrow read APIs for CSV, Arrow IPC, JSON, and JSONL payloads.
+* JSONL payloads decode through PyArrow's native JSON reader. Its type inference
+  applies to the result, so date-like strings now decode as Arrow timestamps
+  rather than strings.
 
 **Fixed:**
 
+* Arrow batch streaming reads one Parquet row group at a time across the local,
+  fsspec, and obstore backends, and accepts a ``batch_size`` bounding each record
+  batch. The obstore backend streams through its seekable reader instead of
+  buffering the whole object in memory, resolves cloud ``base_path`` only once,
+  and closes readers deterministically when a stream is closed early.
+* Decoding a JSONL payload containing a row larger than 1 MiB no longer fails
+  with ``ArrowInvalid: straddling object straddles two block boundaries``.
+* ADBC PostgreSQL connections recover after a failed statement. The aborted
+  transaction is now cleared through the connection rather than by sending a
+  ``ROLLBACK`` statement on a cursor, which the driver rejects on a connection
+  that is already in an error state. Every later statement on that connection
+  previously failed with ``INVALID_STATE: [libpq] cannot start transaction``.
+  Uncommitted work in the aborted transaction is discarded, as PostgreSQL
+  requires; commit before a statement whose failure you intend to recover from.
 * Pointing ``--config`` at a module rather than a configuration object now
   reports the ``module:attribute`` references that module exports, instead of
   failing later with ``AttributeError: module has no attribute 'bind_key'``.
@@ -34,6 +61,16 @@ Unreleased
   paths that contain colons.
 * ``author`` is declared on :class:`~sqlspec.config.MigrationConfig`. The
   migration generator already read it, but type checkers rejected it.
+* psycopg record loads preserve JSON and JSONB object shapes through Arrow COPY.
+  JSON mappings previously reached psycopg COPY as Python dictionaries, which it
+  cannot adapt in text COPY mode.
+
+**Performance:**
+
+* ``AsyncpgDriver.load_from_records()`` writes records directly with one binary
+  COPY call instead of round-tripping them through Arrow. The removed conversion
+  dominated small and medium batches; large batches also overtake
+  ``executemany()`` throughput.
 
 v0.57.0
 ------------------------------------------------------------------------------

docs/reference/storage.rst

@@ -11,15 +11,33 @@ Write and read formats
 
 Row-oriented writes accept only ``json`` and newline-delimited ``jsonl``.
 Arrow-table writes accept only ``parquet``, ``arrow-ipc``, and ``csv``. The
-pipeline rejects mismatched formats before encoding or storage I/O, so it
-cannot write one payload type under another format label. Read APIs retain the
-full format set because they decode all five formats into Arrow tables.
+pipeline raises :class:`~sqlspec.exceptions.StorageCapabilityError` for a
+mismatched format before encoding or storage I/O, so it cannot write one payload
+type under another format label. Read APIs retain the full format set because
+they decode all five formats into Arrow tables.
 
 JSONL reads use PyArrow's native JSON reader. Its type inference applies to the
 result, including conversion of date-like strings to Arrow timestamps. This
-avoids Python per-line decoding and ``Table.from_pylist()`` copies. It does not
-make ``load_from_storage()`` bounded-memory: that API reads the complete object
-payload before decoding it.
+avoids Python per-line decoding and ``Table.from_pylist()`` copies. The reader
+block is sized to the payload, so individual rows may exceed PyArrow's default
+1 MiB block. It does not make ``load_from_storage()`` bounded-memory: that API
+reads the complete object payload before decoding it.
+
+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"`` and raise
+:class:`~sqlspec.exceptions.StorageCapabilityError` for any other format. 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
 =========

pyproject.toml

@@ -24,7 +24,7 @@ maintainers = [{ name = "Litestar Developers", email = "hello@litestar.dev" }]
 name = "sqlspec"
 readme = "README.md"
 requires-python = ">=3.10, <4.0"
-version = "0.57.0"
+version = "0.58.0"
 
 [project.urls]
 Discord = "https://discord.gg/litestar"
@@ -240,6 +240,7 @@ include = [
   "sqlspec/storage/registry.py",                   # Safe storage registry/runtime routing
   "sqlspec/storage/errors.py",                     # Safe storage error normalization
   "sqlspec/storage/_paths.py",                     # Pure storage path handling
+  "sqlspec/storage/_arrow_stream.py",              # Pure Parquet streaming validation and row-group iteration
   "sqlspec/storage/pipeline.py",                   # Storage bridge orchestration with Arrow boundary split out
   "sqlspec/storage/backends/base.py",              # Storage backend runtime base classes
   "sqlspec/storage/backends/fsspec.py",            # fsspec backend import surface
@@ -308,7 +309,7 @@ opt_level = "3"   # Maximum optimization (0-3)
 allow_dirty = true
 commit = false
 commit_args = "--no-verify"
-current_version = "0.57.0"
+current_version = "0.58.0"
 ignore_missing_files = false
 ignore_missing_version = false
 message = "chore(release): bump to v{new_version}"

sqlspec/adapters/adbc/core.py

@@ -3,7 +3,7 @@
 import datetime
 import decimal
 from collections.abc import Sized
-from functools import lru_cache
+from functools import lru_cache, partial
 from typing import TYPE_CHECKING, Any, Final, cast
 from uuid import UUID
 
@@ -464,19 +464,51 @@ def is_postgres_dialect(dialect_name: str) -> bool:
 
 
 def handle_postgres_rollback(dialect: str, cursor: Any, logger: Any | None = None) -> None:
-    """Execute rollback for PostgreSQL after transaction failure.
+    """Clear an aborted PostgreSQL transaction after a statement failure.
+
+    The connection-level rollback is preferred because the driver refuses to
+    start a transaction on an errored connection, which a cursor needs in order
+    to issue a ``ROLLBACK`` statement.
 
     Args:
         dialect: Active dialect identifier.
-        cursor: Database cursor to execute rollback.
+        cursor: Database cursor whose connection should be rolled back.
         logger: Optional logger for diagnostics.
     """
     if not is_postgres_dialect(dialect):
         return
+
+    connection = getattr(cursor, "connection", None)
+    if connection is not None and _try_rollback(connection.rollback):
+        _log_rollback(logger)
+        return
+
+    if _try_rollback(partial(cursor.execute, "ROLLBACK")):
+        _log_rollback(logger)
+
+
+def _try_rollback(rollback: "Callable[[], Any]") -> bool:
+    """Run a rollback callable and report whether it succeeded.
+
+    Args:
+        rollback: Zero-argument callable performing the rollback.
+
+    Returns:
+        True if the rollback completed without raising.
+    """
     try:
-        cursor.execute("ROLLBACK")
+        rollback()
     except Exception:
-        return
+        return False
+    return True
+
+
+def _log_rollback(logger: "Any | None") -> None:
+    """Record that a PostgreSQL rollback completed.
+
+    Args:
+        logger: Optional logger for diagnostics.
+    """
     if logger is not None:
         logger.debug("PostgreSQL rollback executed after transaction failure")
 

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)

sqlspec/storage/_arrow_payload.py

@@ -13,6 +13,8 @@
 
 StorageFormat = Literal["jsonl", "json", "parquet", "arrow-ipc", "csv"]
 
+_PYARROW_JSON_BLOCK_SIZE = 1 << 20
+
 
 def encode_arrow_payload(
     table: "ArrowTable",
@@ -64,6 +66,7 @@ def decode_arrow_payload(payload: bytes, format_choice: StorageFormat) -> "Arrow
         if payload == b"":
             return cast("ArrowTable", pa.table({}))
         pa_json = import_pyarrow_json()
-        return cast("ArrowTable", pa_json.read_json(pa.BufferReader(payload)))
+        read_options = pa_json.ReadOptions(block_size=max(_PYARROW_JSON_BLOCK_SIZE, len(payload)))
+        return cast("ArrowTable", pa_json.read_json(pa.BufferReader(payload), read_options=read_options))
     msg = f"Unsupported storage format for Arrow decoding: {format_choice}"
     raise ValueError(msg)

sqlspec/storage/_arrow_stream.py

@@ -0,0 +1,48 @@
+"""Shared helpers for bounded Parquet batch streaming."""
+
+from pathlib import PurePath
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.exceptions import StorageCapabilityError
+
+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"})
+
+_STREAM_REMEDIATION = "Read this format with the Arrow read APIs instead of batch streaming."
+
+
+def validate_parquet_stream_options(pattern: str, file_format: str, batch_size: int) -> None:
+    """Validate a Parquet streaming request before storage is accessed.
+
+    Args:
+        pattern: Glob pattern selecting objects to stream.
+        file_format: Requested storage format.
+        batch_size: Maximum number of rows in each yielded record batch.
+
+    Raises:
+        StorageCapabilityError: If a non-Parquet format is requested.
+        ValueError: If ``batch_size`` is not greater than zero.
+    """
+    if file_format != "parquet":
+        msg = f"Arrow batch streaming supports only Parquet files; received file_format={file_format!r}"
+        raise StorageCapabilityError(msg, capability="arrow_batch_streaming", remediation=_STREAM_REMEDIATION)
+    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 StorageCapabilityError(msg, capability="arrow_batch_streaming", remediation=_STREAM_REMEDIATION)
+
+
+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)

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