polars-checkpoint 0.1.0__tar.gz

polars_checkpoint-0.1.0/PKG-INFO

@@ -0,0 +1,88 @@
+Metadata-Version: 2.4
+Name: polars-checkpoint
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.11.1
+Description-Content-Type: text/markdown
+Requires-Dist: polars>=1.39.3
+
+# polars-checkpoint
+
+Materialise Polars LazyFrames to parquet files and scan them back lazily. Defaults to the streaming engine for sink/scan. Useful for managing & reducing memory pressure due to expensive intermediate results in complex multi-step transforms.
+
+## Installation
+
+Requires `polars` and Python 3.10+.
+
+## Quick Start
+
+```python
+import polars as pl
+from checkpoint import checkpoint
+
+lf = pl.LazyFrame({"x": range(1_000_000)}).with_columns(y=pl.col("x") * 2)
+
+# Materialises to a temp parquet file; returns a lazy re-scan
+lf = checkpoint(lf)
+
+lf.filter(pl.col("y") > 100).collect()
+```
+
+A process-wide default session manages the temp directory and cleans it up at exit.
+
+## Session API
+
+For explicit control over storage location and lifecycle, use `CheckpointSession`:
+
+```python
+from checkpoint import CheckpointSession
+
+# As a context manager — cleans up on exit from the block
+with CheckpointSession(root_dir="./my_checkpoints") as sess:
+    lf = pl.scan_csv("big.csv")
+    lf = sess.checkpoint(lf, name="after-parse")
+    lf = lf.filter(pl.col("status") == "active")
+    lf = sess.checkpoint(lf, name="filtered")
+```
+
+```python
+# Without a context manager — cleans up at GC or interpreter shutdown,
+# or when you call close() explicitly
+sess = CheckpointSession(root_dir="./my_checkpoints")
+lf = sess.checkpoint(pl.scan_csv("big.csv"), name="raw")
+reloaded = sess["raw"]
+print(sess.summary())
+
+sess.close()  # optional; triggers early cleanup
+```
+
+### `CheckpointSession` constructor
+
+| Parameter | Default | Description |
+|---|---|---|
+| `root_dir` | `None` (auto temp dir) | Parent directory for checkpoint folders. |
+| `cleanup` | `True` | Delete checkpoint files on close / GC / interpreter exit. |
+| `default_sink_kwargs` | `{"compression": "zstd"}` | Defaults passed to `sink_parquet` / `write_parquet`. |
+| `default_scan_kwargs` | `{}` | Defaults passed to `scan_parquet`. |
+
+### Key methods & features
+
+- **`checkpoint(lf, *, name=None, streaming=True, ...)`** — Materialise a LazyFrame to parquet. Auto-generates a name if none given. Falls back to `collect().write_parquet()` when `streaming=False`.
+- **`session[name]`** — Retrieve a checkpoint as a `LazyFrame`.
+- **`name in session`** — Check existence.
+- **`len(session)`** / **`iter(session)`** — Count / list checkpoints.
+- **`summary()`** — Returns a Polars DataFrame with name, size (MB), and path of each checkpoint.
+- **`close(timeout=None)`** — Waits for in-flight writes, then cleans up. Also usable as a context manager.
+
+## Thread Safety
+
+Sessions are internally locked. Concurrent `checkpoint()` calls from multiple threads are safe; `close()` waits for all in-flight materialisations before removing files.
+
+## Cleanup Behaviour
+
+| Scenario | `cleanup=True` (default) | `cleanup=False` |
+|---|---|---|
+| `close()` / `__exit__` | Files deleted | Files retained |
+| GC / interpreter shutdown | Files deleted (via `weakref.finalize`) | Files retained |
+
+When `root_dir` is auto-generated, the entire temp directory is removed. When user-supplied, only the individual checkpoint subdirectories created by the session are removed.

polars_checkpoint-0.1.0/README.md

