polars-map 0.1.0__tar.gz

polars_map-0.1.0/.github/workflows/build.yml

@@ -0,0 +1,31 @@
+name: build
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install uv
+        run: python -m pip install uv
+      - name: Check ruff format
+        run: uv run ruff format --check
+      - name: Check with ruff
+        run: uv run ruff check
+      - name: Run pyright
+        run: uv run pyright
+      - name: Test with pytest
+        run: uv run pytest

polars_map-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,121 @@
+name: release
+
+on:
+  workflow_dispatch:
+    inputs:
+      bump:
+        description: "Version bump type"
+        required: true
+        type: choice
+        options:
+          - patch
+          - minor
+          - major
+
+permissions:
+  contents: write
+  id-token: write
+
+jobs:
+  gate:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install uv
+        run: python -m pip install uv
+      - name: Check ruff format
+        run: uv run ruff format --check
+      - name: Check with ruff
+        run: uv run ruff check
+      - name: Run pyright
+        run: uv run pyright
+      - name: Test with pytest
+        run: uv run pytest
+
+  prepare:
+    runs-on: ubuntu-latest
+    needs: gate
+    outputs:
+      version: ${{ steps.bump.outputs.version }}
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install uv
+        run: python -m pip install uv
+      - name: Compute new version
+        id: bump
+        run: |
+          current=$(uv version | cut -d' ' -f2)
+          IFS='.' read -r major minor patch <<< "$current"
+          case "${{ inputs.bump }}" in
+            major) major=$((major + 1)); minor=0; patch=0 ;;
+            minor) minor=$((minor + 1)); patch=0 ;;
+            patch) patch=$((patch + 1)) ;;
+          esac
+          version="${major}.${minor}.${patch}"
+          echo "version=${version}" >> "$GITHUB_OUTPUT"
+          echo "New version: ${version}"
+
+  build:
+    needs: prepare
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install uv
+        run: python -m pip install uv
+      - name: Set version
+        run: uv version ${{ needs.prepare.outputs.version }} --frozen
+      - name: Build
+        run: uv build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          packages-dir: dist/
+
+  release:
+    needs: [prepare, publish]
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install uv
+        run: python -m pip install uv
+      - name: Set version and commit
+        env:
+          VERSION: ${{ needs.prepare.outputs.version }}
+        run: |
+          uv version "$VERSION" --frozen
+          uv lock
+          git config user.name "github-actions[bot]"
+          git config user.email "github-actions[bot]@users.noreply.github.com"
+          git add pyproject.toml uv.lock
+          git commit -m "v${VERSION}"
+          git tag -a "v${VERSION}" -m "v${VERSION}"
+          git push --follow-tags
+      - name: Create GitHub release
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: gh release create "v${{ needs.prepare.outputs.version }}" --generate-notes

polars_map-0.1.0/.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+.coverage
+.venv
+*.egg-info
+*.pyo
+*.pyc
+*.so
+*.dylib
+build/
+dist/
+out/
+target/
+wheels/

polars_map-0.1.0/.python-version

@@ -0,0 +1 @@
+3.13

polars_map-0.1.0/PKG-INFO

@@ -0,0 +1,9 @@
+Metadata-Version: 2.4
+Name: polars-map
+Version: 0.1.0
+Summary: Polars plugin providing Map operations on List(Struct({key, value})) columns
+Project-URL: Homepage, https://github.com/hafaio/polars-map
+Project-URL: Repository, https://github.com/hafaio/polars-map
+Project-URL: Issues, https://github.com/hafaio/polars-map/issues
+Requires-Python: >=3.10
+Requires-Dist: polars>=1.13.0

polars_map-0.1.0/README.md

