fix(artifacts): make stored DataFrames servable a page at a time

deepnote_toolkit/artifacts/_dataframe_storage.py

@@ -7,27 +7,37 @@
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import pandas as pd
-from packaging.version import Version
 from typing_extensions import TypeAlias, assert_never
 
 from deepnote_toolkit.config import resolve_notebook_root
 from deepnote_toolkit.ocelots.types import PandasDF, PolarsEagerDF
 
 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,120 @@ 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)
 
-            # 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
 
-        # 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())
 
-        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
+    assert_never(frame)
 
-        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 _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:
+        # 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():
+                    for part in _split_oversized(batch):
+                        writer.write_batch(part)
+                        batch_rows.append(part.num_rows)
+    return batch_rows
 
 
 def _delete_replaced_frames(
@@ -125,15 +213,18 @@ 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)
         if isinstance(exc, OSError) and exc.errno == errno.EROFS:
             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

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):

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."
+        )