Tusk DFS Implementation Plan¶

# tests/test_packaging.py
"""Packaging and exception hierarchy smoke tests."""

import pytest

import tusk
from tusk.exceptions import (
    MissingPrimaryKeyWarning,
    PrimitiveError,
    SchemaError,
    TuskError,
)


def test_package_imports():
    assert tusk.__version__


@pytest.mark.parametrize("exc", [SchemaError, PrimitiveError])
def test_errors_share_a_base(exc):
    assert issubclass(exc, TuskError)


def test_missing_primary_key_warning_is_filterable():
    assert issubclass(MissingPrimaryKeyWarning, UserWarning)

# src/tusk/exceptions.py
"""Exception and warning types raised by tusk."""

from __future__ import annotations


class TuskError(Exception):
    """Base class for all errors raised by tusk."""


class SchemaError(TuskError):
    """Raised when an entity set's schema is invalid or inconsistent."""


class PrimitiveError(TuskError):
    """Raised when a primitive is unknown or cannot be applied."""


class MissingPrimaryKeyWarning(UserWarning):
    """Warns that a table without a primary key has reduced capabilities."""

# src/tusk/__init__.py
"""Deep feature synthesis for narwhals lazy dataframes."""

from __future__ import annotations

__version__ = "0.1.0"

__all__ = ["__version__"]

# tests/test_dtypes.py
import narwhals as nw
import pytest

from tusk.dtypes import DtypeFamily, matches

F = DtypeFamily


@pytest.mark.parametrize(
    ("dtype", "family", "expected"),
    [
        (nw.Int64(), F.NUMERIC, True),
        (nw.Float64(), F.NUMERIC, True),
        (nw.Boolean(), F.NUMERIC, False),
        (nw.String(), F.NUMERIC, False),
        (nw.Datetime(), F.TEMPORAL, True),
        (nw.Date(), F.TEMPORAL, True),
        (nw.Int64(), F.TEMPORAL, False),
        (nw.String(), F.STRING, True),
        (nw.Int64(), F.STRING, False),
        (nw.Boolean(), F.BOOLEAN, True),
        (nw.Int64(), F.BOOLEAN, False),
        (nw.String(), F.ANY, True),
        (nw.Datetime(), F.ANY, True),
    ],
)
def test_matches(dtype, family, expected):
    assert matches(dtype, family) is expected

# src/tusk/dtypes.py
"""Dtype families used to decide which primitives apply to which columns.

Matching is done on narwhals dtypes alone. There are no logical types and no
semantic tags: the cost is losing the distinction between an integer that is a
quantity and one that is a category; the benefit is identical behaviour on
every backend with no extra dependencies.
"""

from __future__ import annotations

from enum import Enum

import narwhals as nw


class DtypeFamily(Enum):
    """A group of narwhals dtypes that a primitive can accept."""

    NUMERIC = "numeric"
    TEMPORAL = "temporal"
    STRING = "string"
    BOOLEAN = "boolean"
    ANY = "any"


def matches(dtype: nw.dtypes.DType, family: DtypeFamily) -> bool:
    """Report whether a dtype belongs to a family.

    Args:
        dtype: The narwhals dtype to test.
        family: The family to test membership of.

    Returns:
        True if the dtype belongs to the family.
    """
    if family is DtypeFamily.ANY:
        return True
    if family is DtypeFamily.NUMERIC:
        return bool(dtype.is_numeric())
    if family is DtypeFamily.TEMPORAL:
        return bool(dtype.is_temporal())
    if family is DtypeFamily.STRING:
        return dtype == nw.String
    return dtype == nw.Boolean

# tests/conftest.py
"""Shared fixtures. The backend list is the portability knob: adding a backend
here is intended to be a one-line change."""

import datetime as dt

import polars as pl
import pytest

import tusk

BACKENDS = ["polars"]


@pytest.fixture(params=BACKENDS)
def backend(request):
    """Name of the dataframe backend under test."""
    return request.param


