Class Count

python/ray/data/aggregate.py:392–458

Defines count aggregation. Example: .. testcode:: import ray from ray.data.aggregate import Count ds = ray.data.range(100) # Schema: {'id': int64} ds = ds.add_column( "group_key", lambda batch: batch["id"].as

@PublicAPI
class Count(AggregateFnV2[int, int]):
    """Defines count aggregation.

    Example:

        .. testcode::

            import ray
            from ray.data.aggregate import Count

            ds = ray.data.range(100)
            # Schema: {'id': int64}
            ds = ds.add_column(
                "group_key", lambda batch: batch["id"].astype("int64") % 3
            )
            # Schema: {'id': int64, 'group_key': int64}

            # Counting all rows:
            result = ds.aggregate(Count())
            # result: {'count()': 100}


            # Counting all rows per group:
            result = ds.groupby("group_key").aggregate(Count(on="id")).take_all()
            # result: [{'group_key': 0, 'count(id)': 34},
            #          {'group_key': 1, 'count(id)': 33},
            #          {'group_key': 2, 'count(id)': 33}]


    Args:
        on: Optional name of the column to count values on. If None, counts rows.
        ignore_nulls: Whether to ignore null values when counting. Only applies if
            `on` is specified. Default is `False` which means `Count()` on a column
            will count nulls by default. To match pandas default behavior of not counting nulls,
            set `ignore_nulls=True`.
        alias_name: Optional name for the resulting column.
    """

    def __init__(
        self,
        on: Optional[str] = None,
        ignore_nulls: bool = False,
        alias_name: Optional[str] = None,
    ):
        super().__init__(
            alias_name if alias_name else f"count({on or ''})",
            on=on,
            ignore_nulls=ignore_nulls,
            zero_factory=lambda: 0,
        )

    def aggregate_block(self, block: Block) -> int:
        block_accessor = BlockAccessor.for_block(block)

        if self._target_col_name is None:
            # In case of global count, simply fetch number of rows
            return block_accessor.num_rows()