Python Adapter Wiring¶

# Users import everything from the adapter package
from my_adapter import MySource, CastTypes, FilterRows

# my_adapter/source.py
from __future__ import annotations
from typing import Any

import _rypipe_myfmt
from rypipe import Source


class MySource(Source):
    """Pipeline-capable source for MyFormat files."""

    def _read_arrow(self, plan_overrides: dict[str, Any] | None = None) -> Any:
        # Start with construction-time kwargs (field_mapping, drop_fields, etc.)
        plan = self._build_plan_kwargs()
        # Fused pipeline stages override construction-time kwargs
        if plan_overrides:
            plan.update(plan_overrides)
        # Pass the merged plan to the Rust reader
        return _rypipe_myfmt.read(str(self._path), **plan)

{
    "field_mapping": {"old_name": "new_name"},
    "drop_fields": ["internal_id"],
    "filter": {"field": "status", "op": "==", "value": "active"},
    "field_types": {"amount": "float64"},
}

# my_adapter/rypipe_adapter.py
from __future__ import annotations
from typing import Any

from .source import MySource


class MyAdapter:
    """rypipe-compatible adapter for MyFormat files."""

    def read(self, path: str, **kwargs: Any) -> Any:
        """Parse ``path`` and return a ``pyarrow.Table``."""
        return MySource(path, **kwargs).to_arrow()

    def iter_record_batches(
        self, path: str, memory: str | int = "64MiB",
        batch_size: int | None = None, **kwargs: Any,
    ):
        """Yield ``pyarrow.RecordBatch`` objects with constant memory."""
        yield from MySource(path, **kwargs).iter_record_batches(
            memory=memory, batch_size=batch_size
        )

# my_adapter/rypipe_adapter.py (continued)

def _register() -> None:
    try:
        import rypipe
    except Exception:  # pragma: no cover — rypipe is optional
        return
    rypipe.register_adapter("myfmt", MyAdapter(), extensions=[".myfmt"])


_register()  # runs on import

# my_adapter/__init__.py
import importlib

# Side-effect import: registers the adapter with rypipe on import
from . import rypipe_adapter  # noqa: F401

__all__ = [
    "MySource",
    "MyAdapter",
    "CastTypes",
    "FilterRows",
    "RenameFields",
    "DropFields",
]

_modules = {
    "MySource": ".source",
    "CastTypes": ".stages",
    "FilterRows": ".stages",
    "RenameFields": ".stages",
    "DropFields": ".stages",
}


def __getattr__(name):
    if name in _modules:
        mod = importlib.import_module(_modules[name], __package__)
        return getattr(mod, name)
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


def __dir__():
    return __all__

from typing import Callable

_PY_TO_RUST_TYPE = {
    int: "int64",
    float: "float64",
    str: None,
    bool: "bool",
}


class CastTypes:
    __slots__ = ("_mapping",)

    def __init__(self, mapping: dict[str, Callable]):
        self._mapping = mapping

    def apply(self, record: dict) -> dict:
        for field, cast_fn in self._mapping.items():
            try:
                record[field] = cast_fn(record[field])
            except KeyError:
                pass  # field not in this row — skip silently
            except (ValueError, TypeError) as e:
                raise ValueError(
                    f"CastTypes: cannot cast field '{field}' "
                    f"value {record[field]!r}: {e}"
                ) from e
        return record

    def __call__(self, stream):
        return map(self.apply, stream)

    def _plan_kwargs(self) -> dict | None:
        ft = {}
        for field, fn in self._mapping.items():
            rust_type = _PY_TO_RUST_TYPE.get(fn)
            if rust_type is None:
                if fn is str:
                    continue  # str cast is a no-op
                return None  # unsupported type — can't fuse
            ft[field] = rust_type
        return {"field_types": ft} if ft else None

class FilterRows:
    __slots__ = ("_predicate", "_filter_spec")

    def __init__(self, predicate=None, *, field=None, op=None, value=None,
                 field_a=None, field_b=None):
        if predicate is not None:
            self._predicate = predicate
            self._filter_spec = None
        elif field is not None and op is not None and value is not None:
            self._filter_spec = {"field": field, "op": op, "value": value}
            self._predicate = lambda r: (
                r.get(field) == value if op in ("==", "eq")
                else r.get(field) != value
            )
        elif field_a is not None and op is not None and field_b is not None:
            self._filter_spec = {"field_a": field_a, "op": op, "field_b": field_b}
            ops = {
                ">": lambda a, b: a > b, "<": lambda a, b: a < b,
                ">=": lambda a, b: a >= b, "<=": lambda a, b: a <= b,
                "==": lambda a, b: a == b, "!=": lambda a, b: a != b,
            }
            fn = ops[op]
            self._predicate = lambda r: bool(fn(r.get(field_a), r.get(field_b)))
        else:
            raise ValueError(
                "FilterRows requires a callable predicate or "
                "keyword arguments (field+op+value or field_a+op+field_b)"
            )

    def apply(self, record: dict) -> dict | None:
        return record if self._predicate(record) else None

    def __call__(self, stream):
        return (r for r in map(self.apply, stream) if r is not None)

    def _plan_kwargs(self) -> dict | None:
        return {"filter": self._filter_spec} if self._filter_spec else None

class RenameFields:
    __slots__ = ("_mapping",)

    def __init__(self, mapping: dict[str, str]):
        self._mapping = mapping

    def apply(self, record: dict) -> dict:
        return {self._mapping.get(k, k): v for k, v in record.items()}

    def __call__(self, stream):
        return map(self.apply, stream)

    def _plan_kwargs(self) -> dict | None:
        return {"field_mapping": self._mapping}


class DropFields:
    __slots__ = ("_fields_set",)

    def __init__(self, fields: list[str]):
        if isinstance(fields, str):
            raise TypeError(
                f"DropFields expects a list, got a string; use DropFields([{fields!r}])"
            )
        self._fields_set = frozenset(fields)

    def apply(self, record: dict) -> dict:
        return {k: v for k, v in record.items() if k not in self._fields_set}

    def __call__(self, stream):
        return map(self.apply, stream)

    def _plan_kwargs(self) -> dict | None:
        return {"drop_fields": sorted(self._fields_set)}

class MySource(Source):
    def _read_arrow(self, plan_overrides=None):
        plan = self._build_plan_kwargs()
        if plan_overrides:
            plan.update(plan_overrides)
        return _rypipe_myfmt.read(str(self._path), **plan)

    def iter_record_batches(self, memory="64MiB", batch_size=None, **kwargs):
        plan = self._build_plan_kwargs()
        return _rypipe_myfmt.iter_batches(
            str(self._path), memory=memory, batch_size=batch_size, **plan
        )

from my_adapter import MySource

src = MySource("huge_file.myfmt")
for batch in src.iter_record_batches(memory="256MiB"):
    process(batch)