def _frames():
    """Three related tables with hand-checkable values.

    customers 1 and 2 have sessions; customer 3 has none (empty-group case).
    session 30 has no transactions (nested empty-group case).
    """
    customers = pl.LazyFrame(
        {
            "id": [1, 2, 3],
            "age": [30, 40, 50],
            "signed_up_at": [dt.datetime(2024, 1, 1)] * 3,
        }
    )
    sessions = pl.LazyFrame(
        {
            "id": [10, 20, 30],
            "customer_id": [1, 1, 2],
            "started_at": [
                dt.datetime(2024, 3, 4),  # Monday
                dt.datetime(2024, 3, 5),
                dt.datetime(2024, 3, 6),
            ],
        }
    )
    transactions = pl.LazyFrame(
        {
            "id": [100, 101, 102, 103],
            "session_id": [10, 10, 20, 20],
            "amount": [1.0, 3.0, 10.0, 20.0],
            "occurred_at": [
                dt.datetime(2024, 3, 4, 1),
                dt.datetime(2024, 3, 4, 2),
                dt.datetime(2024, 3, 5, 1),
                dt.datetime(2024, 3, 5, 2),
            ],
        }
    )
    return customers, sessions, transactions


@pytest.fixture
def es():
    """A three-table retail entity set."""
    customers, sessions, transactions = _frames()
    return (
        tusk.EntitySet("retail")
        .add_dataframe("customers", customers, primary_key="id",
                       row_creation_time="signed_up_at")
        .add_dataframe("sessions", sessions, primary_key="id",
                       row_creation_time="started_at")
        .add_dataframe("transactions", transactions, primary_key="id",
                       row_creation_time="occurred_at")
        .add_relationship(parent="customers", child="sessions",
                          foreign_key="customer_id")
        .add_relationship(parent="sessions", child="transactions",
                          foreign_key="session_id")
    )

# tests/test_entityset.py
import narwhals as nw
import polars as pl
import pytest

import tusk
from tusk.entityset import Relationship
from tusk.exceptions import MissingPrimaryKeyWarning, SchemaError


def test_schema_is_read_from_the_frame(es):
    schema = es.schema("transactions")
    assert schema.primary_key == "id"
    assert schema.row_creation_time == "occurred_at"
    assert schema.dtypes["amount"] == nw.Float64


def test_relationship_accessors(es):
    assert es.children_of("customers") == [
        Relationship(parent="customers", child="sessions", foreign_key="customer_id")
    ]
    assert es.parents_of("transactions") == [
        Relationship(parent="sessions", child="transactions", foreign_key="session_id")
    ]
    assert es.children_of("transactions") == []


def test_key_columns_include_pk_fk_and_time(es):
    assert es.key_columns("sessions") == frozenset({"id", "customer_id", "started_at"})


def test_missing_primary_key_warns():
    with pytest.warns(MissingPrimaryKeyWarning, match="cannot be used as a relationship parent"):
        tusk.EntitySet("x").add_dataframe("t", pl.LazyFrame({"a": [1]}))


def test_unknown_column_raises():
    with pytest.raises(SchemaError, match="nope"):
        tusk.EntitySet("x").add_dataframe("t", pl.LazyFrame({"a": [1]}), primary_key="nope")


def test_composite_key_raises():
    with pytest.raises(SchemaError, match="[Cc]omposite"):
        tusk.EntitySet("x").add_dataframe("t", pl.LazyFrame({"a": [1]}), primary_key=["a"])


def test_parent_without_primary_key_raises():
    es = tusk.EntitySet("x").add_dataframe("p", pl.LazyFrame({"a": [1]}), primary_key="a")
    with pytest.warns(MissingPrimaryKeyWarning):
        es.add_dataframe("c", pl.LazyFrame({"a": [1], "p_a": [1]}))
    with pytest.raises(SchemaError, match="primary_key"):
        es.add_relationship(parent="c", child="p", foreign_key="a")


def test_unknown_foreign_key_raises(es):
    with pytest.raises(SchemaError, match="missing"):
        es.add_relationship(parent="customers", child="sessions", foreign_key="nope")


