aggregations.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import abc
import dataclasses
import typing
from typing import Callable, ClassVar, Iterable, Optional, TYPE_CHECKING

import numpy as np
import pandas as pd
import pyarrow as pa

from bigframes.core import agg_expressions
import bigframes.dtypes as dtypes
import bigframes.operations.type as signatures

if TYPE_CHECKING:
    from bigframes.core import expression


@dataclasses.dataclass(frozen=True)
class WindowOp:
    @property
    def skips_nulls(self):
        """Whether the window op skips null rows."""
        return True

    @property
    def nulls_count_for_min_values(self) -> bool:
        """Whether null values count for min_values."""
        return not self.skips_nulls

    @property
    def implicitly_inherits_order(self):
        """
        Whether the operator implicitly inherits the underlying array order, should it exist.

        Notably, rank operations do not want to inherit ordering. Even order-independent operations
        may inherit order when needed for row bounds.
        """
        return True

    @property
    def order_independent(self):
        """
        True if the output of the operator does not depend on the ordering of input rows.

        Aggregation functions are usually order independent, except array_agg, string_agg.

        Navigation functions are a notable case that are not order independent.
        """
        return False

    @abc.abstractmethod
    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        ...


@dataclasses.dataclass(frozen=True)
class NullaryWindowOp(WindowOp):
    @property
    def arguments(self) -> int:
        return 0


@dataclasses.dataclass(frozen=True)
class UnaryWindowOp(WindowOp):
    @property
    def arguments(self) -> int:
        return 1

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return input_types[0]


@dataclasses.dataclass(frozen=True)
class AggregateOp(WindowOp):
    """Aggregate ops can be applied with or without a window clause."""

    @property
    @abc.abstractmethod
    def name(self) -> str:
        ...

    @property
    @abc.abstractmethod
    def arguments(self) -> int:
        ...

    @property
    def order_independent(self):
        return True

    @property
    def uses_total_row_ordering(self):
        return False


@dataclasses.dataclass(frozen=True)
class NullaryAggregateOp(AggregateOp, NullaryWindowOp):
    @property
    def arguments(self) -> int:
        return 0

    def as_expr(
        self,
        *exprs: typing.Union[str, expression.Expression],
    ) -> agg_expressions.NullaryAggregation:
        from bigframes.core import agg_expressions

        return agg_expressions.NullaryAggregation(self)


@dataclasses.dataclass(frozen=True)
class UnaryAggregateOp(AggregateOp, UnaryWindowOp):
    @property
    def arguments(self) -> int:
        return 1

    def as_expr(
        self,
        *exprs: typing.Union[str, expression.Expression],
    ) -> agg_expressions.UnaryAggregation:
        from bigframes.core import agg_expressions
        from bigframes.operations.base_ops import _convert_expr_input

        # Keep this in sync with output_type and compilers
        inputs: list[expression.Expression] = []

        for expr in exprs:
            inputs.append(_convert_expr_input(expr))
        return agg_expressions.UnaryAggregation(
            self,
            inputs[0],
        )


@dataclasses.dataclass(frozen=True)
class BinaryAggregateOp(AggregateOp):
    @property
    def arguments(self) -> int:
        return 2

    def as_expr(
        self,
        *exprs: typing.Union[str, expression.Expression],
    ) -> agg_expressions.BinaryAggregation:
        from bigframes.core import agg_expressions
        from bigframes.operations.base_ops import _convert_expr_input

        # Keep this in sync with output_type and compilers
        inputs: list[expression.Expression] = []

        for expr in exprs:
            inputs.append(_convert_expr_input(expr))

        return agg_expressions.BinaryAggregation(self, inputs[0], inputs[1])


@dataclasses.dataclass(frozen=True)
class SizeOp(NullaryAggregateOp):
    name: ClassVar[str] = "size"

    def output_type(self, *input_types: dtypes.ExpressionType):
        return dtypes.INT_DTYPE


# TODO: Remove this temporary hack once nullary ops are better supported in APIs
@dataclasses.dataclass(frozen=True)
class SizeUnaryOp(UnaryAggregateOp):
    name: ClassVar[str] = "size"

    def output_type(self, *input_types: dtypes.ExpressionType):
        return dtypes.INT_DTYPE