@@ -0,0 +1,86 @@
+# polars-map
+
+[![build](https://github.com/hafaio/polars-map/actions/workflows/build.yml/badge.svg)](https://github.com/hafaio/polars-map/actions/workflows/build.yml)
+[![pypi](https://img.shields.io/pypi/v/polars-map)](https://pypi.org/project/polars-map/)
+
+Polars plugin providing a Map extension type and functions.
+Maps represent a mapping from unique keys of any type to values, and are stored as `List(Struct({key, value}))` columns.
+All function in the `.map` namespace can be used on the extension type or on the
+underlying list.
+
+## Installation
+
+```bash
+pip install polars-map
+```
+
+## Supported operations (`.map.*`)
+
+| Category   | Methods                                                    |
+| ---------- | ---------------------------------------------------------- |
+| Accessors  | `entries`, `keys`, `values`, `len`, `get`, `contains_key`  |
+| Filtering  | `filter`, `filter_keys`, `filter_values`                   |
+| Transform  | `eval`, `eval_keys`, `eval_values`                         |
+| Set ops    | `merge`, `intersection`, `difference`                      |
+| Conversion | `from_entries`                                             |
+| Iteration  | `__iter__`, `to_list` (Series only)                        |
+
+## Usage
+
+```python
+import polars as pl
+from polars_map import Map
+
+# Construction
+ser = pl.Series(
+    "m",
+    [
+        [{"key": "a", "value": 1}, {"key": "b", "value": 2}],
+        [{"key": "x", "value": 10}],
+    ],
+    dtype=Map(pl.String(), pl.Int64()),
+)
+df = pl.DataFrame([ser])
+
+# Accessors
+df.select(pl.col("m").map.keys())        # [["a", "b"], ["x"]]
+df.select(pl.col("m").map.values())       # [[1, 2], [10]]
+df.select(pl.col("m").map.len())          # [2, 1]
+
+# Lookup
+df.select(pl.col("m").map.get("a"))           # [1, None]
+df.select(pl.col("m").map.contains_key("a"))  # [True, False]
+
+# Filtering
+df.select(pl.col("m").map.filter(pl.element().struct["value"] > 1))
+df.select(pl.col("m").map.filter_keys(pl.element() > "a"))
+df.select(pl.col("m").map.filter_values(pl.element() >= 2))
+
+# Transform keys or values
+df.select(pl.col("m").map.eval_keys(pl.element().str.to_uppercase()))
+df.select(pl.col("m").map.eval_values(pl.element() * 2))
+
+# Merge (right-side wins on key conflict)
+left = pl.Series("l", [[{"key": "a", "value": 1}, {"key": "b", "value": 2}]], dtype=Map(pl.String(), pl.Int64()))
+right = pl.Series("r", [[{"key": "a", "value": 99}, {"key": "c", "value": 3}]], dtype=Map(pl.String(), pl.Int64()))
+pl.DataFrame([left, right]).select(pl.col("l").map.merge(pl.col("r")))
+# [{"a": 99, "b": 2, "c": 3}]
+
+# Set operations
+pl.DataFrame([left, right]).select(pl.col("l").map.intersection(pl.col("r")))  # keys in both
+pl.DataFrame([left, right]).select(pl.col("l").map.difference(pl.col("r")))    # keys only in left
+
+# Convert to/from plain List(Struct)
+df.select(pl.col("m").map.entries())        # strip Map type
+df.select(pl.col("m").map.from_entries())   # wrap as Map (with deduplication)
+
+# Series iteration yields Python dicts
+for d in ser.map:
+    print(d)  # {"a": 1, "b": 2}, {"x": 10}
+```
+
+## Caveats
+
+- **Extension types** — used to wrap the underlying `List(Struct)` storage with a semantic `Map` dtype is not yet stabilized and may change across Polars releases.
+- **`pl.dtype_of`** — used to efficiently cast to the extension type after _some_ operations is also unstable.
+- **GIL** - is required to automatically wrap an expression as the extension type, and so operations which could change the underlying key or value types will briefly lock the GIL to do the cast. This may also prevent the polars engine from reasoning about the type.

polars_map-0.1.0/polars_map/__init__.py

@@ -0,0 +1,7 @@
+"""Polars plugin providing Map operations on List(Struct({key, value})) columns."""
+
+from polars_map._dtype import Map
+from polars_map._expr import MapExpr
+from polars_map._series import MapSeries
+
+__all__ = ("Map", "MapExpr", "MapSeries")

polars_map-0.1.0/polars_map/_dtype.py

@@ -0,0 +1,40 @@
+"""Map extension data type for Polars."""
+
+from __future__ import annotations
+
+import polars as pl
+
+
+def _ensure_instance(dt: pl.DataType | type[pl.DataType]) -> pl.DataType:
+    return dt() if isinstance(dt, type) else dt
+
+
+class Map(pl.BaseExtension):
+    """Map extension type backed by List(Struct({key, value})).
+
+    Usage as a dtype for Series construction::
+
+        dtype = Map(pl.String(), pl.Int64())
+    """
+
+    def __init__(self, key: pl.DataType, value: pl.DataType) -> None:
+        storage = pl.List(pl.Struct({"key": key, "value": value}))
+        super().__init__("polars_map.map", storage)
+
+    @property
+    def key(self) -> pl.DataType:
+        """Key data type."""
+        [key, _] = self.ext_storage().inner.fields  # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType]
+        return _ensure_instance(key.dtype)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType]
+
+    @property
+    def value(self) -> pl.DataType:
+        """Value data type."""
+        [_, value] = self.ext_storage().inner.fields  # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownVariableType]
+        return _ensure_instance(value.dtype)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType]
+
+    def _string_repr(self) -> str:
+        return f"map[{self.key._string_repr()},{self.value._string_repr()}]"  # pyright: ignore[reportUnknownMemberType]
+
+
+pl.register_extension_type("polars_map.map", Map)

polars_map-0.1.0/polars_map/_expr.py

@@ -0,0 +1,215 @@
+"""Expr namespace for Map operations."""
+
+from __future__ import annotations
+
+import functools
+from dataclasses import dataclass
+
+import polars as pl
+
+from ._utils import expr_eval, infer_map, validate
+
+
+@pl.api.register_expr_namespace("map")
+@dataclass(frozen=True)
+class MapExpr:
+    """Expression namespace for Map operations on List(Struct({key, value})) columns."""
+
+    _expr: pl.Expr
+
+    def _as_self(self, expr: pl.Expr) -> pl.Expr:
+        """Wrap a List(Struct) result back as Map, preserving the original dtype."""
+        return expr.ext.to(pl.dtype_of(self._expr))
+
+    def from_entries(
+        self,
+        *,
+        validate_fields: bool = True,
+        deduplicate: bool = True,
+        parallel: bool = False,
+    ) -> pl.Expr:
+        """Wrap a List(Struct({key, value})) expression as a Map extension type.
+
+        Parameters
+        ----------
+        deduplicate
+            If True, deduplicate by key, keeping the first occurrence.
+        parallel
+            Run list evaluations in parallel.
+        """
+        return infer_map(
+            self._expr.list.eval(
+                validate(
+                    pl.element(),
+                    validate_fields=validate_fields,
+                    deduplicate=deduplicate,
+                ),
+                parallel=parallel,
+            )
+        )
+
+    @functools.cached_property
+    def _entries(self) -> pl.Expr:
+        return self._expr.ext.storage()
+
+    def entries(self) -> pl.Expr:
+        """Strip the Map extension type, returning raw List(Struct({key, value}))."""
+        return self._entries
+
+    def keys(self) -> pl.Expr:
+        """Extract all keys as a List column."""
+        return self._entries.list.eval(pl.element().struct["key"])
+
+    def values(self) -> pl.Expr:
+        """Extract all values as a List column."""
+        return self._entries.list.eval(pl.element().struct["value"])
+
+    def len(self) -> pl.Expr:
+        """Return the number of entries in the map."""
+        return self._entries.list.len()
+
+    def _get(self, key: object) -> pl.Expr:
+        """Look up a value by key. Returns scalar per row."""
+        return (
+            self._entries.list.eval(
+                pl.element()
+                .filter(pl.element().struct["key"] == key)
+                .struct["value"]
+                .first()
+            )
+            .list.first()
+            .alias(str(key))
+        )
+
+    def get(self, key: object, *keys: object) -> pl.Expr:
+        """Look up a value by key. Returns scalar per row."""
+        if keys:
+            return pl.struct(  # pyright: ignore[reportUnknownMemberType]
+                self._get(key), *(self._get(k) for k in keys)
+            ).struct.unnest()
+        else:
+            return self._get(key)
+
+    def contains_key(self, key: object) -> pl.Expr:
+        """Check if a key exists in the map."""
+        return self._entries.list.eval(pl.element().struct["key"] == key).list.any()
+
+    def eval(
+        self,
+        expr: pl.Expr,
+        *,
+        validate_fields: bool = True,
+        deduplicate: bool = True,
+        parallel: bool = False,
+    ) -> pl.Expr:
+        """Evaluate an expression on entries, returning a Map.
+
+        The expression operates on the struct elements via ``pl.element()``.
+
+        Example
+        -------
+        >>> col.map.eval(pl.element().struct.with_fields(pl.element().struct["value"] * 2))
+        """
+        inner = validate(expr, validate_fields=validate_fields, deduplicate=deduplicate)
+        return infer_map(self._entries.list.eval(inner, parallel=parallel))
+
+    def eval_keys(
+        self, expr: pl.Expr, *, deduplicate: bool = True, parallel: bool = False
+    ) -> pl.Expr:
+        """Transform keys, returning a Map with new key type.
+
+        The expression operates on each key via ``pl.element()``.
+
+        Example
+        -------
+        >>> col.map.eval_keys(pl.element().str.to_uppercase())
+        """
+        inner: pl.Expr = pl.element().struct.with_fields(  # pyright: ignore[reportUnknownMemberType]
+            key=expr_eval(pl.element().struct["key"], expr)
+        )
+        if deduplicate:
+            inner = inner.filter(pl.element().struct["key"].is_first_distinct())
+        return infer_map(self._entries.list.eval(inner, parallel=parallel))
+
+    def eval_values(self, expr: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Transform values, returning a Map with new value type.
+
+        The expression operates on each value via ``pl.element()``.
+
+        Example
+        -------
+        >>> col.map.eval_values(pl.element() * 2)
+        """
+        inner = pl.element().struct.with_fields(  # pyright: ignore[reportUnknownMemberType]
+            value=expr_eval(pl.element().struct["value"], expr)
+        )
+        return infer_map(self._entries.list.eval(inner, parallel=parallel))
+
+    def filter(self, predicate: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Filter entries by a predicate on the struct entry.
+
+        Example
+        -------
+        >>> col.map.filter(pl.element().struct["key"] > "b")
+        """
+        return self._as_self(
+            self._entries.list.eval(pl.element().filter(predicate), parallel=parallel)
+        )
+
+    def filter_keys(self, predicate: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Filter entries where the key satisfies the predicate.
+
+        Example
+        -------
+        >>> col.map.filter_keys(pl.element() > "b")
+        """
+        inner = pl.element().filter(
+            expr_eval(pl.element().struct["key"], predicate)  # pyright: ignore[reportUnknownMemberType]
+        )
+        return self._as_self(self._entries.list.eval(inner, parallel=parallel))
+
+    def filter_values(self, predicate: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Filter entries where the value satisfies the predicate.
+
+        Example
+        -------
+        >>> col.map.filter_values(pl.element() > 5)
+        """
+        inner = pl.element().filter(
+            expr_eval(pl.element().struct["value"], predicate)  # pyright: ignore[reportUnknownMemberType]
+        )
+        return self._as_self(self._entries.list.eval(inner, parallel=parallel))
+
+    def merge(self, other: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Merge two maps. Right-side values win on key conflict."""
+        combined = pl.concat_list([self._entries, other.map.entries()])  # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue,reportUnknownVariableType]
+        return self._as_self(
+            combined.list.eval(
+                pl.element().filter(pl.element().struct["key"].is_last_distinct()),
+                parallel=parallel,
+            )
+        )
+
+    def intersection(self, other: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Keep entries from self where the key also exists in other."""
+        combined = pl.concat_list([self._entries, other.map.entries()])  # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue,reportUnknownVariableType]
+        return self._as_self(
+            combined.list.eval(
+                pl.element().filter(
+                    pl.element().struct["key"].is_duplicated()
+                    & pl.element().struct["key"].is_first_distinct()
+                ),
+                parallel=parallel,
+            )
+        )
+
+    def difference(self, other: pl.Expr, *, parallel: bool = False) -> pl.Expr:
+        """Keep entries from self where the key does NOT exist in other."""
+        other_entries = other.map.entries()  # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue,reportUnknownVariableType]
+        combined = pl.concat_list([self._entries, other_entries, other_entries])  # pyright: ignore[reportUnknownMemberType]
+        return self._as_self(
+            combined.list.eval(
+                pl.element().filter(~pl.element().struct["key"].is_duplicated()),
+                parallel=parallel,
+            )
+        )