@@ -0,0 +1,80 @@
+# polars-checkpoint
+
+Materialise Polars LazyFrames to parquet files and scan them back lazily. Defaults to the streaming engine for sink/scan. Useful for managing & reducing memory pressure due to expensive intermediate results in complex multi-step transforms.
+
+## Installation
+
+Requires `polars` and Python 3.10+.
+
+## Quick Start
+
+```python
+import polars as pl
+from checkpoint import checkpoint
+
+lf = pl.LazyFrame({"x": range(1_000_000)}).with_columns(y=pl.col("x") * 2)
+
+# Materialises to a temp parquet file; returns a lazy re-scan
+lf = checkpoint(lf)
+
+lf.filter(pl.col("y") > 100).collect()
+```
+
+A process-wide default session manages the temp directory and cleans it up at exit.
+
+## Session API
+
+For explicit control over storage location and lifecycle, use `CheckpointSession`:
+
+```python
+from checkpoint import CheckpointSession
+
+# As a context manager — cleans up on exit from the block
+with CheckpointSession(root_dir="./my_checkpoints") as sess:
+    lf = pl.scan_csv("big.csv")
+    lf = sess.checkpoint(lf, name="after-parse")
+    lf = lf.filter(pl.col("status") == "active")
+    lf = sess.checkpoint(lf, name="filtered")
+```
+
+```python
+# Without a context manager — cleans up at GC or interpreter shutdown,
+# or when you call close() explicitly
+sess = CheckpointSession(root_dir="./my_checkpoints")
+lf = sess.checkpoint(pl.scan_csv("big.csv"), name="raw")
+reloaded = sess["raw"]
+print(sess.summary())
+
+sess.close()  # optional; triggers early cleanup
+```
+
+### `CheckpointSession` constructor
+
+| Parameter | Default | Description |
+|---|---|---|
+| `root_dir` | `None` (auto temp dir) | Parent directory for checkpoint folders. |
+| `cleanup` | `True` | Delete checkpoint files on close / GC / interpreter exit. |
+| `default_sink_kwargs` | `{"compression": "zstd"}` | Defaults passed to `sink_parquet` / `write_parquet`. |
+| `default_scan_kwargs` | `{}` | Defaults passed to `scan_parquet`. |
+
+### Key methods & features
+
+- **`checkpoint(lf, *, name=None, streaming=True, ...)`** — Materialise a LazyFrame to parquet. Auto-generates a name if none given. Falls back to `collect().write_parquet()` when `streaming=False`.
+- **`session[name]`** — Retrieve a checkpoint as a `LazyFrame`.
+- **`name in session`** — Check existence.
+- **`len(session)`** / **`iter(session)`** — Count / list checkpoints.
+- **`summary()`** — Returns a Polars DataFrame with name, size (MB), and path of each checkpoint.
+- **`close(timeout=None)`** — Waits for in-flight writes, then cleans up. Also usable as a context manager.
+
+## Thread Safety
+
+Sessions are internally locked. Concurrent `checkpoint()` calls from multiple threads are safe; `close()` waits for all in-flight materialisations before removing files.
+
+## Cleanup Behaviour
+
+| Scenario | `cleanup=True` (default) | `cleanup=False` |
+|---|---|---|
+| `close()` / `__exit__` | Files deleted | Files retained |
+| GC / interpreter shutdown | Files deleted (via `weakref.finalize`) | Files retained |
+
+When `root_dir` is auto-generated, the entire temp directory is removed. When user-supplied, only the individual checkpoint subdirectories created by the session are removed.

polars_checkpoint-0.1.0/pyproject.toml

@@ -0,0 +1,9 @@
+[project]
+name = "polars-checkpoint"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11.1"
+dependencies = [
+    "polars>=1.39.3",
+]

polars_checkpoint-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

polars_checkpoint-0.1.0/src/polars_checkpoint/polars_checkpoint.py

