From f1c6735c64734771aed7b5bdfc4696d1c093a27f Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 24 Oct 2026 08:14:09 +0000 Subject: [PATCH 01/10] fix(artifacts): record batch row counts and cap batch sizes of stored DataFrames store_dataframe wrote frames that could be loaded whole but not served a page at a time. The manifest had no per-batch row counts, batches were capped by rows only, and compression could fall back to lz4. - Write pandas and polars frames through one pyarrow IPC path that cuts the record batches itself and records their row counts in the manifest as batch_rows. - Cap batches at about 8 MiB as well as 65,536 rows, sizing them from the average row, and merge small chunks into whole batches. - Compress with zstd, else not at all. - Raise the manifest cap to 64 KiB, and refuse a frame whose manifest would exceed it. - Refuse polars Object columns, nested ones included. Rechunk chunked Categoricals on polars before 1.35.1. Fixes #141 Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage.py | 167 ++++++++--- .../artifacts/_dataframe_storage_manifest.py | 43 ++- deepnote_toolkit/artifacts/dataframes.py | 14 +- tests/unit/artifacts/test_arguments.py | 8 +- tests/unit/artifacts/test_file_format.py | 274 ++++++++++++++---- .../artifacts/test_manifest_and_cleanup.py | 31 +- tests/unit/artifacts/test_read_only_mount.py | 19 +- tests/unit/artifacts/test_storing.py | 10 +- tests/unit/artifacts/test_write_failures.py | 71 ++++- tests/unit/helpers/dataframe_storage.py | 4 +- 10 files changed, 505 insertions(+), 136 deletions(-) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 95efac1..7647f88 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -7,7 +7,7 @@ import time import uuid from pathlib import Path -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import pandas as pd from packaging.version import Version @@ -18,16 +18,23 @@ from ._dataframe_storage_manifest import ( Manifest, + encode_manifest, is_frame_name, read_manifest, write_manifest, ) +if TYPE_CHECKING: + import pyarrow as pa + FRAMES_DIR = Path("deepnote_dataframes") STALE_FRAME_AGE_SECONDS = 24 * 60 * 60 -POLARS_BATCH_ROWS = 65_536 -POLARS_BATCH_MIN_VERSION = Version("1.37") +MAX_BATCH_ROWS = 65_536 +# A page spanning two batches reads at most half the reader's 32 MiB limit, +# leaving room for dictionaries, IPC framing and uneven rows. +TARGET_BATCH_BYTES = 8 * 1024 * 1024 POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") +POLARS_CHUNK_DICTIONARIES_MIN_VERSION = Version("1.35.1") StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -44,42 +51,131 @@ def resolve_project_root() -> Path: return resolve_notebook_root() -def _select_compression() -> str: - """Return the best codec this pyarrow build has: zstd, lz4, else uncompressed.""" - import pyarrow as pa +def _select_compression() -> Optional[str]: + """Return zstd if this pyarrow build has it, else None for uncompressed. - for codec in ("zstd", "lz4"): - if pa.Codec.is_available(codec): - return codec - return "uncompressed" + Never lz4: the dataframe reader cannot decode it. + """ + import pyarrow as pa + return "zstd" if pa.Codec.is_available("zstd") else None + + +def _contains_polars_dtype(dtype: Any, *targets: Any) -> bool: + """Return whether `dtype` is one of `targets` or nests one at any depth.""" + import polars + + if any(dtype == target for target in targets): + return True + if isinstance(dtype, (polars.List, polars.Array)): + return _contains_polars_dtype(dtype.inner, *targets) + if isinstance(dtype, polars.Struct): + return any( + _contains_polars_dtype(field.dtype, *targets) for field in dtype.fields + ) + return False + + +def _to_arrow_table(frame: StorableFrame) -> "pa.Table": + """Return the frame as an Arrow table; raise TypeError for polars Object data.""" + if isinstance(frame, pd.DataFrame): + import pyarrow as pa + + # As write_feather does: RangeIndex to metadata, other indexes to columns. + return pa.Table.from_pandas(frame, preserve_index=None) + + # polars is an optional dependency, so import it only for a non-pandas frame. + import polars + + if isinstance(frame, polars.DataFrame): + # to_arrow() exports an Object value as its 8-byte memory address. + objects = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Object) + ] + if objects: + raise TypeError(f"Polars Object columns cannot be stored, got {objects}") + polars_version = Version(polars.__version__) + # Before 1.35.1, to_arrow() copies a column's whole dictionary into every + # chunk, so many chunks of a large Categorical could exhaust memory. Only + # the first column's chunks are counted: to_arrow() itself rechunks a frame + # whose columns are chunked differently. + if ( + polars_version < POLARS_CHUNK_DICTIONARIES_MIN_VERSION + and frame.n_chunks() > 1 + and any( + _contains_polars_dtype(dtype, polars.Categorical, polars.Enum) + for dtype in frame.schema.values() + ) + ): + frame = frame.rechunk() + # pyarrow < 15, apache-arrow < 21.2 can't read 1.0+'s string/binary views. + if polars_version >= POLARS_COMPAT_LEVEL_MIN_VERSION: + return frame.to_arrow(compat_level=polars.CompatLevel.oldest()) + return frame.to_arrow() -def _write_frame_file(frame: StorableFrame, path: Path, compression: str) -> None: - """Write the frame to `path` as an Arrow IPC file.""" - with path.open("wb") as sink: - if isinstance(frame, pd.DataFrame): - from pyarrow import feather + assert_never(frame) - # pandas < 2.1's to_feather needs a default index and string column names. - feather.write_feather(frame, sink, compression=compression) - return - # polars is an optional dependency, so import it only for a non-pandas frame. - import polars +def _has_dictionary(data_type: "pa.DataType") -> bool: + """Return whether `data_type` is a dictionary type or nests one.""" + import pyarrow as pa - if isinstance(frame, polars.DataFrame): - polars_version = Version(polars.__version__) - options: dict[str, Any] = {} - # Cap batches like write_feather so a page read stays cheap. - if polars_version >= POLARS_BATCH_MIN_VERSION: - options["record_batch_size"] = POLARS_BATCH_ROWS - # pyarrow < 15, apache-arrow < 21.2 can't read 1.0+'s string/binary views. - if polars_version >= POLARS_COMPAT_LEVEL_MIN_VERSION: - options["compat_level"] = polars.CompatLevel.oldest() - frame.write_ipc(sink, compression=compression, **options) - return + if pa.types.is_dictionary(data_type): + return True + return any( + _has_dictionary(data_type.field(i).type) for i in range(data_type.num_fields) + ) + + +def _rows_per_batch(table: "pa.Table") -> int: + """Return how many rows make a batch of about TARGET_BATCH_BYTES.""" + # Every batch writes its own copy of each column and chunk, as nbytes counts + # them, even where they share buffers. A dictionary is the exception: the file + # holds it once, while nbytes counts it once per chunk that shares it. + # get_total_buffer_size() counts it once, but counts a slice's whole buffer. + nbytes = sum( + ( + min(column.nbytes, column.get_total_buffer_size()) + if _has_dictionary(column.type) + else column.nbytes + ) + for column in table.columns + ) + if table.num_rows == 0 or nbytes == 0: + return MAX_BATCH_ROWS + row_bytes = nbytes / table.num_rows + return max(1, min(MAX_BATCH_ROWS, int(TARGET_BATCH_BYTES // row_bytes))) + + +def _write_frame_file( + frame: StorableFrame, path: Path, compression: Optional[str] +) -> list[int]: + """Write the frame to `path` as an Arrow IPC file; return its batches' row counts. + + The dataframe reader needs every record batch's row count, so the batches are + cut here rather than by pyarrow. + """ + import pyarrow as pa - assert_never(frame) + with path.open("wb") as sink: + # A file with more than one dictionary for a column is refused by pyarrow. + table = _to_arrow_table(frame).unify_dictionaries() + rows_per_batch = _rows_per_batch(table) + options = pa.ipc.IpcWriteOptions(compression=compression, allow_64bit=True) + batch_rows: list[int] = [] + with pa.ipc.new_file(sink, table.schema, options=options) as writer: + for offset in range(0, table.num_rows, rows_per_batch): + # slice() does not clamp the length of a table without columns. + length = min(rows_per_batch, table.num_rows - offset) + # to_batches() alone never merges chunks, so small chunks would + # make small batches; combining a slice copies at most one batch. + batch_table = table.slice(offset, length).combine_chunks() + for batch in batch_table.to_batches(): + writer.write_batch(batch) + batch_rows.append(batch.num_rows) + return batch_rows def _delete_replaced_frames( @@ -125,8 +221,11 @@ def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: new_file = f"{reference_id}.arrow" try: directory.mkdir(parents=True, exist_ok=True) - _write_frame_file(frame, directory / new_file, _select_compression()) + batch_rows = _write_frame_file( + frame, directory / new_file, _select_compression() + ) old = read_manifest(directory) + manifest = encode_manifest(new_file, old.file if old else None, batch_rows) except BaseException as exc: with contextlib.suppress(OSError): (directory / new_file).unlink(missing_ok=True) @@ -134,6 +233,6 @@ def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]: return None raise - write_manifest(directory, new_file, old.file if old else None, len(frame)) + write_manifest(directory, manifest) _delete_replaced_frames(directory, old, new_file) return reference_id diff --git a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py index bdf94c3..b01d1e5 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py @@ -19,7 +19,9 @@ MANIFEST_NAME = "manifest.json" MANIFEST_VERSION = 1 -MANIFEST_MAX_BYTES = 4096 +# The dataframe reader has the same cap. A batch costs about 7 bytes in +# batch_rows, so this fits about 9,000 batches. +MANIFEST_MAX_BYTES = 64 * 1024 _FRAME_NAME_REGEX = re.compile( r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.arrow", @@ -33,13 +35,17 @@ def is_frame_name(name: str) -> bool: class Manifest(BaseModel): - """The manifest file; `previous` is the frame `file` replaced.""" + """The manifest file; `previous` is the frame `file` replaced. + + `batch_rows` holds the row count of each record batch in `file`, in order. + """ version: StrictInt file: StrictStr previous: Optional[StrictStr] = None written_at: StrictStr rows: StrictInt + batch_rows: list[StrictInt] toolkit_version: StrictStr @@ -68,22 +74,39 @@ def read_manifest(directory: Path) -> Optional[Manifest]: return manifest -def write_manifest( - directory: Path, new_file: str, previous: Optional[str], rows: int -) -> None: - """Point the manifest at `new_file`. +def encode_manifest( + new_file: str, previous: Optional[str], batch_rows: list[int] +) -> str: + """Return the manifest naming `new_file`, whose record batches hold `batch_rows`. - On s3fs a rename is a copy and then a delete, so a failed rename may already - have replaced the manifest. + Raises: + ValueError: The manifest would exceed MANIFEST_MAX_BYTES, so neither the + toolkit nor the dataframe reader could read it. """ manifest = Manifest( version=MANIFEST_VERSION, file=new_file, previous=previous, written_at=datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"), - rows=rows, + rows=sum(batch_rows), + batch_rows=batch_rows, toolkit_version=__version__, ) + encoded = json.dumps(_model_dump_compat(manifest)) + if len(encoded.encode("utf-8")) > MANIFEST_MAX_BYTES: + raise ValueError( + f"The DataFrame needs {len(batch_rows)} record batches, more than its " + f"manifest can list in {MANIFEST_MAX_BYTES} bytes" + ) + return encoded + + +def write_manifest(directory: Path, encoded: str) -> None: + """Replace the directory's manifest with `encoded`. + + On s3fs a rename is a copy and then a delete, so a failed rename may already + have replaced the manifest. + """ # delete=False: s3fs uploads at close(), so the file must be closed before the # rename, and deleting it only when the context ends needs Python 3.12's # delete_on_close=False. @@ -97,7 +120,7 @@ def write_manifest( ) try: with temporary: - temporary.write(json.dumps(_model_dump_compat(manifest))) + temporary.write(encoded) os.replace(temporary.name, directory / MANIFEST_NAME) except BaseException: with contextlib.suppress(OSError): diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index aa11fd8..f716f33 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -24,7 +24,7 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: """Store `df` as the current frame under `name`, replacing the one stored before. Args: - df: A pandas or polars DataFrame, written whole with no row or size limit. + df: A pandas or polars DataFrame, written whole. name: The frame's name, unique within the project, and its directory under `deepnote_dataframes/`. Must match `[A-Za-z0-9_-]{1,128}`. @@ -33,13 +33,17 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: read-only, in which case nothing is stored. Raises: - ValueError: `name` is not valid. + ValueError: `name` is not valid, or `df` needs more record batches than + its manifest can list, which takes some 600M rows or 70 GiB of Arrow + data. TypeError: `df` is a SQL query preview or a frame derived from one, or is not a pandas or polars DataFrame. That includes PySpark and - pandas-on-Spark DataFrames and polars LazyFrames. + pandas-on-Spark DataFrames and polars LazyFrames. Also raised for a + polars DataFrame with an Object column, or Object in a List, Array or + Struct. OSError: The write failed. - Exception: `df` cannot be serialized to Arrow: whatever pyarrow's - `write_feather()` or polars' `df.write_ipc()` raises. + Exception: `df` cannot be converted to Arrow: whatever pyarrow's + `Table.from_pandas()` or polars' `df.to_arrow()` raises. A failed or interrupted call leaves the name's previous frame current, unless it stops after the manifest names the new frame: an interrupt can arrive then, diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index 9d1c28e..eeee7d7 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -2,12 +2,10 @@ from pathlib import Path from typing import Any, Callable -from unittest import mock import pandas as pd import polars as pl import pytest -from pyarrow import feather from deepnote_toolkit import artifacts from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview @@ -46,12 +44,10 @@ def test_valid_names_are_accepted(tmp_project_root: Path, name: str) -> None: def test_invalid_name_raises_before_anything_is_written( tmp_project_root: Path, name: Any ) -> None: - with mock.patch.object(feather, "write_feather", autospec=True) as write: - with pytest.raises(ValueError) as raised: - artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) + with pytest.raises(ValueError) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": [1]}), name) assert str(raised.value).endswith(f"got {name!r}") - write.assert_not_called() assert list(tmp_project_root.iterdir()) == [] diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index 96f1eb1..a66acfd 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -1,9 +1,10 @@ """Tests for the compression, record batches and column types of stored DataFrames.""" from pathlib import Path -from typing import Any, Dict +from typing import Any, Callable, Dict, List, Optional from unittest import mock +import numpy as np import pandas as pd import polars as pl import pyarrow as pa @@ -43,80 +44,184 @@ def is_available(self, name: str) -> bool: "availability, expected", [ pytest.param({"zstd": True, "lz4": True}, "zstd", id="zstd"), - pytest.param({"zstd": False, "lz4": True}, "lz4", id="zstd-missing"), - pytest.param({"zstd": False, "lz4": False}, "uncompressed", id="none"), + pytest.param({"zstd": False, "lz4": True}, None, id="zstd-missing"), ], ) -def test_compression_falls_back_from_zstd_to_lz4_to_none( - monkeypatch: pytest.MonkeyPatch, availability: Dict[str, bool], expected: str +def test_compression_is_zstd_or_none_never_lz4( + monkeypatch: pytest.MonkeyPatch, + availability: Dict[str, bool], + expected: Optional[str], ) -> None: monkeypatch.setattr(pa, "Codec", _FakeCodec(availability)) assert _dataframe_storage._select_compression() == expected -@pytest.mark.parametrize( - "make_frame", [pd.DataFrame, pl.DataFrame], ids=["pandas", "polars"] -) -def test_frame_is_stored_uncompressed_when_no_codec_is_available( - tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch, make_frame: Any +def test_frame_is_stored_uncompressed_without_zstd_even_with_lz4( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(pa, "Codec", _FakeCodec({"zstd": False, "lz4": False})) + monkeypatch.setattr(pa, "Codec", _FakeCodec({"zstd": False, "lz4": True})) - artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) + artifacts.store_dataframe(pd.DataFrame({"a": ["x" * 100] * 10_000}), NAME) stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file assert pa.ipc.open_file(stored).read_all().num_rows == 10_000 assert stored.stat().st_size > 1_000_000 -def test_pandas_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> None: - artifacts.store_dataframe(pd.DataFrame({"a": range(70_000)}), NAME) +def _file_batch_rows(root: Path) -> List[int]: + reader = pa.ipc.open_file(frame_dir(root) / stored_manifest(root).file) + return [reader.get_batch(i).num_rows for i in range(reader.num_record_batches)] + + +def _categories_in_many_chunks(in_lists: bool = False) -> pl.DataFrame: + """2,000 chunks that share one unified dictionary of 200,000 values.""" + with pl.StringCache(): + parts = [] + for i in range(2_000): + ids = [f"id-{i:05d}-{j:03d}" for j in range(100)] + column = ( + pl.Series([[id_] for id_ in ids], dtype=pl.List(pl.Categorical)) + if in_lists + else pl.Series(ids, dtype=pl.Categorical) + ) + parts.append(pl.DataFrame({"id": column})) + return pl.concat(parts, rechunk=False) - reader = pa.ipc.open_file( - frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file - ) - assert reader.num_record_batches == 2 - assert reader.read_all().num_rows == 70_000 +def _categories_in_two_batches() -> pl.DataFrame: + """Two chunks whose dictionaries differ and fall into different batches.""" + with pl.StringCache(): + first = pl.DataFrame( + {"c": pl.Series(["x", "y"] * 32_768, dtype=pl.Categorical)} + ) + second = pl.DataFrame({"c": pl.Series(["z", "x"], dtype=pl.Categorical)}) + return pl.concat([first, second], rechunk=False) -def test_polars_frame_is_written_in_64k_row_batches(tmp_project_root: Path) -> None: - artifacts.store_dataframe(pl.DataFrame({"a": range(200_000)}), NAME) + +@pytest.mark.parametrize( + "make_frame, expected", + [ + pytest.param( + lambda: pd.DataFrame({"a": range(70_000)}), + [65_536, 4_464], + id="pandas", + ), + pytest.param( + lambda: pl.concat( + [pl.DataFrame({"a": range(10)}), pl.DataFrame({"a": range(70_000)})], + rechunk=False, + ), + [65_536, 4_474], + id="polars-small-chunk-merged", + ), + pytest.param( + lambda: pd.DataFrame( + {"a": pd.array(range(2_000_000), dtype="int64[pyarrow]")} + ).head(70_000), + [65_536, 4_464], + id="pandas-arrow-backed-head-of-a-larger-frame", + ), + pytest.param( + _categories_in_many_chunks, + [65_536, 65_536, 65_536, 3_392], + id="polars-categorical-many-chunks", + ), + pytest.param( + lambda: _categories_in_many_chunks(in_lists=True), + [65_536, 65_536, 65_536, 3_392], + id="polars-list-of-categoricals-many-chunks", + ), + pytest.param(lambda: pd.DataFrame({"a": []}), [], id="pandas-empty"), + pytest.param(lambda: pd.DataFrame(index=range(5)), [5], id="pandas-no-columns"), + ], +) +def test_manifest_batch_rows_match_the_files_record_batches( + tmp_project_root: Path, make_frame: Callable[[], Any], expected: List[int] +) -> None: + frame = make_frame() + + artifacts.store_dataframe(frame, NAME) + + manifest = stored_manifest(tmp_project_root) + assert manifest.batch_rows == expected + assert _file_batch_rows(tmp_project_root) == expected + assert manifest.rows == len(frame) + + +def test_dictionaries_of_chunks_are_unified(tmp_project_root: Path) -> None: + artifacts.store_dataframe(_categories_in_two_batches(), NAME) + + values = stored_table(tmp_project_root).column("c").to_pylist() + assert values == ["x", "y"] * 32_768 + ["z", "x"] + + +def test_long_text_is_split_into_batches_of_at_most_8_mib( + tmp_project_root: Path, +) -> None: + # 20,000 rows of 2,000 bytes: about 40 MB, which needs five 8 MiB batches. + artifacts.store_dataframe(pd.DataFrame({"text": ["x" * 2_000] * 20_000}), NAME) reader = pa.ipc.open_file( frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file ) - assert reader.num_record_batches == 4 - assert reader.read_all().num_rows == 200_000 + batches = [reader.get_batch(i) for i in range(reader.num_record_batches)] + assert len(batches) == 5 + assert max(batch.nbytes for batch in batches) <= 8 * 1024 * 1024 + assert stored_manifest(tmp_project_root).batch_rows == [ + batch.num_rows for batch in batches + ] -@pytest.mark.parametrize( - "polars_version, expected_options", - [("1.36.9", {}), ("1.37.0", {"record_batch_size": 65_536})], - ids=["before-1.37", "from-1.37"], -) -def test_polars_batch_size_is_passed_to_polars_from_1_37( +def test_row_larger_than_a_batch_gets_a_batch_of_its_own( tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - expected_options: Dict[str, int], ) -> None: - monkeypatch.setattr(pl, "__version__", polars_version) - original = pl.DataFrame.write_ipc + artifacts.store_dataframe(pd.DataFrame({"text": ["x" * 10_000_000] * 3}), NAME) - with mock.patch.object( - pl.DataFrame, "write_ipc", autospec=True, side_effect=original - ) as write_ipc: - artifacts.store_dataframe(pl.DataFrame({"a": [1, 2, 3]}), NAME) + assert stored_manifest(tmp_project_root).batch_rows == [1, 1, 1] + assert _file_batch_rows(tmp_project_root) == [1, 1, 1] - write_ipc.assert_called_once() - options = { - name: value - for name, value in write_ipc.call_args.kwargs.items() - if name == "record_batch_size" - } - assert options == expected_options - assert stored_table(tmp_project_root).to_pydict() == {"a": [1, 2, 3]} + +def _arrow_text(rows: int) -> pd.Series: + return pd.Series(["x" * 500] * rows, dtype="string[pyarrow]") + + +@pytest.mark.parametrize( + "make_frame", + [ + pytest.param( + lambda: pd.concat( + [pd.DataFrame({"t": _arrow_text(20_000)})] * 4, ignore_index=True + ), + id="pandas-frame-concatenated-with-itself", + ), + pytest.param( + lambda: pd.DataFrame({"t": _arrow_text(20_000)}).pipe( + lambda frame: frame.assign(a=frame["t"], b=frame["t"], c=frame["t"]) + ), + id="pandas-column-and-its-copies", + ), + pytest.param( + lambda: pl.DataFrame( + {"e": pl.Series(np.ones((20_000, 128), dtype=np.float32))} + ).with_columns( + pl.col("e").alias("a"), pl.col("e").alias("b"), pl.col("e").alias("c") + ), + id="polars-column-and-its-aliases", + ), + ], +) +def test_buffers_shared_by_columns_or_chunks_count_once_per_use( + tmp_project_root: Path, make_frame: Callable[[], Any] +) -> None: + # Every batch writes its own copy of each column and chunk. + artifacts.store_dataframe(make_frame(), NAME) + + reader = pa.ipc.open_file( + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file + ) + batches = [reader.get_batch(i) for i in range(reader.num_record_batches)] + assert max(batch.nbytes for batch in batches) <= 8 * 1024 * 1024 def test_polars_strings_are_stored_without_view_types(tmp_project_root: Path) -> None: @@ -150,12 +255,81 @@ def test_polars_compat_level_is_passed_to_polars_from_1_2_1( passed: bool, ) -> None: monkeypatch.setattr(pl, "__version__", polars_version) - original = pl.DataFrame.write_ipc + original = pl.DataFrame.to_arrow with mock.patch.object( - pl.DataFrame, "write_ipc", autospec=True, side_effect=original - ) as write_ipc: + pl.DataFrame, "to_arrow", autospec=True, side_effect=original + ) as to_arrow: artifacts.store_dataframe(pl.DataFrame({"a": ["x"]}), NAME) - assert ("compat_level" in write_ipc.call_args.kwargs) is passed + assert ("compat_level" in to_arrow.call_args.kwargs) is passed assert stored_table(tmp_project_root).to_pydict() == {"a": ["x"]} + + +def _two_chunks(make_column: Callable[[List[str]], pl.Series]) -> pl.DataFrame: + """Two chunks of column `c` beside an Int64 column, which needs no rechunk.""" + return pl.concat( + [ + pl.DataFrame({"n": [1, 2], "c": make_column(["x", "y"])}), + pl.DataFrame({"n": [3], "c": make_column(["z"])}), + ], + rechunk=False, + ) + + +@pytest.mark.parametrize( + "polars_version, make_column, converted_chunks", + [ + pytest.param( + "1.34.0", + lambda values: pl.Series(values, dtype=pl.Categorical), + 1, + id="categorical-before-1.35.1", + ), + pytest.param( + "1.34.0", + lambda values: pl.Series(values, dtype=pl.Enum(["x", "y", "z"])), + 1, + id="enum-before-1.35.1", + ), + pytest.param( + "1.34.0", + lambda values: pl.Series( + [[value] for value in values], dtype=pl.List(pl.Categorical) + ), + 1, + id="list-of-categoricals-before-1.35.1", + ), + pytest.param( + "1.35.1", + lambda values: pl.Series(values, dtype=pl.Categorical), + 2, + id="categorical-from-1.35.1", + ), + pytest.param( + "1.34.0", + lambda values: pl.Series(values), + 2, + id="strings-before-1.35.1", + ), + ], +) +def test_chunked_polars_categoricals_are_rechunked_before_1_35_1( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + make_column: Callable[[List[str]], pl.Series], + converted_chunks: int, +) -> None: + with pl.StringCache(): + frame = _two_chunks(make_column) + monkeypatch.setattr(pl, "__version__", polars_version) + original = pl.DataFrame.to_arrow + + with mock.patch.object( + pl.DataFrame, "to_arrow", autospec=True, side_effect=original + ) as to_arrow: + artifacts.store_dataframe(frame, NAME) + + assert to_arrow.call_args.args[0].get_column("c").n_chunks() == converted_chunks + assert stored_manifest(tmp_project_root).rows == len(frame) diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index 9a35a33..e0295e5 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -93,6 +93,7 @@ def add(name: str, age_seconds: float) -> Path: "previous": FRAME_B, "written_at": "2026-10-15T09:30:00Z", "rows": 1, + "batch_rows": [1], "toolkit_version": "1.0.0", } @@ -114,11 +115,13 @@ def _manifest_bytes(without: str = "", **fields: Any) -> bytes: pytest.param(_manifest_bytes(version=True), id="version-a-bool"), pytest.param(_manifest_bytes(without="file"), id="no-file"), pytest.param(_manifest_bytes(rows="1"), id="rows-a-string"), + pytest.param(_manifest_bytes(without="batch_rows"), id="no-batch-rows"), + pytest.param(_manifest_bytes(batch_rows=["1"]), id="batch-rows-strings"), pytest.param( - _manifest_bytes(pad="x" * 5000), id="oversized-truncated-mid-string" + _manifest_bytes(pad="x" * 70_000), id="oversized-truncated-mid-string" ), pytest.param( - _manifest_bytes() + b" " * 5000, id="oversized-valid-json-then-spaces" + _manifest_bytes().ljust(65_537), id="oversized-valid-json-then-spaces" ), ], ) @@ -144,7 +147,7 @@ def test_unusable_manifest_counts_as_no_manifest( @pytest.mark.parametrize( "manifest_bytes", [ - pytest.param(_manifest_bytes().ljust(4096), id="exactly-4096-bytes"), + pytest.param(_manifest_bytes().ljust(65_536), id="exactly-64-kib"), pytest.param( _manifest_bytes(written_by_a_newer_toolkit=True), id="unknown-key" ), @@ -155,7 +158,7 @@ def test_usable_manifest_is_honoured( ) -> None: directory = frame_dir(tmp_project_root) directory.mkdir(parents=True) - assert len(manifest_bytes) <= 4096 + assert len(manifest_bytes) <= 65_536 (directory / "manifest.json").write_bytes(manifest_bytes) for name in (FRAME_A, FRAME_B): (directory / name).write_bytes(b"named by the manifest") @@ -167,6 +170,19 @@ def test_usable_manifest_is_honoured( assert not (directory / FRAME_B).exists() +def test_manifest_is_written_up_to_the_size_the_reader_accepts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # written_at has a fixed width, so the manifest's size does not change. + size = len(_dataframe_storage_manifest.encode_manifest(FRAME_A, None, [1])) + + monkeypatch.setattr(_dataframe_storage_manifest, "MANIFEST_MAX_BYTES", size) + _dataframe_storage_manifest.encode_manifest(FRAME_A, None, [1]) + monkeypatch.setattr(_dataframe_storage_manifest, "MANIFEST_MAX_BYTES", size - 1) + with pytest.raises(ValueError): + _dataframe_storage_manifest.encode_manifest(FRAME_A, None, [1]) + + @pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) def test_previous_in_a_parent_directory_is_not_deleted( tmp_project_root: Path, absolute: bool @@ -187,7 +203,7 @@ def test_previous_in_a_parent_directory_is_not_deleted( def test_manifest_nested_too_deeply_to_parse_counts_as_no_manifest( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - """Before 3.12, json raises RecursionError on the nesting 4 KiB allows.""" + """json raises RecursionError on the nesting 64 KiB allows.""" artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) def too_deep(raw: bytes) -> Any: @@ -213,10 +229,11 @@ def test_manifest_is_written_after_the_frame_file_is_complete( seen: List[Dict[str, Any]] = [] original = _dataframe_storage.write_manifest - def spy(directory: Path, new_file: str, *args: Any) -> None: + def spy(directory: Path, encoded: str) -> None: + new_file = json.loads(encoded)["file"] table = pa.ipc.open_file(directory / new_file).read_all() seen.append({"closed": arrow_sinks[0].closed, "rows": table.num_rows}) - original(directory, new_file, *args) + original(directory, encoded) monkeypatch.setattr(_dataframe_storage, "write_manifest", spy) diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index 608823b..f89ce45 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -6,10 +6,11 @@ from unittest import mock import pandas as pd +import pyarrow as pa import pytest -from pyarrow import feather from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import _dataframe_storage from tests.unit.helpers.dataframe_storage import ( NAME, arrow_files, @@ -24,11 +25,13 @@ def test_read_only_mount_stores_nothing_and_does_not_convert( ) -> None: read_only(monkeypatch) - with mock.patch.object(feather, "write_feather", autospec=True) as write: + with mock.patch.object( + _dataframe_storage, "_to_arrow_table", autospec=True + ) as convert: result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) assert result is None - write.assert_not_called() + convert.assert_not_called() assert list(tmp_project_root.iterdir()) == [] @@ -39,11 +42,13 @@ def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() read_only(monkeypatch) - with mock.patch.object(feather, "write_feather", autospec=True) as write: + with mock.patch.object( + _dataframe_storage, "_to_arrow_table", autospec=True + ) as convert: result = artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) assert result is None - write.assert_not_called() + convert.assert_not_called() assert ( frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes @@ -53,10 +58,10 @@ def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( def test_erofs_from_the_write_stores_nothing( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def refuse(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + def refuse(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: raise OSError(errno.EROFS, "Read-only file system") - monkeypatch.setattr(feather, "write_feather", refuse) + monkeypatch.setattr(pa.ipc, "new_file", refuse) result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index c6eebc9..7ca17b0 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -40,11 +40,19 @@ def test_stores_pandas_frame_and_manifest(tmp_project_root: Path) -> None: } # Manifest ignores unknown keys and defaults previous, so the keys are checked raw. written = json.loads((frame_dir(tmp_project_root) / "manifest.json").read_text()) - assert set(written) == {"version", "file", "written_at", "rows", "toolkit_version"} + assert set(written) == { + "version", + "file", + "written_at", + "rows", + "batch_rows", + "toolkit_version", + } manifest = stored_manifest(tmp_project_root) assert manifest.version == 1 assert manifest.file == f"{reference_id}.arrow" assert manifest.rows == 3 + assert manifest.batch_rows == [3] assert manifest.toolkit_version == deepnote_toolkit.__version__ assert re.fullmatch(r"\d{4}-\d\d-\d\dT\d\d:\d\d:\d\dZ", manifest.written_at) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 9137c5f..936a2f4 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -10,8 +10,8 @@ import pandas as pd import polars as pl +import pyarrow as pa import pytest -from pyarrow import feather from deepnote_toolkit import artifacts from deepnote_toolkit.artifacts import _dataframe_storage, _dataframe_storage_manifest @@ -74,18 +74,65 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} +@pytest.mark.parametrize( + "dtype", + [ + pytest.param(pl.Object, id="object"), + pytest.param( + pl.List(pl.Object), + id="list-of-objects", + # polars panics, unraisably, while it builds a List(Object) column. + marks=pytest.mark.filterwarnings( + "ignore::pytest.PytestUnraisableExceptionWarning" + ), + ), + pytest.param(pl.Array(pl.Object, 1), id="array-of-objects"), + pytest.param( + pl.List(pl.List(pl.Object)), + id="list-of-lists-of-objects", + marks=pytest.mark.filterwarnings( + "ignore::pytest.PytestUnraisableExceptionWarning" + ), + ), + ], +) def test_polars_object_column_raises_and_keeps_the_previous_frame( - tmp_project_root: Path, + tmp_project_root: Path, dtype: Any ) -> None: artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) - with pytest.raises(pl.exceptions.ComputeError) as raised: - artifacts.store_dataframe( - pl.DataFrame({"a": pl.Series([object()], dtype=pl.Object)}), NAME - ) + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(frame, NAME) - raised.match("'Object' datatype") + raised.match(r"Object columns cannot be stored, got \['a'\]") + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 + + +def test_object_in_a_struct_is_refused() -> None: + # polars 1.2 and later refuse to build such a frame, but older versions store + # the Object field's memory addresses. + assert _dataframe_storage._contains_polars_dtype( + pl.Struct({"a": pl.Int64, "o": pl.Object}), pl.Object + ) + + +def test_frame_with_too_many_batches_for_the_manifest_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + artifacts.store_dataframe(pd.DataFrame({"ok": [1]}), NAME) + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + # 25,000 two-row batches need a manifest of about 75 KB. + monkeypatch.setattr(_dataframe_storage, "MAX_BATCH_ROWS", 2) + + with pytest.raises(ValueError) as raised: + artifacts.store_dataframe(pd.DataFrame({"a": range(50_000)}), NAME) + + raised.match("25000 record batches") assert ( frame_dir(tmp_project_root) / "manifest.json" ).read_bytes() == manifest_bytes @@ -95,11 +142,11 @@ def test_polars_object_column_raises_and_keeps_the_previous_frame( def test_partial_file_is_deleted_when_the_write_fails( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def write_some_then_fail(frame: pd.DataFrame, sink: Any, **kwargs: Any) -> None: + def write_some_then_fail(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: sink.write(b"ARROW1 partial") raise OSError(errno.EIO, "Input/output error") - monkeypatch.setattr(feather, "write_feather", write_some_then_fail) + monkeypatch.setattr(pa.ipc, "new_file", write_some_then_fail) with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -300,15 +347,13 @@ def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( ) -> None: interrupt = KeyboardInterrupt() - def write_some_then_interrupt( - frame: pd.DataFrame, sink: Any, **kwargs: Any - ) -> None: + def write_some_then_interrupt(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: sink.write(b"ARROW1 partial") raise interrupt artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() - monkeypatch.setattr(feather, "write_feather", write_some_then_interrupt) + monkeypatch.setattr(pa.ipc, "new_file", write_some_then_interrupt) with pytest.raises(KeyboardInterrupt) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) diff --git a/tests/unit/helpers/dataframe_storage.py b/tests/unit/helpers/dataframe_storage.py index 9a7ecdc..2d0c4eb 100644 --- a/tests/unit/helpers/dataframe_storage.py +++ b/tests/unit/helpers/dataframe_storage.py @@ -1,7 +1,6 @@ """Helpers shared by the DataFrame storage tests.""" import errno -import io import json import os import time @@ -10,7 +9,6 @@ import pyarrow as pa import pytest -from pyarrow import feather from deepnote_core.pydantic_compat_helpers import model_validate_compat from deepnote_toolkit.artifacts._dataframe_storage_manifest import Manifest @@ -80,7 +78,7 @@ def mkdir( def probe_failure(frame: Any) -> Optional[type]: """The exception type pyarrow itself raises for this frame, if any.""" try: - feather.write_feather(frame, io.BytesIO()) + pa.Table.from_pandas(frame, preserve_index=None) except Exception as exc: # pylint: disable=broad-exception-caught return type(exc) return None From c7cd85a57ee757d93631551b42b3423301308d55 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 26 Oct 2026 19:33:50 +0000 Subject: [PATCH 02/10] fix(artifacts): require polars 1.35.1, and 1.43.2 for Array columns Before 1.43.2, polars' to_arrow() applies the offset of a sliced Array column with nulls twice (pola-rs/polars#28583), so store_dataframe stored shifted values, or a file that aborts pyarrow on read. On those versions, a frame with an Array column at any depth now raises TypeError with an upgrade hint. Struct had the same bug before 1.43.0, but pyarrow rejects that export, so storing fails instead of writing wrong values; Struct frames are not refused. polars older than 1.35.1 is refused outright. That removes the rechunk of chunked Categoricals before 1.35.1, the plain to_arrow() branch before 1.2.1, and the polars.Enum lookup that raised AttributeError on polars older than 0.20. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage.py | 58 ++++++---- deepnote_toolkit/artifacts/dataframes.py | 6 +- tests/unit/artifacts/test_file_format.py | 102 ------------------ tests/unit/artifacts/test_write_failures.py | 77 ++++++++++++- 4 files changed, 116 insertions(+), 127 deletions(-) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 7647f88..3ce7019 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -33,8 +33,8 @@ # A page spanning two batches reads at most half the reader's 32 MiB limit, # leaving room for dictionaries, IPC framing and uneven rows. TARGET_BATCH_BYTES = 8 * 1024 * 1024 -POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") -POLARS_CHUNK_DICTIONARIES_MIN_VERSION = Version("1.35.1") +POLARS_MIN_VERSION = Version("1.35.1") +POLARS_ARRAY_MIN_VERSION = Version("1.43.2") StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -77,7 +77,12 @@ def _contains_polars_dtype(dtype: Any, *targets: Any) -> bool: def _to_arrow_table(frame: StorableFrame) -> "pa.Table": - """Return the frame as an Arrow table; raise TypeError for polars Object data.""" + """Return the frame as an Arrow table. + + Raise TypeError for polars Object data, on polars older than + POLARS_MIN_VERSION, and for Array data on polars older than + POLARS_ARRAY_MIN_VERSION. + """ if isinstance(frame, pd.DataFrame): import pyarrow as pa @@ -88,6 +93,15 @@ def _to_arrow_table(frame: StorableFrame) -> "pa.Table": import polars if isinstance(frame, polars.DataFrame): + polars_version = Version(polars.__version__) + # Before 1.35.1, to_arrow() copies a column's whole dictionary into every + # chunk, so many chunks of a large Categorical could exhaust memory. + if polars_version < POLARS_MIN_VERSION: + raise TypeError( + f"Storing polars DataFrames needs polars {POLARS_MIN_VERSION} or " + f"newer, found {polars.__version__}. Upgrade it (pip install -U " + "polars) and restart the kernel." + ) # to_arrow() exports an Object value as its 8-byte memory address. objects = [ name @@ -96,24 +110,26 @@ def _to_arrow_table(frame: StorableFrame) -> "pa.Table": ] if objects: raise TypeError(f"Polars Object columns cannot be stored, got {objects}") - polars_version = Version(polars.__version__) - # Before 1.35.1, to_arrow() copies a column's whole dictionary into every - # chunk, so many chunks of a large Categorical could exhaust memory. Only - # the first column's chunks are counted: to_arrow() itself rechunks a frame - # whose columns are chunked differently. - if ( - polars_version < POLARS_CHUNK_DICTIONARIES_MIN_VERSION - and frame.n_chunks() > 1 - and any( - _contains_polars_dtype(dtype, polars.Categorical, polars.Enum) - for dtype in frame.schema.values() - ) - ): - frame = frame.rechunk() - # pyarrow < 15, apache-arrow < 21.2 can't read 1.0+'s string/binary views. - if polars_version >= POLARS_COMPAT_LEVEL_MIN_VERSION: - return frame.to_arrow(compat_level=polars.CompatLevel.oldest()) - return frame.to_arrow() + # Before 1.43.2, to_arrow() applies the offset of a sliced Array column with + # nulls twice: the file holds shifted values or aborts pyarrow on read. + # Struct had the same bug before 1.43.0, but pyarrow rejects that export, + # so storing fails instead of writing wrong values. + if polars_version < POLARS_ARRAY_MIN_VERSION: + arrays = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Array) + ] + if arrays: + raise TypeError( + "Storing polars Array columns needs polars " + f"{POLARS_ARRAY_MIN_VERSION} or newer, found " + f"{polars.__version__}, got {arrays}. Upgrade it (pip install " + "-U polars) and restart the kernel." + ) + # Pinned, though it's the default: pyarrow < 15 and apache-arrow < 21.2 + # can't read polars' string/binary views. + return frame.to_arrow(compat_level=polars.CompatLevel.oldest()) assert_never(frame) diff --git a/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index f716f33..63bae58 100644 --- a/deepnote_toolkit/artifacts/dataframes.py +++ b/deepnote_toolkit/artifacts/dataframes.py @@ -38,8 +38,10 @@ def store_dataframe(df: Any, name: str) -> Optional[str]: data. TypeError: `df` is a SQL query preview or a frame derived from one, or is not a pandas or polars DataFrame. That includes PySpark and - pandas-on-Spark DataFrames and polars LazyFrames. Also raised for a - polars DataFrame with an Object column, or Object in a List, Array or + pandas-on-Spark DataFrames and polars LazyFrames. Also raised for any + polars DataFrame on polars older than 1.35.1; for one with an Object + column, or Object in a List, Array or Struct; and, on polars older than + 1.43.2, for one with an Array column, or Array in a List, Array or Struct. OSError: The write failed. Exception: `df` cannot be converted to Arrow: whatever pyarrow's diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index a66acfd..d55e04d 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -2,9 +2,7 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Optional -from unittest import mock -import numpy as np import pandas as pd import polars as pl import pyarrow as pa @@ -201,14 +199,6 @@ def _arrow_text(rows: int) -> pd.Series: ), id="pandas-column-and-its-copies", ), - pytest.param( - lambda: pl.DataFrame( - {"e": pl.Series(np.ones((20_000, 128), dtype=np.float32))} - ).with_columns( - pl.col("e").alias("a"), pl.col("e").alias("b"), pl.col("e").alias("c") - ), - id="polars-column-and-its-aliases", - ), ], ) def test_buffers_shared_by_columns_or_chunks_count_once_per_use( @@ -241,95 +231,3 @@ def test_polars_strings_are_stored_without_view_types(tmp_project_root: Path) -> pa.large_binary(), pa.large_list(pa.large_string()), ] - - -@pytest.mark.parametrize( - "polars_version, passed", - [("1.2.0", False), ("1.2.1", True)], - ids=["before-1.2.1", "from-1.2.1"], -) -def test_polars_compat_level_is_passed_to_polars_from_1_2_1( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - passed: bool, -) -> None: - monkeypatch.setattr(pl, "__version__", polars_version) - original = pl.DataFrame.to_arrow - - with mock.patch.object( - pl.DataFrame, "to_arrow", autospec=True, side_effect=original - ) as to_arrow: - artifacts.store_dataframe(pl.DataFrame({"a": ["x"]}), NAME) - - assert ("compat_level" in to_arrow.call_args.kwargs) is passed - assert stored_table(tmp_project_root).to_pydict() == {"a": ["x"]} - - -def _two_chunks(make_column: Callable[[List[str]], pl.Series]) -> pl.DataFrame: - """Two chunks of column `c` beside an Int64 column, which needs no rechunk.""" - return pl.concat( - [ - pl.DataFrame({"n": [1, 2], "c": make_column(["x", "y"])}), - pl.DataFrame({"n": [3], "c": make_column(["z"])}), - ], - rechunk=False, - ) - - -@pytest.mark.parametrize( - "polars_version, make_column, converted_chunks", - [ - pytest.param( - "1.34.0", - lambda values: pl.Series(values, dtype=pl.Categorical), - 1, - id="categorical-before-1.35.1", - ), - pytest.param( - "1.34.0", - lambda values: pl.Series(values, dtype=pl.Enum(["x", "y", "z"])), - 1, - id="enum-before-1.35.1", - ), - pytest.param( - "1.34.0", - lambda values: pl.Series( - [[value] for value in values], dtype=pl.List(pl.Categorical) - ), - 1, - id="list-of-categoricals-before-1.35.1", - ), - pytest.param( - "1.35.1", - lambda values: pl.Series(values, dtype=pl.Categorical), - 2, - id="categorical-from-1.35.1", - ), - pytest.param( - "1.34.0", - lambda values: pl.Series(values), - 2, - id="strings-before-1.35.1", - ), - ], -) -def test_chunked_polars_categoricals_are_rechunked_before_1_35_1( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - make_column: Callable[[List[str]], pl.Series], - converted_chunks: int, -) -> None: - with pl.StringCache(): - frame = _two_chunks(make_column) - monkeypatch.setattr(pl, "__version__", polars_version) - original = pl.DataFrame.to_arrow - - with mock.patch.object( - pl.DataFrame, "to_arrow", autospec=True, side_effect=original - ) as to_arrow: - artifacts.store_dataframe(frame, NAME) - - assert to_arrow.call_args.args[0].get_column("c").n_chunks() == converted_chunks - assert stored_manifest(tmp_project_root).rows == len(frame) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 936a2f4..96e1616 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -114,13 +114,86 @@ def test_polars_object_column_raises_and_keeps_the_previous_frame( def test_object_in_a_struct_is_refused() -> None: - # polars 1.2 and later refuse to build such a frame, but older versions store - # the Object field's memory addresses. + # polars refuses to build such a frame, so only the dtype walk can be checked. assert _dataframe_storage._contains_polars_dtype( pl.Struct({"a": pl.Int64, "o": pl.Object}), pl.Object ) +@pytest.mark.parametrize( + "polars_version, dtype, message", + [ + pytest.param( + "1.35.0", + pl.Int64, + r"needs polars 1\.35\.1 or newer, found 1\.35\.0", + id="before-1.35.1", + ), + pytest.param( + "1.43.1", + pl.Array(pl.Int64, 2), + r"Array columns needs polars 1\.43\.2 or newer, found 1\.43\.1, " + r"got \['a'\]", + id="array-before-1.43.2", + ), + pytest.param( + "1.43.1", + pl.List(pl.Array(pl.Int64, 2)), + r"Array columns needs polars 1\.43\.2", + id="list-of-arrays-before-1.43.2", + ), + pytest.param( + "1.43.1", + pl.Struct({"x": pl.Array(pl.Int64, 2)}), + r"Array columns needs polars 1\.43\.2", + id="struct-of-arrays-before-1.43.2", + ), + ], +) +def test_polars_too_old_for_the_frame_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + dtype: Any, + message: str, +) -> None: + artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) + monkeypatch.setattr(pl, "__version__", polars_version) + + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(frame, NAME) + + raised.match(message) + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 + + +@pytest.mark.parametrize( + "polars_version, dtype", + [ + pytest.param("1.35.1", pl.Int64, id="1.35.1"), + pytest.param("1.43.2", pl.Array(pl.Int64, 2), id="array-from-1.43.2"), + pytest.param("1.43.1", pl.Struct({"x": pl.Int64}), id="struct-before-1.43.2"), + ], +) +def test_polars_frame_the_version_checks_allow_is_stored( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + dtype: Any, +) -> None: + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) + monkeypatch.setattr(pl, "__version__", polars_version) + + artifacts.store_dataframe(frame, NAME) + + assert stored_table(tmp_project_root).to_pydict() == {"ok": [1], "a": [None]} + + def test_frame_with_too_many_batches_for_the_manifest_raises_and_keeps_the_previous_frame( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From 5c92afa2055af569e29eb234352f5b23f06d3f3c Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 27 Oct 2026 12:42:53 +0000 Subject: [PATCH 03/10] refactor(artifacts): move polars checks into their own functions Split the polars version, Object and Array checks out of _to_arrow_table, and reword the lz4 note in _select_compression. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage.py | 91 ++++++++++++------- .../artifacts/_dataframe_storage_manifest.py | 2 - 2 files changed, 56 insertions(+), 37 deletions(-) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 3ce7019..6515140 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -54,7 +54,7 @@ def resolve_project_root() -> Path: def _select_compression() -> Optional[str]: """Return zstd if this pyarrow build has it, else None for uncompressed. - Never lz4: the dataframe reader cannot decode it. + Never lz4: Deepnote Cloud doesn't support it yet. """ import pyarrow as pa @@ -76,6 +76,58 @@ def _contains_polars_dtype(dtype: Any, *targets: Any) -> bool: return False +def _refuse_old_polars() -> None: + """Raise TypeError on polars older than POLARS_MIN_VERSION.""" + import polars + + # Before 1.35.1, to_arrow() copies a column's whole dictionary into every + # chunk, so many chunks of a large Categorical could exhaust memory. + if Version(polars.__version__) < POLARS_MIN_VERSION: + raise TypeError( + f"Storing polars DataFrames needs polars {POLARS_MIN_VERSION} or " + f"newer, found {polars.__version__}. Upgrade it (pip install -U " + "polars) and restart the kernel." + ) + + +def _refuse_object_columns(frame: PolarsEagerDF) -> None: + """Raise TypeError for Object data.""" + import polars + + # to_arrow() exports an Object value as its 8-byte memory address. + objects = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Object) + ] + if objects: + raise TypeError(f"Polars Object columns cannot be stored, got {objects}") + + +def _refuse_arrays_on_old_polars(frame: PolarsEagerDF) -> None: + """Raise TypeError for Array data on polars older than POLARS_ARRAY_MIN_VERSION.""" + import polars + + # Before 1.43.2, to_arrow() applies the offset of a sliced Array column with + # nulls twice: the file holds shifted values or aborts pyarrow on read. + # Struct had the same bug before 1.43.0, but pyarrow rejects that export, + # so storing fails instead of writing wrong values. + if Version(polars.__version__) >= POLARS_ARRAY_MIN_VERSION: + return + arrays = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Array) + ] + if arrays: + raise TypeError( + "Storing polars Array columns needs polars " + f"{POLARS_ARRAY_MIN_VERSION} or newer, found " + f"{polars.__version__}, got {arrays}. Upgrade it (pip install " + "-U polars) and restart the kernel." + ) + + def _to_arrow_table(frame: StorableFrame) -> "pa.Table": """Return the frame as an Arrow table. @@ -93,40 +145,9 @@ def _to_arrow_table(frame: StorableFrame) -> "pa.Table": import polars if isinstance(frame, polars.DataFrame): - polars_version = Version(polars.__version__) - # Before 1.35.1, to_arrow() copies a column's whole dictionary into every - # chunk, so many chunks of a large Categorical could exhaust memory. - if polars_version < POLARS_MIN_VERSION: - raise TypeError( - f"Storing polars DataFrames needs polars {POLARS_MIN_VERSION} or " - f"newer, found {polars.__version__}. Upgrade it (pip install -U " - "polars) and restart the kernel." - ) - # to_arrow() exports an Object value as its 8-byte memory address. - objects = [ - name - for name, dtype in frame.schema.items() - if _contains_polars_dtype(dtype, polars.Object) - ] - if objects: - raise TypeError(f"Polars Object columns cannot be stored, got {objects}") - # Before 1.43.2, to_arrow() applies the offset of a sliced Array column with - # nulls twice: the file holds shifted values or aborts pyarrow on read. - # Struct had the same bug before 1.43.0, but pyarrow rejects that export, - # so storing fails instead of writing wrong values. - if polars_version < POLARS_ARRAY_MIN_VERSION: - arrays = [ - name - for name, dtype in frame.schema.items() - if _contains_polars_dtype(dtype, polars.Array) - ] - if arrays: - raise TypeError( - "Storing polars Array columns needs polars " - f"{POLARS_ARRAY_MIN_VERSION} or newer, found " - f"{polars.__version__}, got {arrays}. Upgrade it (pip install " - "-U polars) and restart the kernel." - ) + _refuse_old_polars() + _refuse_object_columns(frame) + _refuse_arrays_on_old_polars(frame) # Pinned, though it's the default: pyarrow < 15 and apache-arrow < 21.2 # can't read polars' string/binary views. return frame.to_arrow(compat_level=polars.CompatLevel.oldest()) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py index b01d1e5..95e75eb 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py @@ -19,8 +19,6 @@ MANIFEST_NAME = "manifest.json" MANIFEST_VERSION = 1 -# The dataframe reader has the same cap. A batch costs about 7 bytes in -# batch_rows, so this fits about 9,000 batches. MANIFEST_MAX_BYTES = 64 * 1024 _FRAME_NAME_REGEX = re.compile( From 2480edcdd84f00498b3e91cf38fe284abaa90a03 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 14:15:23 +0000 Subject: [PATCH 04/10] refactor(artifacts): move the polars checks and their tests into their own modules Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage.py | 90 ++---------- .../artifacts/_polars_compatibility_checks.py | 77 ++++++++++ .../test_polars_compatibility_checks.py | 136 ++++++++++++++++++ tests/unit/artifacts/test_write_failures.py | 121 ---------------- 4 files changed, 223 insertions(+), 201 deletions(-) create mode 100644 deepnote_toolkit/artifacts/_polars_compatibility_checks.py create mode 100644 tests/unit/artifacts/test_polars_compatibility_checks.py diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 6515140..42d1c1b 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -7,10 +7,9 @@ import time import uuid from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Optional, Union import pandas as pd -from packaging.version import Version from typing_extensions import TypeAlias, assert_never from deepnote_toolkit.config import resolve_notebook_root @@ -23,6 +22,11 @@ read_manifest, write_manifest, ) +from ._polars_compatibility_checks import ( + refuse_arrays_on_old_polars, + refuse_object_columns, + refuse_old_polars, +) if TYPE_CHECKING: import pyarrow as pa @@ -33,8 +37,6 @@ # A page spanning two batches reads at most half the reader's 32 MiB limit, # leaving room for dictionaries, IPC framing and uneven rows. TARGET_BATCH_BYTES = 8 * 1024 * 1024 -POLARS_MIN_VERSION = Version("1.35.1") -POLARS_ARRAY_MIN_VERSION = Version("1.43.2") StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -61,80 +63,8 @@ def _select_compression() -> Optional[str]: return "zstd" if pa.Codec.is_available("zstd") else None -def _contains_polars_dtype(dtype: Any, *targets: Any) -> bool: - """Return whether `dtype` is one of `targets` or nests one at any depth.""" - import polars - - if any(dtype == target for target in targets): - return True - if isinstance(dtype, (polars.List, polars.Array)): - return _contains_polars_dtype(dtype.inner, *targets) - if isinstance(dtype, polars.Struct): - return any( - _contains_polars_dtype(field.dtype, *targets) for field in dtype.fields - ) - return False - - -def _refuse_old_polars() -> None: - """Raise TypeError on polars older than POLARS_MIN_VERSION.""" - import polars - - # Before 1.35.1, to_arrow() copies a column's whole dictionary into every - # chunk, so many chunks of a large Categorical could exhaust memory. - if Version(polars.__version__) < POLARS_MIN_VERSION: - raise TypeError( - f"Storing polars DataFrames needs polars {POLARS_MIN_VERSION} or " - f"newer, found {polars.__version__}. Upgrade it (pip install -U " - "polars) and restart the kernel." - ) - - -def _refuse_object_columns(frame: PolarsEagerDF) -> None: - """Raise TypeError for Object data.""" - import polars - - # to_arrow() exports an Object value as its 8-byte memory address. - objects = [ - name - for name, dtype in frame.schema.items() - if _contains_polars_dtype(dtype, polars.Object) - ] - if objects: - raise TypeError(f"Polars Object columns cannot be stored, got {objects}") - - -def _refuse_arrays_on_old_polars(frame: PolarsEagerDF) -> None: - """Raise TypeError for Array data on polars older than POLARS_ARRAY_MIN_VERSION.""" - import polars - - # Before 1.43.2, to_arrow() applies the offset of a sliced Array column with - # nulls twice: the file holds shifted values or aborts pyarrow on read. - # Struct had the same bug before 1.43.0, but pyarrow rejects that export, - # so storing fails instead of writing wrong values. - if Version(polars.__version__) >= POLARS_ARRAY_MIN_VERSION: - return - arrays = [ - name - for name, dtype in frame.schema.items() - if _contains_polars_dtype(dtype, polars.Array) - ] - if arrays: - raise TypeError( - "Storing polars Array columns needs polars " - f"{POLARS_ARRAY_MIN_VERSION} or newer, found " - f"{polars.__version__}, got {arrays}. Upgrade it (pip install " - "-U polars) and restart the kernel." - ) - - def _to_arrow_table(frame: StorableFrame) -> "pa.Table": - """Return the frame as an Arrow table. - - Raise TypeError for polars Object data, on polars older than - POLARS_MIN_VERSION, and for Array data on polars older than - POLARS_ARRAY_MIN_VERSION. - """ + """Return the frame as an Arrow table.""" if isinstance(frame, pd.DataFrame): import pyarrow as pa @@ -145,9 +75,9 @@ def _to_arrow_table(frame: StorableFrame) -> "pa.Table": import polars if isinstance(frame, polars.DataFrame): - _refuse_old_polars() - _refuse_object_columns(frame) - _refuse_arrays_on_old_polars(frame) + refuse_old_polars() + refuse_object_columns(frame) + refuse_arrays_on_old_polars(frame) # Pinned, though it's the default: pyarrow < 15 and apache-arrow < 21.2 # can't read polars' string/binary views. return frame.to_arrow(compat_level=polars.CompatLevel.oldest()) diff --git a/deepnote_toolkit/artifacts/_polars_compatibility_checks.py b/deepnote_toolkit/artifacts/_polars_compatibility_checks.py new file mode 100644 index 0000000..e830f63 --- /dev/null +++ b/deepnote_toolkit/artifacts/_polars_compatibility_checks.py @@ -0,0 +1,77 @@ +"""Refuses polars versions and data types that DataFrame storage can't handle.""" + +from typing import Any + +from packaging.version import Version + +from deepnote_toolkit.ocelots.types import PolarsEagerDF + +POLARS_MIN_VERSION = Version("1.35.1") +POLARS_ARRAY_MIN_VERSION = Version("1.43.2") + + +def _contains_polars_dtype(dtype: Any, *targets: Any) -> bool: + """Return whether `dtype` is one of `targets` or nests one at any depth.""" + import polars + + if any(dtype == target for target in targets): + return True + if isinstance(dtype, (polars.List, polars.Array)): + return _contains_polars_dtype(dtype.inner, *targets) + if isinstance(dtype, polars.Struct): + return any( + _contains_polars_dtype(field.dtype, *targets) for field in dtype.fields + ) + return False + + +def refuse_old_polars() -> None: + """Raise TypeError on polars older than POLARS_MIN_VERSION.""" + import polars + + # Before 1.35.1, to_arrow() copies a column's whole dictionary into every + # chunk, so many chunks of a large Categorical could exhaust memory. + if Version(polars.__version__) < POLARS_MIN_VERSION: + raise TypeError( + f"Storing polars DataFrames needs polars {POLARS_MIN_VERSION} or " + f"newer, found {polars.__version__}. Upgrade it (pip install -U " + "polars) and restart the kernel." + ) + + +def refuse_object_columns(frame: PolarsEagerDF) -> None: + """Raise TypeError for Object data.""" + import polars + + # to_arrow() exports an Object value as its 8-byte memory address. + objects = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Object) + ] + if objects: + raise TypeError(f"Polars Object columns cannot be stored, got {objects}") + + +def refuse_arrays_on_old_polars(frame: PolarsEagerDF) -> None: + """Raise TypeError for Array data on polars older than POLARS_ARRAY_MIN_VERSION.""" + import polars + + # Before 1.43.2, to_arrow() applies the offset of a sliced Array column with + # nulls twice: the file holds shifted values or aborts pyarrow on read. + # Struct had the same bug before 1.43.0, but pyarrow rejects that export, + # so storing fails instead of writing wrong values. + if Version(polars.__version__) >= POLARS_ARRAY_MIN_VERSION: + return + arrays = [ + name + for name, dtype in frame.schema.items() + if _contains_polars_dtype(dtype, polars.Array) + ] + if arrays: + raise TypeError( + "Storing polars Array columns needs polars " + f"{POLARS_ARRAY_MIN_VERSION} or newer, found " + f"{polars.__version__}, got {arrays}. Upgrade it (pip install " + "-U polars) and restart the kernel." + ) diff --git a/tests/unit/artifacts/test_polars_compatibility_checks.py b/tests/unit/artifacts/test_polars_compatibility_checks.py new file mode 100644 index 0000000..8c0e818 --- /dev/null +++ b/tests/unit/artifacts/test_polars_compatibility_checks.py @@ -0,0 +1,136 @@ +"""Tests for the polars versions and data types that DataFrame storage refuses.""" + +from pathlib import Path +from typing import Any + +import polars as pl +import pytest + +from deepnote_toolkit import artifacts +from deepnote_toolkit.artifacts import _polars_compatibility_checks +from tests.unit.helpers.dataframe_storage import ( + NAME, + arrow_files, + frame_dir, + stored_table, +) + + +@pytest.mark.parametrize( + "dtype", + [ + pytest.param(pl.Object, id="object"), + pytest.param( + pl.List(pl.Object), + id="list-of-objects", + # polars panics, unraisably, while it builds a List(Object) column. + marks=pytest.mark.filterwarnings( + "ignore::pytest.PytestUnraisableExceptionWarning" + ), + ), + pytest.param(pl.Array(pl.Object, 1), id="array-of-objects"), + pytest.param( + pl.List(pl.List(pl.Object)), + id="list-of-lists-of-objects", + marks=pytest.mark.filterwarnings( + "ignore::pytest.PytestUnraisableExceptionWarning" + ), + ), + ], +) +def test_polars_object_column_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, dtype: Any +) -> None: + artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) + + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(frame, NAME) + + raised.match(r"Object columns cannot be stored, got \['a'\]") + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 + + +def test_object_in_a_struct_is_refused() -> None: + # polars refuses to build such a frame, so only the dtype walk can be checked. + assert _polars_compatibility_checks._contains_polars_dtype( + pl.Struct({"a": pl.Int64, "o": pl.Object}), pl.Object + ) + + +@pytest.mark.parametrize( + "polars_version, dtype, message", + [ + pytest.param( + "1.35.0", + pl.Int64, + r"needs polars 1\.35\.1 or newer, found 1\.35\.0", + id="before-1.35.1", + ), + pytest.param( + "1.43.1", + pl.Array(pl.Int64, 2), + r"Array columns needs polars 1\.43\.2 or newer, found 1\.43\.1, " + r"got \['a'\]", + id="array-before-1.43.2", + ), + pytest.param( + "1.43.1", + pl.List(pl.Array(pl.Int64, 2)), + r"Array columns needs polars 1\.43\.2", + id="list-of-arrays-before-1.43.2", + ), + pytest.param( + "1.43.1", + pl.Struct({"x": pl.Array(pl.Int64, 2)}), + r"Array columns needs polars 1\.43\.2", + id="struct-of-arrays-before-1.43.2", + ), + ], +) +def test_polars_too_old_for_the_frame_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + dtype: Any, + message: str, +) -> None: + artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) + manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) + monkeypatch.setattr(pl, "__version__", polars_version) + + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(frame, NAME) + + raised.match(message) + assert ( + frame_dir(tmp_project_root) / "manifest.json" + ).read_bytes() == manifest_bytes + assert len(arrow_files(tmp_project_root)) == 1 + + +@pytest.mark.parametrize( + "polars_version, dtype", + [ + pytest.param("1.35.1", pl.Int64, id="1.35.1"), + pytest.param("1.43.2", pl.Array(pl.Int64, 2), id="array-from-1.43.2"), + pytest.param("1.43.1", pl.Struct({"x": pl.Int64}), id="struct-before-1.43.2"), + ], +) +def test_polars_frame_the_version_checks_allow_is_stored( + tmp_project_root: Path, + monkeypatch: pytest.MonkeyPatch, + polars_version: str, + dtype: Any, +) -> None: + frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) + monkeypatch.setattr(pl, "__version__", polars_version) + + artifacts.store_dataframe(frame, NAME) + + assert stored_table(tmp_project_root).to_pydict() == {"ok": [1], "a": [None]} diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 96e1616..8665a95 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -9,7 +9,6 @@ from typing import Any, Callable, Dict, Optional import pandas as pd -import polars as pl import pyarrow as pa import pytest @@ -74,126 +73,6 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} -@pytest.mark.parametrize( - "dtype", - [ - pytest.param(pl.Object, id="object"), - pytest.param( - pl.List(pl.Object), - id="list-of-objects", - # polars panics, unraisably, while it builds a List(Object) column. - marks=pytest.mark.filterwarnings( - "ignore::pytest.PytestUnraisableExceptionWarning" - ), - ), - pytest.param(pl.Array(pl.Object, 1), id="array-of-objects"), - pytest.param( - pl.List(pl.List(pl.Object)), - id="list-of-lists-of-objects", - marks=pytest.mark.filterwarnings( - "ignore::pytest.PytestUnraisableExceptionWarning" - ), - ), - ], -) -def test_polars_object_column_raises_and_keeps_the_previous_frame( - tmp_project_root: Path, dtype: Any -) -> None: - artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) - manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() - frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) - - with pytest.raises(TypeError) as raised: - artifacts.store_dataframe(frame, NAME) - - raised.match(r"Object columns cannot be stored, got \['a'\]") - assert ( - frame_dir(tmp_project_root) / "manifest.json" - ).read_bytes() == manifest_bytes - assert len(arrow_files(tmp_project_root)) == 1 - - -def test_object_in_a_struct_is_refused() -> None: - # polars refuses to build such a frame, so only the dtype walk can be checked. - assert _dataframe_storage._contains_polars_dtype( - pl.Struct({"a": pl.Int64, "o": pl.Object}), pl.Object - ) - - -@pytest.mark.parametrize( - "polars_version, dtype, message", - [ - pytest.param( - "1.35.0", - pl.Int64, - r"needs polars 1\.35\.1 or newer, found 1\.35\.0", - id="before-1.35.1", - ), - pytest.param( - "1.43.1", - pl.Array(pl.Int64, 2), - r"Array columns needs polars 1\.43\.2 or newer, found 1\.43\.1, " - r"got \['a'\]", - id="array-before-1.43.2", - ), - pytest.param( - "1.43.1", - pl.List(pl.Array(pl.Int64, 2)), - r"Array columns needs polars 1\.43\.2", - id="list-of-arrays-before-1.43.2", - ), - pytest.param( - "1.43.1", - pl.Struct({"x": pl.Array(pl.Int64, 2)}), - r"Array columns needs polars 1\.43\.2", - id="struct-of-arrays-before-1.43.2", - ), - ], -) -def test_polars_too_old_for_the_frame_raises_and_keeps_the_previous_frame( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - dtype: Any, - message: str, -) -> None: - artifacts.store_dataframe(pl.DataFrame({"ok": [1]}), NAME) - manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() - frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) - monkeypatch.setattr(pl, "__version__", polars_version) - - with pytest.raises(TypeError) as raised: - artifacts.store_dataframe(frame, NAME) - - raised.match(message) - assert ( - frame_dir(tmp_project_root) / "manifest.json" - ).read_bytes() == manifest_bytes - assert len(arrow_files(tmp_project_root)) == 1 - - -@pytest.mark.parametrize( - "polars_version, dtype", - [ - pytest.param("1.35.1", pl.Int64, id="1.35.1"), - pytest.param("1.43.2", pl.Array(pl.Int64, 2), id="array-from-1.43.2"), - pytest.param("1.43.1", pl.Struct({"x": pl.Int64}), id="struct-before-1.43.2"), - ], -) -def test_polars_frame_the_version_checks_allow_is_stored( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - dtype: Any, -) -> None: - frame = pl.DataFrame({"ok": [1]}).with_columns(pl.lit(None, dtype=dtype).alias("a")) - monkeypatch.setattr(pl, "__version__", polars_version) - - artifacts.store_dataframe(frame, NAME) - - assert stored_table(tmp_project_root).to_pydict() == {"ok": [1], "a": [None]} - - def test_frame_with_too_many_batches_for_the_manifest_raises_and_keeps_the_previous_frame( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: From c7a56408420c3508a597e513249210f9ebbf869a Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 14:15:23 +0000 Subject: [PATCH 05/10] refactor(artifacts): describe the manifest fields with Field descriptions Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage_manifest.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py index 95e75eb..e03cd52 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage_manifest.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Optional -from pydantic import BaseModel, StrictInt, StrictStr, ValidationError +from pydantic import BaseModel, Field, StrictInt, StrictStr, ValidationError from deepnote_core.pydantic_compat_helpers import ( _model_dump_compat, @@ -33,17 +33,18 @@ def is_frame_name(name: str) -> bool: class Manifest(BaseModel): - """The manifest file; `previous` is the frame `file` replaced. - - `batch_rows` holds the row count of each record batch in `file`, in order. - """ + """The manifest file""" version: StrictInt file: StrictStr - previous: Optional[StrictStr] = None + previous: Optional[StrictStr] = Field( + default=None, description="The previous frame file, if any" + ) written_at: StrictStr - rows: StrictInt - batch_rows: list[StrictInt] + rows: StrictInt = Field(description="The total number of rows in the DataFrame.") + batch_rows: list[StrictInt] = Field( + description="The number of rows in each record batch, in order." + ) toolkit_version: StrictStr From a76d0a1284f31f15730ac49354de9671be73759c Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 14:35:17 +0000 Subject: [PATCH 06/10] test(artifacts): type the patched writers' sink as BufferedWriter Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- tests/unit/artifacts/test_read_only_mount.py | 3 ++- tests/unit/artifacts/test_write_failures.py | 9 +++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index f89ce45..f3e1e0b 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -1,6 +1,7 @@ """Tests for storing and deleting DataFrames on a read-only mount.""" import errno +from io import BufferedWriter from pathlib import Path from typing import Any from unittest import mock @@ -58,7 +59,7 @@ def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( def test_erofs_from_the_write_stores_nothing( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def refuse(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: + def refuse(sink: BufferedWriter, schema: pa.Schema, **kwargs: Any) -> None: raise OSError(errno.EROFS, "Read-only file system") monkeypatch.setattr(pa.ipc, "new_file", refuse) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 8665a95..22dde53 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -4,6 +4,7 @@ import os import shutil import tempfile +from io import BufferedWriter from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Dict, Optional @@ -94,7 +95,9 @@ def test_frame_with_too_many_batches_for_the_manifest_raises_and_keeps_the_previ def test_partial_file_is_deleted_when_the_write_fails( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def write_some_then_fail(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: + def write_some_then_fail( + sink: BufferedWriter, schema: pa.Schema, **kwargs: Any + ) -> None: sink.write(b"ARROW1 partial") raise OSError(errno.EIO, "Input/output error") @@ -299,7 +302,9 @@ def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( ) -> None: interrupt = KeyboardInterrupt() - def write_some_then_interrupt(sink: Any, schema: pa.Schema, **kwargs: Any) -> None: + def write_some_then_interrupt( + sink: BufferedWriter, schema: pa.Schema, **kwargs: Any + ) -> None: sink.write(b"ARROW1 partial") raise interrupt From 096b95b0b3e1b70f066418921c0eed5523263157 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 15:06:12 +0000 Subject: [PATCH 07/10] test(artifacts): fail the write in write_batch rather than new_file new_file writes nothing; the first bytes reach the file in write_batch. The failing mocks now let the real write_batch write a batch, then raise. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- tests/unit/artifacts/test_read_only_mount.py | 7 +++--- tests/unit/artifacts/test_write_failures.py | 24 ++++++++++++-------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/tests/unit/artifacts/test_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index f3e1e0b..b2446c6 100644 --- a/tests/unit/artifacts/test_read_only_mount.py +++ b/tests/unit/artifacts/test_read_only_mount.py @@ -1,7 +1,6 @@ """Tests for storing and deleting DataFrames on a read-only mount.""" import errno -from io import BufferedWriter from pathlib import Path from typing import Any from unittest import mock @@ -59,10 +58,12 @@ def test_read_only_mount_keeps_the_stored_frame_and_does_not_convert( def test_erofs_from_the_write_stores_nothing( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def refuse(sink: BufferedWriter, schema: pa.Schema, **kwargs: Any) -> None: + def refuse( + self: pa.ipc.RecordBatchFileWriter, batch: pa.RecordBatch, **kwargs: Any + ) -> None: raise OSError(errno.EROFS, "Read-only file system") - monkeypatch.setattr(pa.ipc, "new_file", refuse) + monkeypatch.setattr(pa.ipc.RecordBatchFileWriter, "write_batch", refuse) result = artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 22dde53..2f9606c 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -4,7 +4,6 @@ import os import shutil import tempfile -from io import BufferedWriter from pathlib import Path from types import SimpleNamespace from typing import Any, Callable, Dict, Optional @@ -95,13 +94,17 @@ def test_frame_with_too_many_batches_for_the_manifest_raises_and_keeps_the_previ def test_partial_file_is_deleted_when_the_write_fails( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: - def write_some_then_fail( - sink: BufferedWriter, schema: pa.Schema, **kwargs: Any + real_write_batch = pa.ipc.RecordBatchFileWriter.write_batch + + def write_a_batch_then_fail( + self: pa.ipc.RecordBatchFileWriter, batch: pa.RecordBatch, **kwargs: Any ) -> None: - sink.write(b"ARROW1 partial") + real_write_batch(self, batch, **kwargs) raise OSError(errno.EIO, "Input/output error") - monkeypatch.setattr(pa.ipc, "new_file", write_some_then_fail) + monkeypatch.setattr( + pa.ipc.RecordBatchFileWriter, "write_batch", write_a_batch_then_fail + ) with pytest.raises(OSError) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) @@ -301,16 +304,19 @@ def test_interrupt_during_the_write_raises_and_keeps_the_previous_frame( tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch ) -> None: interrupt = KeyboardInterrupt() + real_write_batch = pa.ipc.RecordBatchFileWriter.write_batch - def write_some_then_interrupt( - sink: BufferedWriter, schema: pa.Schema, **kwargs: Any + def write_a_batch_then_interrupt( + self: pa.ipc.RecordBatchFileWriter, batch: pa.RecordBatch, **kwargs: Any ) -> None: - sink.write(b"ARROW1 partial") + real_write_batch(self, batch, **kwargs) raise interrupt artifacts.store_dataframe(pd.DataFrame({"a": [1]}), NAME) manifest_bytes = (frame_dir(tmp_project_root) / "manifest.json").read_bytes() - monkeypatch.setattr(pa.ipc, "new_file", write_some_then_interrupt) + monkeypatch.setattr( + pa.ipc.RecordBatchFileWriter, "write_batch", write_a_batch_then_interrupt + ) with pytest.raises(KeyboardInterrupt) as raised: artifacts.store_dataframe(pd.DataFrame({"a": [2]}), NAME) From 7dbd1f2fd169cd2378daa7674ed47ba2455f96ab Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 17:48:53 +0000 Subject: [PATCH 08/10] test(artifacts): replace Any with concrete types in the artifact tests Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- tests/unit/artifacts/conftest.py | 8 +++++--- tests/unit/artifacts/test_arguments.py | 2 +- tests/unit/artifacts/test_file_format.py | 10 ++++++---- tests/unit/artifacts/test_manifest_and_cleanup.py | 10 +++++----- tests/unit/artifacts/test_storing.py | 6 ++++-- tests/unit/artifacts/test_write_failures.py | 2 +- 6 files changed, 22 insertions(+), 16 deletions(-) diff --git a/tests/unit/artifacts/conftest.py b/tests/unit/artifacts/conftest.py index acab15c..e6081ee 100644 --- a/tests/unit/artifacts/conftest.py +++ b/tests/unit/artifacts/conftest.py @@ -54,12 +54,14 @@ def spark_without_ansi_mode(spark: SparkSession) -> Iterator[SparkSession]: @pytest.fixture -def arrow_sinks(monkeypatch: pytest.MonkeyPatch) -> List[Any]: +def arrow_sinks(monkeypatch: pytest.MonkeyPatch) -> List[IO[Any]]: """Every file object opened for writing a frame, in order.""" - sinks: List[Any] = [] + sinks: List[IO[Any]] = [] real_open = Path.open - def open_and_record(self: Path, mode: str = "r", *args: Any, **kwargs: Any) -> Any: + def open_and_record( + self: Path, mode: str = "r", *args: Any, **kwargs: Any + ) -> IO[Any]: handle = real_open(self, mode, *args, **kwargs) if self.suffix == ".arrow" and "w" in mode: sinks.append(handle) diff --git a/tests/unit/artifacts/test_arguments.py b/tests/unit/artifacts/test_arguments.py index eeee7d7..f1bad71 100644 --- a/tests/unit/artifacts/test_arguments.py +++ b/tests/unit/artifacts/test_arguments.py @@ -89,7 +89,7 @@ def test_unsupported_value_raises_type_error( @pytest.mark.parametrize("make_preview", QUERY_PREVIEWS) def test_query_preview_raises_type_error( - tmp_project_root: Path, make_preview: Callable[[], Any] + tmp_project_root: Path, make_preview: Callable[[], DeepnoteQueryPreview] ) -> None: with pytest.raises(TypeError, match="SQL query preview"): artifacts.store_dataframe(make_preview(), NAME) diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index d55e04d..c1d0aee 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -1,7 +1,7 @@ """Tests for the compression, record batches and column types of stored DataFrames.""" from pathlib import Path -from typing import Any, Callable, Dict, List, Optional +from typing import Callable, Dict, List, Optional, Type, Union import pandas as pd import polars as pl @@ -21,7 +21,9 @@ @pytest.mark.parametrize( "make_frame", [pd.DataFrame, pl.DataFrame], ids=["pandas", "polars"] ) -def test_frames_are_compressed(tmp_project_root: Path, make_frame: Any) -> None: +def test_frames_are_compressed( + tmp_project_root: Path, make_frame: Union[Type[pd.DataFrame], Type[pl.DataFrame]] +) -> None: artifacts.store_dataframe(make_frame({"a": ["x" * 100] * 10_000}), NAME) stored = frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file @@ -135,7 +137,7 @@ def _categories_in_two_batches() -> pl.DataFrame: ], ) def test_manifest_batch_rows_match_the_files_record_batches( - tmp_project_root: Path, make_frame: Callable[[], Any], expected: List[int] + tmp_project_root: Path, make_frame: Callable[[], pd.DataFrame], expected: List[int] ) -> None: frame = make_frame() @@ -202,7 +204,7 @@ def _arrow_text(rows: int) -> pd.Series: ], ) def test_buffers_shared_by_columns_or_chunks_count_once_per_use( - tmp_project_root: Path, make_frame: Callable[[], Any] + tmp_project_root: Path, make_frame: Callable[[], pd.DataFrame] ) -> None: # Every batch writes its own copy of each column and chunk. artifacts.store_dataframe(make_frame(), NAME) diff --git a/tests/unit/artifacts/test_manifest_and_cleanup.py b/tests/unit/artifacts/test_manifest_and_cleanup.py index e0295e5..dcf36cb 100644 --- a/tests/unit/artifacts/test_manifest_and_cleanup.py +++ b/tests/unit/artifacts/test_manifest_and_cleanup.py @@ -6,7 +6,7 @@ import threading from pathlib import Path from types import SimpleNamespace -from typing import Any, Dict, Iterator, List +from typing import IO, Any, Dict, Iterator, List, Tuple import pandas as pd import pyarrow as pa @@ -224,22 +224,22 @@ def too_deep(raw: bytes) -> Any: def test_manifest_is_written_after_the_frame_file_is_complete( - tmp_project_root: Path, arrow_sinks: List[Any], monkeypatch: pytest.MonkeyPatch + tmp_project_root: Path, arrow_sinks: List[IO[Any]], monkeypatch: pytest.MonkeyPatch ) -> None: - seen: List[Dict[str, Any]] = [] + seen: List[Tuple[bool, int]] = [] original = _dataframe_storage.write_manifest def spy(directory: Path, encoded: str) -> None: new_file = json.loads(encoded)["file"] table = pa.ipc.open_file(directory / new_file).read_all() - seen.append({"closed": arrow_sinks[0].closed, "rows": table.num_rows}) + seen.append((arrow_sinks[0].closed, table.num_rows)) original(directory, encoded) monkeypatch.setattr(_dataframe_storage, "write_manifest", spy) artifacts.store_dataframe(pd.DataFrame({"a": [1, 2, 3]}), NAME) - assert seen == [{"closed": True, "rows": 3}] + assert seen == [(True, 3)] def test_concurrent_writers_leave_a_complete_current_frame( diff --git a/tests/unit/artifacts/test_storing.py b/tests/unit/artifacts/test_storing.py index 7ca17b0..c06ec8b 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -5,7 +5,7 @@ import re from datetime import datetime, timezone from pathlib import Path -from typing import Any +from typing import Any, Union import pandas as pd import polars as pl @@ -134,7 +134,9 @@ def test_transposed_frame_is_stored_with_string_names(tmp_project_root: Path) -> [pd.DataFrame({"a": []}), pl.DataFrame({"a": []})], ids=["pandas", "polars"], ) -def test_empty_frame_is_stored(tmp_project_root: Path, frame: Any) -> None: +def test_empty_frame_is_stored( + tmp_project_root: Path, frame: Union[pd.DataFrame, pl.DataFrame] +) -> None: artifacts.store_dataframe(frame, NAME) assert stored_table(tmp_project_root).num_rows == 0 diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 2f9606c..14eb9e1 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -34,7 +34,7 @@ def _store_a_good_frame(tmp_project_root: Path) -> bytes: def _store_bad_frame_after_a_good_one( - tmp_project_root: Path, make_frame: Callable[[], Any], expected_error: type + tmp_project_root: Path, make_frame: Callable[[], pd.DataFrame], expected_error: type ) -> Dict[str, Any]: """Store a good frame, then a bad one. From e7ac53c04600abfc5e8da422de7fd55c7aba04d5 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 28 Oct 2026 18:00:52 +0000 Subject: [PATCH 09/10] fix(tests): update type hint for make_frame to support both pandas and polars DataFrames --- tests/unit/artifacts/test_file_format.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index c1d0aee..9884aff 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -137,7 +137,7 @@ def _categories_in_two_batches() -> pl.DataFrame: ], ) def test_manifest_batch_rows_match_the_files_record_batches( - tmp_project_root: Path, make_frame: Callable[[], pd.DataFrame], expected: List[int] + tmp_project_root: Path, make_frame: Callable[[], Union[pd.DataFrame, pl.DataFrame]], expected: List[int] ) -> None: frame = make_frame() From bac0397f55efb0b422efbeee1279bcce53369ae9 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 29 Oct 2026 18:18:16 +0000 Subject: [PATCH 10/10] fix(artifacts): split record batches that long rows push past 12 MiB Batches are cut by the frame's average row size, so long rows that sit together could make one batch several times the 8 MiB target, and the dataframe reader refuses any page that reads more than 32 MiB. A batch over 12 MiB is now split into the longest prefixes that fit. Halving it instead would leave a trail of ever-smaller batches around the long rows, and the reader refuses a page that spans more than 8. Also wraps a test annotation that black rejected. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx --- .../artifacts/_dataframe_storage.py | 29 ++++++++++- tests/unit/artifacts/test_file_format.py | 48 ++++++++++++++++++- 2 files changed, 73 insertions(+), 4 deletions(-) diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 42d1c1b..d661b90 100644 --- a/deepnote_toolkit/artifacts/_dataframe_storage.py +++ b/deepnote_toolkit/artifacts/_dataframe_storage.py @@ -37,6 +37,7 @@ # A page spanning two batches reads at most half the reader's 32 MiB limit, # leaving room for dictionaries, IPC framing and uneven rows. TARGET_BATCH_BYTES = 8 * 1024 * 1024 +MAX_BATCH_BYTES = 12 * 1024 * 1024 StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -116,6 +117,29 @@ def _rows_per_batch(table: "pa.Table") -> int: return max(1, min(MAX_BATCH_ROWS, int(TARGET_BATCH_BYTES // row_bytes))) +def _split_oversized(batch: "pa.RecordBatch") -> list["pa.RecordBatch"]: + """Split `batch` into parts of at most MAX_BATCH_BYTES, or one row each. + + Each part but the last is the longest prefix that fits, so long rows don't + leave a trail of small batches that one page would have to span. + """ + import pyarrow as pa + + parts = [] + while batch.num_rows > 1 and pa.ipc.get_record_batch_size(batch) > MAX_BATCH_BYTES: + fits, too_many = 1, batch.num_rows + while too_many - fits > 1: + middle = (fits + too_many) // 2 + if pa.ipc.get_record_batch_size(batch.slice(0, middle)) <= MAX_BATCH_BYTES: + fits = middle + else: + too_many = middle + parts.append(batch.slice(0, fits)) + batch = batch.slice(fits) + parts.append(batch) + return parts + + def _write_frame_file( frame: StorableFrame, path: Path, compression: Optional[str] ) -> list[int]: @@ -140,8 +164,9 @@ def _write_frame_file( # make small batches; combining a slice copies at most one batch. batch_table = table.slice(offset, length).combine_chunks() for batch in batch_table.to_batches(): - writer.write_batch(batch) - batch_rows.append(batch.num_rows) + for part in _split_oversized(batch): + writer.write_batch(part) + batch_rows.append(part.num_rows) return batch_rows diff --git a/tests/unit/artifacts/test_file_format.py b/tests/unit/artifacts/test_file_format.py index 9884aff..d173820 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -132,12 +132,35 @@ def _categories_in_two_batches() -> pl.DataFrame: [65_536, 65_536, 65_536, 3_392], id="polars-list-of-categoricals-many-chunks", ), + pytest.param( + lambda: pd.DataFrame( + {"c": pd.Categorical([f"{i:0100d}" for i in range(200_000)])} + ), + [65_536, 65_536, 65_536, 3_392], + id="pandas-categorical-dictionary-larger-than-a-batch", + ), + pytest.param( + lambda: pd.DataFrame( + {"t": ["y"] * 150 + ["x" * 7_000_000] * 2 + ["y"] * 19_848} + ), + [151, 11_747, 8_102], + id="pandas-long-rows-cut-where-a-batch-is-full", + ), + pytest.param( + lambda: pd.DataFrame( + {"t": ["y"] * 150 + ["x" * 13_000_000] + ["y"] * 19_849} + ), + [150, 1, 12_656, 7_193], + id="pandas-row-over-the-limit-among-short-ones", + ), pytest.param(lambda: pd.DataFrame({"a": []}), [], id="pandas-empty"), pytest.param(lambda: pd.DataFrame(index=range(5)), [5], id="pandas-no-columns"), ], ) def test_manifest_batch_rows_match_the_files_record_batches( - tmp_project_root: Path, make_frame: Callable[[], Union[pd.DataFrame, pl.DataFrame]], expected: List[int] + tmp_project_root: Path, + make_frame: Callable[[], Union[pd.DataFrame, pl.DataFrame]], + expected: List[int], ) -> None: frame = make_frame() @@ -176,12 +199,33 @@ def test_long_text_is_split_into_batches_of_at_most_8_mib( def test_row_larger_than_a_batch_gets_a_batch_of_its_own( tmp_project_root: Path, ) -> None: - artifacts.store_dataframe(pd.DataFrame({"text": ["x" * 10_000_000] * 3}), NAME) + artifacts.store_dataframe(pd.DataFrame({"text": ["x" * 13_000_000] * 3}), NAME) assert stored_manifest(tmp_project_root).batch_rows == [1, 1, 1] assert _file_batch_rows(tmp_project_root) == [1, 1, 1] +def test_long_rows_together_are_split_into_batches_of_at_most_12_mib( + tmp_project_root: Path, +) -> None: + # 20 MB in the first and in the last 400 rows, which the frame's average row size + # puts in one batch each. + long_rows = ["x" * 50_000] * 400 + frame = pd.DataFrame({"t": long_rows + ["y"] * 99_200 + long_rows}) + + artifacts.store_dataframe(frame, NAME) + + reader = pa.ipc.open_file( + frame_dir(tmp_project_root) / stored_manifest(tmp_project_root).file + ) + batches = [reader.get_batch(i) for i in range(reader.num_record_batches)] + assert max(batch.nbytes for batch in batches) <= 12 * 1024 * 1024 + assert stored_manifest(tmp_project_root).batch_rows == [ + batch.num_rows for batch in batches + ] + assert pa.Table.from_batches(batches).column("t").to_pylist() == list(frame["t"]) + + def _arrow_text(rows: int) -> pd.Series: return pd.Series(["x" * 500] * rows, dtype="string[pyarrow]")