shackleton 2.0.1__tar.gz → 2.1.0__tar.gz

shackleton-2.1.0/PERF.md

@@ -0,0 +1,41 @@
+# Performance
+
+## Running Benchmarks
+
+```bash
+uv run bench.py
+uv run bench.py --rows <n>
+```
+
+Rows default to `DEFAULT_ROWS` in `bench.py`; columns are an int id, a float and two strings.
+
+## Metrics
+
+All measured in rows/second.
+
+| Operation | Baseline (`DEFAULT_ROWS`, polars 2.0) |
+| --- | --: |
+| Single-file extend | ~50M rows/s |
+| Extend into stored rows (id_col merge) | ~17M rows/s |
+| Partitioned write (`N_CATEGORIES` partitions) | ~9.4M rows/s |
+| Hash-partitioned write (`N_BUCKETS` buckets) | ~5.5M rows/s |
+| Extend with `max_records=MAX_RECORDS` | ~11M rows/s |
+| IPC extend | ~30M rows/s |
+| Full read (collect) | ~127M rows/s |
+| Lazy read + filter | ~158M rows/s |
+| Single partition read | ~19M rows/s |
+| IPC full read | ~122M rows/s |
+| Replace records (`N_UPSERTS` upserts) | ~248k rows/s |
+
+## Known Bottlenecks
+
+- `extend` into a stored file and `replace_records` read-merge-rewrite the partition, so their cost scales with the stored partition, not the batch: an append-heavy store partitions so that each partition stays small or is written once
+- With `max_records` unset, a merge with stored data streams through polars' sink; with `max_records > 0`, `replace_records` and `compact` collect the whole partition in memory
+- Hash-partitioned writes scale sub-linearly with bucket count due to per-partition file overhead
+- Partition reads are limited by parquet file open latency when partitions are small
+
+## What to Watch
+
+- Polars version upgrades can significantly affect parquet I/O throughput
+- Row group size interacts with partition count — too many small files hurts read performance
+- The GIL-releasing path through polars I/O makes thread-pool scaling important for parallel reads

{shackleton-2.0.1 → shackleton-2.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: shackleton
-Version: 2.0.1
+Version: 2.1.0
 Summary: Polars-based persistent tabular store backed by Parquet or Arrow IPC files
 Project-URL: Homepage, https://github.com/endremborza/shackleton
 Author-email: Endre Márk Borza <endremborza@gmail.com>
@@ -131,7 +131,7 @@ repo.get_full_df().shape  # (4, 2) — event_id 2 kept once
 
 ## Sorted merge with `id_col`
 
-When `id_col` is set, each extend keeps the file sorted by it (when `max_records == 0`) and `replace_records` does an upsert: matching IDs are updated, new IDs appended. `id_col` is a column or a list of columns, sorted in that order. On equal keys stored rows come before incoming ones, and within one `replace_records` batch the last row for a key wins.
+When `id_col` is set, each extend keeps the file sorted by it (when `max_records == 0`) and `replace_records` does an upsert: matching IDs are updated, new IDs appended. `id_col` is a column or a list of columns, sorted in that order. On equal keys stored rows come before incoming ones, and within one `replace_records` batch the last row for a key wins. Id columns must not be null: a write with a null id raises `NullIdError`.
 
 ```python
 repo = TableRepo(Path("/tmp/prices"), id_col="ticker")
@@ -202,8 +202,7 @@ with repo.get_replacing_record_writer(batch_size=10_000) as writer:
 
 ## Schema evolution
 
-Columns added in later writes fill missing values with `null` in earlier files.
-No migration step needed.
+Columns added in later writes fill missing values with `null` in earlier files. No migration step needed. The `id_col` merges match columns by name, so a batch may also carry the stored columns in another order.
 
 ```python
 repo = TableRepo(Path("/tmp/evolving"))
@@ -334,22 +333,4 @@ result = pl.DataFrame(repo.map_partitions(summarise, workers=0))
 
 ## Performance
 
-Benchmarked on a single core, NVMe SSD, 1 million rows (4 columns: i64, f64, str, str).
-Run `uv run bench.py` to reproduce.
-
-| Scenario | Throughput |
-|---|---|
-| extend — single file | ~33M rows/s |
-| extend — sorted merge (`id_col`) | ~49M rows/s |
-| extend — 26 column partitions | ~11M rows/s |
-| extend — 64 hash partitions | ~7M rows/s |
-| extend — `max_records=10k` | ~7M rows/s |
-| extend — IPC format | ~23M rows/s |
-| read full (collect) | ~131M rows/s |
-| read full (lazy + filter) | ~170M rows/s |
-| read single partition (1 of 26) | ~18M rows/s |
-| read full — IPC format | ~102M rows/s |
-| replace_records (10k upserts) | ~107k rows/s |
-
-`replace_records` does a full read-merge-rewrite, so throughput scales with
-existing file size, not batch size. Use partitioned repos to bound the merge scope.
+Throughput per scenario and the known bottlenecks are in [PERF.md](PERF.md); `uv run bench.py` reproduces it.

{shackleton-2.0.1 → shackleton-2.1.0}/README.md

@@ -113,7 +113,7 @@ repo.get_full_df().shape  # (4, 2) — event_id 2 kept once
 
 ## Sorted merge with `id_col`
 
-When `id_col` is set, each extend keeps the file sorted by it (when `max_records == 0`) and `replace_records` does an upsert: matching IDs are updated, new IDs appended. `id_col` is a column or a list of columns, sorted in that order. On equal keys stored rows come before incoming ones, and within one `replace_records` batch the last row for a key wins.
+When `id_col` is set, each extend keeps the file sorted by it (when `max_records == 0`) and `replace_records` does an upsert: matching IDs are updated, new IDs appended. `id_col` is a column or a list of columns, sorted in that order. On equal keys stored rows come before incoming ones, and within one `replace_records` batch the last row for a key wins. Id columns must not be null: a write with a null id raises `NullIdError`.
 
 ```python
 repo = TableRepo(Path("/tmp/prices"), id_col="ticker")
@@ -184,8 +184,7 @@ with repo.get_replacing_record_writer(batch_size=10_000) as writer:
 
 ## Schema evolution
 
-Columns added in later writes fill missing values with `null` in earlier files.
-No migration step needed.
+Columns added in later writes fill missing values with `null` in earlier files. No migration step needed. The `id_col` merges match columns by name, so a batch may also carry the stored columns in another order.
 
 ```python
 repo = TableRepo(Path("/tmp/evolving"))
@@ -316,22 +315,4 @@ result = pl.DataFrame(repo.map_partitions(summarise, workers=0))
 
 ## Performance
 
-Benchmarked on a single core, NVMe SSD, 1 million rows (4 columns: i64, f64, str, str).
-Run `uv run bench.py` to reproduce.
-
-| Scenario | Throughput |
-|---|---|
-| extend — single file | ~33M rows/s |
-| extend — sorted merge (`id_col`) | ~49M rows/s |
-| extend — 26 column partitions | ~11M rows/s |
-| extend — 64 hash partitions | ~7M rows/s |
-| extend — `max_records=10k` | ~7M rows/s |
-| extend — IPC format | ~23M rows/s |
-| read full (collect) | ~131M rows/s |
-| read full (lazy + filter) | ~170M rows/s |
-| read single partition (1 of 26) | ~18M rows/s |
-| read full — IPC format | ~102M rows/s |
-| replace_records (10k upserts) | ~107k rows/s |
-
-`replace_records` does a full read-merge-rewrite, so throughput scales with
-existing file size, not batch size. Use partitioned repos to bound the merge scope.
+Throughput per scenario and the known bottlenecks are in [PERF.md](PERF.md); `uv run bench.py` reproduces it.

{shackleton-2.0.1 → shackleton-2.1.0}/bench.py

@@ -14,6 +14,12 @@ import polars as pl
 
 from shackleton import HashPartitioner, TableRepo
 
+DEFAULT_ROWS = 1_000_000
+N_CATEGORIES = 26
+N_BUCKETS = 64
+MAX_RECORDS = 10_000
+N_UPSERTS = 10_000
+
 # ---------------------------------------------------------------------------
 # Helpers
 # ---------------------------------------------------------------------------
@@ -23,7 +29,7 @@ def make_df(n: int) -> pl.DataFrame:
         {
             "id": pl.arange(0, n, eager=True),
             "value": pl.Series([float(i) * 1.1 for i in range(n)]),
-            "category": pl.Series([chr(65 + i % 26) for i in range(n)]),
+            "category": pl.Series([chr(65 + i % N_CATEGORIES) for i in range(n)]),
             "tag": pl.Series([f"tag-{i % 100}" for i in range(n)]),
         }
     )