@@ -0,0 +1,341 @@
+from __future__ import annotations
+
+import logging
+import shutil
+import tempfile
+import threading
+import time
+import uuid
+import weakref
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+import polars as pl
+
+logger = logging.getLogger(__name__)
+
+_default_session_lock = threading.Lock()
+_default_session: CheckpointSession | None = None
+
+
+def _cleanup_files(
+    root_dir: Path,
+    auto_root: bool,
+    owned_dirs: list[Path],
+) -> None:
+    """Safety-net cleanup invoked by weakref.finalize."""
+    if auto_root:
+        shutil.rmtree(root_dir, ignore_errors=True)
+    else:
+        for d in owned_dirs:
+            shutil.rmtree(d, ignore_errors=True)
+
+
+class CheckpointSession:
+    """Manages a directory of parquet checkpoints for Polars LazyFrames."""
+
+    def __init__(
+        self,
+        root_dir: str | Path | None = None,
+        *,
+        cleanup: bool = True,
+        default_sink_kwargs: dict[str, Any] | None = None,
+        default_scan_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+        self._auto_root = root_dir is None
+        if self._auto_root:
+            self.root_dir = Path(tempfile.mkdtemp(prefix="pl-ckpt-"))
+        else:
+            self.root_dir = Path(root_dir)
+            self.root_dir.mkdir(parents=True, exist_ok=True)
+
+        self.cleanup = cleanup
+        self.default_sink_kwargs = {
+            "compression": "zstd",
+            **(default_sink_kwargs or {}),
+        }
+        self.default_scan_kwargs = dict(default_scan_kwargs or {})
+
+        self._owned_dirs: list[Path] = []
+        self._active_ops = 0
+        self._closed = False
+
+        # finalize fires at GC or interpreter shutdown (atexit=True),
+        # whichever comes first.  close() calls detach() to prevent
+        # double-cleanup.
+        if self.cleanup:
+            self._finaliser = weakref.finalize(
+                self,
+                _cleanup_files,
+                self.root_dir,
+                self._auto_root,
+                self._owned_dirs,  # same list object, mutated in place
+            )
+            self._finaliser.atexit = True
+        else:
+            self._finaliser = None
+
+    def __repr__(self) -> str:
+        if self._closed:
+            state = "closed"
+        else:
+            n = len(self._owned_dirs)
+            state = f"open, {n} checkpoint{'s' if n != 1 else ''}"
+        return f"<CheckpointSession root_dir={str(self.root_dir)!r} {state}>"
+
+    # -- context manager (optional, for early deterministic cleanup) -----------
+
+    def __enter__(self) -> CheckpointSession:
+        with self._cond:
+            self._ensure_open()
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    # -- lifecycle -------------------------------------------------------------
+
+    def close(self, *, timeout: float | None = None) -> None:
+        """Close the session, wait for in-flight checkpoints, clean up files."""
+        with self._cond:
+            if self._closed:
+                return
+            self._closed = True
+
+            deadline = None if timeout is None else time.monotonic() + timeout
+            while self._active_ops > 0:
+                remaining = (
+                    None if deadline is None else max(0.0, deadline - time.monotonic())
+                )
+                if remaining is not None and remaining <= 0:
+                    logger.warning(
+                        "Timed out waiting for %d in-flight checkpoint(s); "
+                        "proceeding with cleanup",
+                        self._active_ops,
+                    )
+                    break
+                self._cond.wait(timeout=remaining)
+
+            owned = list(self._owned_dirs)
+
+        # Deactivate the finaliser so it won't double-clean.
+        # If the finaliser has already fired (e.g. at GC), detach() is a
+        # no-op, and the rmtree calls below are idempotent.
+        if self._finaliser is not None:
+            self._finaliser.detach()
+
+        if not self.cleanup:
+            return
+
+        if self._auto_root:
+            shutil.rmtree(self.root_dir, ignore_errors=True)
+        else:
+            for d in owned:
+                shutil.rmtree(d, ignore_errors=True)
+
+    # -- collection protocol ---------------------------------------------------
+
+    def __getitem__(self, name: str) -> pl.LazyFrame:
+        with self._cond:
+            self._ensure_open()
+        slug = _normalise_name(name)
+        path = self.root_dir / slug / "data.parquet"
+        if not path.exists():
+            with self._cond:
+                available = [
+                    d.name for d in self._owned_dirs if (d / "data.parquet").exists()
+                ]
+            raise KeyError(f"No checkpoint named {name!r}. Available: {available}")
+        return pl.scan_parquet(path, **self.default_scan_kwargs)
+
+    def __contains__(self, name: object) -> bool:
+        with self._cond:
+            self._ensure_open()
+        if not isinstance(name, str):
+            return False
+        try:
+            slug = _normalise_name(name)
+        except ValueError:
+            return False
+        return (self.root_dir / slug / "data.parquet").exists()
+
+    def __len__(self) -> int:
+        with self._cond:
+            self._ensure_open()
+            return len(self._owned_dirs)
+
+    def __iter__(self) -> Iterator[str]:
+        with self._cond:
+            self._ensure_open()
+            dirs = list(self._owned_dirs)
+        for d in dirs:
+            yield d.name
+
+    # -- introspection ---------------------------------------------------------
+
+    def summary(self) -> pl.DataFrame:
+        """Return a DataFrame listing all live checkpoints and their sizes."""
+        # Should maybe make it a .show() method that prints a table?
+        with self._cond:
+            self._ensure_open()
+            dirs = list(self._owned_dirs)
+        rows: list[dict[str, Any]] = []
+        for d in dirs:
+            p = d / "data.parquet"
+            if p.exists():
+                rows.append(
+                    {
+                        "name": d.name,
+                        "size_mb": round(p.stat().st_size / 1_048_576, 2),
+                        "path": str(p),
+                    }
+                )
+        if not rows:
+            return pl.DataFrame(
+                schema={
+                    "name": pl.Utf8,
+                    "size_mb": pl.Float64,
+                    "path": pl.Utf8,
+                },
+            )
+        return pl.DataFrame(rows)
+
+    # -- checkpointing ---------------------------------------------------------
+
+    def checkpoint(
+        self,
+        lf: pl.LazyFrame,
+        *,
+        name: str | None = None,
+        sink_kwargs: dict[str, Any] | None = None,
+        scan_kwargs: dict[str, Any] | None = None,
+        streaming: bool = True,
+    ) -> pl.LazyFrame:
+        """Materialise LazyFrame to a parquet checkpoint and return a lazy re-scan."""
+        with self._cond:
+            self._ensure_open()
+            checkpoint_dir = self._new_checkpoint_dir(name)
+            self._active_ops += 1
+
+        checkpoint_path = checkpoint_dir / "data.parquet"
+        sink_opts = {**self.default_sink_kwargs, **(sink_kwargs or {})}
+        scan_opts = {**self.default_scan_kwargs, **(scan_kwargs or {})}
+
+        ok = False
+        try:
+            self._materialise(lf, checkpoint_path, sink_opts, streaming)
+            ok = True
+        finally:
+            with self._cond:
+                if not ok:
+                    try:
+                        self._owned_dirs.remove(checkpoint_dir)
+                    except ValueError:
+                        pass
+                self._active_ops -= 1
+                self._cond.notify_all()
+            if not ok:
+                shutil.rmtree(checkpoint_dir, ignore_errors=True)
+
+        return pl.scan_parquet(checkpoint_path, **scan_opts)
+
+    # -- internals -------------------------------------------------------------
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("CheckpointSession is closed")
+
+    def _new_checkpoint_dir(self, name: str | None) -> Path:
+        # Must be called under self._cond / self._lock.
+        slug = _normalise_name(name) if name is not None else uuid.uuid4().hex
+        path = self.root_dir / slug
+        try:
+            path.mkdir(parents=True, exist_ok=False)
+        except FileExistsError:
+            msg = f"Checkpoint directory {slug!r} already exists"
+            if name is not None and slug != name:
+                msg += f" (normalised from {name!r})"
+            raise FileExistsError(msg) from None
+        self._owned_dirs.append(path)
+        return path
+
+    @staticmethod
+    def _materialise(
+        lf: pl.LazyFrame,
+        path: Path,
+        sink_opts: dict[str, Any],
+        streaming: bool,
+    ) -> None:
+        mode = "streaming" if streaming else "collect"
+        logger.debug("Materialising checkpoint to %s (%s)", path, mode)
+        t0 = time.perf_counter()
+
+        if streaming:
+            try:
+                lf.sink_parquet(path, **sink_opts)
+            except Exception as exc:
+                raise RuntimeError(
+                    f"Checkpoint materialisation failed: {exc}\n\n"
+                    "If the streaming engine does not support this query "
+                    "plan, retry with streaming=False."
+                ) from exc
+        else:
+            lf.collect().write_parquet(path, **sink_opts)
+
+        elapsed = time.perf_counter() - t0
+        size_mb = path.stat().st_size / 1_048_576
+        logger.debug("Checkpoint written in %.2fs (%.1f MB)", elapsed, size_mb)
+
+
+# -- standalone function --------------------------------------------------------
+
+
+def checkpoint(
+    lf: pl.LazyFrame,
+    *,
+    session: CheckpointSession | None = None,
+    name: str | None = None,
+    sink_kwargs: dict[str, Any] | None = None,
+    scan_kwargs: dict[str, Any] | None = None,
+    streaming: bool = True,
+) -> pl.LazyFrame:
+    """Materialise LazyFrame to a parquet checkpoint."""
+    sess = session if session is not None else _get_default_session()
+    return sess.checkpoint(
+        lf,
+        name=name,
+        sink_kwargs=sink_kwargs,
+        scan_kwargs=scan_kwargs,
+        streaming=streaming,
+    )
+
+
+# -- default session -----------------------------------------------------------
+
+
+def _get_default_session() -> CheckpointSession:
+    """Lazily create a process-wide default session."""
+    global _default_session
+    with _default_session_lock:
+        if _default_session is None or _default_session._closed:
+            _default_session = CheckpointSession()
+        return _default_session
+
+
+# -- utilities -----------------------------------------------------------------
+
+
+def _normalise_name(name: str) -> str:
+    out = []
+    for ch in name:
+        if ch.isalnum() or ch in {"-", "_", "."}:
+            out.append(ch)
+        else:
+            out.append("_")
+    normalised = "".join(out).strip("._")
+    if not normalised:
+        raise ValueError(f"Checkpoint name {name!r} normalised to an empty string")
+    return normalised

polars_checkpoint-0.1.0/src/polars_checkpoint.egg-info/PKG-INFO

@@ -0,0 +1,88 @@
+Metadata-Version: 2.4
+Name: polars-checkpoint
+Version: 0.1.0
+Summary: Add your description here
+Requires-Python: >=3.11.1
+Description-Content-Type: text/markdown
+Requires-Dist: polars>=1.39.3
+
+# polars-checkpoint
+
+Materialise Polars LazyFrames to parquet files and scan them back lazily. Defaults to the streaming engine for sink/scan. Useful for managing & reducing memory pressure due to expensive intermediate results in complex multi-step transforms.
+
+## Installation
+
+Requires `polars` and Python 3.10+.
+
+## Quick Start
+
+```python
+import polars as pl
+from checkpoint import checkpoint
+
+lf = pl.LazyFrame({"x": range(1_000_000)}).with_columns(y=pl.col("x") * 2)
+
+# Materialises to a temp parquet file; returns a lazy re-scan
+lf = checkpoint(lf)
+
+lf.filter(pl.col("y") > 100).collect()
+```
+
+A process-wide default session manages the temp directory and cleans it up at exit.
+
+## Session API
+
+For explicit control over storage location and lifecycle, use `CheckpointSession`:
+
+```python
+from checkpoint import CheckpointSession
+
+# As a context manager — cleans up on exit from the block
+with CheckpointSession(root_dir="./my_checkpoints") as sess:
+    lf = pl.scan_csv("big.csv")
+    lf = sess.checkpoint(lf, name="after-parse")
+    lf = lf.filter(pl.col("status") == "active")
+    lf = sess.checkpoint(lf, name="filtered")
+```
+
+```python
+# Without a context manager — cleans up at GC or interpreter shutdown,
+# or when you call close() explicitly
+sess = CheckpointSession(root_dir="./my_checkpoints")
+lf = sess.checkpoint(pl.scan_csv("big.csv"), name="raw")
+reloaded = sess["raw"]
+print(sess.summary())
+
+sess.close()  # optional; triggers early cleanup
+```
+
+### `CheckpointSession` constructor
+
+| Parameter | Default | Description |
+|---|---|---|
+| `root_dir` | `None` (auto temp dir) | Parent directory for checkpoint folders. |
+| `cleanup` | `True` | Delete checkpoint files on close / GC / interpreter exit. |
+| `default_sink_kwargs` | `{"compression": "zstd"}` | Defaults passed to `sink_parquet` / `write_parquet`. |
+| `default_scan_kwargs` | `{}` | Defaults passed to `scan_parquet`. |
+
+### Key methods & features
+
+- **`checkpoint(lf, *, name=None, streaming=True, ...)`** — Materialise a LazyFrame to parquet. Auto-generates a name if none given. Falls back to `collect().write_parquet()` when `streaming=False`.
+- **`session[name]`** — Retrieve a checkpoint as a `LazyFrame`.
+- **`name in session`** — Check existence.
+- **`len(session)`** / **`iter(session)`** — Count / list checkpoints.
+- **`summary()`** — Returns a Polars DataFrame with name, size (MB), and path of each checkpoint.
+- **`close(timeout=None)`** — Waits for in-flight writes, then cleans up. Also usable as a context manager.
+
+## Thread Safety
+
+Sessions are internally locked. Concurrent `checkpoint()` calls from multiple threads are safe; `close()` waits for all in-flight materialisations before removing files.
+
+## Cleanup Behaviour
+
+| Scenario | `cleanup=True` (default) | `cleanup=False` |
+|---|---|---|
+| `close()` / `__exit__` | Files deleted | Files retained |
+| GC / interpreter shutdown | Files deleted (via `weakref.finalize`) | Files retained |
+
+When `root_dir` is auto-generated, the entire temp directory is removed. When user-supplied, only the individual checkpoint subdirectories created by the session are removed.

polars_checkpoint-0.1.0/src/polars_checkpoint.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+pyproject.toml
+src/polars_checkpoint/__init__.py
+src/polars_checkpoint/polars_checkpoint.py
+src/polars_checkpoint.egg-info/PKG-INFO
+src/polars_checkpoint.egg-info/SOURCES.txt
+src/polars_checkpoint.egg-info/dependency_links.txt
+src/polars_checkpoint.egg-info/requires.txt
+src/polars_checkpoint.egg-info/top_level.txt

polars_checkpoint-0.1.0/src/polars_checkpoint.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

polars_checkpoint-0.1.0/src/polars_checkpoint.egg-info/requires.txt

@@ -0,0 +1 @@
+polars>=1.39.3

polars_checkpoint-0.1.0/src/polars_checkpoint.egg-info/top_level.txt

@@ -0,0 +1 @@
+polars_checkpoint