feat(artifacts): add store_dataframe and delete_dataframe

deepnote_core/pydantic_compat_helpers.py

@@ -1,4 +1,8 @@
-from typing import Any, Mapping
+from typing import Any, Mapping, TypeVar
+
+from pydantic import BaseModel
+
+Model = TypeVar("Model", bound=BaseModel)
 
 
 def _model_dump_compat(model_cls) -> dict:
@@ -45,7 +49,7 @@ def _get_description_compat(field: Any) -> str:
     return getattr(field_info, "description", "") if field_info else ""
 
 
-def model_validate_compat(model_cls, data: Mapping[str, Any]) -> Any:
+def model_validate_compat(model_cls: type[Model], data: Mapping[str, Any]) -> Model:
     """Validate using Pydantic v2 if available, else fall back to v1.
 
     Pydantic v2 exposes `model_validate`, while v1 uses `parse_obj` on the class.

deepnote_toolkit/__init__.py

@@ -38,6 +38,7 @@
     (".variable_explorer", "deepnote_get_data_preview_json"),
     (".variable_explorer", "get_var_list"),
     (".ocelots", None),
+    (".artifacts", None),
 ]
 
 # Perform safe imports

deepnote_toolkit/artifacts/__init__.py

@@ -0,0 +1,5 @@
+"""Stores and deletes named DataFrames in the project's filesystem."""
+
+from .dataframes import delete_dataframe, store_dataframe
+
+__all__ = ["delete_dataframe", "store_dataframe"]

deepnote_toolkit/artifacts/dataframe_storage.py

@@ -0,0 +1,135 @@
+"""Writes named DataFrames into the project's filesystem as Arrow files."""
+
+import contextlib
+import errno
+import os
+import re
+import time
+import uuid
+from pathlib import Path
+from typing import Any, 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,
+    is_frame_name,
+    read_manifest,
+    write_manifest,
+)
+
+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")
+
+StorableFrame: TypeAlias = Union[PandasDF, PolarsEagerDF]
+
+_SAFE_ID_REGEX = re.compile(r"[A-Za-z0-9_-]{1,128}")
+
+
+def is_safe_id(value: object) -> bool:
+    """Return whether `value` can be used as a single path segment."""
+    return isinstance(value, str) and _SAFE_ID_REGEX.fullmatch(value) is not None
+
+
+def resolve_project_root() -> Path:
+    """Return the project's root directory, the one set_notebook_path() uses."""
+    return resolve_notebook_root()
+
+
+def _select_compression() -> str:
+    """Return the best codec this pyarrow build has: zstd, lz4, else uncompressed."""
+    import pyarrow as pa
+
+    for codec in ("zstd", "lz4"):
+        if pa.Codec.is_available(codec):
+            return codec
+    return "uncompressed"
+
+
+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
+
+            # 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
+
+        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)
+
+
+def _delete_replaced_frames(
+    directory: Path, old: Optional[Manifest], new_file: str
+) -> None:
+    """Delete the frame the replaced manifest called `previous`, and stale strays.
+
+    Another kernel may still be writing an unnamed frame here, so a stray must be a
+    day old before it is deleted.
+    """
+    named = {new_file}
+    if old is not None:
+        named.add(old.file)
+        if old.previous is not None and old.previous not in named:
+            (directory / old.previous).unlink(missing_ok=True)
+
+    now = time.time()
+    with os.scandir(directory) as entries:
+        for entry in entries:
+            if entry.name in named or not is_frame_name(entry.name):
+                continue
+            with contextlib.suppress(FileNotFoundError):
+                if (
+                    entry.is_file(follow_symlinks=False)
+                    and now - entry.stat(follow_symlinks=False).st_mtime
+                    > STALE_FRAME_AGE_SECONDS
+                ):
+                    (directory / entry.name).unlink()
+
+
+def store_frame(frame: StorableFrame, name: str, root: Path) -> Optional[str]:
+    """Store the frame as the current one under `name` and return its reference ID.
+
+    Returns None, having stored nothing, if the filesystem is read-only. EROFS from
+    the manifest step on raises like any other error: the new frame may be current.
+    """
+    directory = root / FRAMES_DIR / name
+    reference_id = str(uuid.uuid4())
+    new_file = f"{reference_id}.arrow"
+    try:
+        directory.mkdir(parents=True, exist_ok=True)
+        _write_frame_file(frame, directory / new_file, _select_compression())
+        old = read_manifest(directory)
+    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))
+    _delete_replaced_frames(directory, old, new_file)
+    return reference_id

deepnote_toolkit/artifacts/dataframe_storage_manifest.py