def test_eager_input_is_recorded_and_lazified():
    es = tusk.EntitySet("x").add_dataframe("t", pl.DataFrame({"a": [1]}), primary_key="a")
    assert es.is_eager is True
    assert isinstance(es.frame("t"), nw.LazyFrame)


def test_self_reference_is_allowed():
    es = tusk.EntitySet("x").add_dataframe(
        "employees", pl.LazyFrame({"id": [1, 2], "manager_id": [None, 1]}), primary_key="id"
    )
    es.add_relationship(parent="employees", child="employees", foreign_key="manager_id")
    assert es.children_of("employees")[0].foreign_key == "manager_id"

# src/tusk/entityset.py
"""The schema model: tables, relationships, and the entity set that holds them."""

from __future__ import annotations

import warnings
from dataclasses import dataclass
from typing import Any, Mapping

import narwhals as nw

from tusk.exceptions import MissingPrimaryKeyWarning, SchemaError


@dataclass(frozen=True)
class TableSchema:
    """Everything phase 1 knows about a table.

    Attributes:
        name: Table name within the entity set.
        primary_key: Column uniquely identifying a row, if declared.
        row_creation_time: Column recording when a row became knowable.
        dtypes: Mapping of column name to narwhals dtype.
    """

    name: str
    primary_key: str | None
    row_creation_time: str | None
    dtypes: Mapping[str, Any]


@dataclass(frozen=True)
class Relationship:
    """A one-to-many link from a parent table to a child table.

    The parent side is always the parent's ``primary_key``; ``foreign_key``
    names the child's column.

    Attributes:
        parent: Name of the parent table.
        child: Name of the child table.
        foreign_key: Column on the child pointing at the parent's primary key.
    """

    parent: str
    child: str
    foreign_key: str


