fix(#668): harden storage payload codecs

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.
 
+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.
+
+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.
+
 Pipelines
 =========
 

sqlspec/storage/_arrow_payload.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, Any, Literal, cast
 
-from sqlspec.storage._utils import import_pyarrow, import_pyarrow_csv, import_pyarrow_parquet
+from sqlspec.storage._utils import import_pyarrow, import_pyarrow_csv, import_pyarrow_json, import_pyarrow_parquet
 from sqlspec.utils.serializers import from_json
 
 if TYPE_CHECKING:
@@ -56,13 +56,14 @@ def decode_arrow_payload(payload: bytes, format_choice: StorageFormat) -> "Arrow
     if format_choice == "csv":
         pa_csv = import_pyarrow_csv()
         return cast("ArrowTable", pa_csv.read_csv(pa.BufferReader(payload)))
-    text_payload = payload.decode()
     if format_choice == "json":
-        data = from_json(text_payload)
+        data = from_json(payload.decode())
         rows = data if isinstance(data, list) else [data]
         return cast("ArrowTable", pa.Table.from_pylist(rows))
     if format_choice == "jsonl":
-        rows = [from_json(line) for line in text_payload.splitlines() if line.strip()]
-        return cast("ArrowTable", pa.Table.from_pylist(rows))
+        if payload == b"":
+            return cast("ArrowTable", pa.table({}))
+        pa_json = import_pyarrow_json()
+        return cast("ArrowTable", pa_json.read_json(pa.BufferReader(payload)))
     msg = f"Unsupported storage format for Arrow decoding: {format_choice}"
     raise ValueError(msg)

sqlspec/storage/_utils.py

@@ -6,7 +6,13 @@
 from sqlspec.utils.logging import get_logger, log_with_context
 from sqlspec.utils.module_loader import ensure_pyarrow
 
-__all__ = ("_log_storage_event", "import_pyarrow", "import_pyarrow_csv", "import_pyarrow_parquet")
+__all__ = (
+    "_log_storage_event",
+    "import_pyarrow",
+    "import_pyarrow_csv",
+    "import_pyarrow_json",
+    "import_pyarrow_parquet",
+)
 
 logger = get_logger(__name__)
 
@@ -50,6 +56,19 @@ def import_pyarrow_csv() -> "Any":
     return pa_csv
 
 
+def import_pyarrow_json() -> "Any":
+    """Import PyArrow JSON module with optional dependency guard.
+
+    Returns:
+        PyArrow JSON module.
+    """
+
+    ensure_pyarrow()
+    import pyarrow.json as pa_json
+
+    return pa_json
+
+
 def _log_storage_event(
     event: str,
     *,

sqlspec/storage/pipeline.py

@@ -144,6 +144,8 @@ def reset(self) -> None:
 _METRICS = _StorageBridgeMetrics()
 _RECENT_STORAGE_EVENTS: "deque[StorageTelemetry]" = deque(maxlen=25)
 _EMPTY_STORAGE_OPTIONS: dict[str, Any] = {}
+_ARROW_WRITE_FORMATS = frozenset({"parquet", "arrow-ipc", "csv"})
+_ROW_WRITE_FORMATS = frozenset({"json", "jsonl"})
 
 
 def _storage_options(default_options: "dict[str, Any]", storage_options: "dict[str, Any] | None") -> "dict[str, Any]":
@@ -225,6 +227,18 @@ def _encode_row_payload(rows: "list[Any]", format_hint: StorageFormat) -> bytes:
     return bytes(buffer)
 
 
+def _validate_arrow_write_format(format_choice: StorageFormat) -> None:
+    if format_choice not in _ARROW_WRITE_FORMATS:
+        msg = "Arrow storage writes support only Parquet, Arrow IPC, and CSV formats"
+        raise ValueError(msg)
+
+
+def _validate_row_write_format(format_choice: StorageFormat) -> None:
+    if format_choice not in _ROW_WRITE_FORMATS:
+        msg = "Row storage writes support only JSON and JSONL formats"
+        raise ValueError(msg)
+
+
 def _encode_arrow_payload(
     table: "ArrowTable",
     format_choice: StorageFormat,
@@ -353,8 +367,9 @@ def write_rows(
     ) -> StorageTelemetry:
         """Write dictionary rows to storage using cached serializers."""
 
-        serialized = serialize_collection(rows)
         format_choice = format_hint or "jsonl"
+        _validate_row_write_format(format_choice)
+        serialized = serialize_collection(rows)
         payload = _encode_row_payload(serialized, format_choice)
         resolved_options = _storage_options(self._storage_options, storage_options)
         return self._write_bytes(
@@ -373,6 +388,7 @@ def write_arrow(
         """Write an Arrow table to storage using zero-copy buffers."""
 
         format_choice = format_hint or "parquet"
+        _validate_arrow_write_format(format_choice)
         resolved_options = _storage_options(self._storage_options, storage_options)
         format_write_options = _csv_write_options(
             format_choice, resolved_options, self._storage_options, self._csv_write_options
@@ -485,8 +501,9 @@ async def write_rows(
         format_hint: StorageFormat | None = None,
         storage_options: "dict[str, Any] | None" = None,
     ) -> StorageTelemetry:
-        serialized = serialize_collection(rows)
         format_choice = format_hint or "jsonl"
+        _validate_row_write_format(format_choice)
+        serialized = serialize_collection(rows)
         payload = await async_(_encode_row_payload)(serialized, format_choice)
         resolved_options = _storage_options(self._storage_options, storage_options)
         return await self._write_bytes_async(
@@ -503,6 +520,7 @@ async def write_arrow(
         compression: str | None = None,
     ) -> StorageTelemetry:
         format_choice = format_hint or "parquet"
+        _validate_arrow_write_format(format_choice)
         resolved_options = _storage_options(self._storage_options, storage_options)
         format_write_options = _csv_write_options(
             format_choice, resolved_options, self._storage_options, self._csv_write_options

tests/unit/storage/test_payload_codecs.py

@@ -0,0 +1,140 @@
+# pyright: reportPrivateUsage=false
+"""Focused storage payload codec and write-format contract tests."""
+
+from typing import Any
+
+import pyarrow as pa
+import pytest
+
+import sqlspec.storage.pipeline as storage_pipeline
+from sqlspec.storage._arrow_payload import decode_arrow_payload, encode_arrow_payload
+from sqlspec.storage.pipeline import (
+    AsyncStoragePipeline,
+    SyncStoragePipeline,
+    get_recent_storage_events,
+    reset_storage_bridge_events,
+)
+
+
+class _TrackingBackend:
+    backend_type = "tracking"
+
+    def __init__(self) -> None:
+        self.writes: list[tuple[str, bytes]] = []
+
+    def write_bytes_sync(self, path: str, payload: bytes) -> None:
+        self.writes.append((path, payload))
+
+    async def write_bytes_async(self, path: str, payload: bytes) -> None:
+        self.writes.append((path, payload))
+
+
+def _install_backend(monkeypatch: pytest.MonkeyPatch, pipeline_type: type[Any]) -> _TrackingBackend:
+    backend = _TrackingBackend()
+    monkeypatch.setattr(pipeline_type, "_backend", lambda *_args, **_kwargs: (backend, "payload", backend.backend_type))
+    return backend
+
+
+def test_decode_jsonl_uses_native_timestamp_inference() -> None:
+    payload = b'{"id":1,"created_at":"2026-08-01T12:34:56"}\n'
+
+    table = decode_arrow_payload(payload, "jsonl")
+
+    assert table.schema.field("created_at").type == pa.timestamp("s")
+    assert table.to_pylist() == [{"id": 1, "created_at": table.column("created_at")[0].as_py()}]
+
+
+def test_decode_empty_jsonl_preserves_empty_table() -> None:
+    table = decode_arrow_payload(b"", "jsonl")
+
+    assert table.equals(pa.table({}))
+
+
+def test_decode_whitespace_only_jsonl_is_delegated_to_pyarrow() -> None:
+    table = decode_arrow_payload(b" \n\t", "jsonl")
+
+    assert table.equals(pa.table({}))
+
+
+def test_json_array_decode_keeps_existing_shape() -> None:
+    table = decode_arrow_payload(b'[{"id":1},{"id":2}]', "json")
+
+    assert table.to_pylist() == [{"id": 1}, {"id": 2}]
+
+
+@pytest.mark.parametrize("format_choice", ["parquet", "arrow-ipc", "csv"])
+def test_arrow_payload_round_trip_has_expected_signature(format_choice: str) -> None:
+    table = pa.table({"id": [1, 2], "name": ["alpha", "beta"]})
+
+    payload = encode_arrow_payload(table, format_choice, compression=None)  # type: ignore[arg-type]
+    restored = decode_arrow_payload(payload, format_choice)  # type: ignore[arg-type]
+
+    signatures = {"parquet": b"PAR1", "arrow-ipc": b"ARROW1", "csv": b'"id","name"'}
+    assert payload.startswith(signatures[format_choice])
+    assert restored.to_pylist() == table.to_pylist()
+
+
+@pytest.mark.parametrize("format_hint", ["parquet", "arrow-ipc", "csv"])
+def test_sync_row_write_rejects_arrow_formats_before_encoding_or_io(
+    monkeypatch: pytest.MonkeyPatch, format_hint: str
+) -> None:
+    pipeline = SyncStoragePipeline()
+    backend = _install_backend(monkeypatch, SyncStoragePipeline)
+    monkeypatch.setattr(storage_pipeline, "serialize_collection", pytest.fail)
+    monkeypatch.setattr(storage_pipeline, "_encode_row_payload", pytest.fail)
+    reset_storage_bridge_events()
+
+    with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"):
+        pipeline.write_rows([{"id": 1}], "payload", format_hint=format_hint)  # type: ignore[arg-type]
+
+    assert backend.writes == []
+    assert get_recent_storage_events() == []
+
+
+@pytest.mark.parametrize("format_hint", ["json", "jsonl"])
+def test_sync_arrow_write_rejects_row_formats_before_encoding_or_io(
+    monkeypatch: pytest.MonkeyPatch, format_hint: str
+) -> None:
+    pipeline = SyncStoragePipeline()
+    backend = _install_backend(monkeypatch, SyncStoragePipeline)
+    monkeypatch.setattr(storage_pipeline, "_encode_arrow_payload", pytest.fail)
+    reset_storage_bridge_events()
+
+    with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"):
+        pipeline.write_arrow(pa.table({"id": [1]}), "payload", format_hint=format_hint)  # type: ignore[arg-type]
+
+    assert backend.writes == []
+    assert get_recent_storage_events() == []
+
+
+@pytest.mark.parametrize("format_hint", ["parquet", "arrow-ipc", "csv"])
+async def test_async_row_write_rejects_arrow_formats_before_io(
+    monkeypatch: pytest.MonkeyPatch, format_hint: str
+) -> None:
+    pipeline = AsyncStoragePipeline()
+    backend = _install_backend(monkeypatch, AsyncStoragePipeline)
+    monkeypatch.setattr(storage_pipeline, "serialize_collection", pytest.fail)
+    monkeypatch.setattr(storage_pipeline, "_encode_row_payload", pytest.fail)
+    reset_storage_bridge_events()
+
+    with pytest.raises(ValueError, match="Row storage writes support only JSON and JSONL"):
+        await pipeline.write_rows([{"id": 1}], "payload", format_hint=format_hint)  # type: ignore[arg-type]
+
+    assert backend.writes == []
+    assert get_recent_storage_events() == []
+
+
+@pytest.mark.parametrize("format_hint", ["json", "jsonl"])
+async def test_async_arrow_write_rejects_row_formats_before_io(
+    monkeypatch: pytest.MonkeyPatch, format_hint: str
+) -> None:
+    pipeline = AsyncStoragePipeline()
+    backend = _install_backend(monkeypatch, AsyncStoragePipeline)
+    monkeypatch.setattr(storage_pipeline, "_encode_arrow_payload", pytest.fail)
+    reset_storage_bridge_events()
+
+    with pytest.raises(ValueError, match="Arrow storage writes support only Parquet, Arrow IPC, and CSV"):
+        await pipeline.write_arrow(pa.table({"id": [1]}), "payload", format_hint=format_hint)  # type: ignore[arg-type]
+
+    assert backend.writes == []
+    assert get_recent_storage_events() == []