@@ -0,0 +1,105 @@
+"""Reads and writes the manifest that names a stored DataFrame's current file."""
+
+import contextlib
+import json
+import os
+import re
+import tempfile
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Optional
+
+from pydantic import BaseModel, StrictInt, StrictStr, ValidationError
+
+from deepnote_core.pydantic_compat_helpers import (
+    _model_dump_compat,
+    model_validate_compat,
+)
+from deepnote_toolkit import __version__
+
+MANIFEST_NAME = "manifest.json"
+MANIFEST_VERSION = 1
+MANIFEST_MAX_BYTES = 4096
+
+_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",
+    re.IGNORECASE,
+)
+
+
+def is_frame_name(name: str) -> bool:
+    """Return whether `name` is `<uuid>.arrow`, the only form of frame file name."""
+    return _FRAME_NAME_REGEX.fullmatch(name) is not None
+
+
+class Manifest(BaseModel):
+    """The manifest file; `previous` is the frame `file` replaced."""
+
+    version: StrictInt
+    file: StrictStr
+    previous: Optional[StrictStr] = None
+    written_at: StrictStr
+    rows: StrictInt
+    toolkit_version: StrictStr
+
+
+def read_manifest(directory: Path) -> Optional[Manifest]:
+    """Return the directory's manifest, or None if there is no usable one.
+
+    Any code in the project can write the manifest, so a malformed one counts as
+    absent.
+    """
+    try:
+        with (directory / MANIFEST_NAME).open("rb") as handle:
+            raw = handle.read(MANIFEST_MAX_BYTES + 1)
+    except FileNotFoundError:
+        return None
+    if len(raw) > MANIFEST_MAX_BYTES:
+        return None
+    try:
+        data = json.loads(raw)
+    except (ValueError, RecursionError):
+        return None
+    try:
+        manifest = model_validate_compat(Manifest, data)
+    except (ValidationError, TypeError):
+        return None
+
+    return manifest
+
+
+def write_manifest(
+    directory: Path, new_file: str, previous: Optional[str], rows: int
+) -> None:
+    """Point the manifest at `new_file`.
+
+    On s3fs a rename is a copy and then a delete, so a failed rename may already
+    have replaced the manifest.
+    """
+    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,
+        toolkit_version=__version__,
+    )
+    # 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.
+    temporary = tempfile.NamedTemporaryFile(
+        "w",
+        encoding="utf-8",
+        dir=directory,
+        prefix=f".{MANIFEST_NAME}.",
+        suffix=".tmp",
+        delete=False,
+    )
+    try:
+        with temporary:
+            temporary.write(json.dumps(_model_dump_compat(manifest)))
+        os.replace(temporary.name, directory / MANIFEST_NAME)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.unlink(temporary.name)
+        raise

deepnote_toolkit/artifacts/dataframes.py

@@ -0,0 +1,90 @@
+"""Stores and deletes named DataFrames in the project's filesystem."""
+
+import errno
+import shutil
+from typing import Any, Optional
+
+from deepnote_toolkit.ocelots.utils import (
+    is_pandas_dataframe,
+    is_polars_eager_dataframe,
+)
+from deepnote_toolkit.sql.query_preview import DeepnoteQueryPreview
+
+from . import dataframe_storage
+
+
+def _check_name(name: object) -> None:
+    if not dataframe_storage.is_safe_id(name):
+        raise ValueError(
+            f"A DataFrame name must match [A-Za-z0-9_-]{{1,128}}, got {name!r}"
+        )
+
+
+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.
+        name: The frame's name, unique within the project, and its directory
+            under `deepnote_dataframes/`. Must match `[A-Za-z0-9_-]{1,128}`.
+
+    Returns:
+        The new frame's reference ID, or None if the project's filesystem is
+        read-only, in which case nothing is stored.
+
+    Raises:
+        ValueError: `name` is not valid.
+        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.
+        OSError: The write failed.
+        Exception: `df` cannot be serialized to Arrow: whatever pyarrow's
+            `write_feather()` or polars' `df.write_ipc()` 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,
+    deleting the frames it replaced can fail, and so can an s3fs rename that has
+    already put the new manifest in place.
+    """
+    _check_name(name)
+    if isinstance(df, DeepnoteQueryPreview):
+        raise TypeError(
+            "store_dataframe refuses a SQL query preview or a frame derived from "
+            "one, which holds at most 100 rows of the query's result. Set the SQL "
+            "block's output mode to DataFrame, or pass pd.DataFrame(df) to store "
+            "these rows"
+        )
+    if not (is_pandas_dataframe(df) or is_polars_eager_dataframe(df)):
+        raise TypeError(
+            "store_dataframe needs a pandas or polars DataFrame, "
+            f"got {type(df).__qualname__}"
+        )
+
+    root = dataframe_storage.resolve_project_root()
+    return dataframe_storage.store_frame(df, name, root)
+
+
+def delete_dataframe(name: str) -> bool:
+    """Delete the frames stored under `name`.
+
+    Returns:
+        Whether there was anything to delete. False on a read-only filesystem, where
+        nothing is deleted.
+
+    Raises:
+        ValueError: `name` is not valid.
+        OSError: The delete failed, possibly part-way, leaving some of the name's
+            files. A symbolic link in place of the name's directory raises too.
+    """
+    _check_name(name)
+    root = dataframe_storage.resolve_project_root()
+    directory = root / dataframe_storage.FRAMES_DIR / name
+    if not directory.is_dir():
+        return False
+    try:
+        shutil.rmtree(directory)
+    except OSError as exc:
+        if exc.errno == errno.EROFS:
+            return False
+        raise
+    return True

deepnote_toolkit/config.py

@@ -15,6 +15,8 @@
 from deepnote_core.config.loader import ConfigurationLoader
 from deepnote_core.config.models import DeepnoteConfig
 
+from . import env
+
 
 @lru_cache(maxsize=128)
 def _load_config_cached(config_path: Optional[Path]) -> DeepnoteConfig:
@@ -63,3 +65,20 @@ def get_config(
 
     # Call the cached function
     return _load_config_cached(cfg_path)
+
+
+def resolve_notebook_root() -> Path:
+    """Return the project's root directory, which notebook paths are relative to."""
+    try:
+        cfg = get_config()
+        if cfg.paths.notebook_root:
+            return Path(cfg.paths.notebook_root)
+        if cfg.paths.home_dir:
+            return Path(cfg.paths.home_dir)
+        if cfg.runtime.dev_mode:
+            return Path("/work")
+        return Path(os.environ.get("HOME", str(Path.home()))) / "work"
+    except Exception:  # pylint: disable=broad-exception-caught
+        if env.get_env("DEEPNOTE_RUNNING_IN_DEV_MODE"):
+            return Path("/work")
+        return Path(os.environ.get("HOME", str(Path.home()))) / "work"