From f1c6735c64734771aed7b5bdfc4696d1c093a27f Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 24 Oct 2026 08:14:09 +0000 Subject: [PATCH 1/2] 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 2/2] 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: