shackleton 1.1.0__tar.gz → 1.2.0__tar.gz

{shackleton-1.1.0 → shackleton-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shackleton
-Version: 1.1.0
+Version: 1.2.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>
@@ -110,10 +110,30 @@ df = repo.get_full_df()
 
 ---
 
+## Deduplication with `dedup_cols`
+
+When `dedup_cols` is set, each `extend` drops duplicate rows on those columns
+after merging with existing data — only the first occurrence is kept. Combine
+with `id_col` for sorted-merge dedup; use alone for unordered dedup.
+Not supported with `max_records > 0` (raises at construction time).
+
+```python
+repo = TableRepo(Path("/tmp/events"), dedup_cols=["event_id"])
+
+repo.extend(pl.DataFrame({"event_id": [1, 2, 3], "val": [10, 20, 30]}))
+# Re-ingest with overlapping IDs — existing rows are kept, new ones appended
+repo.extend(pl.DataFrame({"event_id": [2, 4], "val": [99, 40]}))
+
+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 that column and
-`replace_records` does an upsert: matching IDs are updated, new IDs appended.
+When `id_col` is set, each extend keeps the file sorted by that column
+(when `max_records == 0`) and `replace_records` does an upsert: matching IDs
+are updated, new IDs appended.
 
 ```python
 repo = TableRepo(Path("/tmp/prices"), id_col="ticker")
@@ -198,6 +218,41 @@ df = repo.get_full_df()
 
 ---
 
+## Partition-level overwrite
+
+`replace_partition` atomically clears and rewrites only the partitions covered
+by the DataFrame. `purge_partition` deletes matching files without rewriting.
+
+```python
+repo = TableRepo(Path("/tmp/events"), partition_cols=["date"])
+repo.extend(pl.DataFrame({
+    "date": ["2024-01-01", "2024-01-01", "2024-01-02"],
+    "val":  [1, 2, 3],
+}))
+
+# Replace just the 2024-01-01 partition with corrected data
+repo.replace_partition(pl.DataFrame({
+    "date": ["2024-01-01"],
+    "val":  [99],
+}))
+
+# Or wipe it without rewriting
+repo.purge_partition({"date": "2024-01-01"})
+```
+
+---
+
+## Compression
+
+Control the codec and, for Parquet, the compression level.
+
+```python
+repo = TableRepo(Path("/tmp/data"), compression="zstd", compression_level=3)
+repo_ipc = TableRepo(Path("/tmp/fast"), ipc=True, compression="lz4")
+```
+
+---
+
 ## Arrow IPC format
 
 Slightly faster reads for local workloads that don't need Parquet compatibility.

{shackleton-1.1.0 → shackleton-1.2.0}/README.md

@@ -92,10 +92,30 @@ df = repo.get_full_df()
 
 ---
 
+## Deduplication with `dedup_cols`
+
+When `dedup_cols` is set, each `extend` drops duplicate rows on those columns
+after merging with existing data — only the first occurrence is kept. Combine
+with `id_col` for sorted-merge dedup; use alone for unordered dedup.
+Not supported with `max_records > 0` (raises at construction time).
+
+```python
+repo = TableRepo(Path("/tmp/events"), dedup_cols=["event_id"])
+
+repo.extend(pl.DataFrame({"event_id": [1, 2, 3], "val": [10, 20, 30]}))
+# Re-ingest with overlapping IDs — existing rows are kept, new ones appended
+repo.extend(pl.DataFrame({"event_id": [2, 4], "val": [99, 40]}))
+
+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 that column and
-`replace_records` does an upsert: matching IDs are updated, new IDs appended.
+When `id_col` is set, each extend keeps the file sorted by that column
+(when `max_records == 0`) and `replace_records` does an upsert: matching IDs
+are updated, new IDs appended.
 
 ```python
 repo = TableRepo(Path("/tmp/prices"), id_col="ticker")
@@ -180,6 +200,41 @@ df = repo.get_full_df()
 
 ---
 
+## Partition-level overwrite
+
+`replace_partition` atomically clears and rewrites only the partitions covered
+by the DataFrame. `purge_partition` deletes matching files without rewriting.
+
+```python
+repo = TableRepo(Path("/tmp/events"), partition_cols=["date"])
+repo.extend(pl.DataFrame({
+    "date": ["2024-01-01", "2024-01-01", "2024-01-02"],
+    "val":  [1, 2, 3],
+}))
+
+# Replace just the 2024-01-01 partition with corrected data
+repo.replace_partition(pl.DataFrame({
+    "date": ["2024-01-01"],
+    "val":  [99],
+}))
+
+# Or wipe it without rewriting
+repo.purge_partition({"date": "2024-01-01"})
+```
+
+---
+
+## Compression
+
+Control the codec and, for Parquet, the compression level.
+
+```python
+repo = TableRepo(Path("/tmp/data"), compression="zstd", compression_level=3)
+repo_ipc = TableRepo(Path("/tmp/fast"), ipc=True, compression="lz4")
+```
+
+---
+
 ## Arrow IPC format
 
 Slightly faster reads for local workloads that don't need Parquet compatibility.

shackleton-1.2.0/coverage/report.md

@@ -0,0 +1,6 @@
+| Name                           |    Stmts |     Miss |    Cover |
+|------------------------------- | -------: | -------: | -------: |
+| shackleton/\_\_init\_\_.py     |        3 |        0 |     100% |
+| shackleton/core.py             |      289 |        0 |     100% |
+| shackleton/tests/test\_core.py |      349 |        0 |     100% |
+| **TOTAL**                      |  **641** |    **0** | **100%** |

{shackleton-1.1.0 → shackleton-1.2.0}/docs/release_notes.md

@@ -1,3 +1,9 @@
+## v1.2.0
+
+deduped
+
+---
+
 ## v1.1.0
 
 realign, simplify

{shackleton-1.1.0 → shackleton-1.2.0}/shackleton/__init__.py

@@ -4,4 +4,4 @@ from .core import DfBatchWriter, HashPartitioner, RecordWriter, TableRepo  # noq
 
 TableShack = TableRepo  # backward-compat alias
 
-__version__ = "1.1.0"
+__version__ = "1.2.0"

{shackleton-1.1.0 → shackleton-1.2.0}/shackleton/core.py

@@ -3,7 +3,7 @@ import fcntl
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass, field
-from functools import cached_property, partial
+from functools import cached_property
 from pathlib import Path
 from typing import Any, Callable, Iterable, TypeVar
 
@@ -76,7 +76,12 @@ class TableRepo:
     Args:
         root_path: Root directory for all data files.
         id_col: Column used as a unique row identifier. Required for
-            ``replace_records``; enables sorted merge on extend.
+            ``replace_records``; enables sorted merge on ``extend`` when
+            ``max_records == 0`` (no effect on the per-file append path).
+        dedup_cols: Columns forming a composite dedup key. When set, ``extend``
+            drops duplicates after merging with existing data. Only effective
+            when ``max_records == 0``; raises ``ValueError`` otherwise.
+            ``replace_records`` still requires a scalar ``id_col``.
         partition_cols: Column name(s) to partition by, or a
             ``HashPartitioner`` for hash-based bucketing. Partition values
             are encoded as hive-style directories (``col=val/``).
@@ -86,14 +91,22 @@ class TableRepo:
         ipc: Use Arrow IPC format instead of Parquet.
         compression: Compression codec passed to the writer
             (e.g. ``"zstd"``, ``"gzip"``). ``None`` uses the writer default.
+        compression_level: Compression level for Parquet writers. ``None`` uses
+            the codec default. Ignored for Arrow IPC.
     """
 
     root_path: Path
     id_col: str | None = None
+    dedup_cols: list[str] | None = None
     partition_cols: list[str] | HashPartitioner = field(default_factory=list)
     max_records: int = 0
     ipc: bool = False
     compression: str | None = None
+    compression_level: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.dedup_cols and self.max_records:
+            raise ValueError("dedup_cols is not supported with max_records > 0")
 
     @cached_property
     def extension(self) -> str:
@@ -119,6 +132,20 @@ class TableRepo:
         self.purge()
         self.extend(df)
 
+    def purge_partition(self, partition_values: dict[str, Any]) -> None:
+        """Delete all files in the matching partition(s)."""
+        for d in {p.parent for p in self._partition_paths(partition_values)}:
+            self._purge_dir(d)
+
+    def replace_partition(self, df: pl.DataFrame) -> None:
+        """Overwrite the partition(s) covered by ``df``, deleting existing files first.
+
+        ``df`` must contain the partition column(s).
+        """
+        if not self._part_keys:
+            raise ValueError("replace_partition requires partition_cols to be set")
+        self._for_each_partition(df, self._replace_partition_dir)
+
     def purge(self) -> None:
         """Delete all data files."""
         for p in self.paths:
@@ -265,15 +292,30 @@ class TableRepo:
 
     def _write(self, df: pl.DataFrame, path: Path) -> None:
         path.parent.mkdir(parents=True, exist_ok=True)
-        write_fn = pl.DataFrame.write_ipc if self.ipc else pl.DataFrame.write_parquet
+        kw: dict = {}
         if self.compression:
-            write_fn = partial(write_fn, compression=self.compression)
-        write_fn(df, path)
+            kw["compression"] = self.compression
+        if self.ipc:
+            df.write_ipc(path, **kw)
+        else:
+            if self.compression_level is not None:
+                kw["compression_level"] = self.compression_level
+            df.write_parquet(path, **kw)
 
     def _merge_extend(self, new: pl.LazyFrame, old: pl.LazyFrame) -> pl.LazyFrame:
         if self.id_col:
-            return old.merge_sorted(new.sort(self.id_col), key=self.id_col)
-        return pl.concat([old, new], how="diagonal")
+            out = old.merge_sorted(new.sort(self.id_col), key=self.id_col)
+        else:
+            out = pl.concat([old, new], how="diagonal")
+        if self.dedup_cols:
+            # when id_col is set, merge_sorted makes dupes adjacent; keep="first"
+            # with maintain_order=True then runs in O(n) and preserves sort
+            out = out.unique(
+                subset=self.dedup_cols,
+                keep="first" if self.id_col else "any",
+                maintain_order=bool(self.id_col),
+            )
+        return out
 
     def _merge_replace(self, new: pl.LazyFrame, old: pl.LazyFrame) -> pl.LazyFrame:
         return (
@@ -308,11 +350,24 @@ class TableRepo:
             else:
                 path = self._single_path(root)
                 if not path.exists():
-                    out = df.sort(self.id_col) if self.id_col else df
+                    out = df
+                    if self.dedup_cols:
+                        out = out.unique(subset=self.dedup_cols, maintain_order=False)
+                    if self.id_col:
+                        out = out.sort(self.id_col)
                 else:
                     out = self._merge_extend(df.lazy(), self._lazy_read(path)).collect()
                 self._write(out, path)
 
+    def _purge_dir(self, root: Path) -> None:
+        with _dir_lock(str(root)):
+            for p in root.glob("*" + self.extension):
+                p.unlink()
+
+    def _replace_partition_dir(self, df: pl.DataFrame, root: Path) -> None:
+        self._purge_dir(root)
+        self._extend_dir(df, root)
+
     def _replace_dir(self, df: pl.DataFrame, root: Path) -> None:
         with _dir_lock(str(root)):
             existing = sorted(root.glob("*" + self.extension))

{shackleton-1.1.0 → shackleton-1.2.0}/shackleton/tests/test_core.py

@@ -7,21 +7,16 @@ import pytest
 
 from shackleton import HashPartitioner, TableRepo
 
-# ---------------------------------------------------------------------------
-# Fixtures / helpers
-# ---------------------------------------------------------------------------
-
 
-def make_df(n: int, seed: int = 0, with_id: bool = False) -> pl.DataFrame:
+def make_df(n: int, seed: int = 0) -> pl.DataFrame:
     rng = random.Random(seed)
-    data: dict = {
-        "A": [rng.random() for _ in range(n)],
-        "B": [rng.randint(0, 9) for _ in range(n)],
-        "C": [rng.choice(list("xyz")) for _ in range(n)],
-    }
-    if with_id:
-        data["id"] = list(range(n))
-    return pl.DataFrame(data)
+    return pl.DataFrame(
+        {
+            "A": [rng.random() for _ in range(n)],
+            "B": [rng.randint(0, 9) for _ in range(n)],
+            "C": [rng.choice(list("xyz")) for _ in range(n)],
+        }
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -31,10 +26,8 @@ def make_df(n: int, seed: int = 0, with_id: bool = False) -> pl.DataFrame:
 
 def test_extend_single_file(tmp_path):
     repo = TableRepo(tmp_path / "data")
-    df1 = make_df(10, seed=1)
-    df2 = make_df(5, seed=2)
-    repo.extend(df1)
-    repo.extend(df2)
+    repo.extend(make_df(10, seed=1))
+    repo.extend(make_df(5, seed=2))
     assert repo.get_full_df().shape[0] == 15
     assert repo.n_files == 1
 
@@ -42,10 +35,8 @@ def test_extend_single_file(tmp_path):
 def test_replace_all(tmp_path):
     repo = TableRepo(tmp_path / "data")
     repo.extend(make_df(20))
-    df = make_df(7, seed=99)
-    repo.replace_all(df)
-    result = repo.get_full_df()
-    assert result.shape[0] == 7
+    repo.replace_all(make_df(7, seed=99))
+    assert repo.get_full_df().shape[0] == 7
 
 
 def test_purge(tmp_path):
@@ -90,15 +81,13 @@ def test_max_records_multiple_extends(tmp_path):
     for seed in range(4):
         repo.extend(make_df(4, seed=seed))
     assert repo.get_full_df().shape[0] == 16
-    # 16 rows at 5/file → 4 files (5, 5, 5, 1)
-    assert repo.n_files == 4
+    assert repo.n_files == 4  # 16 rows at 5/file → (5,5,5,1)
 
 
 def test_max_records_fills_last_file(tmp_path):
     repo = TableRepo(tmp_path / "data", max_records=10)
     repo.extend(make_df(7))
     repo.extend(make_df(6))
-    # 13 rows → file-0: 10, file-1: 3
     assert repo.n_files == 2
     assert repo.get_full_df().shape[0] == 13
 
@@ -108,20 +97,11 @@ def test_max_records_fills_last_file(tmp_path):
 # ---------------------------------------------------------------------------
 
 
-def test_partition_single_col(tmp_path):
-    repo = TableRepo(tmp_path / "data", partition_cols=["B"])
-    df = make_df(100)
-    repo.extend(df)
-    full = repo.get_full_df()
-    assert full.shape[0] == 100
-    assert set(full["B"].to_list()).issubset(set(range(10)))
-
-
-def test_partition_multi_col(tmp_path):
-    repo = TableRepo(tmp_path / "data", partition_cols=["B", "C"])
-    df = make_df(200)
-    repo.extend(df)
-    assert repo.get_full_df().shape[0] == 200
+@pytest.mark.parametrize("partition_cols", [["B"], ["B", "C"]])
+def test_partition_extend(tmp_path, partition_cols):
+    repo = TableRepo(tmp_path / "data", partition_cols=partition_cols)
+    repo.extend(make_df(100))
+    assert repo.get_full_df().shape[0] == 100
 
 
 def test_get_partition(tmp_path):
@@ -136,11 +116,13 @@ def test_get_partition(tmp_path):
         }
     )
     repo.extend(df)
-    part = repo.get_partition_df({"B": 4})
-    assert part.shape[0] == df.filter(pl.col("B") == 4).shape[0]
-
-    part2 = repo.get_partition_df({"C": byte_set[0]})
-    assert part2.shape[0] == df.filter(pl.col("C") == byte_set[0]).shape[0]
+    assert (
+        repo.get_partition_df({"B": 4}).shape[0] == df.filter(pl.col("B") == 4).shape[0]
+    )
+    assert (
+        repo.get_partition_df({"C": byte_set[0]}).shape[0]
+        == df.filter(pl.col("C") == byte_set[0]).shape[0]
+    )
 
 
 def test_partition_extend_twice(tmp_path):
@@ -154,8 +136,7 @@ def test_partition_extend_twice(tmp_path):
 def test_partition_replace_all(tmp_path):
     repo = TableRepo(tmp_path / "data", partition_cols=["C"])
     repo.extend(make_df(30))
-    df2 = make_df(10, seed=99)
-    repo.replace_all(df2)
+    repo.replace_all(make_df(10, seed=99))
     assert repo.get_full_df().shape[0] == 10
 
 
@@ -163,7 +144,6 @@ def test_partition_max_records(tmp_path):
     repo = TableRepo(tmp_path / "data", partition_cols=["C"], max_records=3)
     repo.extend(make_df(30))
     assert repo.get_full_df().shape[0] == 30
-    # Each partition ("x", "y", "z") should have multiple files
     assert repo.n_files > 3
 
 
@@ -172,14 +152,10 @@ def test_partition_max_records(tmp_path):
 # ---------------------------------------------------------------------------
 
 
-def test_replace_records_basic(tmp_path):
+def test_replace_records_insert_and_update(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="id")
-    df = pl.DataFrame({"id": [1, 2, 3], "val": [10, 20, 30]})
-    repo.extend(df)
-
-    updates = pl.DataFrame({"id": [2, 4], "val": [200, 400]})
-    repo.replace_records(updates)
-
+    repo.extend(pl.DataFrame({"id": [1, 2, 3], "val": [10, 20, 30]}))
+    repo.replace_records(pl.DataFrame({"id": [2, 4], "val": [200, 400]}))
     result = repo.get_full_df().sort("id")
     assert result["id"].to_list() == [1, 2, 3, 4]
     assert result["val"].to_list() == [10, 200, 30, 400]
@@ -193,25 +169,18 @@ def test_replace_records_no_id_raises(tmp_path):
 
 def test_replace_records_dedup(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="id")
-    df = pl.DataFrame({"id": [1, 2, 3], "val": [1, 2, 3]})
-    repo.extend(df)
-    # Replace with duplicate ids in the input — last wins
-    dups = pl.DataFrame({"id": [2, 2], "val": [20, 21]})
-    repo.replace_records(dups)
+    repo.extend(pl.DataFrame({"id": [1, 2, 3], "val": [1, 2, 3]}))
+    repo.replace_records(pl.DataFrame({"id": [2, 2], "val": [20, 21]}))
     result = repo.get_full_df().sort("id")
-    assert result.filter(pl.col("id") == 2)["val"][0] in (20, 21)
     assert result.shape[0] == 3
+    assert result.filter(pl.col("id") == 2)["val"][0] in (20, 21)
 
 
 def test_replace_records_with_max_records(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="id", max_records=2)
-    df = pl.DataFrame({"id": [1, 2, 3], "val": [10, 20, 30]})
-    repo.extend(df)
-    assert repo.n_files == 2  # 2 + 1
-
-    updates = pl.DataFrame({"id": [2, 3, 4], "val": [200, 300, 400]})
-    repo.replace_records(updates)
-
+    repo.extend(pl.DataFrame({"id": [1, 2, 3], "val": [10, 20, 30]}))
+    assert repo.n_files == 2  # 2+1
+    repo.replace_records(pl.DataFrame({"id": [2, 3, 4], "val": [200, 300, 400]}))
     result = repo.get_full_df().sort("id")
     assert result["id"].to_list() == [1, 2, 3, 4]
     assert result["val"].to_list() == [10, 200, 300, 400]
@@ -220,27 +189,77 @@ def test_replace_records_with_max_records(tmp_path):
 
 def test_replace_records_empty_repo(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="id")
-    df = pl.DataFrame({"id": [1, 2], "val": [1, 2]})
-    repo.replace_records(df)
+    repo.replace_records(pl.DataFrame({"id": [1, 2], "val": [1, 2]}))
     assert repo.get_full_df().shape[0] == 2
 
 
+def test_id_col_sorted_merge_on_extend(tmp_path):
+    repo = TableRepo(tmp_path / "data", id_col="id")
+    repo.extend(pl.DataFrame({"id": [1, 3], "val": [10, 30]}))
+    repo.extend(pl.DataFrame({"id": [2, 4], "val": [20, 40]}))
+    result = repo.get_full_df()
+    assert result["id"].to_list() == [1, 2, 3, 4]
+    assert result["val"].to_list() == [10, 20, 30, 40]
+
+
 # ---------------------------------------------------------------------------
-# IPC format
+# dedup_cols
 # ---------------------------------------------------------------------------
 
 
-def test_ipc_format(tmp_path):
-    repo = TableRepo(tmp_path / "data", ipc=True)
-    df = make_df(20)
-    repo.extend(df)
+def test_dedup_cols_max_records_raises(tmp_path):
+    with pytest.raises(ValueError, match="dedup_cols"):
+        TableRepo(tmp_path / "data", dedup_cols=["id"], max_records=100)
+
+
+def test_dedup_cols_initial_write(tmp_path):
+    repo = TableRepo(tmp_path / "data", dedup_cols=["id"])
+    repo.extend(pl.DataFrame({"id": [1, 1, 2], "val": [10, 99, 20]}))
+    result = repo.get_full_df()
+    assert result.shape[0] == 2
+    assert set(result["id"].to_list()) == {1, 2}
+
+
+def test_dedup_cols_extend(tmp_path):
+    repo = TableRepo(tmp_path / "data", dedup_cols=["id"])
+    repo.extend(pl.DataFrame({"id": [1, 2, 3], "val": [10, 20, 30]}))
+    repo.extend(pl.DataFrame({"id": [2, 4], "val": [99, 40]}))
+    result = repo.get_full_df().sort("id")
+    assert result.shape[0] == 4
+    assert set(result["id"].to_list()) == {1, 2, 3, 4}
+    # id=4 is new, so its value is unambiguous
+    assert result.filter(pl.col("id") == 4)["val"][0] == 40
+
+
+def test_dedup_cols_with_id_col_old_wins(tmp_path):
+    # id_col + dedup_cols: merge_sorted places old rows first, keep="first" retains them
+    repo = TableRepo(tmp_path / "data", id_col="id", dedup_cols=["id"])
+    repo.extend(pl.DataFrame({"id": [1, 2], "val": [10, 20]}))
+    repo.extend(pl.DataFrame({"id": [2, 3], "val": [99, 30]}))
+    result = repo.get_full_df().sort("id")
+    assert result.shape[0] == 3
+    assert result["id"].to_list() == [1, 2, 3]
+    assert result.filter(pl.col("id") == 2)["val"][0] == 20  # old value kept
+    assert result.filter(pl.col("id") == 3)["val"][0] == 30  # new row inserted
+
+
+# ---------------------------------------------------------------------------
+# IPC / Parquet format
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("ipc", [False, True])
+def test_format_basic(tmp_path, ipc):
+    repo = TableRepo(tmp_path / "data", ipc=ipc)
+    repo.extend(make_df(20))
     assert repo.n_files == 1
-    assert list(repo.paths)[0].suffix == ".arrow"
+    assert list(repo.paths)[0].suffix == (".arrow" if ipc else ".parquet")
     assert repo.get_full_df().shape[0] == 20
 
 
-def test_ipc_with_partitions(tmp_path):
-    repo = TableRepo(tmp_path / "data", ipc=True, partition_cols=["C"])
+@pytest.mark.parametrize("ipc", [False, True])
+def test_format_with_partitions(tmp_path, ipc):
+    repo = TableRepo(tmp_path / "data", ipc=ipc, partition_cols=["C"])
     repo.extend(make_df(30))
     assert repo.get_full_df().shape[0] == 30
 
@@ -250,8 +269,15 @@ def test_ipc_with_partitions(tmp_path):
 # ---------------------------------------------------------------------------
 
 
-def test_compression(tmp_path):
-    repo = TableRepo(tmp_path / "data", compression="zstd")
+@pytest.mark.parametrize(
+    "kw",
+    [
+        {"compression": "zstd"},
+        {"compression": "zstd", "compression_level": 1},
+    ],
+)
+def test_compression(tmp_path, kw):
+    repo = TableRepo(tmp_path / "data", **kw)
     repo.extend(make_df(50))
     assert repo.get_full_df().shape[0] == 50
 
@@ -264,12 +290,8 @@ def test_compression(tmp_path):
 def test_hash_partitioner_distributes(tmp_path):
     hp = HashPartitioner("B", num_groups=4)
     repo = TableRepo(tmp_path / "data", partition_cols=hp)
-    df = make_df(200)
-    repo.extend(df)
-
-    total = repo.get_full_df().shape[0]
-    assert total == 200
-    # Buckets: 1–4 files (collisions allowed); at least 1
+    repo.extend(make_df(200))
+    assert repo.get_full_df().shape[0] == 200
     assert 1 <= repo.n_files <= 4
 
 
@@ -277,8 +299,7 @@ def test_hash_partitioner_key_stripped(tmp_path):
     hp = HashPartitioner("B", num_groups=4)
     repo = TableRepo(tmp_path / "data", partition_cols=hp)
     repo.extend(make_df(20))
-    df = repo.get_full_df()
-    assert hp.key not in df.columns
+    assert hp.key not in repo.get_full_df().columns
 
 
 def test_hash_partitioner_get_partition(tmp_path):
@@ -286,13 +307,11 @@ def test_hash_partitioner_get_partition(tmp_path):
     repo = TableRepo(tmp_path / "data", partition_cols=hp)
     df = make_df(100)
     repo.extend(df)
-
-    # All rows with B=3 land in the same bucket
     b3_bucket = str(int(pl.Series([3]).cast(pl.String).hash()[0]) % 8).zfill(1)
     part = repo.get_partition_df({hp.key: b3_bucket})
-    b3_count = df.filter(pl.col("B") == 3).shape[0]
-    # Partition contains all B=3 rows (and possibly other B values in same bucket)
-    assert part.filter(pl.col("B") == 3).shape[0] == b3_count
+    assert (
+        part.filter(pl.col("B") == 3).shape[0] == df.filter(pl.col("B") == 3).shape[0]
+    )
 
 
 def test_hash_partitioner_extend_twice(tmp_path):
@@ -314,29 +333,30 @@ def test_record_writer_extend(tmp_path):
     with repo.get_extending_record_writer(7) as writer:
         for i in range(30):
             writer.add_to_batch({"i": i, "val": i * 2})
-    assert repo.get_full_df().shape[0] == 30
+    result = repo.get_full_df().sort("i")
+    assert result.shape[0] == 30
+    assert result["val"].to_list() == [i * 2 for i in range(30)]
 
 
 def test_record_writer_replace(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="i")
-    df = pl.DataFrame({"i": list(range(10)), "val": list(range(10))})
-    repo.extend(df)
-
+    repo.extend(pl.DataFrame({"i": list(range(10)), "val": list(range(10))}))
     with repo.get_replacing_record_writer(3) as writer:
         for i in range(5):
             writer.add_to_batch({"i": i, "val": i * 100})
-
     result = repo.get_full_df().sort("i")
     assert result.shape[0] == 10
-    assert result.filter(pl.col("i") == 0)["val"][0] == 0
+    assert result.filter(pl.col("i") < 5)["val"].to_list() == [
+        i * 100 for i in range(5)
+    ]
+    assert result.filter(pl.col("i") >= 5)["val"].to_list() == list(range(5, 10))
 
 
 def test_record_writer_context_manager_flushes(tmp_path):
     repo = TableRepo(tmp_path / "data")
     with repo.get_extending_record_writer(100) as writer:
-        for i in range(5):  # below batch_size
+        for i in range(5):
             writer.add_to_batch({"i": i})
-    # context exit should have flushed
     assert repo.get_full_df().shape[0] == 5
 
 
@@ -347,29 +367,22 @@ def test_record_writer_context_manager_flushes(tmp_path):
 
 def test_df_writer_extend(tmp_path):
     repo = TableRepo(tmp_path / "data")
-    dfs = [make_df(10, seed=i) for i in range(5)]
     with repo.get_extending_df_writer(25) as writer:
-        for df in dfs:
-            writer.add_to_batch(df)
+        for i in range(5):
+            writer.add_to_batch(make_df(10, seed=i))
     assert repo.get_full_df().shape[0] == 50
 
 
 def test_df_writer_replaces(tmp_path):
     repo = TableRepo(tmp_path / "data", id_col="id")
-    base = pl.DataFrame({"id": list(range(20)), "val": list(range(20))})
-    repo.extend(base)
-
-    updates = [
-        pl.DataFrame({"id": [0, 1, 2], "val": [100, 101, 102]}),
-        pl.DataFrame({"id": [3, 4, 5], "val": [103, 104, 105]}),
-    ]
+    repo.extend(pl.DataFrame({"id": list(range(20)), "val": list(range(20))}))
     with repo.get_replacing_df_writer(10) as writer:
-        for df in updates:
-            writer.add_to_batch(df)
-
+        writer.add_to_batch(pl.DataFrame({"id": [0, 1, 2], "val": [100, 101, 102]}))
+        writer.add_to_batch(pl.DataFrame({"id": [3, 4, 5], "val": [103, 104, 105]}))
     result = repo.get_full_df().sort("id")
     assert result.shape[0] == 20
-    assert result.filter(pl.col("id") == 0)["val"][0] == 100
+    assert result.filter(pl.col("id") <= 5)["val"].to_list() == list(range(100, 106))
+    assert result.filter(pl.col("id") > 5)["val"].to_list() == list(range(6, 20))
 
 
 def test_df_writer_context_flushes(tmp_path):
@@ -384,27 +397,15 @@ def test_df_writer_context_flushes(tmp_path):
 # ---------------------------------------------------------------------------
 
 
-def test_concurrent_extend(tmp_path):
-    repo = TableRepo(tmp_path / "data")
-    n_threads = 8
-    rows_per_thread = 100
-
-    def worker(seed):
-        repo.extend(make_df(rows_per_thread, seed=seed))
-
-    threads = [threading.Thread(target=worker, args=(i,)) for i in range(n_threads)]
-    for t in threads:
-        t.start()
-    for t in threads:
-        t.join()
-
-    assert repo.get_full_df().shape[0] == n_threads * rows_per_thread
-
-
-def test_concurrent_extend_partitioned(tmp_path):
-    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
-    n_threads = 6
-    rows_per_thread = 50
+@pytest.mark.parametrize(
+    "partition_cols,n_threads,rows_per_thread",
+    [
+        ([], 8, 100),
+        (["C"], 6, 50),
+    ],
+)
+def test_concurrent_extend(tmp_path, partition_cols, n_threads, rows_per_thread):
+    repo = TableRepo(tmp_path / "data", partition_cols=partition_cols)
 
     def worker(seed):
         repo.extend(make_df(rows_per_thread, seed=seed))
@@ -414,7 +415,6 @@ def test_concurrent_extend_partitioned(tmp_path):
         t.start()
     for t in threads:
         t.join()
-
     assert repo.get_full_df().shape[0] == n_threads * rows_per_thread
 
 
@@ -426,11 +426,9 @@ def test_concurrent_extend_partitioned(tmp_path):
 def test_schema_mismatch_extend(tmp_path):
     repo = TableRepo(tmp_path / "data")
     repo.extend(pl.DataFrame({"A": [1, 2], "B": [3, 4]}))
-    # Add column C that didn't exist before
     repo.extend(pl.DataFrame({"A": [5], "C": [6]}))
     full = repo.get_full_df()
     assert full.shape[0] == 3
-    # B and C both present (nulls for missing values)
     assert "B" in full.columns
     assert "C" in full.columns
 
@@ -443,9 +441,6 @@ def test_schema_mismatch_extend(tmp_path):
 def test_empty_extend(tmp_path):
     repo = TableRepo(tmp_path / "data")
     repo.extend(pl.DataFrame({"A": pl.Series([], dtype=pl.Int64)}))
-    # Empty extend should not create a file (nothing to write in _append_parts)
-    # or create an empty file - behaviour: _extend_dir writes an empty file
-    # The important thing is get_full_df works
     assert repo.get_full_df().shape[0] == 0
 
 
@@ -454,18 +449,17 @@ def test_empty_extend(tmp_path):
 # ---------------------------------------------------------------------------
 
 
-def test_lock_file_created(tmp_path):
-    repo = TableRepo(tmp_path / "data")
-    repo.extend(make_df(10))
-    assert (tmp_path / "data" / ".lock").exists()
-
-
-def test_lock_file_created_partitioned(tmp_path):
-    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
-    repo.extend(make_df(30))
-    # Each partition dir should have a .lock file
-    lock_files = list((tmp_path / "data").glob("**/.lock"))
-    assert len(lock_files) == 3  # x, y, z partitions
+@pytest.mark.parametrize(
+    "partition_cols,n,expected_locks",
+    [
+        ([], 10, 1),
+        (["C"], 30, 3),  # x, y, z partitions
+    ],
+)
+def test_lock_file_created(tmp_path, partition_cols, n, expected_locks):
+    repo = TableRepo(tmp_path / "data", partition_cols=partition_cols)
+    repo.extend(make_df(n))
+    assert len(list((tmp_path / "data").glob("**/.lock"))) == expected_locks
 
 
 # ---------------------------------------------------------------------------
@@ -478,15 +472,15 @@ def test_map_partitions_basic(tmp_path):
     repo.extend(make_df(90))
     counts = repo.map_partitions(lambda df: len(df))
     assert sum(counts) == 90
-    assert len(counts) == 3  # x, y, z
+    assert len(counts) == 3
 
 
 def test_map_partitions_parallel(tmp_path):
     repo = TableRepo(tmp_path / "data", partition_cols=["C"])
     repo.extend(make_df(90))
-    counts_seq = sorted(repo.map_partitions(lambda df: len(df)))
-    counts_par = sorted(repo.map_partitions(lambda df: len(df), workers=2))
-    assert counts_seq == counts_par
+    assert sorted(repo.map_partitions(lambda df: len(df))) == sorted(
+        repo.map_partitions(lambda df: len(df), workers=2)
+    )
 
 
 def test_map_partitions_empty(tmp_path):
@@ -498,9 +492,7 @@ def test_map_partitions_no_partition_combines_files(tmp_path):
     repo = TableRepo(tmp_path / "data", max_records=5)
     repo.extend(make_df(15))
     assert repo.n_files == 3
-    # map_partitions should combine all 3 files into one call
-    result = repo.map_partitions(lambda df: df.shape[0])
-    assert result == [15]
+    assert repo.map_partitions(lambda df: df.shape[0]) == [15]
 
 
 # ---------------------------------------------------------------------------
@@ -544,10 +536,75 @@ def test_compact_ipc(tmp_path):
 
 
 # ---------------------------------------------------------------------------
+# purge_partition / replace_partition
+# ---------------------------------------------------------------------------
+
+
+def test_purge_partition(tmp_path):
+    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
+    repo.extend(make_df(90))
+    x_count = repo.get_partition_df({"C": "x"}).shape[0]
+    repo.purge_partition({"C": "x"})
+    result = repo.get_full_df()
+    assert result.filter(pl.col("C") == "x").shape[0] == 0
+    assert result.shape[0] == 90 - x_count
+
+
+def test_purge_partition_no_match(tmp_path):
+    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
+    repo.extend(make_df(30))
+    n = repo.get_full_df().shape[0]
+    repo.purge_partition({"C": "nonexistent"})
+    assert repo.get_full_df().shape[0] == n
+
+
+def test_replace_partition_basic(tmp_path):
+    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
+    repo.extend(make_df(90))
+    repo.replace_partition(pl.DataFrame({"A": [9.0], "B": [5], "C": ["x"]}))
+    x_rows = repo.get_full_df().filter(pl.col("C") == "x")
+    assert x_rows.shape[0] == 1
+    assert x_rows["A"][0] == 9.0
+
+
+def test_replace_partition_no_partition_raises(tmp_path):
+    repo = TableRepo(tmp_path / "data")
+    with pytest.raises(ValueError, match="partition_cols"):
+        repo.replace_partition(pl.DataFrame({"A": [1]}))
+
+
+def test_replace_partition_with_max_records(tmp_path):
+    repo = TableRepo(tmp_path / "data", partition_cols=["C"], max_records=5)
+    repo.extend(make_df(90))
+    repo.replace_partition(
+        pl.DataFrame({"A": [1.0, 2.0], "B": [1, 2], "C": ["x", "x"]})
+    )
+    assert repo.get_full_df().filter(pl.col("C") == "x").shape[0] == 2
+
+
+# ---------------------------------------------------------------------------
+# get_partition_lf edge cases
+# ---------------------------------------------------------------------------
+
+
+def test_get_partition_lf_no_match(tmp_path):
+    repo = TableRepo(tmp_path / "data", partition_cols=["C"])
+    repo.extend(make_df(30))
+    assert repo.get_partition_lf({"C": "nonexistent"}).collect().shape[0] == 0
+
+
+def test_get_partition_lf_no_partition_cols(tmp_path):
+    repo = TableRepo(tmp_path / "data")
+    repo.extend(make_df(10))
+    assert repo.get_partition_lf({}).collect().shape[0] == 0
+
+
+# ---------------------------------------------------------------------------
+# Performance smoke test
+# ---------------------------------------------------------------------------
 
 
 def test_large_extend_performance(tmp_path):
-    """Smoke test: 40 batches of 10k rows should complete without error."""
     repo = TableRepo(tmp_path / "data")
     rng = random.Random(7)
     for _ in range(40):

shackleton-1.1.0/coverage/report.md

@@ -1,6 +0,0 @@
-| Name                           |    Stmts |     Miss |   Cover |
-|------------------------------- | -------: | -------: | ------: |
-| shackleton/\_\_init\_\_.py     |        3 |        0 |    100% |
-| shackleton/core.py             |      259 |        3 |     99% |
-| shackleton/tests/test\_core.py |      324 |        1 |     99% |
-| **TOTAL**                      |  **586** |    **4** | **99%** |