class EntitySet:
    """A collection of related tables that DFS can synthesize features over."""

    def __init__(self, id: str) -> None:
        """Create an empty entity set.

        Args:
            id: Human-readable identifier for this entity set.
        """
        self.id = id
        self._frames: dict[str, nw.LazyFrame] = {}
        self._schemas: dict[str, TableSchema] = {}
        self._relationships: list[Relationship] = []
        self._backend: Any = None
        self._is_eager: bool | None = None

    @property
    def is_eager(self) -> bool:
        """Whether the caller supplied eager frames."""
        return bool(self._is_eager)

    @property
    def table_names(self) -> tuple[str, ...]:
        """Names of every table in the entity set."""
        return tuple(self._schemas)

    def add_dataframe(
        self,
        name: str,
        dataframe: Any,
        primary_key: str | None = None,
        row_creation_time: str | None = None,
    ) -> EntitySet:
        """Add a table to the entity set.

        Args:
            name: Name to register the table under.
            dataframe: A native frame or a narwhals frame.
            primary_key: Column uniquely identifying a row. Required for a
                table used as a relationship parent or as the DFS target.
            row_creation_time: Column recording when a row became knowable.
                Required for order-dependent primitives on this table.

        Returns:
            This entity set, to allow chaining.

        Raises:
            SchemaError: If the name is taken, a declared column is missing,
                a key is composite, or the backend differs from earlier tables.

        Warns:
            MissingPrimaryKeyWarning: If ``primary_key`` is omitted.
        """
        if name in self._schemas:
            raise SchemaError(f"table {name!r} is already in this entity set")
        _reject_composite(primary_key, "primary_key")
        _reject_composite(row_creation_time, "row_creation_time")

        frame = dataframe if isinstance(dataframe, (nw.DataFrame, nw.LazyFrame)) else nw.from_native(dataframe)
        is_eager = isinstance(frame, nw.DataFrame)
        lazy = frame.lazy() if is_eager else frame

        if self._backend is None:
            self._backend = lazy.implementation
            self._is_eager = is_eager
        elif lazy.implementation != self._backend:
            raise SchemaError(
                f"table {name!r} uses backend {lazy.implementation}, but this entity set "
                f"uses {self._backend}; narwhals cannot join across backends"
            )

        dtypes = dict(lazy.collect_schema())
        for column, label in ((primary_key, "primary_key"), (row_creation_time, "row_creation_time")):
            if column is not None and column not in dtypes:
                raise SchemaError(f"{label} {column!r} is not a column of {name!r}")

        if primary_key is None:
            warnings.warn(
                f"{name!r} has no primary_key: it cannot be used as a relationship parent "
                f"or as a DFS target, and order-dependent primitives on it will have "
                f"non-deterministic tiebreaks",
                MissingPrimaryKeyWarning,
                stacklevel=2,
            )

        self._frames[name] = lazy
        self._schemas[name] = TableSchema(name, primary_key, row_creation_time, dtypes)
        return self

    def add_relationship(self, parent: str, child: str, foreign_key: str) -> EntitySet:
        """Link a parent table to a child table.

        Args:
            parent: Name of the parent table. Must have a ``primary_key``.
            child: Name of the child table.
            foreign_key: The child's column pointing at the parent's primary key.

        Returns:
            This entity set, to allow chaining.

        Raises:
            SchemaError: If a table is unknown, the parent has no primary key,
                the foreign key is composite, or the child lacks that column.
        """
        _reject_composite(foreign_key, "foreign_key")
        for table in (parent, child):
            if table not in self._schemas:
                raise SchemaError(f"unknown table {table!r}")
        if self._schemas[parent].primary_key is None:
            raise SchemaError(f"parent table {parent!r} needs a primary_key to be a relationship parent")
        if foreign_key not in self._schemas[child].dtypes:
            raise SchemaError(f"child table {child!r} is missing foreign_key column {foreign_key!r}")
        self._relationships.append(Relationship(parent, child, foreign_key))
        return self

    def schema(self, name: str) -> TableSchema:
        """Return a table's schema.

        Args:
            name: Table name.

        Returns:
            The table's schema.

        Raises:
            SchemaError: If the table is unknown.
        """
        try:
            return self._schemas[name]
        except KeyError:
            raise SchemaError(f"unknown table {name!r}") from None

    def frame(self, name: str) -> nw.LazyFrame:
        """Return a table's lazy frame.

        Args:
            name: Table name.

        Returns:
            The table's narwhals LazyFrame.

        Raises:
            SchemaError: If the table is unknown.
        """
        try:
            return self._frames[name]
        except KeyError:
            raise SchemaError(f"unknown table {name!r}") from None

    def children_of(self, name: str) -> list[Relationship]:
        """Return relationships where this table is the parent.

        Args:
            name: Table name.

        Returns:
            Matching relationships, in insertion order.
        """
        return [r for r in self._relationships if r.parent == name]

    def parents_of(self, name: str) -> list[Relationship]:
        """Return relationships where this table is the child.

        Args:
            name: Table name.

        Returns:
            Matching relationships, in insertion order.
        """
        return [r for r in self._relationships if r.child == name]

    def key_columns(self, name: str) -> frozenset[str]:
        """Return columns excluded from use as primitive inputs.

        Primary keys, foreign keys, and the row creation time are structural,
        not measurements: ``MEAN(customer_id)`` is noise.

        Args:
            name: Table name.

        Returns:
            The table's structural column names.
        """
        schema = self.schema(name)
        keys = {schema.primary_key, schema.row_creation_time}
        keys.update(r.foreign_key for r in self.parents_of(name))
        keys.discard(None)
        return frozenset(keys)  # type: ignore[arg-type]


def _reject_composite(value: Any, label: str) -> None:
    """Raise if a key was given as a sequence.

    Args:
        value: The declared key.
        label: Parameter name, used in the message.

    Raises:
        SchemaError: If the value is a list or tuple.
    """
    if isinstance(value, (list, tuple)):
        raise SchemaError(f"composite {label} is not supported; got {value!r}")

from tusk.entityset import EntitySet, Relationship, TableSchema

__all__ = ["__version__", "EntitySet", "Relationship", "TableSchema"]

# tests/test_primitives_base.py
import pickle
from dataclasses import FrozenInstanceError, dataclass, is_dataclass

import narwhals as nw
import polars as pl
import pytest