@dataclasses.dataclass(frozen=True)
class SumOp(UnaryAggregateOp):
    name: ClassVar[str] = "sum"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] == dtypes.TIMEDELTA_DTYPE:
            return dtypes.TIMEDELTA_DTYPE

        if dtypes.is_numeric(input_types[0]):
            if pd.api.types.is_bool_dtype(input_types[0]):
                return dtypes.INT_DTYPE
            return input_types[0]

        raise TypeError(f"Type {input_types[0]} is not numeric or timedelta")


@dataclasses.dataclass(frozen=True)
class MedianOp(UnaryAggregateOp):
    name: ClassVar[str] = "median"

    @property
    def order_independent(self) -> bool:
        return True

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        # These will change if median is changed to exact implementation.
        if not dtypes.is_orderable(input_types[0]):
            raise TypeError(f"Type {input_types[0]} is not orderable")
        if pd.api.types.is_bool_dtype(input_types[0]):
            return dtypes.INT_DTYPE
        else:
            return input_types[0]


@dataclasses.dataclass(frozen=True)
class QuantileOp(UnaryAggregateOp):
    q: float
    should_floor_result: bool = False

    @property
    def name(self):
        return f"{int(self.q * 100)}%"

    @property
    def order_independent(self) -> bool:
        return True

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] == dtypes.TIMEDELTA_DTYPE:
            return dtypes.TIMEDELTA_DTYPE
        return signatures.UNARY_REAL_NUMERIC.output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class ApproxQuartilesOp(UnaryAggregateOp):
    quartile: int

    @property
    def name(self):
        return f"{self.quartile * 25}%"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if not dtypes.is_orderable(input_types[0]):
            raise TypeError(f"Type {input_types[0]} is not orderable")
        return input_types[0]


@dataclasses.dataclass(frozen=True)
class ApproxTopCountOp(UnaryAggregateOp):
    name: typing.ClassVar[str] = "approx_top_count"
    number: int

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if not dtypes.is_orderable(input_types[0]):
            raise TypeError(f"Type {input_types[0]} is not orderable")

        input_type = input_types[0]
        fields = [
            pa.field("value", dtypes.bigframes_dtype_to_arrow_dtype(input_type)),
            pa.field("count", pa.int64()),
        ]
        return pd.ArrowDtype(pa.list_(pa.struct(fields)))


@dataclasses.dataclass(frozen=True)
class MeanOp(UnaryAggregateOp):
    name: ClassVar[str] = "mean"

    should_floor_result: bool = False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] == dtypes.TIMEDELTA_DTYPE:
            return dtypes.TIMEDELTA_DTYPE
        return signatures.UNARY_REAL_NUMERIC.output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class ProductOp(UnaryAggregateOp):
    name: ClassVar[str] = "product"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_numeric, dtypes.FLOAT_DTYPE, "numeric"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class MaxOp(UnaryAggregateOp):
    name: ClassVar[str] = "max"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.TypePreserving(dtypes.is_orderable, "orderable").output_type(
            input_types[0]
        )


@dataclasses.dataclass(frozen=True)
class MinOp(UnaryAggregateOp):
    name: ClassVar[str] = "min"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.TypePreserving(dtypes.is_orderable, "orderable").output_type(
            input_types[0]
        )


@dataclasses.dataclass(frozen=True)
class StdOp(UnaryAggregateOp):
    name: ClassVar[str] = "std"

    should_floor_result: bool = False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] == dtypes.TIMEDELTA_DTYPE:
            return dtypes.TIMEDELTA_DTYPE

        return signatures.FixedOutputType(
            dtypes.is_numeric, dtypes.FLOAT_DTYPE, "numeric"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class VarOp(UnaryAggregateOp):
    name: ClassVar[str] = "var"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_numeric, dtypes.FLOAT_DTYPE, "numeric"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class PopVarOp(UnaryAggregateOp):
    name: ClassVar[str] = "popvar"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_numeric, dtypes.FLOAT_DTYPE, "numeric"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class CountOp(UnaryAggregateOp):
    name: ClassVar[str] = "count"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            lambda x: True, dtypes.INT_DTYPE, ""
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class ArrayAggOp(UnaryAggregateOp):
    name: ClassVar[str] = "arrayagg"

    @property
    def order_independent(self):
        return False

    @property
    def skips_nulls(self):
        return True

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return dtypes.list_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class StringAggOp(UnaryAggregateOp):
    name: ClassVar[str] = "string_agg"
    sep: str = ","

    @property
    def order_independent(self):
        return False

    @property
    def skips_nulls(self):
        return True

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] != dtypes.STRING_DTYPE:
            raise TypeError(f"Type {input_types[0]} is not string-like")
        return dtypes.STRING_DTYPE


