Aggregations#

def agg(self, *to_agg: Union[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(
        ...     col("age").min().alias("min_age"),
        ...     col("age").max().alias("max_age"),
        ...     col("pet").count().alias("count"),
        ...     col("name").any_value(),
        ... )
        >>> grouped_df = grouped_df.sort("pet")
        >>> grouped_df.show()
        ╭──────┬─────────┬─────────┬────────┬────────╮
        │ pet  ┆ min_age ┆ max_age ┆ count  ┆ name   │
        │ ---  ┆ ---     ┆ ---     ┆ ---    ┆ ---    │
        │ Utf8 ┆ Int64   ┆ Int64   ┆ UInt64 ┆ Utf8   │
        ╞══════╪═════════╪═════════╪════════╪════════╡
        │ 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 agg_concat(self, *cols: ColumnInputType) -> DataFrame:
    """Performs grouped concat on this GroupedDataFrame.

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

def agg_list(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.agg_list, cols, self.group_by)

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

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

    Returns:
        DataFrame: DataFrame with grouped set per column.
    """
    return self.df._apply_agg_fn(Expression.agg_set, cols, 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 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.udf(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               │
        │ ---   ┆ ---                │
        │ Utf8  ┆ 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 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) -> DataFrame:
    """Performs grouped standard deviation on this GroupedDataFrame.

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

    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             │
        │ ---  ┆ ---               │
        │ Utf8 ┆ Float64           │
        ╞══════╪═══════════════════╡
        │ a    ┆ 0.816496580927726 │
        ├╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
        │ b    ┆ 0                 │
        ╰──────┴───────────────────╯
        <BLANKLINE>
        (Showing first 2 of 2 rows)

    """
    return self.df._apply_agg_fn(Expression.stddev, 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)