Aggregations#

def agg(self, *to_agg: Expression | Iterable[Expression]) -> DataFrame:
    """Perform aggregations on this GroupedDataFrame. Allows for mixed aggregations.

    Args:
        *to_agg (Union[Expression, Iterable[Expression]]): aggregation expressions

    Returns:
        DataFrame: DataFrame with grouped aggregations

    Examples:
        >>> import daft
        >>> from daft import col
        >>> df = daft.from_pydict(
        ...     {
        ...         "pet": ["cat", "dog", "dog", "cat"],
        ...         "age": [1, 2, 3, 4],
        ...         "name": ["Alex", "Jordan", "Sam", "Riley"],
        ...     }
        ... )
        >>> grouped_df = df.groupby("pet").agg(
        ...     df["age"].min().alias("min_age"),
        ...     df["age"].max().alias("max_age"),
        ...     df["pet"].count().alias("count"),
        ...     df["name"].any_value(),
        ... )
        >>> grouped_df = grouped_df.sort("pet")
        >>> grouped_df.show()
        ╭────────┬─────────┬─────────┬────────┬────────╮
        │ pet    ┆ min_age ┆ max_age ┆ count  ┆ name   │
        │ ---    ┆ ---     ┆ ---     ┆ ---    ┆ ---    │
        │ String ┆ Int64   ┆ Int64   ┆ UInt64 ┆ String │
        ╞════════╪═════════╪═════════╪════════╪════════╡
        │ cat    ┆ 1       ┆ 4       ┆ 2      ┆ Alex   │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
        │ dog    ┆ 2       ┆ 3       ┆ 2      ┆ Jordan │
        ╰────────┴─────────┴─────────┴────────┴────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    to_agg_list = (
        list(to_agg[0])
        if (len(to_agg) == 1 and not isinstance(to_agg[0], Expression))
        else list(typing.cast("tuple[Expression]", to_agg))
    )

    for expr in to_agg_list:
        if not isinstance(expr, Expression):
            raise ValueError(f"GroupedDataFrame.agg() only accepts expression type, received: {type(expr)}")

    return self.df._agg(to_agg_list, group_by=self.group_by)

def any_value(self, *cols: ColumnInputType) -> DataFrame:
    """Returns an arbitrary value on this GroupedDataFrame.

    Values for each column are not guaranteed to be from the same row.

    Args:
        *cols (Union[str, Expression]): columns to get

    Returns:
        DataFrame: DataFrame with any values.
    """
    return self.df._apply_agg_fn(Expression.any_value, cols, self.group_by)

def count(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped count on this GroupedDataFrame.

    Returns:
        DataFrame: DataFrame with grouped count per column.
    """
    return self.df._apply_agg_fn(Expression.count, cols, self.group_by)

def count_distinct(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped count of distinct values on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to count distinct values

    Returns:
        DataFrame: DataFrame with grouped count of distinct values per column.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"keys": ["a", "a", "a", "b", "b", "b"], "vals": [1, 1, 2, 3, 3, 3]})
        >>> df = df.groupby("keys").count_distinct("vals")
        >>> df = df.sort("keys")
        >>> df.show()
        ╭────────┬────────╮
        │ keys   ┆ vals   │
        │ ---    ┆ ---    │
        │ String ┆ UInt64 │
        ╞════════╪════════╡
        │ a      ┆ 2      │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
        │ b      ┆ 1      │
        ╰────────┴────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._apply_agg_fn(Expression.count_distinct, cols, self.group_by)