@@ -49,22 +55,42 @@ def bench_extend_single(df: pl.DataFrame, root: Path) -> None:
 
 def bench_extend_with_id(df: pl.DataFrame, root: Path) -> None:
     repo = TableRepo(root / "extend_id", id_col="id")
-    bench("extend + sorted merge (id_col)", lambda: repo.extend(df), scale=len(df))
+    repo.extend(df.with_columns(pl.col("id") * 2))
+    interleaved = df.with_columns(pl.col("id") * 2 + 1)
+    bench(
+        "extend into stored rows (id_col merge)",
+        lambda: repo.extend(interleaved),
+        scale=len(df),
+    )
 
 
 def bench_extend_partitioned(df: pl.DataFrame, root: Path) -> None:
     repo = TableRepo(root / "extend_part", partition_cols=["category"])
-    bench("extend (26 column partitions)", lambda: repo.extend(df), scale=len(df))
+    bench(
+        f"extend ({N_CATEGORIES} column partitions)",
+        lambda: repo.extend(df),
+        scale=len(df),
+    )
 
 
 def bench_extend_hash(df: pl.DataFrame, root: Path) -> None:
-    repo = TableRepo(root / "extend_hash", partition_cols=HashPartitioner("tag", 64))
-    bench("extend (64 hash partitions)", lambda: repo.extend(df), scale=len(df))
+    repo = TableRepo(
+        root / "extend_hash", partition_cols=HashPartitioner("tag", N_BUCKETS)
+    )
+    bench(
+        f"extend ({N_BUCKETS} hash partitions)",
+        lambda: repo.extend(df),
+        scale=len(df),
+    )
 
 
 def bench_extend_max_records(df: pl.DataFrame, root: Path) -> None:
-    repo = TableRepo(root / "extend_max", max_records=10_000)
-    bench("extend (max_records=10k files)", lambda: repo.extend(df), scale=len(df))
+    repo = TableRepo(root / "extend_max", max_records=MAX_RECORDS)
+    bench(
+        f"extend (max_records={MAX_RECORDS:,})",
+        lambda: repo.extend(df),
+        scale=len(df),
+    )
 
 
 def bench_read_full(df: pl.DataFrame, root: Path) -> None:
@@ -87,19 +113,23 @@ def bench_read_partition(df: pl.DataFrame, root: Path) -> None:
     repo = TableRepo(root / "read_part", partition_cols=["category"])
     repo.extend(df)
     bench(
-        "get_partition_df (1 of 26 partitions)",
+        f"get_partition_df (1 of {N_CATEGORIES} partitions)",
         lambda: repo.get_partition_df({"category": "A"}),
-        scale=len(df) // 26,
+        scale=len(df) // N_CATEGORIES,
     )
 
 
 def bench_replace_records(df: pl.DataFrame, root: Path) -> None:
     repo = TableRepo(root / "replace_rec", id_col="id")
     repo.extend(df)