import tusk.primitives  # noqa: F401  -- registers the built-in primitives
from tusk.dtypes import DtypeFamily as F
from tusk.exceptions import PrimitiveError
from tusk.primitives.base import AggregationPrimitive, TransformPrimitive
from tusk.primitives.registry import _REGISTRY, register, resolve, resolve_all


@register
@dataclass(frozen=True)
class Doubled(TransformPrimitive):
    name = "doubled"
    input_dtypes = (F.NUMERIC,)

    def build(self, expr):
        return expr * 2


@register
@dataclass(frozen=True)
class Spread(AggregationPrimitive):
    name = "spread"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64

    def build(self, expr):
        return expr.max() - expr.min()


@dataclass(frozen=True)
class Pair(AggregationPrimitive):
    name = "pair"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64
    scale: float = 1.0

    @property
    def number_of_outputs(self):
        return 2

    def build(self, expr):
        return [expr.min() * self.scale, expr.max() * self.scale]


def test_name_generation():
    assert Doubled().generate_name(("amount",)) == "DOUBLED(amount)"


def test_single_output_names():
    assert Spread().output_names("SPREAD(amount)") == ("SPREAD(amount)",)


def test_multi_output_names_are_indexed():
    assert Pair().output_names("PAIR(amount)") == ("PAIR(amount)[0]", "PAIR(amount)[1]")


def test_outputs_is_always_a_tuple():
    assert len(Spread().outputs(nw.col("a"))) == 1
    assert len(Pair().outputs(nw.col("a"))) == 2


def test_return_dtype_preserves_input_by_default():
    assert Doubled().return_dtype((nw.Int64(),)) == nw.Int64()
    assert Spread().return_dtype((nw.Int64(),)) == nw.Float64


def test_parameters_are_dataclass_fields_and_affect_equality():
    assert Pair(scale=2.0) != Pair(scale=1.0)
    assert Pair(scale=2.0) == Pair(scale=2.0)


def test_registry_resolves_names_and_passes_instances_through():
    assert resolve("doubled") == Doubled()
    instance = Pair(scale=3.0)
    assert resolve(instance) is instance
    assert resolve_all(["doubled", "spread"]) == (Doubled(), Spread())


def test_unknown_primitive_lists_alternatives():
    with pytest.raises(PrimitiveError, match="doubled"):
        resolve("dubbled")


def test_expressions_actually_evaluate():
    lf = nw.from_native(pl.LazyFrame({"g": [1, 1], "a": [2.0, 5.0]}))
    got = lf.group_by("g").agg(Spread().outputs(nw.col("a"))[0].alias("s")).collect()
    assert got.to_native()["s"][0] == 3.0


def test_every_registered_primitive_is_a_frozen_dataclass():
    for name in _REGISTRY:
        primitive = resolve(name)
        assert is_dataclass(primitive), f"{name} is not a dataclass"
        with pytest.raises(FrozenInstanceError):
            setattr(primitive, "name", "mutated")  # noqa: B010


def test_every_registered_primitive_round_trips_through_pickle():
    # Primitives built by a metaprogramming helper rather than a class body
    # land in the wrong module and silently break any process-parallel or
    # cached use. Instantiating via the registry keeps this honest.
    for name in _REGISTRY:
        primitive = resolve(name)
        assert pickle.loads(pickle.dumps(primitive)) == primitive