@dataclasses.dataclass(frozen=True)
class CutOp(UnaryWindowOp):
    # TODO: Unintuitive, refactor into multiple ops?
    bins: typing.Union[int, Iterable]
    right: Optional[bool]
    labels: typing.Union[bool, Iterable[str], None]

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if self.labels is False:
            return dtypes.INT_DTYPE
        elif isinstance(self.labels, Iterable):
            return dtypes.STRING_DTYPE
        else:
            # Assumption: buckets use same numeric type
            if isinstance(self.bins, int):
                interval_dtype = pa.float64()
            elif len(list(self.bins)) == 0:
                interval_dtype = pa.int64()
            else:
                interval_dtype = dtypes.infer_literal_arrow_type(list(self.bins)[0][0])
            pa_type = pa.struct(
                [
                    pa.field(
                        "left_exclusive" if self.right else "left_inclusive",
                        interval_dtype,
                        nullable=True,
                    ),
                    pa.field(
                        "right_inclusive" if self.right else "right_exclusive",
                        interval_dtype,
                        nullable=True,
                    ),
                ]
            )

            return pd.ArrowDtype(pa_type)

    @property
    def order_independent(self):
        return True


@dataclasses.dataclass(frozen=True)
class QcutOp(UnaryWindowOp):  # bucket op
    quantiles: typing.Union[int, typing.Tuple[float, ...]]

    @property
    def name(self):
        return f"qcut-{self.quantiles}"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_orderable, dtypes.INT_DTYPE, "orderable"
        ).output_type(input_types[0])

    @property
    def order_independent(self):
        return True


@dataclasses.dataclass(frozen=True)
class NuniqueOp(UnaryAggregateOp):
    name: ClassVar[str] = "nunique"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return dtypes.INT_DTYPE


@dataclasses.dataclass(frozen=True)
class AnyValueOp(UnaryAggregateOp):
    # Warning: only use if all values are equal. Non-deterministic otherwise.
    # Do not expose to users. For special cases only (e.g. pivot).
    name: ClassVar[str] = "any_value"

    @property
    def skips_nulls(self):
        return True


# This should really by a NullaryWindowOp, but APIs don't support that yet.
@dataclasses.dataclass(frozen=True)
class RowNumberOp(NullaryWindowOp):
    name: ClassVar[str] = "rownumber"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return dtypes.INT_DTYPE


@dataclasses.dataclass(frozen=True)
class RankOp(UnaryWindowOp):
    name: ClassVar[str] = "rank"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return dtypes.INT_DTYPE

    @property
    def implicitly_inherits_order(self):
        return False


@dataclasses.dataclass(frozen=True)
class DenseRankOp(UnaryWindowOp):
    name: ClassVar[str] = "dense_rank"

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return dtypes.INT_DTYPE

    @property
    def implicitly_inherits_order(self):
        return False


@dataclasses.dataclass(frozen=True)
class FirstOp(UnaryWindowOp):
    name: ClassVar[str] = "first"


@dataclasses.dataclass(frozen=True)
class FirstNonNullOp(UnaryWindowOp):
    @property
    def skips_nulls(self):
        return False

    @property
    def nulls_count_for_min_values(self) -> bool:
        return False


@dataclasses.dataclass(frozen=True)
class LastOp(UnaryWindowOp):
    name: ClassVar[str] = "last"


@dataclasses.dataclass(frozen=True)
class LastNonNullOp(UnaryWindowOp):
    @property
    def skips_nulls(self):
        return False

    @property
    def nulls_count_for_min_values(self) -> bool:
        return False


@dataclasses.dataclass(frozen=True)
class ShiftOp(UnaryWindowOp):
    periods: int

    @property
    def skips_nulls(self):
        return False


