xolars 0.2.0__py3-none-any.whl

xolars/__init__.py

@@ -0,0 +1,12 @@
+"""xolars — an xarray Dataset paired with per-dimension Polars frames.
+
+The :class:`Xolars` container keeps an :class:`xarray.Dataset` and one Polars
+frame per dimension aligned to the Dataset's coordinate order, including under
+``isel``/``sel`` selection and zarr+parquet round-trips.
+"""
+
+from __future__ import annotations
+
+from xolars._core import Xolars
+
+__all__ = ["Xolars"]

xolars/_core.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from collections.abc import Hashable, Mapping
+from pathlib import Path
+from typing import Any, Generic, Iterable, TypeVar, cast
+
+import numpy as np
+import polars as pl
+import xarray as xr
+from attrs import define, evolve
+from typing_extensions import Self
+from xarray.core.types import ErrorOptionsWithWarn, ZarrWriteModes
+
+F = TypeVar("F", bound=pl.DataFrame | pl.LazyFrame)
+
+
+@define(frozen=True)
+class Xolars(Generic[F]):
+    ds: xr.Dataset
+    df: Mapping[Hashable, F]
+
+    def __attrs_post_init__(self):
+        new_df = {}
+        for dim, frame in self.df.items():
+            dim_str = str(dim)
+            if dim_str not in self.ds.sizes:
+                raise ValueError(
+                    f"'{dim}' is not a dimension in ds. Available: {list(self.ds.dims)}"
+                )
+            col_names = (
+                frame.collect_schema().names()
+                if isinstance(frame, pl.LazyFrame)
+                else frame.columns
+            )
+            if dim_str not in col_names:
+                raise ValueError(
+                    f"DataFrame for dim '{dim}' is missing column '{dim_str}'"
+                )
+            # Collect only the dim column for validation (cheap — just IDs)
+            if isinstance(frame, pl.LazyFrame):
+                id_col: pl.Series = frame.select(dim_str).collect()[dim_str]
+            else:
+                id_col = cast(pl.DataFrame, frame)[dim_str]
+            ds_dim_values = self.ds[dim_str]
+            df_dim_values: pl.Series = id_col
+            common = np.intersect1d(ds_dim_values.to_numpy(), df_dim_values.to_numpy())
+            if len(common) != len(ds_dim_values) or len(common) != len(df_dim_values):
+                raise ValueError(
+                    f"DataFrame for dim '{dim}' values don't match ds['{dim_str}'] coordinates. "
+                    f"Expected {len(ds_dim_values)} number of shared values, got {len(common)}"
+                )
+            new_df[dim] = _reorder(frame, ds_dim_values)
+        object.__setattr__(self, "df", new_df)
+
+    @classmethod
+    def open(cls, path: Path) -> Xolars[pl.LazyFrame]:
+        ds = xr.open_zarr(path / "dataset.zarr")
+        df: dict[Hashable, pl.LazyFrame] = {}
+        for parquet_path in sorted(path.glob("*.parquet")):
+            df[parquet_path.stem] = pl.scan_parquet(parquet_path)
+        return Xolars(ds=ds, df=df)
+
+    def write(self, path: Path, mode: ZarrWriteModes):
+        path.mkdir(parents=True, exist_ok=True)
+        self.ds.to_zarr(path / "dataset.zarr", mode=mode)
+        for dim, frame in self.df.items():
+            out = path / f"{dim}.parquet"
+            if isinstance(frame, pl.LazyFrame):
+                frame.sink_parquet(out)
+            else:
+                cast(pl.DataFrame, frame).write_parquet(out)
+
+    def isel(
+        self,
+        indexers: Mapping[Any, Any] | None = None,
+        drop: bool = False,
+        missing_dims: ErrorOptionsWithWarn = "raise",
+        **indexers_kwargs: Any,
+    ) -> Self:
+        merged = dict(indexers or {})
+        merged.update(indexers_kwargs)
+        new_ds = self.ds.isel(merged, drop=drop, missing_dims=missing_dims)
+        return evolve(self, ds=new_ds, df=_filter_df(self.df, new_ds, merged))
+
+    def sel(
+        self,
+        indexers: Mapping[Any, Any] | None = None,
+        method: str | None = None,
+        tolerance: int | float | Iterable[int | float] | None = None,
+        drop: bool = False,
+        **indexers_kwargs: Any,
+    ) -> Self:
+        merged = dict(indexers or {})
+        merged.update(indexers_kwargs)
+        new_ds = self.ds.sel(merged, method=method, tolerance=tolerance, drop=drop)
+        return evolve(self, ds=new_ds, df=_filter_df(self.df, new_ds, merged))
+
+    def collect(self) -> Xolars[pl.DataFrame]:
+        df: dict[Hashable, pl.DataFrame] = {
+            k: v.collect() if isinstance(v, pl.LazyFrame) else cast(pl.DataFrame, v)
+            for k, v in self.df.items()
+        }
+        return Xolars(self.ds, df)
+
+
+def _reorder(frame: F, coords: xr.DataArray) -> F:
+    """Reorder frame rows to match the given coordinate order."""
+    dim_name = str(coords.name)
+    order_df = pl.DataFrame({dim_name: coords.to_numpy()}).with_row_index("__i__")
+    if isinstance(frame, pl.LazyFrame):
+        joined = frame.join(order_df.lazy(), on=dim_name, how="right")
+    else:
+        joined = cast(pl.DataFrame, frame).join(order_df, on=dim_name, how="right")
+    return cast(F, joined.sort("__i__").drop("__i__"))
+
+
+def _filter_df(
+    df: Mapping[Hashable, F],
+    new_ds: xr.Dataset,
+    merged: dict[str, Any],
+) -> dict[Hashable, F]:
+    """Filter df entries whose dim appears in merged indexers to match new_ds coords."""
+    new_df: dict[Hashable, F] = {}
+    for dim, frame in df.items():
+        dim_str = str(dim)
+        if dim_str in merged:
+            if dim_str not in new_ds.dims:
+                continue  # scalar index dropped this dimension
+            keep = new_ds[dim_str].to_numpy()
+            new_df[dim] = cast(F, frame.filter(pl.col(dim_str).is_in(keep)))
+        else:
+            new_df[dim] = frame
+    return new_df