# src/tusk/primitives/base.py
"""Primitive protocol.

A primitive builds narwhals expressions and never sees a value. That is what
lets feature computation push down into the backend engine. Featuretools' model
-- a callable over a materialized pandas Series -- is deliberately not used.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any, ClassVar, Sequence

import narwhals as nw

from tusk.dtypes import DtypeFamily


class Primitive(ABC):
    """Base class for every primitive.

    Attributes:
        name: Registry key, also the upper-cased stem of generated names.
        input_dtypes: One dtype family per input. Empty means the primitive
            takes no column input, e.g. ``count``.
        output_dtype: Fixed output dtype, or None to preserve the first input's.
        commutative: Whether argument order is irrelevant, so that only one of
            ``f(a, b)`` and ``f(b, a)`` is generated.
        stack_on_self: Whether this primitive may be applied to its own output.
        default_value: Value substituted for empty groups after a left join.
    """

    name: ClassVar[str]
    input_dtypes: ClassVar[tuple[DtypeFamily, ...]] = ()
    output_dtype: ClassVar[Any] = None
    commutative: ClassVar[bool] = False
    stack_on_self: ClassVar[bool] = True
    default_value: ClassVar[Any] = None

    @property
    def number_of_outputs(self) -> int:
        """How many columns this primitive produces."""
        return 1

    def return_dtype(self, input_dtypes: tuple[Any, ...]) -> Any:
        """Compute the output dtype without touching data.

        Args:
            input_dtypes: Dtypes of the input features, in order.

        Returns:
            The dtype of this primitive's output.
        """
        if self.output_dtype is not None:
            return self.output_dtype
        return input_dtypes[0]

    def generate_name(self, arg_names: Sequence[str]) -> str:
        """Build the display name for an application of this primitive.

        Args:
            arg_names: Names of the inputs. For a zero-input aggregation this
                is the child table's name, giving e.g. ``COUNT(transactions)``.

        Returns:
            The feature name.
        """
        return f"{self.name.upper()}({', '.join(arg_names)})"

    def output_names(self, base_name: str) -> tuple[str, ...]:
        """Expand a feature name into one name per output column.

        Args:
            base_name: The name from :meth:`generate_name`.

        Returns:
            One name per output column; indexed when there is more than one.
        """
        if self.number_of_outputs == 1:
            return (base_name,)
        return tuple(f"{base_name}[{i}]" for i in range(self.number_of_outputs))

    def outputs(self, *inputs: nw.Expr) -> tuple[nw.Expr, ...]:
        """Normalize :meth:`build` to a tuple of expressions.

        Args:
            *inputs: One expression per declared input.

        Returns:
            One expression per output column.
        """
        built = self.build(*inputs)
        if isinstance(built, (list, tuple)):
            return tuple(built)
        return (built,)

    @abstractmethod
    def build(self, *inputs: nw.Expr) -> nw.Expr | Sequence[nw.Expr]:
        """Build this primitive's narwhals expression.

        Args:
            *inputs: One expression per declared input.

        Returns:
            A single expression, or a sequence for multi-output primitives.
        """


class AggregationPrimitive(Primitive):
    """A primitive applied to a child table's rows, grouped by foreign key."""


class TransformPrimitive(Primitive):
    """A primitive applied row-wise within a single table.

    Attributes:
        order_dependent: Whether the expression needs an explicit ordering.
            Narwhals requires ``.over(order_by=...)`` for these on lazy
            backends, so tusk requires a ``row_creation_time`` on the table.
    """

    order_dependent: ClassVar[bool] = False

# src/tusk/primitives/registry.py
"""Name-to-primitive registry."""

from __future__ import annotations

from collections.abc import Iterable
from typing import TypeVar

from tusk.exceptions import PrimitiveError
from tusk.primitives.base import Primitive

_REGISTRY: dict[str, type[Primitive]] = {}


def register(cls: type[Primitive]) -> type[Primitive]:
    """Register a primitive class under its ``name``.

    Args:
        cls: The primitive class to register.

    Returns:
        The class unchanged, so this works as a decorator.

    Raises:
        PrimitiveError: If the name is already registered.
    """
    if cls.name in _REGISTRY and _REGISTRY[cls.name] is not cls:
        raise PrimitiveError(f"primitive name {cls.name!r} is already registered")
    _REGISTRY[cls.name] = cls
    return cls


def resolve(spec: str | Primitive) -> Primitive:
    """Turn a name or instance into a primitive instance.

    Args:
        spec: A registered primitive name, or an already-built instance.

    Returns:
        A primitive instance.

    Raises:
        PrimitiveError: If the name is not registered.
    """
    if isinstance(spec, Primitive):
        return spec
    try:
        return _REGISTRY[spec]()
    except KeyError:
        known = ", ".join(sorted(_REGISTRY))
        raise PrimitiveError(f"unknown primitive {spec!r}; available: {known}") from None