@dataclasses.dataclass(frozen=True)
class DiffOp(UnaryWindowOp):
    name: ClassVar[str] = "diff"
    periods: int

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if dtypes.is_date_like(input_types[0]):
            return dtypes.TIMEDELTA_DTYPE
        return super().output_type(*input_types)


@dataclasses.dataclass(frozen=True)
class TimeSeriesDiffOp(UnaryWindowOp):
    periods: int

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if dtypes.is_datetime_like(input_types[0]):
            return dtypes.TIMEDELTA_DTYPE
        raise TypeError(f"expect datetime-like types, but got {input_types[0]}")


@dataclasses.dataclass(frozen=True)
class DateSeriesDiffOp(UnaryWindowOp):
    periods: int

    @property
    def skips_nulls(self):
        return False

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        if input_types[0] == dtypes.DATE_DTYPE:
            return dtypes.TIMEDELTA_DTYPE
        raise TypeError(f"expect date type, but got {input_types[0]}")


@dataclasses.dataclass(frozen=True)
class AllOp(UnaryAggregateOp):
    name: ClassVar[str] = "all"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_bool_coercable, dtypes.BOOL_DTYPE, "convertible to boolean"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class AnyOp(UnaryAggregateOp):
    name: ClassVar[str] = "any"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.FixedOutputType(
            dtypes.is_bool_coercable, dtypes.BOOL_DTYPE, "convertible to boolean"
        ).output_type(input_types[0])


@dataclasses.dataclass(frozen=True)
class CorrOp(BinaryAggregateOp):
    name: ClassVar[str] = "corr"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.BINARY_REAL_NUMERIC.output_type(
            input_types[0], input_types[1]
        )


@dataclasses.dataclass(frozen=True)
class CovOp(BinaryAggregateOp):
    name: ClassVar[str] = "cov"

    def output_type(self, *input_types: dtypes.ExpressionType) -> dtypes.ExpressionType:
        return signatures.BINARY_REAL_NUMERIC.output_type(
            input_types[0], input_types[1]
        )


size_op = SizeOp()
sum_op = SumOp()
mean_op = MeanOp()
median_op = MedianOp()
product_op = ProductOp()
max_op = MaxOp()
min_op = MinOp()
std_op = StdOp()
var_op = VarOp()
count_op = CountOp()
nunique_op = NuniqueOp()
rank_op = RankOp()
dense_rank_op = DenseRankOp()
all_op = AllOp()
any_op = AnyOp()
first_op = FirstOp()


# TODO: Alternative names and lookup from numpy function objects
_STRING_TO_AGG_OP: typing.Dict[
    str, typing.Union[UnaryAggregateOp, NullaryAggregateOp]
] = {
    op.name: op
    for op in [
        sum_op,
        mean_op,
        median_op,
        product_op,
        max_op,
        min_op,
        std_op,
        var_op,
        count_op,
        all_op,
        any_op,
        nunique_op,
        ApproxQuartilesOp(1),
        ApproxQuartilesOp(2),
        ApproxQuartilesOp(3),
    ]
    + [
        # Add size_op separately to avoid Mypy type inference errors.
        size_op,
    ]
}

_CALLABLE_TO_AGG_OP: typing.Dict[
    Callable, typing.Union[UnaryAggregateOp, NullaryAggregateOp]
] = {
    np.sum: sum_op,
    np.mean: mean_op,
    np.median: median_op,
    np.prod: product_op,
    np.max: max_op,
    np.min: min_op,
    np.std: std_op,
    np.var: var_op,
    np.all: all_op,
    np.any: any_op,
    np.unique: nunique_op,
    # TODO(b/443252872): Solve
    # list: ArrayAggOp(),
    np.size: size_op,
}


def lookup_agg_func(
    key,
) -> tuple[typing.Union[UnaryAggregateOp, NullaryAggregateOp], str]:
    if key in _STRING_TO_AGG_OP:
        return (_STRING_TO_AGG_OP[key], key)
    if key in _CALLABLE_TO_AGG_OP:
        return (_CALLABLE_TO_AGG_OP[key], key.__name__)
    else:
        raise ValueError(f"Unrecognize aggregate function: {key}")