-    updates = df.sample(n=min(10_000, len(df) // 10), seed=42).with_columns(
+    updates = df.sample(n=min(N_UPSERTS, len(df) // 10), seed=42).with_columns(
         pl.col("value") * 2
     )
-    bench("replace_records (10k updates)", lambda: repo.replace_records(updates), scale=len(updates))
+    bench(
+        f"replace_records ({len(updates):,} updates)",
+        lambda: repo.replace_records(updates),
+        scale=len(updates),
+    )
 
 
 def bench_ipc_extend(df: pl.DataFrame, root: Path) -> None:
@@ -119,7 +149,7 @@ def bench_ipc_read(df: pl.DataFrame, root: Path) -> None:
 
 def main() -> None:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--rows", type=int, default=1_000_000)
+    parser.add_argument("--rows", type=int, default=DEFAULT_ROWS)
     args = parser.parse_args()
 
     n = args.rows

{shackleton-2.0.1 → shackleton-2.1.0}/docs/release_notes.md

@@ -1,3 +1,11 @@
+## v2.1.0
+
+- `id_col` merges now match columns by name instead of position, so a batch whose columns arrive reordered, or that adds or drops columns, merges correctly instead of merging wrong or failing.
+- Writing a null value into an `id_col` column now raises the new exported `NullIdError` instead of silently misordering the merge.
+- A batch that lands entirely in one partition is no longer re-sorted on write, fixing a case where it lost its already-sorted order.
+
+---
+
 ## v2.0.1
 
 - Fixed composite-key merges to be order-exact: stored rows now reliably come before incoming ones on equal keys, where polars' streaming engine could previously reorder ties.

{shackleton-2.0.1 → shackleton-2.1.0}/shackleton/__init__.py

@@ -4,10 +4,11 @@ from .core import (  # noqa: F401
     DfBatchWriter,
     HashDriftError,
     HashPartitioner,
+    NullIdError,
     RecordWriter,
     TableRepo,
 )
 
 TableShack = TableRepo  # backward-compat alias
 
-__version__ = "2.0.1"
+__version__ = "2.1.0"

{shackleton-2.0.1 → shackleton-2.1.0}/shackleton/core.py

@@ -22,6 +22,10 @@ class HashDriftError(Exception):
     """Stored rows sit in buckets the installed polars no longer hashes them to."""
 
 
+class NullIdError(ValueError):
+    """A written batch has nulls in an ``id_col`` column."""
+
+
 _PATH_LOCKS: dict[str, threading.Lock] = {}
 _LOCKS_LOCK = threading.Lock()
 
@@ -93,6 +97,7 @@ class TableRepo:
             ``max_records == 0`` (no effect on the per-file append path). On
             equal keys, stored rows precede incoming ones and a batch keeps
             its own order: ``replace_records`` keeps the batch's last row.
+            Must not be null: writes raise ``NullIdError``.
         dedup_cols: Columns forming a composite dedup key. When set, ``extend``
             drops duplicates after merging with existing data, keeping the
             stored row when ``id_col`` is set. Only effective when
@@ -351,11 +356,20 @@ class TableRepo:
     def _sort(self, lf: pl.LazyFrame) -> pl.LazyFrame:
         return lf.sort(self.id_cols, maintain_order=True)
 
+    def _merge_sorted(self, old: pl.LazyFrame, new: pl.LazyFrame) -> pl.LazyFrame:
+        """``old`` and the sorted ``new`` merged on ``id_cols``. ``merge_sorted``
+        matches columns by position, so each side first takes the diagonal union's
+        schema, the one the diagonal concat path writes."""
+        new = self._sort(new)
+        return pl.concat([old, new.clear()], how="diagonal").merge_sorted(
+            pl.concat([old.clear(), new], how="diagonal"),
+            key=self.id_cols,
+            maintain_order=True,
+        )
+
     def _merge_extend(self, new: pl.LazyFrame, old: pl.LazyFrame) -> pl.LazyFrame:
         if self.id_cols:
-            out = old.merge_sorted(
-                self._sort(new), key=self.id_cols, maintain_order=True
-            )
+            out = self._merge_sorted(old, new)
         else:
             out = pl.concat([old, new], how="diagonal")
         if self.dedup_cols:
@@ -371,13 +385,7 @@ class TableRepo:
     def _merge_replace(
         self, new: pl.LazyFrame, old: pl.LazyFrame | None
     ) -> pl.LazyFrame:
-        merged = (
-            self._sort(new)
-            if old is None
-            else old.merge_sorted(
-                self._sort(new), key=self.id_cols, maintain_order=True
-            )
-        )
+        merged = self._sort(new) if old is None else self._merge_sorted(old, new)
         # sorted, so a key's rows are adjacent: keep the last row of each run
         last_of_run = pl.any_horizontal(
             pl.col(c).ne_missing(pl.col(c).shift(-1)) for c in self.id_cols
@@ -385,6 +393,9 @@ class TableRepo:
         return merged.filter(last_of_run)
 
     def _for_each_partition(self, df: pl.DataFrame, handler: Callable) -> None:
+        # polars' streaming merge_sorted misorders null keys on a single key column
+        if null_ids := [c for c in self.id_cols if df[c].null_count()]:
+            raise NullIdError(f"{self.root_path}: nulls in id column(s) {null_ids}")
         keys = self._part_keys
         if not keys:
             handler(df, self.root_path)
@@ -395,8 +406,11 @@ class TableRepo:
             self._verify_buckets(hasher)
             df = hasher.add_to_df(df)
 
-        for group_vals, group_df in df.group_by(keys):
-            group_vals = group_vals if isinstance(group_vals, tuple) else (group_vals,)
+        # a frame of one partition goes through whole: group_by would gather a copy
+        # and drop its sorted flags, so the write would sort it again
+        single = df.select(pl.struct(keys).n_unique()).item() == 1
+        groups = [(df.select(keys).row(0), df)] if single else df.group_by(keys)
+        for group_vals, group_df in groups:
             hive_parts = [f"{k}={v}" for k, v in zip(keys, group_vals)]
             part_dir = Path(self.root_path, *hive_parts)
             if hasher is not None:

{shackleton-2.0.1 → shackleton-2.1.0}/shackleton/tests/test_core.py

@@ -5,7 +5,7 @@ from base64 import b16encode
 import polars as pl
 import pytest
 
-from shackleton import HashDriftError, HashPartitioner, TableRepo
+from shackleton import HashDriftError, HashPartitioner, NullIdError, TableRepo
 
 
 def make_df(n: int, seed: int = 0) -> pl.DataFrame:
@@ -503,6 +503,25 @@ def test_schema_mismatch_extend(tmp_path):
     assert "C" in full.columns
 
 
+def test_id_col_merges_match_columns_by_name(tmp_path):
+    repo = TableRepo(tmp_path / "data", id_col="id")
+    repo.extend(pl.DataFrame({"id": [1, 3], "a": ["x", "y"], "b": [1.0, 2.0]}))
+    repo.extend(pl.DataFrame({"id": [2], "b": [9.0], "a": ["z"], "c": [7]}))
+    repo.replace_records(pl.DataFrame({"a": ["w"], "id": [3], "b": [None]}))
+    assert repo.get_full_df().to_dicts() == [
+        {"id": 1, "a": "x", "b": 1.0, "c": None},
+        {"id": 2, "a": "z", "b": 9.0, "c": 7},
+        {"id": 3, "a": "w", "b": None, "c": None},
+    ]
+
+
+@pytest.mark.parametrize("method", ["extend", "replace_records"])
+def test_null_ids_raise(tmp_path, method):
+    repo = TableRepo(tmp_path / "data", id_col=["k", "id"])
+    with pytest.raises(NullIdError):
+        getattr(repo, method)(pl.DataFrame({"k": [1, 1], "id": [2, None]}))
+
+
 # ---------------------------------------------------------------------------
 # Edge cases
 # ---------------------------------------------------------------------------

shackleton-2.0.1/PERF.md

@@ -1,37 +0,0 @@
-# Performance
-
-## Running Benchmarks
-
-```bash
-uv run bench.py
-uv run bench.py --rows 5000000
-```
-
-Default: 1M rows, 5 columns (int, float, str, date, bool).
-
-## Metrics
-
-All measured in rows/second.
-
-| Operation | Baseline (1M rows) |
-|-----------|-------------------|
-| Single-file extend | ~33M rows/s |
-| Extend with sorted merge (id_col) | ~49M rows/s |
-| Partitioned write (26 partitions) | ~11M rows/s |
-| Hash-partitioned write (64 buckets) | ~7M rows/s |
-| Full read (collect) | ~131M rows/s |
-| Lazy read + filter | ~170M rows/s |
-| Single partition read | ~18M rows/s |
-| Replace records (10k upserts) | ~107k rows/s |
-
-## Known Bottlenecks
-
-- `replace_records` requires a full merge — O(n) in total rows, not O(k) in replacement size
-- Hash-partitioned writes scale sub-linearly with bucket count due to per-partition file overhead
-- Partition reads are limited by parquet file open latency when partitions are small
-
-## What to Watch
-
-- Polars version upgrades can significantly affect parquet I/O throughput
-- Row group size interacts with partition count — too many small files hurts read performance
-- The GIL-releasing path through polars I/O makes thread-pool scaling important for parallel reads