def resolve_all(specs: Iterable[str | Primitive]) -> tuple[Primitive, ...]:
    """Resolve a collection of names or instances.

    Args:
        specs: Names or instances.

    Returns:
        Primitive instances in the given order.
    """
    return tuple(resolve(spec) for spec in specs)

# src/tusk/primitives/__init__.py
"""Primitives: the expression builders DFS composes into features."""

from __future__ import annotations

from tusk.primitives.base import AggregationPrimitive, Primitive, TransformPrimitive
from tusk.primitives.registry import register, resolve, resolve_all

__all__ = [
    "AggregationPrimitive",
    "Primitive",
    "TransformPrimitive",
    "register",
    "resolve",
    "resolve_all",
]

# tests/test_primitives_aggregation.py
import narwhals as nw
import polars as pl
import pytest

from tusk.primitives.aggregation import AGG_DEFAULTS, Count, Quantiles
from tusk.primitives.registry import resolve


@pytest.fixture
def lf():
    return nw.from_native(
        pl.LazyFrame(
            {
                "g": [1, 1, 1, 2],
                "v": [1.0, 2.0, 6.0, 4.0],
                "b": [True, False, True, True],
                "s": ["a", "a", "b", "c"],
            }
        )
    )


def _agg(lf, primitive, column):
    exprs = primitive.outputs(nw.col(column))
    named = [e.alias(f"o{i}") for i, e in enumerate(exprs)]
    got = lf.group_by("g").agg(*named).sort("g").collect().to_native()
    return got


@pytest.mark.parametrize(
    ("name", "column", "expected"),
    [
        ("sum", "v", 9.0),
        ("mean", "v", 3.0),
        ("min", "v", 1.0),
        ("max", "v", 6.0),
        ("median", "v", 2.0),
        ("n_unique", "s", 2),
        ("percent_true", "b", pytest.approx(2 / 3)),
    ],
)
def test_aggregations_over_group_one(lf, name, column, expected):
    got = _agg(lf, resolve(name), column)
    assert got["o0"][0] == expected


def test_count_takes_no_column_input(lf):
    assert Count().input_dtypes == ()
    got = lf.group_by("g").agg(Count().outputs()[0].alias("n")).sort("g").collect().to_native()
    assert got["n"].to_list() == [3, 1]


def test_count_defaults_empty_groups_to_zero():
    assert Count().default_value == 0
    assert Count().stack_on_self is False


def test_count_and_n_unique_declare_int64(lf):
    got = lf.group_by("g").agg(
        Count().outputs()[0].alias("n"),
        resolve("n_unique").outputs(nw.col("s"))[0].alias("u"),
    ).collect()
    assert got.collect_schema()["n"] == nw.Int64
    assert got.collect_schema()["u"] == nw.Int64


def test_quantiles_is_multi_output(lf):
    q = Quantiles(qs=(0.0, 0.5, 1.0))
    assert q.number_of_outputs == 3
    assert q.output_names("QUANTILES(v)") == (
        "QUANTILES(v)[0]",
        "QUANTILES(v)[1]",
        "QUANTILES(v)[2]",
    )
    got = _agg(lf, q, "v")
    assert [got["o0"][0], got["o1"][0], got["o2"][0]] == [1.0, 2.0, 6.0]


def test_defaults_are_the_documented_set():
    assert AGG_DEFAULTS == ("count", "sum", "mean", "min", "max", "std", "n_unique")