def list_agg(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped list on this GroupedDataFrame.

    Returns:
        DataFrame: DataFrame with grouped list per column.
    """
    return self.df._apply_agg_fn(Expression.list_agg, cols, self.group_by)

def list_agg_distinct(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped list distinct on this GroupedDataFrame (ignoring nulls).

    Args:
        *cols (Union[str, Expression]): columns to form into a set

    Returns:
        DataFrame: DataFrame with grouped list distinct per column.
    """
    return self.df._apply_agg_fn(Expression.list_agg_distinct, cols, self.group_by)

def map_groups(self, udf: Expression) -> DataFrame:
    """Apply a user-defined function to each group. The name of the resultant column will default to the name of the first input column.

    Args:
        udf (Expression): User-defined function to apply to each group.

    Returns:
        DataFrame: DataFrame with grouped aggregations

    Examples:
        >>> import daft, statistics
        >>>
        >>> df = daft.from_pydict({"group": ["a", "a", "a", "b", "b", "b"], "data": [1, 20, 30, 4, 50, 600]})
        >>>
        >>> @daft.func.batch(return_dtype=daft.DataType.float64())
        ... def std_dev(data):
        ...     return [statistics.stdev(data)]
        >>>
        >>> df = df.groupby("group").map_groups(std_dev(df["data"]))
        >>> df = df.sort("group")
        >>> df.show()
        ╭────────┬────────────────────╮
        │ group  ┆ data               │
        │ ---    ┆ ---                │
        │ String ┆ Float64            │
        ╞════════╪════════════════════╡
        │ a      ┆ 14.730919862656235 │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ b      ┆ 331.62026476076517 │
        ╰────────┴────────────────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._map_groups(udf, group_by=self.group_by)

def max(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped max on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to max

    Returns:
        DataFrame: DataFrame with grouped max.
    """
    return self.df._apply_agg_fn(Expression.max, cols, self.group_by)

def mean(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped mean on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to mean

    Returns:
        DataFrame: DataFrame with grouped mean.
    """
    return self.df._apply_agg_fn(Expression.mean, cols, self.group_by)

def min(self, *cols: ColumnInputType) -> DataFrame:
    """Perform grouped min on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to min

    Returns:
        DataFrame: DataFrame with grouped min.
    """
    return self.df._apply_agg_fn(Expression.min, cols, self.group_by)

def product(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped product on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to product

    Returns:
        DataFrame: DataFrame with grouped products.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"keys": ["a", "a", "a", "b"], "col_a": [1, 2, 3, 100]})
        >>> df = df.groupby("keys").product()
        >>> df = df.sort("keys")
        >>> df.show()
        ╭────────┬───────╮
        │ keys   ┆ col_a │
        │ ---    ┆ ---   │
        │ String ┆ Int64 │
        ╞════════╪═══════╡
        │ a      ┆ 6     │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
        │ b      ┆ 100   │
        ╰────────┴───────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._apply_agg_fn(Expression.product, cols, self.group_by)

def skew(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped skew on this GroupedDataFrame.

    Returns:
        DataFrame: DataFrame with the grouped skew per column.
    """
    return self.df._apply_agg_fn(Expression.skew, cols, self.group_by)

def stddev(self, *cols: ColumnInputType, ddof: int = 1) -> DataFrame:
    """Performs grouped standard deviation on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to stddev
        ddof (int): Delta degrees of freedom used in the denominator `N - ddof`.
            Defaults to 1 (sample standard deviation).

    Returns:
        DataFrame: DataFrame with grouped standard deviation.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"keys": ["a", "a", "a", "b"], "col_a": [0, 1, 2, 100]})
        >>> df = df.groupby("keys").stddev()
        >>> df = df.sort("keys")
        >>> df.show()
        ╭────────┬─────────╮
        │ keys   ┆ col_a   │
        │ ---    ┆ ---     │
        │ String ┆ Float64 │
        ╞════════╪═════════╡
        │ a      ┆ 1       │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ b      ┆ None    │
        ╰────────┴─────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._apply_agg_fn(lambda expr: Expression.stddev(expr, ddof), cols, self.group_by)

def string_agg(self, *cols: ColumnInputType, delimiter: str | None = None) -> DataFrame:
    """Performs grouped string concat on this GroupedDataFrame.

    Returns:
        DataFrame: DataFrame with grouped string concatenated per column.
    """
    return self.df._apply_agg_fn(lambda expr: Expression.string_agg(expr, delimiter=delimiter), cols, self.group_by)

def sum(self, *cols: ColumnInputType) -> DataFrame:
    """Perform grouped sum on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to sum

    Returns:
        DataFrame: DataFrame with grouped sums.
    """
    return self.df._apply_agg_fn(Expression.sum, cols, self.group_by)

def var(self, *cols: ColumnInputType, ddof: int = 1) -> DataFrame:
    """Performs grouped variance on this GroupedDataFrame.

    Args:
        *cols (Union[str, Expression]): columns to compute variance for
        ddof (int): Delta degrees of freedom used in the denominator `N - ddof`.
            Defaults to 1 (sample variance).

    Returns:
        DataFrame: DataFrame with grouped variance.

    Examples:
        >>> import daft
        >>> df = daft.from_pydict({"keys": ["a", "a", "a", "b"], "col_a": [0, 1, 2, 100]})
        >>> df = df.groupby("keys").var()
        >>> df = df.sort("keys")
        >>> df.show()
        ╭────────┬─────────╮
        │ keys   ┆ col_a   │
        │ ---    ┆ ---     │
        │ String ┆ Float64 │
        ╞════════╪═════════╡
        │ a      ┆ 1       │
        ├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
        │ b      ┆ None    │
        ╰────────┴─────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._apply_agg_fn(lambda expr: Expression.var(expr, ddof), cols, self.group_by)