xolars-0.2.0.dist-info/METADATA

@@ -0,0 +1,96 @@
+Metadata-Version: 2.4
+Name: xolars
+Version: 0.2.0
+Summary: Xarray + Polars: an xarray Dataset paired with per-dimension Polars frames that stay aligned under selection.
+Keywords: xarray,polars,zarr,dataframe,dataset,bioinformatics
+Author: d-laub
+Author-email: d-laub <dlaub@ucsd.edu>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Typing :: Typed
+Requires-Dist: numpy<3
+Requires-Dist: polars>=1.39,<2
+Requires-Dist: xarray>=2026.2.0,<2027
+Requires-Dist: attrs>=24
+Requires-Dist: typing-extensions>=4.10
+Requires-Dist: zarr>=3.1.5,<4
+Requires-Python: >=3.11
+Project-URL: Homepage, https://github.com/d-laub/xolars
+Project-URL: Repository, https://github.com/d-laub/xolars
+Project-URL: Issues, https://github.com/d-laub/xolars/issues
+Description-Content-Type: text/markdown
+
+# xolars
+
+**X**array + p**olars**: an [xarray](https://xarray.dev) `Dataset` paired with one
+[Polars](https://pola.rs) frame per dimension, kept aligned to the Dataset's
+coordinate order — including under `isel`/`sel` selection and zarr + parquet
+round-trips.
+
+## Install
+
+```bash
+uv add xolars      # or: pip install xolars
+```
+
+## Usage
+
+```python
+import numpy as np
+import polars as pl
+import xarray as xr
+from xolars import Xolars
+
+ds = xr.Dataset(
+    {"expr": (["gene_id", "sample_id"], np.arange(12.0).reshape(3, 4))},
+    coords={"gene_id": ["G1", "G2", "G3"], "sample_id": ["S1", "S2", "S3", "S4"]},
+)
+genes = pl.DataFrame({"gene_id": ["G1", "G2", "G3"], "chrom": ["c1", "c2", "c3"]})
+
+xol = Xolars(ds=ds, df={"gene_id": genes})
+
+# Selection filters the Dataset AND every per-dimension frame, together:
+sub = xol.sel(gene_id=["G3", "G1"])
+assert list(sub.df["gene_id"]["gene_id"]) == list(sub.ds["gene_id"].values)
+
+# Persist to zarr (Dataset) + parquet (per-dim frames), then reopen lazily:
+xol.write("mydata.xolars", mode="w")
+reloaded = Xolars.open("mydata.xolars")        # frames are pl.LazyFrame
+eager = reloaded.collect()                      # -> pl.DataFrame
+```
+
+`Xolars` is a frozen, generic container: `Xolars[pl.LazyFrame]` after `open`,
+`Xolars[pl.DataFrame]` after `collect`. Construction validates that each frame's
+dim column exactly matches the Dataset coordinate (same set, same multiplicity)
+and reorders rows to the Dataset's order.
+
+## Development
+
+```bash
+uv sync                          # create .venv with runtime + dev deps
+uv run pytest -q                 # tests
+uv run prek run --all-files      # ruff + pyrefly + hygiene hooks
+uv run prek install              # enable git hooks locally
+```
+
+## Releasing
+
+Releases run via the manual `Release` GitHub Actions workflow
+(`workflow_dispatch`), which uses [commitizen](https://commitizen-tools.github.io/commitizen/)
+to bump the version from conventional commits, tag, create a GitHub release, and
+publish to PyPI via OIDC trusted publishing. Before the workflow can succeed,
+the following one-time setup is required on GitHub (out of scope for the initial
+extraction):
+
+1. Push this repository to `https://github.com/d-laub/xolars`.
+2. Add a `GH_ACTIONS` repository secret (a PAT able to push to `main`) so the
+   bump commit + tag can be pushed past branch protection.
+3. Configure a `pypi` environment and enable
+   [PyPI trusted publishing](https://docs.pypi.org/trusted-publishers/) for the
+   `xolars` project pointing at the `publish` job.

xolars-0.2.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+xolars/__init__.py,sha256=4A92AJpojB0dlDPjZxoNIkwJf104NYo03SufwsGHvMI,383
+xolars/_core.py,sha256=7-Bdds1ZK47JFdymP8u6rpIqLkqGzTzLQ4b5sTmxGEk,5141
+xolars/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xolars-0.2.0.dist-info/WHEEL,sha256=LlB9zUOn921TC3CC5yCeS6O5jsLlxqKqIpg_Zk6XXcQ,81
+xolars-0.2.0.dist-info/METADATA,sha256=CKfBctG6h43ifYFwujqh7jtiJM92dOPYi1j8ExtPkak,3662
+xolars-0.2.0.dist-info/RECORD,,

xolars-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.11.20
+Root-Is-Purelib: true
+Tag: py3-none-any