# src/tusk/primitives/aggregation.py
"""Built-in aggregation primitives.

Every expression here is legal inside a lazy ``group_by().agg()``. Length-changing
expressions such as ``mode()`` are not -- narwhals rejects them on lazy frames --
which is why ``quantiles`` rather than ``n_most_common`` is the multi-output
primitive.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Sequence

import narwhals as nw

from tusk.dtypes import DtypeFamily as F
from tusk.primitives.base import AggregationPrimitive
from tusk.primitives.registry import register

AGG_DEFAULTS: tuple[str, ...] = ("count", "sum", "mean", "min", "max", "std", "n_unique")


@register
@dataclass(frozen=True)
class Count(AggregationPrimitive):
    """Number of child rows in the group."""

    name = "count"
    input_dtypes = ()
    output_dtype = nw.Int64
    default_value = 0
    stack_on_self = False

    def build(self) -> nw.Expr:
        """Build the row-count expression.

        Returns:
            A narwhals expression counting rows.
        """
        return nw.len().cast(nw.Int64)


@register
@dataclass(frozen=True)
class Sum(AggregationPrimitive):
    """Sum of a numeric column."""

    name = "sum"
    input_dtypes = (F.NUMERIC,)

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the sum expression.

        Args:
            expr: The column to sum.

        Returns:
            A narwhals expression.
        """
        return expr.sum()


@register
@dataclass(frozen=True)
class Mean(AggregationPrimitive):
    """Arithmetic mean of a numeric column."""

    name = "mean"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the mean expression.

        Args:
            expr: The column to average.

        Returns:
            A narwhals expression.
        """
        return expr.mean()


@register
@dataclass(frozen=True)
class Min(AggregationPrimitive):
    """Smallest value of a numeric column."""

    name = "min"
    input_dtypes = (F.NUMERIC,)

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the minimum expression.

        Args:
            expr: The column to reduce.

        Returns:
            A narwhals expression.
        """
        return expr.min()


@register
@dataclass(frozen=True)
class Max(AggregationPrimitive):
    """Largest value of a numeric column."""

    name = "max"
    input_dtypes = (F.NUMERIC,)

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the maximum expression.

        Args:
            expr: The column to reduce.

        Returns:
            A narwhals expression.
        """
        return expr.max()


@register
@dataclass(frozen=True)
class Std(AggregationPrimitive):
    """Sample standard deviation of a numeric column."""

    name = "std"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the standard-deviation expression.

        Args:
            expr: The column to reduce.

        Returns:
            A narwhals expression.
        """
        return expr.std()


@register
@dataclass(frozen=True)
class Median(AggregationPrimitive):
    """Median of a numeric column."""

    name = "median"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the median expression.

        Args:
            expr: The column to reduce.

        Returns:
            A narwhals expression.
        """
        return expr.median()


@register
@dataclass(frozen=True)
class NUnique(AggregationPrimitive):
    """Number of distinct values in a column."""

    name = "n_unique"
    input_dtypes = (F.ANY,)
    output_dtype = nw.Int64
    default_value = 0
    stack_on_self = False

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the distinct-count expression.

        Args:
            expr: The column to count distinct values of.

        Returns:
            A narwhals expression.
        """
        return expr.n_unique().cast(nw.Int64)


@register
@dataclass(frozen=True)
class PercentTrue(AggregationPrimitive):
    """Fraction of rows where a boolean column is true."""

    name = "percent_true"
    input_dtypes = (F.BOOLEAN,)
    output_dtype = nw.Float64

    def build(self, expr: nw.Expr) -> nw.Expr:
        """Build the true-fraction expression.

        Args:
            expr: The boolean column.

        Returns:
            A narwhals expression.
        """
        return expr.cast(nw.Int64).mean()


@register
@dataclass(frozen=True)
class Quantiles(AggregationPrimitive):
    """Several quantiles of a numeric column, one output column per quantile.

    Attributes:
        qs: The quantiles to compute, each in [0, 1].
    """

    name = "quantiles"
    input_dtypes = (F.NUMERIC,)
    output_dtype = nw.Float64

    qs: tuple[float, ...] = field(default=(0.25, 0.5, 0.75))

    @property
    def number_of_outputs(self) -> int:
        """One output column per requested quantile."""
        return len(self.qs)

    def build(self, expr: nw.Expr) -> Sequence[nw.Expr]:
        """Build one expression per quantile.

        Args:
            expr: The column to reduce.

        Returns:
            One narwhals expression per quantile.
        """
        return [expr.quantile(q, interpolation="linear") for q in self.qs]