diff --git a/deepnote_toolkit/artifacts/_dataframe_storage.py b/deepnote_toolkit/artifacts/_dataframe_storage.py index 95efac1..c8c6b98 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 Any, Optional, Union +from typing import TYPE_CHECKING, Callable, 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 @@ -18,16 +17,27 @@ from ._dataframe_storage_manifest import ( Manifest, + encode_manifest, is_frame_name, 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 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 +MAX_BATCH_BYTES = 12 * 1024 * 1024 StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF] @@ -44,42 +54,160 @@ 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: Deepnote Cloud doesn't support it yet. + """ 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 _to_arrow_table(frame: StorableFrame) -> "pa.Table": + """Return the frame as an Arrow table.""" + 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): + 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()) + + assert_never(frame) + + +def _contains_type( + data_type: "pa.DataType", matches: Callable[["pa.DataType"], bool] +) -> bool: + """Return whether `matches` holds for `data_type` or any type nested in it.""" + import pyarrow as pa + + if matches(data_type): + return True + if isinstance(data_type, pa.BaseExtensionType): + return _contains_type(data_type.storage_type, matches) + if pa.types.is_dictionary(data_type): + return _contains_type(data_type.value_type, matches) + return any( + _contains_type(data_type.field(i).type, matches) + for i in range(data_type.num_fields) + ) + + +def _is_view_type(data_type: "pa.DataType") -> bool: + """Return whether `data_type` is string_view, binary_view or a list view.""" + import pyarrow as pa - # pandas < 2.1's to_feather needs a default index and string column names. - feather.write_feather(frame, sink, compression=compression) - return + # pyproject allows pyarrow older than 16, which lacks these functions. + return any( + hasattr(pa.types, check) and getattr(pa.types, check)(data_type) + for check in ( + "is_string_view", + "is_binary_view", + "is_list_view", + "is_large_list_view", + ) + ) - # 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__) - 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 +def _refuse_view_types(table: "pa.Table") -> None: + """Raise TypeError naming each column whose type is or nests a view type.""" + # IPC writes a slice of a view column with up to all of the column's data, so + # every batch would repeat it. pyarrow older than 24 can't measure their nbytes. + views = [ + field.name + for field in table.schema + if _contains_type(field.type, _is_view_type) + ] + if views: + raise TypeError(f"Arrow view-type columns cannot be stored, got {views}") - assert_never(frame) + +def _rows_per_batch(table: "pa.Table") -> int: + """Return how many rows make a batch of about TARGET_BATCH_BYTES.""" + import pyarrow as pa + + # 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 _contains_type(column.type, pa.types.is_dictionary) + 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 _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]: + """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 + + with path.open("wb") as sink: + table = _to_arrow_table(frame) + _refuse_view_types(table) + # A file with more than one dictionary for a column is refused by pyarrow. + table = table.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(): + for part in _split_oversized(batch): + writer.write_batch(part) + batch_rows.append(part.num_rows) + return batch_rows def _delete_replaced_frames( @@ -125,8 +253,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 +265,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..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, @@ -19,7 +19,7 @@ MANIFEST_NAME = "manifest.json" MANIFEST_VERSION = 1 -MANIFEST_MAX_BYTES = 4096 +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 +33,18 @@ def is_frame_name(name: str) -> bool: class Manifest(BaseModel): - """The manifest file; `previous` is the frame `file` replaced.""" + """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 + 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 @@ -68,22 +73,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 +119,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/_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/deepnote_toolkit/artifacts/dataframes.py b/deepnote_toolkit/artifacts/dataframes.py index aa11fd8..17492e4 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,21 @@ 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. Also raised for a pandas DataFrame with a column or index + whose Arrow type is or nests a view type (string_view, binary_view, + list_view or large_list_view). 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/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 9d1c28e..f1bad71 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()) == [] @@ -93,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 96f1eb1..d173820 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 Callable, Dict, List, Optional, Type, Union import pandas as pd import polars as pl @@ -22,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 @@ -43,80 +44,220 @@ 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) + + +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( + {"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], +) -> 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 == 2 - assert reader.read_all().num_rows == 70_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 + ] -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 test_row_larger_than_a_batch_gets_a_batch_of_its_own( + tmp_project_root: Path, +) -> None: + 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 ) - 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 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]") @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"], + "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_polars_batch_size_is_passed_to_polars_from_1_37( - tmp_project_root: Path, - monkeypatch: pytest.MonkeyPatch, - polars_version: str, - expected_options: Dict[str, int], +def test_buffers_shared_by_columns_or_chunks_count_once_per_use( + tmp_project_root: Path, make_frame: Callable[[], pd.DataFrame] ) -> 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": [1, 2, 3]}), NAME) + # 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 +277,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..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 @@ -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: @@ -208,21 +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, 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) + 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_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_read_only_mount.py b/tests/unit/artifacts/test_read_only_mount.py index 608823b..b2446c6 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,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(frame: pd.DataFrame, sink: Any, **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(feather, "write_feather", 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_storing.py b/tests/unit/artifacts/test_storing.py index c6eebc9..af390cf 100644 --- a/tests/unit/artifacts/test_storing.py +++ b/tests/unit/artifacts/test_storing.py @@ -5,10 +5,11 @@ 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 +import pyarrow as pa import pytest import deepnote_toolkit @@ -40,11 +41,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) @@ -126,13 +135,32 @@ 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 assert stored_manifest(tmp_project_root).rows == 0 +def test_frame_is_stored_on_pyarrow_older_than_16( + tmp_project_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # pyarrow older than 16 has none of these. + for check in ( + "is_string_view", + "is_binary_view", + "is_list_view", + "is_large_list_view", + ): + monkeypatch.delattr(pa.types, check) + + artifacts.store_dataframe(pd.DataFrame({"a": [1], "s": ["x"]}), NAME) + + assert stored_table(tmp_project_root).to_pydict() == {"a": [1], "s": ["x"]} + + def test_names_are_stored_apart(tmp_project_root: Path) -> None: artifacts.store_dataframe(pd.DataFrame({"a": [1]}), "first") artifacts.store_dataframe(pd.DataFrame({"a": [2]}), "second") diff --git a/tests/unit/artifacts/test_write_failures.py b/tests/unit/artifacts/test_write_failures.py index 9137c5f..3b8c3d5 100644 --- a/tests/unit/artifacts/test_write_failures.py +++ b/tests/unit/artifacts/test_write_failures.py @@ -9,9 +9,8 @@ from typing import Any, Callable, Dict, Optional 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 @@ -35,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. @@ -74,32 +73,108 @@ def test_unserializable_frame_raises_and_keeps_the_previous_frame( assert stored_table(tmp_project_root).to_pydict() == {"ok": [1]} -def test_polars_object_column_raises_and_keeps_the_previous_frame( - tmp_project_root: Path, +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(pl.DataFrame({"ok": [1]}), NAME) + 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 + assert len(arrow_files(tmp_project_root)) == 1 + + +def _arrow_column(values: pa.Array) -> pd.Series: + return pd.Series(pd.arrays.ArrowExtensionArray(values)) + + +@pytest.mark.parametrize( + "make_values", + [ + pytest.param(lambda: pa.array(["a"], pa.string_view()), id="string-view"), + pytest.param(lambda: pa.array([b"a"], pa.binary_view()), id="binary-view"), + pytest.param( + lambda: pa.array([["a"]], pa.list_(pa.string_view())), + id="list-of-string-views", + ), + pytest.param( + lambda: pa.array( + [{"n": 1, "s": "a"}], + pa.struct([("n", pa.int64()), ("s", pa.string_view())]), + ), + id="struct-with-a-string-view-second", + ), + pytest.param( + lambda: pa.DictionaryArray.from_arrays( + pa.array([0], pa.int64()), pa.array(["a"], pa.string_view()) + ), + id="dictionary-of-string-views", + ), + pytest.param( + lambda: pa.array([["a"]], pa.fixed_shape_tensor(pa.string_view(), [1])), + id="tensor-extension-of-string-views", + ), + pytest.param(lambda: pa.array([[1]], pa.list_view(pa.int64())), id="list-view"), + pytest.param( + lambda: pa.array([[1]], pa.large_list_view(pa.int64())), + id="large-list-view", + ), + ], +) +def test_view_type_column_raises_and_keeps_the_previous_frame( + tmp_project_root: Path, make_values: Callable[[], pa.Array] +) -> None: + manifest_bytes = _store_a_good_frame(tmp_project_root) + frame = pd.DataFrame({"ok": [1], "a": _arrow_column(make_values())}) - 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"Arrow view-type 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_view_type_error_names_each_view_type_column_and_no_other( + tmp_project_root: Path, +) -> None: + frame = pd.DataFrame( + { + "s": _arrow_column(pa.array(["a"], pa.string_view())), + "ok": _arrow_column(pa.array(["a"], pa.large_string())), + "l": _arrow_column(pa.array([[1]], pa.list_view(pa.int64()))), + } + ) + + with pytest.raises(TypeError) as raised: + artifacts.store_dataframe(frame, NAME) + + raised.match(r"view-type columns cannot be stored, got \['s', 'l'\]") + + 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: - sink.write(b"ARROW1 partial") + real_write_batch = pa.ipc.RecordBatchFileWriter.write_batch + + def write_a_batch_then_fail( + self: pa.ipc.RecordBatchFileWriter, batch: pa.RecordBatch, **kwargs: Any + ) -> None: + real_write_batch(self, batch, **kwargs) raise OSError(errno.EIO, "Input/output error") - monkeypatch.setattr(feather, "write_feather", 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) @@ -299,16 +374,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( - frame: pd.DataFrame, sink: Any, **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(feather, "write_feather", 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) 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