diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 95efac1..3ce7019 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") -POLARS_COMPAT_LEVEL_MIN_VERSION = Version("1.2.1") +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_MIN_VERSION = Version("1.35.1") +POLARS_ARRAY_MIN_VERSION = Version("1.43.2") StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -44,42 +51,147 @@ 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.""" +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. + """ import pyarrow as pa - for codec in ("zstd", "lz4"): - if pa.Codec.is_available(codec): - return codec - return "uncompressed" + return "zstd" if pa.Codec.is_available("zstd") else None -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 +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 - # 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 _to_arrow_table(frame: StorableFrame) -> "pa.Table": + """Return the frame as an Arrow table. - 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 + 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 + + # 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): + 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." + ) + # 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) + + +def _has_dictionary(data_type: "pa.DataType") -> bool: + """Return whether `data_type` is a dictionary type or nests one.""" + import pyarrow as pa + + 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 +237,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 +249,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..63bae58 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,19 @@ 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 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 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..d55e04d 100644 --- a/tests/unit/artifacts/test_file_format.py +++ b/tests/unit/artifacts/test_file_format.py @@ -1,8 +1,7 @@ """Tests for the compression, record batches and column types of stored DataFrames.""" from pathlib import Path -from typing import Any, Dict -from unittest import mock +from typing import Any, Callable, Dict, List, Optional import pandas as pd import polars as pl @@ -43,80 +42,176 @@ 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)] - 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_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) -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) + +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) + + +@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) + + assert stored_manifest(tmp_project_root).batch_rows == [1, 1, 1] + assert _file_batch_rows(tmp_project_root) == [1, 1, 1] + - 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) +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", + ), + ], +) +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) - 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]} + 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: @@ -136,26 +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.write_ipc - - with mock.patch.object( - pl.DataFrame, "write_ipc", autospec=True, side_effect=original - ) as write_ipc: - artifacts.store_dataframe(pl.DataFrame({"a": ["x"]}), NAME) - - assert ("compat_level" in write_ipc.call_args.kwargs) is passed - assert stored_table(tmp_project_root).to_pydict() == {"a": ["x"]} 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..96e1616 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,138 @@ 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(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(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: + 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 +215,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 +420,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