wehub-resource-sync
1885 lines
71 KiB
Python
1885 lines
71 KiB
Python
import abc
|
|
import enum
|
|
import math
|
|
import pickle
|
|
import re
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
Callable,
|
|
Collection,
|
|
Dict,
|
|
Generic,
|
|
List,
|
|
Optional,
|
|
Protocol,
|
|
Set,
|
|
TypeVar,
|
|
Union,
|
|
)
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
|
|
from ray.data._internal.util import is_null
|
|
from ray.data.block import (
|
|
Block,
|
|
BlockAccessor,
|
|
BlockColumn,
|
|
BlockColumnAccessor,
|
|
KeyType,
|
|
)
|
|
from ray.util.annotations import Deprecated, PublicAPI
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.dataset import Schema
|
|
|
|
|
|
class _SupportsRichComparison(Protocol):
|
|
def __lt__(self, other: Any) -> bool:
|
|
...
|
|
|
|
def __le__(self, other: Any) -> bool:
|
|
...
|
|
|
|
def __gt__(self, other: Any) -> bool:
|
|
...
|
|
|
|
def __ge__(self, other: Any) -> bool:
|
|
...
|
|
|
|
|
|
AccumulatorType = TypeVar("AccumulatorType")
|
|
SupportsRichComparisonType = TypeVar(
|
|
"SupportsRichComparisonType", bound=_SupportsRichComparison
|
|
)
|
|
AggOutputType = TypeVar("AggOutputType")
|
|
|
|
_AGGREGATION_NAME_PATTERN = re.compile(r"^([^(]+)(?:\(.*\))?$")
|
|
|
|
|
|
@Deprecated(message="AggregateFn is deprecated, please use AggregateFnV2")
|
|
@PublicAPI
|
|
class AggregateFn:
|
|
"""NOTE: THIS IS DEPRECATED, PLEASE USE :class:`AggregateFnV2` INSTEAD
|
|
|
|
Defines how to perform a custom aggregation in Ray Data.
|
|
|
|
`AggregateFn` instances are passed to a Dataset's ``.aggregate(...)`` method to
|
|
specify the steps required to transform and combine rows sharing the same key.
|
|
This enables implementing custom aggregators beyond the standard
|
|
built-in options like Sum, Min, Max, Mean, etc.
|
|
|
|
Args:
|
|
init: Function that creates an initial aggregator for each group. Receives a key
|
|
(the group key) and returns the initial accumulator state (commonly 0,
|
|
an empty list, or an empty dictionary).
|
|
merge: Function that merges two accumulators generated by different workers
|
|
into one accumulator.
|
|
name: An optional display name for the aggregator. Useful for debugging.
|
|
accumulate_row: Function that processes an individual row. It receives the current
|
|
accumulator and a row, then returns an updated accumulator. Cannot be
|
|
used if `accumulate_block` is provided.
|
|
accumulate_block: Function that processes an entire block of rows at once. It receives the
|
|
current accumulator and a block of rows, then returns an updated accumulator.
|
|
This allows for vectorized operations. Cannot be used if `accumulate_row`
|
|
is provided.
|
|
finalize: Function that finishes the aggregation by transforming the final
|
|
accumulator state into the desired output. For example, if your
|
|
accumulator is a list of items, you may want to compute a statistic
|
|
from the list. If not provided, the final accumulator state is returned
|
|
as-is.
|
|
|
|
Example:
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import AggregateFn
|
|
|
|
# A simple aggregator that counts how many rows there are per group
|
|
count_agg = AggregateFn(
|
|
init=lambda k: 0,
|
|
accumulate_row=lambda counter, row: counter + 1,
|
|
merge=lambda c1, c2: c1 + c2,
|
|
name="custom_count"
|
|
)
|
|
ds = ray.data.from_items([{"group": "A"}, {"group": "B"}, {"group": "A"}])
|
|
result = ds.groupby("group").aggregate(count_agg).take_all()
|
|
# result: [{'group': 'A', 'custom_count': 2}, {'group': 'B', 'custom_count': 1}]
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
init: Callable[[KeyType], AccumulatorType],
|
|
merge: Callable[[AccumulatorType, AccumulatorType], AccumulatorType],
|
|
name: str,
|
|
accumulate_row: Callable[
|
|
[AccumulatorType, Dict[str, Any]], AccumulatorType
|
|
] = None,
|
|
accumulate_block: Callable[[AccumulatorType, Block], AccumulatorType] = None,
|
|
finalize: Optional[Callable[[AccumulatorType], AggOutputType]] = None,
|
|
):
|
|
if (accumulate_row is None and accumulate_block is None) or (
|
|
accumulate_row is not None and accumulate_block is not None
|
|
):
|
|
raise ValueError(
|
|
"Exactly one of accumulate_row or accumulate_block must be provided."
|
|
)
|
|
|
|
if accumulate_block is None:
|
|
|
|
def accumulate_block(a: AccumulatorType, block: Block) -> AccumulatorType:
|
|
block_acc = BlockAccessor.for_block(block)
|
|
for r in block_acc.iter_rows(public_row_format=False):
|
|
a = accumulate_row(a, r)
|
|
return a
|
|
|
|
if not isinstance(name, str):
|
|
raise TypeError("`name` must be provided.")
|
|
|
|
if finalize is None:
|
|
finalize = lambda a: a # noqa: E731
|
|
|
|
self.name = name
|
|
self.init = init
|
|
self.merge = merge
|
|
self.accumulate_block = accumulate_block
|
|
self.finalize = finalize
|
|
|
|
def _validate(self, schema: Optional["Schema"]) -> None:
|
|
"""Raise an error if this cannot be applied to the given schema."""
|
|
pass
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
"""Return the PyArrow ``Field`` this aggregator produces.
|
|
|
|
Returns ``None`` by default; subclasses that know their output
|
|
type override. Used by ``Aggregate.infer_schema()`` to compute
|
|
the post-aggregation schema without executing the plan. When
|
|
any aggregator returns ``None``, the entire ``Aggregate.infer_schema``
|
|
returns ``None`` (callers fall back to a ``limit(1)`` execution).
|
|
"""
|
|
return None
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class AggregateFnV2(AggregateFn, abc.ABC, Generic[AccumulatorType, AggOutputType]):
|
|
"""Provides an interface to implement efficient aggregations to be applied
|
|
to the dataset.
|
|
|
|
`AggregateFnV2` instances are passed to a Dataset's ``.aggregate(...)`` method to
|
|
perform distributed aggregations. To create a custom aggregation, you should subclass
|
|
`AggregateFnV2` and implement the `aggregate_block` and `combine` methods.
|
|
The `finalize` method can also be overridden if the final accumulated state
|
|
needs further transformation.
|
|
|
|
Aggregation follows these steps:
|
|
|
|
1. **Initialization**: For each group (if grouping) or for the entire dataset,
|
|
an initial accumulator is created using `zero_factory`.
|
|
2. **Block Aggregation**: The `aggregate_block` method is applied to
|
|
each block independently, producing a partial aggregation result for that block.
|
|
3. **Combination**: The `combine` method is used to merge these partial
|
|
results (or an existing accumulated result with a new partial result)
|
|
into a single, combined accumulator.
|
|
4. **Finalization**: Optionally, the `finalize` method transforms the
|
|
final combined accumulator into the desired output format.
|
|
|
|
Generic Type Parameters:
|
|
This class is parameterized by two type variables:
|
|
|
|
- ``AccumulatorType``: The type of the intermediate state (accumulator) used
|
|
during aggregation. This is what `aggregate_block` returns, what `combine`
|
|
takes as inputs and returns, and what `finalize` receives. For simple
|
|
aggregations like `Sum`, this might just be a numeric type. For more complex
|
|
aggregations like `Mean`, this could be a composite type like
|
|
``List[Union[int, float]]`` representing ``[sum, count]``.
|
|
|
|
- ``AggOutputType``: The type of the final result after `finalize` is called.
|
|
This is what gets written to the output dataset. For `Sum`, this is the
|
|
same as the accumulator type (a number). For `Mean`, the accumulator is
|
|
``[sum, count]`` but the output is a single ``float`` (the computed mean).
|
|
|
|
Examples of type parameterization in built-in aggregations::
|
|
|
|
Count(AggregateFnV2[int, int]) # accumulator: int, output: int
|
|
Sum(AggregateFnV2[Union[int, float], ...]) # accumulator: number, output: number
|
|
Mean(AggregateFnV2[List[...], float]) # accumulator: [sum, count], output: float
|
|
Std(AggregateFnV2[List[...], float]) # accumulator: [M2, mean, count], output: float
|
|
|
|
Args:
|
|
name: The name of the aggregation. This will be used as the column name
|
|
in the output, e.g., "sum(my_col)".
|
|
zero_factory: A callable that returns the initial "zero" value for the
|
|
accumulator. For example, for a sum, this would be `lambda: 0`; for
|
|
finding a minimum, `lambda: float("inf")`, for finding a maximum,
|
|
`lambda: float("-inf")`.
|
|
on: The name of the column to perform the aggregation on. If `None`,
|
|
the aggregation is performed over the entire row (e.g., for `Count()`).
|
|
ignore_nulls: Whether to ignore null values during aggregation.
|
|
If `True`, nulls are skipped.
|
|
If `False`, the presence of a null value might result in a null output,
|
|
depending on the aggregation logic.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name: str,
|
|
zero_factory: Callable[[], AccumulatorType],
|
|
*,
|
|
on: Optional[str],
|
|
ignore_nulls: bool,
|
|
):
|
|
if not name:
|
|
raise ValueError(
|
|
f"Non-empty string has to be provided as name (got {name})"
|
|
)
|
|
|
|
self._target_col_name = on
|
|
self._ignore_nulls = ignore_nulls
|
|
|
|
# Extract and store the agg name (e.g., "sum" from "sum(col)")
|
|
# This avoids string parsing later
|
|
match = _AGGREGATION_NAME_PATTERN.match(name)
|
|
if match:
|
|
self._agg_name = match.group(1)
|
|
else:
|
|
self._agg_name = name
|
|
|
|
_safe_combine = _null_safe_combine(self.combine, ignore_nulls)
|
|
_safe_aggregate = _null_safe_aggregate(self.aggregate_block, ignore_nulls)
|
|
_safe_finalize = _null_safe_finalize(self.finalize)
|
|
|
|
_safe_zero_factory = _null_safe_zero_factory(zero_factory, ignore_nulls)
|
|
|
|
super().__init__(
|
|
name=name,
|
|
init=_safe_zero_factory,
|
|
merge=_safe_combine,
|
|
accumulate_block=lambda _, block: _safe_aggregate(block),
|
|
finalize=_safe_finalize,
|
|
)
|
|
|
|
def get_target_column(self) -> Optional[str]:
|
|
return self._target_col_name
|
|
|
|
def get_agg_name(self) -> str:
|
|
"""Return the agg name (e.g., 'sum', 'mean', 'count').
|
|
|
|
Returns the aggregation type extracted from the name during initialization.
|
|
For example, returns 'sum' for an aggregator named 'sum(col)'.
|
|
"""
|
|
return self._agg_name
|
|
|
|
@abc.abstractmethod
|
|
def combine(
|
|
self, current_accumulator: AccumulatorType, new: AccumulatorType
|
|
) -> AccumulatorType:
|
|
"""Combines a new partial aggregation result with the current accumulator.
|
|
|
|
This method defines how two intermediate aggregation states are merged.
|
|
For example, if `aggregate_block` produces partial sums `s1` and `s2` from
|
|
two different blocks, `combine(s1, s2)` should return `s1 + s2`.
|
|
|
|
Args:
|
|
current_accumulator: The current accumulated state (e.g., the result of
|
|
previous `combine` calls or an initial value from `zero_factory`).
|
|
new: A new partially aggregated value, typically the output of
|
|
`aggregate_block` from a new block of data, or another accumulator
|
|
from a parallel task.
|
|
|
|
Returns:
|
|
The updated accumulator after combining it with the new value.
|
|
"""
|
|
...
|
|
|
|
@abc.abstractmethod
|
|
def aggregate_block(self, block: Block) -> AccumulatorType:
|
|
"""Aggregates data within a single block.
|
|
|
|
This method processes all rows in a given `Block` and returns a partial
|
|
aggregation result for that block. For instance, if implementing a sum,
|
|
this method would sum all relevant values within the block.
|
|
|
|
Args:
|
|
block: A `Block` of data to be aggregated.
|
|
|
|
Returns:
|
|
A partial aggregation result for the input block. The type of this
|
|
result (`AggType`) should be consistent with the `current_accumulator`
|
|
and `new` arguments of the `combine` method, and the `accumulator`
|
|
argument of the `finalize` method.
|
|
"""
|
|
...
|
|
|
|
def finalize(self, accumulator: AccumulatorType) -> Optional[AggOutputType]:
|
|
"""Transforms the final accumulated state into the desired output.
|
|
|
|
This method is called once per group after all blocks have been processed
|
|
and all partial results have been combined. It provides an opportunity
|
|
to perform a final transformation on the accumulated data.
|
|
|
|
For many aggregations (e.g., Sum, Count, Min, Max), the accumulated state
|
|
is already the final result, so this method can simply return the
|
|
accumulator as is (which is the default behavior).
|
|
|
|
For other aggregations, like Mean, this method is crucial.
|
|
A Mean aggregation might accumulate `[sum, count]`. The `finalize`
|
|
method would then compute `sum / count` to get the final mean.
|
|
|
|
Args:
|
|
accumulator: The final accumulated state for a group, after all
|
|
`aggregate_block` and `combine` operations.
|
|
|
|
Returns:
|
|
The final result of the aggregation for the group.
|
|
"""
|
|
return accumulator
|
|
|
|
def _validate(self, schema: Optional["Schema"]) -> None:
|
|
if self._target_col_name:
|
|
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
|
|
|
SortKey(self._target_col_name).validate_schema(schema)
|
|
|
|
|
|
def _agg_output_field(
|
|
name: str,
|
|
input_schema: "pa.Schema",
|
|
target_col: Optional[str],
|
|
pyarrow_kernel: Callable[[pa.Array], pa.Scalar],
|
|
) -> Optional["pa.Field"]:
|
|
"""Compute the output field of a scalar reduction aggregator by running
|
|
its PyArrow compute kernel on an empty array of the target column's type.
|
|
|
|
Args:
|
|
name: Name of the *output* column the aggregation produces (the
|
|
aggregator's alias, e.g. ``"sum(a)"`` or a user-supplied
|
|
``alias_name``). This becomes the name of the returned field.
|
|
input_schema: Schema of the aggregator's input (pre-aggregation)
|
|
blocks, used to look up the target column's type.
|
|
target_col: Name of the *input* column being aggregated (the ``on=``
|
|
argument). This is the column the kernel reads, not the group-by
|
|
key. ``None`` for aggregations that don't read a column (e.g.
|
|
``Count()`` over rows), in which case the type can't be inferred.
|
|
pyarrow_kernel: The PyArrow compute kernel implementing the reduction
|
|
(e.g. ``pc.sum``), run on an empty array to derive the output type.
|
|
|
|
Returns:
|
|
The output ``pa.Field`` (``name`` with the inferred type), or ``None``
|
|
if ``target_col`` is ``None``/missing or the pyarrow_kernel rejects the
|
|
column's type (e.g., ``pc.sum`` on a string column).
|
|
"""
|
|
if target_col is None:
|
|
return None
|
|
try:
|
|
in_type = input_schema.field(target_col).type
|
|
except (KeyError, ValueError):
|
|
return None
|
|
try:
|
|
result = pyarrow_kernel(pa.array([], type=in_type))
|
|
except (pa.ArrowNotImplementedError, pa.ArrowInvalid, pa.ArrowTypeError):
|
|
# The kernel has no implementation for this column's type
|
|
# (e.g. ``pc.sum`` on a string/struct/list column). Fall back to
|
|
# an unresolved schema rather than masking a real error.
|
|
return None
|
|
out_type = result.type
|
|
return pa.field(name, out_type, nullable=True)
|
|
|
|
|
|
@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()
|
|
|
|
return block_accessor.count(
|
|
self._target_col_name, ignore_nulls=self._ignore_nulls
|
|
)
|
|
|
|
def combine(self, current_accumulator: int, new: int) -> int:
|
|
return current_accumulator + new
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return pa.field(self.name, pa.int64(), nullable=False)
|
|
|
|
|
|
@PublicAPI
|
|
class AsList(AggregateFnV2[List, List]):
|
|
"""Listing aggregation combining all values within the group into a single
|
|
list element.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
:skipif: True
|
|
|
|
# Skip testing b/c this example require proper ordering of the output
|
|
# to be robust and not flaky
|
|
|
|
import ray
|
|
from ray.data.aggregate import AsList
|
|
|
|
ds = ray.data.range(10)
|
|
# Schema: {'id': int64}
|
|
ds = ds.add_column(
|
|
"group_key", lambda batch: batch["id"].astype("int64") % 3
|
|
)
|
|
# Schema: {'id': int64, 'group_key': int64}
|
|
|
|
# Listing all elements per group:
|
|
result = ds.groupby("group_key").aggregate(AsList(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'list(id)': [0, 3, 6, 9]},
|
|
# {'group_key': 1, 'list(id)': [1, 4, 7]},
|
|
# {'group_key': 2, 'list(id)': [2, 5, 8]}
|
|
|
|
Args:
|
|
on: The name of the column to collect values from. Must be provided.
|
|
alias_name: Optional name for the resulting column.
|
|
ignore_nulls: Whether to ignore null values when collecting. If `True`,
|
|
nulls are skipped. If `False` (default), nulls are included in the list.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
alias_name: Optional[str] = None,
|
|
ignore_nulls: bool = False,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"list({on or ''})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=lambda: [],
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> AccumulatorType:
|
|
column_accessor = BlockColumnAccessor.for_column(
|
|
block[self.get_target_column()]
|
|
)
|
|
|
|
if self._ignore_nulls:
|
|
column_accessor = BlockColumnAccessor.for_column(column_accessor.dropna())
|
|
|
|
return column_accessor.to_pylist()
|
|
|
|
def combine(
|
|
self, current_accumulator: AccumulatorType, new: AccumulatorType
|
|
) -> AccumulatorType:
|
|
return current_accumulator + new
|
|
|
|
|
|
@PublicAPI
|
|
class Sum(AggregateFnV2[Union[int, float], Union[int, float]]):
|
|
"""Defines sum aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Sum
|
|
|
|
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}
|
|
|
|
# Summing all rows per group:
|
|
result = ds.aggregate(Sum(on="id"))
|
|
# result: {'sum(id)': 4950}
|
|
|
|
Args:
|
|
on: The name of the numerical column to sum. Must be provided.
|
|
ignore_nulls: Whether to ignore null values during summation. If `True` (default),
|
|
nulls are skipped. If `False`, the sum will be null if any
|
|
value in the group is null.
|
|
alias_name: Optional name for the resulting column.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"sum({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=lambda: 0,
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> Union[int, float]:
|
|
return BlockAccessor.for_block(block).sum(
|
|
self._target_col_name, self._ignore_nulls
|
|
)
|
|
|
|
def combine(
|
|
self, current_accumulator: Union[int, float], new: Union[int, float]
|
|
) -> Union[int, float]:
|
|
return current_accumulator + new
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return _agg_output_field(self.name, input_schema, self._target_col_name, pc.sum)
|
|
|
|
|
|
@PublicAPI
|
|
class Min(AggregateFnV2[SupportsRichComparisonType, SupportsRichComparisonType]):
|
|
"""Defines min aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Min
|
|
|
|
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}
|
|
|
|
# Finding the minimum value per group:
|
|
result = ds.groupby("group_key").aggregate(Min(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'min(id)': 0},
|
|
# {'group_key': 1, 'min(id)': 1},
|
|
# {'group_key': 2, 'min(id)': 2}]
|
|
|
|
Args:
|
|
on: The name of the column to find the minimum value from. Must be provided.
|
|
ignore_nulls: Whether to ignore null values. If `True` (default), nulls are
|
|
skipped. If `False`, the minimum will be null if any value in
|
|
the group is null (for most data types, or follow type-specific
|
|
comparison rules with nulls).
|
|
alias_name: Optional name for the resulting column.
|
|
zero_factory: A callable that returns the initial "zero" value for the
|
|
accumulator. For example, for a float column, this would be
|
|
`lambda: float("+inf")`. Default is `lambda: float("+inf")`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
zero_factory: Callable[[], SupportsRichComparisonType] = lambda: float("+inf"),
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"min({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=zero_factory,
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> SupportsRichComparisonType:
|
|
return BlockAccessor.for_block(block).min(
|
|
self._target_col_name, self._ignore_nulls
|
|
)
|
|
|
|
def combine(
|
|
self,
|
|
current_accumulator: SupportsRichComparisonType,
|
|
new: SupportsRichComparisonType,
|
|
) -> SupportsRichComparisonType:
|
|
return min(current_accumulator, new)
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return _agg_output_field(self.name, input_schema, self._target_col_name, pc.min)
|
|
|
|
|
|
@PublicAPI
|
|
class Max(AggregateFnV2[SupportsRichComparisonType, SupportsRichComparisonType]):
|
|
"""Defines max aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Max
|
|
|
|
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}
|
|
|
|
# Finding the maximum value per group:
|
|
result = ds.groupby("group_key").aggregate(Max(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'max(id)': ...},
|
|
# {'group_key': 1, 'max(id)': ...},
|
|
# {'group_key': 2, 'max(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the column to find the maximum value from. Must be provided.
|
|
ignore_nulls: Whether to ignore null values. If `True` (default), nulls are
|
|
skipped. If `False`, the maximum will be null if any value in
|
|
the group is null (for most data types, or follow type-specific
|
|
comparison rules with nulls).
|
|
alias_name: Optional name for the resulting column.
|
|
zero_factory: A callable that returns the initial "zero" value for the
|
|
accumulator. For example, for a float column, this would be
|
|
`lambda: float("-inf")`. Default is `lambda: float("-inf")`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
zero_factory: Callable[[], SupportsRichComparisonType] = lambda: float("-inf"),
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"max({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=zero_factory,
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> SupportsRichComparisonType:
|
|
return BlockAccessor.for_block(block).max(
|
|
self._target_col_name, self._ignore_nulls
|
|
)
|
|
|
|
def combine(
|
|
self,
|
|
current_accumulator: SupportsRichComparisonType,
|
|
new: SupportsRichComparisonType,
|
|
) -> SupportsRichComparisonType:
|
|
return max(current_accumulator, new)
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return _agg_output_field(self.name, input_schema, self._target_col_name, pc.max)
|
|
|
|
|
|
@PublicAPI
|
|
class Mean(AggregateFnV2[List[Union[int, float]], float]):
|
|
"""Defines mean (average) aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Mean
|
|
|
|
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}
|
|
|
|
# Calculating the mean value per group:
|
|
result = ds.groupby("group_key").aggregate(Mean(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'mean(id)': ...},
|
|
# {'group_key': 1, 'mean(id)': ...},
|
|
# {'group_key': 2, 'mean(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the numerical column to calculate the mean on. Must be provided.
|
|
ignore_nulls: Whether to ignore null values. If `True` (default), nulls are
|
|
skipped. If `False`, the mean will be null if any value in the
|
|
group is null.
|
|
alias_name: Optional name for the resulting column.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"mean({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
# The accumulator is: [current_sum, current_count].
|
|
# NOTE: We copy the returned list `list([0,0])` as some internal mechanisms
|
|
# might modify accumulators in-place.
|
|
zero_factory=lambda: list([0, 0]), # noqa: C410
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> Optional[List[Union[int, float]]]:
|
|
block_acc = BlockAccessor.for_block(block)
|
|
count = block_acc.count(self._target_col_name, self._ignore_nulls)
|
|
|
|
if count == 0 or count is None:
|
|
# Empty or all null.
|
|
return None
|
|
|
|
sum_ = block_acc.sum(self._target_col_name, self._ignore_nulls)
|
|
|
|
if is_null(sum_):
|
|
# In case of ignore_nulls=False and column containing 'null'
|
|
# return as is (to prevent unnecessary type conversions, when, for ex,
|
|
# using Pandas and returning None)
|
|
return sum_
|
|
|
|
return [sum_, count]
|
|
|
|
def combine(
|
|
self, current_accumulator: List[Union[int, float]], new: List[Union[int, float]]
|
|
) -> List[Union[int, float]]:
|
|
return [current_accumulator[0] + new[0], current_accumulator[1] + new[1]]
|
|
|
|
def finalize(self, accumulator: List[Union[int, float]]) -> Optional[float]:
|
|
# The final accumulator for a group is [total_sum, total_count].
|
|
if accumulator[1] == 0:
|
|
# If total_count is 0 (e.g., group was empty or all nulls ignored),
|
|
# the mean is undefined. Return NaN
|
|
return np.nan
|
|
|
|
return accumulator[0] / accumulator[1]
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return pa.field(self.name, pa.float64(), nullable=True)
|
|
|
|
|
|
@PublicAPI
|
|
class Std(AggregateFnV2[List[Union[int, float]], float]):
|
|
"""Defines standard deviation aggregation.
|
|
|
|
Uses Welford's online algorithm for numerical stability. This method computes
|
|
the standard deviation in a single pass. Results may differ slightly from
|
|
libraries like NumPy or Pandas that use a two-pass algorithm but are generally
|
|
more accurate.
|
|
|
|
See: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Std
|
|
|
|
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}
|
|
|
|
# Calculating the standard deviation per group:
|
|
result = ds.groupby("group_key").aggregate(Std(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'std(id)': ...},
|
|
# {'group_key': 1, 'std(id)': ...},
|
|
# {'group_key': 2, 'std(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the column to calculate standard deviation on.
|
|
ddof: Delta Degrees of Freedom. The divisor used in calculations is `N - ddof`,
|
|
where `N` is the number of elements. Default is 1.
|
|
ignore_nulls: Whether to ignore null values. Default is True.
|
|
alias_name: Optional name for the resulting column.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ddof: int = 1,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"std({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
# Accumulator: [M2, mean, count]
|
|
# M2: sum of squares of differences from the current mean
|
|
# mean: current mean
|
|
# count: current count of non-null elements
|
|
# We need to copy the list as it might be modified in-place by some aggregations.
|
|
zero_factory=lambda: list([0, 0, 0]), # noqa: C410
|
|
)
|
|
|
|
self._ddof = ddof
|
|
|
|
def aggregate_block(self, block: Block) -> List[Union[int, float]]:
|
|
block_acc = BlockAccessor.for_block(block)
|
|
count = block_acc.count(self._target_col_name, ignore_nulls=self._ignore_nulls)
|
|
if count == 0 or count is None:
|
|
# Empty or all null.
|
|
return None
|
|
sum_ = block_acc.sum(self._target_col_name, self._ignore_nulls)
|
|
if is_null(sum_):
|
|
# If sum is null (e.g., ignore_nulls=False and a null was encountered),
|
|
# return as is to prevent type conversions.
|
|
return sum_
|
|
mean = sum_ / count
|
|
M2 = block_acc.sum_of_squared_diffs_from_mean(
|
|
self._target_col_name, self._ignore_nulls, mean
|
|
)
|
|
return [M2, mean, count]
|
|
|
|
def combine(
|
|
self, current_accumulator: List[float], new: List[float]
|
|
) -> List[float]:
|
|
# Merges two accumulators [M2, mean, count] using a parallel algorithm.
|
|
# See: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm
|
|
M2_a, mean_a, count_a = current_accumulator
|
|
M2_b, mean_b, count_b = new
|
|
delta = mean_b - mean_a
|
|
count = count_a + count_b
|
|
# NOTE: We use this mean calculation since it's more numerically
|
|
# stable than mean_a + delta * count_b / count, which actually
|
|
# deviates from Pandas in the ~15th decimal place and causes our
|
|
# exact comparison tests to fail.
|
|
mean = (mean_a * count_a + mean_b * count_b) / count
|
|
# Update the sum of squared differences.
|
|
M2 = M2_a + M2_b + (delta**2) * count_a * count_b / count
|
|
return [M2, mean, count]
|
|
|
|
def finalize(self, accumulator: List[float]) -> Optional[float]:
|
|
# Compute the final standard deviation from the accumulated
|
|
# sum of squared differences from current mean and the count.
|
|
# Final accumulator: [M2, mean, count]
|
|
M2, mean, count = accumulator
|
|
# Denominator for variance calculation is count - ddof
|
|
if count - self._ddof <= 0:
|
|
# If count - ddof is not positive, variance/std is undefined (or zero).
|
|
# Return NaN, consistent with pandas/numpy.
|
|
return np.nan
|
|
# Standard deviation is the square root of variance (M2 / (count - ddof))
|
|
return math.sqrt(M2 / (count - self._ddof))
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
return pa.field(self.name, pa.float64(), nullable=True)
|
|
|
|
|
|
@PublicAPI
|
|
class AbsMax(AggregateFnV2[SupportsRichComparisonType, SupportsRichComparisonType]):
|
|
"""Defines absolute max aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import AbsMax
|
|
|
|
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}
|
|
|
|
# Calculating the absolute maximum value per group:
|
|
result = ds.groupby("group_key").aggregate(AbsMax(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'abs_max(id)': ...},
|
|
# {'group_key': 1, 'abs_max(id)': ...},
|
|
# {'group_key': 2, 'abs_max(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the column to calculate absolute maximum on. Must be provided.
|
|
ignore_nulls: Whether to ignore null values. Default is True.
|
|
alias_name: Optional name for the resulting column.
|
|
zero_factory: A callable that returns the initial "zero" value for the
|
|
accumulator. For example, for a float column, this would be
|
|
`lambda: 0`. Default is `lambda: 0`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
zero_factory: Callable[[], SupportsRichComparisonType] = lambda: 0,
|
|
):
|
|
if on is None or not isinstance(on, str):
|
|
raise ValueError(f"Column to aggregate on has to be provided (got {on})")
|
|
|
|
super().__init__(
|
|
alias_name if alias_name else f"abs_max({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=zero_factory,
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> Optional[SupportsRichComparisonType]:
|
|
block_accessor = BlockAccessor.for_block(block)
|
|
|
|
max_ = block_accessor.max(self._target_col_name, self._ignore_nulls)
|
|
min_ = block_accessor.min(self._target_col_name, self._ignore_nulls)
|
|
|
|
if is_null(max_) or is_null(min_):
|
|
return None
|
|
|
|
return max(abs(max_), abs(min_))
|
|
|
|
def combine(
|
|
self,
|
|
current_accumulator: SupportsRichComparisonType,
|
|
new: SupportsRichComparisonType,
|
|
) -> SupportsRichComparisonType:
|
|
return max(current_accumulator, new)
|
|
|
|
def output_field(self, input_schema: "pa.Schema") -> Optional["pa.Field"]:
|
|
# AbsMax = max(abs(x)). Compose abs + max into a single kernel so the
|
|
# output type (and the type-support check) come from the same pyarrow
|
|
# kernels, returning None for types abs/max reject (e.g. strings).
|
|
return _agg_output_field(
|
|
self.name,
|
|
input_schema,
|
|
self._target_col_name,
|
|
lambda a: pc.max(pc.abs(a)),
|
|
)
|
|
|
|
|
|
@PublicAPI
|
|
class Quantile(AggregateFnV2[List[Any], List[Any]]):
|
|
"""Defines Quantile aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Quantile
|
|
|
|
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}
|
|
|
|
# Calculating the 50th percentile (median) per group:
|
|
result = ds.groupby("group_key").aggregate(Quantile(q=0.5, on="id")).take_all()
|
|
# result: [{'group_key': 0, 'quantile(id)': ...},
|
|
# {'group_key': 1, 'quantile(id)': ...},
|
|
# {'group_key': 2, 'quantile(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the column to calculate the quantile on. Must be provided.
|
|
q: The quantile to compute, which must be between 0 and 1 inclusive.
|
|
For example, q=0.5 computes the median.
|
|
ignore_nulls: Whether to ignore null values. Default is True.
|
|
alias_name: Optional name for the resulting column.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
q: float = 0.5,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
self._q = q
|
|
|
|
super().__init__(
|
|
alias_name if alias_name else f"quantile({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=list,
|
|
)
|
|
|
|
def combine(self, current_accumulator: List[Any], new: List[Any]) -> List[Any]:
|
|
if isinstance(current_accumulator, List) and isinstance(new, List):
|
|
current_accumulator.extend(new)
|
|
return current_accumulator
|
|
|
|
if isinstance(current_accumulator, List) and (not isinstance(new, List)):
|
|
if new is not None and new != "":
|
|
current_accumulator.append(new)
|
|
return current_accumulator
|
|
|
|
if isinstance(new, List) and (not isinstance(current_accumulator, List)):
|
|
if current_accumulator is not None and current_accumulator != "":
|
|
new.append(current_accumulator)
|
|
return new
|
|
|
|
ls = []
|
|
|
|
if current_accumulator is not None and current_accumulator != "":
|
|
ls.append(current_accumulator)
|
|
|
|
if new is not None and new != "":
|
|
ls.append(new)
|
|
|
|
return ls
|
|
|
|
def aggregate_block(self, block: Block) -> List[Any]:
|
|
block_acc = BlockAccessor.for_block(block)
|
|
ls = []
|
|
|
|
for row in block_acc.iter_rows(public_row_format=False):
|
|
ls.append(row.get(self._target_col_name))
|
|
|
|
return ls
|
|
|
|
def finalize(self, accumulator: List[Any]) -> Optional[Any]:
|
|
if self._ignore_nulls:
|
|
accumulator = [v for v in accumulator if not is_null(v)]
|
|
else:
|
|
nulls = [v for v in accumulator if is_null(v)]
|
|
if len(nulls) > 0:
|
|
# If nulls are present and not ignored, the quantile is undefined.
|
|
# Return the first null encountered to preserve column type.
|
|
return nulls[0]
|
|
|
|
if not accumulator:
|
|
# If the list is empty (e.g., all values were null and ignored, or no values),
|
|
# quantile is undefined.
|
|
return None
|
|
|
|
key = lambda x: x # noqa: E731
|
|
input_values = sorted(accumulator)
|
|
k = (len(input_values) - 1) * self._q
|
|
f = math.floor(k)
|
|
c = math.ceil(k)
|
|
|
|
if f == c:
|
|
return key(input_values[int(k)])
|
|
|
|
# Interpolate between the elements at floor and ceil indices.
|
|
d0 = key(input_values[int(f)]) * (c - k)
|
|
d1 = key(input_values[int(c)]) * (k - f)
|
|
|
|
return round(d0 + d1, 5)
|
|
|
|
|
|
@PublicAPI
|
|
class Unique(AggregateFnV2[Set[Any], List[Any]]):
|
|
"""Defines unique aggregation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import Unique
|
|
|
|
ds = ray.data.range(100)
|
|
ds = ds.add_column(
|
|
"group_key", lambda batch: batch["id"].astype("int64") % 3
|
|
)
|
|
|
|
# Calculating the unique values per group:
|
|
result = ds.groupby("group_key").aggregate(Unique(on="id")).take_all()
|
|
# result: [{'group_key': 0, 'unique(id)': ...},
|
|
# {'group_key': 1, 'unique(id)': ...},
|
|
# {'group_key': 2, 'unique(id)': ...}]
|
|
|
|
Args:
|
|
on: The name of the column from which to collect unique values.
|
|
ignore_nulls: Whether to ignore null values when collecting unique items.
|
|
Default is True (nulls are excluded).
|
|
alias_name: Optional name for the resulting column.
|
|
encode_lists: If `True`, encode list elements. If `False`, encode
|
|
whole lists (i.e., the entire list is considered as a single object).
|
|
`False` by default. Note that this is a top-level flatten (not a recursive
|
|
flatten) operation.
|
|
"""
|
|
|
|
class ListEncodingMode(str, enum.Enum):
|
|
"""Controls how to encode individual elements inside the list column:
|
|
|
|
- NONE: no encoding applied, elements (lists) are stored as is and
|
|
unique ones are returned.
|
|
- FLATTEN: column of element lists is flattened into a single list.
|
|
- HASH: each list element is hashed, a list of unique hashes is returned.
|
|
"""
|
|
|
|
FLATTEN = "FLATTEN"
|
|
HASH = "HASH"
|
|
|
|
def __init__(
|
|
self,
|
|
on: Optional[str] = None,
|
|
ignore_nulls: bool = False,
|
|
alias_name: Optional[str] = None,
|
|
encode_lists: Union[bool, ListEncodingMode, None] = None,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"unique({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=set,
|
|
)
|
|
|
|
if isinstance(encode_lists, Unique.ListEncodingMode):
|
|
self._list_encoding_mode = encode_lists
|
|
elif isinstance(encode_lists, bool) and encode_lists:
|
|
self._list_encoding_mode = Unique.ListEncodingMode.FLATTEN
|
|
else:
|
|
self._list_encoding_mode = None
|
|
|
|
def combine(self, current_accumulator: Set[Any], new: Set[Any]) -> Set[Any]:
|
|
return self._to_set(current_accumulator) | self._to_set(new)
|
|
|
|
def _compute_unique(self, block: Block) -> BlockColumn:
|
|
column = block[self._target_col_name]
|
|
column_accessor = BlockColumnAccessor.for_column(column)
|
|
|
|
if (
|
|
column_accessor.is_composed_of_lists()
|
|
and self._list_encoding_mode is not None
|
|
):
|
|
if self._list_encoding_mode == Unique.ListEncodingMode.FLATTEN:
|
|
column_accessor = BlockColumnAccessor.for_column(
|
|
column_accessor.flatten()
|
|
)
|
|
elif self._list_encoding_mode == Unique.ListEncodingMode.HASH:
|
|
column_accessor = BlockColumnAccessor.for_column(column_accessor.hash())
|
|
else:
|
|
raise ValueError(
|
|
f"list encoding mode not supported: {self._list_encoding_mode}"
|
|
)
|
|
|
|
if self._ignore_nulls:
|
|
column_accessor = BlockColumnAccessor.for_column(column_accessor.dropna())
|
|
|
|
return column_accessor.unique()
|
|
|
|
def aggregate_block(self, block: Block) -> List[Any]:
|
|
column = self._compute_unique(block)
|
|
return BlockColumnAccessor.for_column(column).to_pylist()
|
|
|
|
@staticmethod
|
|
def _to_set(x):
|
|
if isinstance(x, set):
|
|
return Unique._normalize_nans(x)
|
|
|
|
elif isinstance(x, list):
|
|
if len(x) > 0 and isinstance(x[0], list):
|
|
# necessary because pyarrow converts all tuples to
|
|
# list internally.
|
|
x = map(lambda v: None if v is None else tuple(v), x)
|
|
|
|
return Unique._normalize_nans(x)
|
|
else:
|
|
return {x}
|
|
|
|
@staticmethod
|
|
def _normalize_nans(x: Collection) -> Set:
|
|
# NOTE: Pandas when converting to Python objects instantiates
|
|
# new `float('nan')` objects which are incomparable b/w each
|
|
# other. Here we canonicalize any nan instances replacing them
|
|
# w/ `np.nan`
|
|
return {v if not (isinstance(v, float) and np.isnan(v)) else np.nan for v in x}
|
|
|
|
|
|
@PublicAPI
|
|
class CountDistinct(Unique):
|
|
"""Defines distinct count aggregation.
|
|
|
|
This aggregation computes the count of distinct values in a column.
|
|
It is similar to SQL's COUNT(DISTINCT column_name) operation.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import CountDistinct
|
|
|
|
# Create a dataset with repeated values
|
|
ds = ray.data.from_items([
|
|
{"category": "A"}, {"category": "B"}, {"category": "A"},
|
|
{"category": "C"}, {"category": "A"}, {"category": "B"}
|
|
])
|
|
|
|
# Count distinct categories
|
|
result = ds.aggregate(CountDistinct(on="category"))
|
|
# result: {'count_distinct(category)': 3}
|
|
|
|
# Using with groupby
|
|
ds = ray.data.from_items([
|
|
{"group": "X", "category": "A"}, {"group": "X", "category": "B"},
|
|
{"group": "Y", "category": "A"}, {"group": "Y", "category": "A"}
|
|
])
|
|
result = ds.groupby("group").aggregate(CountDistinct(on="category")).take_all()
|
|
# result: [{'group': 'X', 'count_distinct(category)': 2},
|
|
# {'group': 'Y', 'count_distinct(category)': 1}]
|
|
|
|
Args:
|
|
on: The name of the column to count distinct values on.
|
|
ignore_nulls: Whether to ignore null values when counting distinct items.
|
|
Default is True (nulls are excluded from the count).
|
|
alias_name: Optional name for the resulting column. If not provided,
|
|
defaults to "count_distinct({on})".
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
alias_name=alias_name if alias_name else f"count_distinct({str(on)})",
|
|
)
|
|
|
|
def finalize(self, accumulator: Set[Any]) -> int:
|
|
"""Return the count of distinct values."""
|
|
return len(accumulator)
|
|
|
|
|
|
@PublicAPI
|
|
class ValueCounter(AggregateFnV2):
|
|
"""Counts the number of times each value appears in a column.
|
|
|
|
This aggregation computes value counts for a specified column, similar to pandas'
|
|
`value_counts()` method. It returns a dictionary with two lists: "values" containing
|
|
the unique values found in the column, and "counts" containing the corresponding
|
|
count for each value.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import ValueCounter
|
|
|
|
# Create a dataset with repeated values
|
|
ds = ray.data.from_items([
|
|
{"category": "A"}, {"category": "B"}, {"category": "A"},
|
|
{"category": "C"}, {"category": "A"}, {"category": "B"}
|
|
])
|
|
|
|
# Count occurrences of each category
|
|
result = ds.aggregate(ValueCounter(on="category"))
|
|
# result: {'value_counter(category)': {'values': ['A', 'B', 'C'], 'counts': [3, 2, 1]}}
|
|
|
|
# Using with groupby
|
|
ds = ray.data.from_items([
|
|
{"group": "X", "category": "A"}, {"group": "X", "category": "B"},
|
|
{"group": "Y", "category": "A"}, {"group": "Y", "category": "A"}
|
|
])
|
|
result = ds.groupby("group").aggregate(ValueCounter(on="category")).take_all()
|
|
# result: [{'group': 'X', 'value_counter(category)': {'values': ['A', 'B'], 'counts': [1, 1]}},
|
|
# {'group': 'Y', 'value_counter(category)': {'values': ['A'], 'counts': [2]}}]
|
|
|
|
Args:
|
|
on: The name of the column to count values in. Must be provided.
|
|
alias_name: Optional name for the resulting column. If not provided,
|
|
defaults to "value_counter({column_name})".
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
super().__init__(
|
|
alias_name if alias_name else f"value_counter({str(on)})",
|
|
on=on,
|
|
ignore_nulls=True,
|
|
zero_factory=lambda: {"values": [], "counts": []},
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> Dict[str, List]:
|
|
|
|
col_accessor = BlockColumnAccessor.for_column(block[self._target_col_name])
|
|
return col_accessor.value_counts()
|
|
|
|
def combine(
|
|
self,
|
|
current_accumulator: Dict[str, List],
|
|
new_accumulator: Dict[str, List],
|
|
) -> Dict[str, List]:
|
|
|
|
values = current_accumulator["values"]
|
|
counts = current_accumulator["counts"]
|
|
|
|
# Build a value → index map once (avoid repeated lookups)
|
|
value_to_index = {v: i for i, v in enumerate(values)}
|
|
|
|
for v_new, c_new in zip(new_accumulator["values"], new_accumulator["counts"]):
|
|
if v_new in value_to_index:
|
|
idx = value_to_index[v_new]
|
|
counts[idx] += c_new
|
|
else:
|
|
value_to_index[v_new] = len(values)
|
|
values.append(v_new)
|
|
counts.append(c_new)
|
|
|
|
return current_accumulator
|
|
|
|
|
|
def _null_safe_zero_factory(zero_factory, ignore_nulls: bool):
|
|
"""NOTE: PLEASE READ CAREFULLY BEFORE CHANGING
|
|
|
|
Null-safe zero factory is crucial for implementing proper aggregation
|
|
protocol (monoid) w/o the need for additional containers.
|
|
|
|
Main hurdle for implementing proper aggregation semantic is to be able to encode
|
|
semantic of an "empty accumulator" and be able to tell it from the case when
|
|
accumulator is actually holding null value:
|
|
|
|
- Empty container can be overridden with any value
|
|
- Container holding null can't be overridden if ignore_nulls=False
|
|
|
|
However, it's possible for us to exploit asymmetry in cases of ignore_nulls being
|
|
True or False:
|
|
|
|
- Case of ignore_nulls=False entails that if there's any "null" in the sequence,
|
|
aggregation is undefined and correspondingly expected to return null
|
|
|
|
- Case of ignore_nulls=True in turn, entails that if aggregation returns "null"
|
|
if and only if the sequence does NOT have any non-null value
|
|
|
|
Therefore, we apply this difference in semantic to zero-factory to make sure that
|
|
our aggregation protocol is adherent to that definition:
|
|
|
|
- If ignore_nulls=True, zero-factory returns null, therefore encoding empty
|
|
container
|
|
- If ignore_nulls=False, couldn't return null as aggregation will incorrectly
|
|
prioritize it, and instead it returns true zero value for the aggregation
|
|
(ie 0 for count/sum, -inf for max, etc).
|
|
"""
|
|
|
|
if ignore_nulls:
|
|
|
|
def _safe_zero_factory(_):
|
|
return None
|
|
|
|
else:
|
|
|
|
def _safe_zero_factory(_):
|
|
return zero_factory()
|
|
|
|
return _safe_zero_factory
|
|
|
|
|
|
def _null_safe_aggregate(
|
|
aggregate: Callable[[Block], AccumulatorType],
|
|
ignore_nulls: bool,
|
|
) -> Callable[[Block], Optional[AccumulatorType]]:
|
|
def _safe_aggregate(block: Block) -> Optional[AccumulatorType]:
|
|
result = aggregate(block)
|
|
# NOTE: If `ignore_nulls=True`, aggregation will only be returning
|
|
# null if the block does NOT contain any non-null elements
|
|
if is_null(result) and ignore_nulls:
|
|
return None
|
|
|
|
return result
|
|
|
|
return _safe_aggregate
|
|
|
|
|
|
def _null_safe_finalize(
|
|
finalize: Callable[[AccumulatorType], AccumulatorType],
|
|
) -> Callable[[Optional[AccumulatorType]], AccumulatorType]:
|
|
def _safe_finalize(acc: Optional[AccumulatorType]) -> AccumulatorType:
|
|
# If accumulator container is not null, finalize.
|
|
# Otherwise, return as is.
|
|
return acc if is_null(acc) else finalize(acc)
|
|
|
|
return _safe_finalize
|
|
|
|
|
|
def _null_safe_combine(
|
|
combine: Callable[[AccumulatorType, AccumulatorType], AccumulatorType],
|
|
ignore_nulls: bool,
|
|
) -> Callable[
|
|
[Optional[AccumulatorType], Optional[AccumulatorType]], Optional[AccumulatorType]
|
|
]:
|
|
"""Null-safe combination have to be an associative operation
|
|
with an identity element (zero) or in other words implement a monoid.
|
|
|
|
To achieve that in the presence of null values following semantic is
|
|
established:
|
|
|
|
- Case of ignore_nulls=True:
|
|
- If current accumulator is null (ie empty), return new accumulator
|
|
- If new accumulator is null (ie empty), return cur
|
|
- Otherwise combine (current and new)
|
|
|
|
- Case of ignore_nulls=False:
|
|
- If new accumulator is null (ie has null in the sequence, b/c we're
|
|
NOT ignoring nulls), return it
|
|
- If current accumulator is null (ie had null in the prior sequence,
|
|
b/c we're NOT ignoring nulls), return it
|
|
- Otherwise combine (current and new)
|
|
"""
|
|
|
|
if ignore_nulls:
|
|
|
|
def _safe_combine(
|
|
cur: Optional[AccumulatorType], new: Optional[AccumulatorType]
|
|
) -> Optional[AccumulatorType]:
|
|
if is_null(cur):
|
|
return new
|
|
elif is_null(new):
|
|
return cur
|
|
else:
|
|
return combine(cur, new)
|
|
|
|
else:
|
|
|
|
def _safe_combine(
|
|
cur: Optional[AccumulatorType], new: Optional[AccumulatorType]
|
|
) -> Optional[AccumulatorType]:
|
|
if is_null(new):
|
|
return new
|
|
elif is_null(cur):
|
|
return cur
|
|
else:
|
|
return combine(cur, new)
|
|
|
|
return _safe_combine
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class MissingValuePercentage(AggregateFnV2[List[int], float]):
|
|
"""Calculates the percentage of null values in a column.
|
|
|
|
This aggregation computes the percentage of null (missing) values in a dataset column.
|
|
It treats both None values and NaN values as null. The result is a percentage value
|
|
between 0.0 and 100.0, where 0.0 means no missing values and 100.0 means all values
|
|
are missing.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import MissingValuePercentage
|
|
|
|
# Create a dataset with some missing values
|
|
ds = ray.data.from_items([
|
|
{"value": 1}, {"value": None}, {"value": 3},
|
|
{"value": None}, {"value": 5}
|
|
])
|
|
|
|
# Calculate missing value percentage
|
|
result = ds.aggregate(MissingValuePercentage(on="value"))
|
|
# result: 40.0 (2 out of 5 values are missing)
|
|
|
|
# Using with groupby
|
|
ds = ray.data.from_items([
|
|
{"group": "A", "value": 1}, {"group": "A", "value": None},
|
|
{"group": "B", "value": 3}, {"group": "B", "value": None}
|
|
])
|
|
result = ds.groupby("group").aggregate(MissingValuePercentage(on="value")).take_all()
|
|
# result: [{'group': 'A', 'missing_pct(value)': 50.0},
|
|
# {'group': 'B', 'missing_pct(value)': 50.0}]
|
|
|
|
Args:
|
|
on: The name of the column to calculate missing value percentage on.
|
|
alias_name: Optional name for the resulting column. If not provided,
|
|
defaults to "missing_pct({column_name})".
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
# Initialize with a list accumulator [null_count, total_count]
|
|
super().__init__(
|
|
alias_name if alias_name else f"missing_pct({str(on)})",
|
|
on=on,
|
|
ignore_nulls=False, # Include nulls for this calculation
|
|
zero_factory=lambda: [0, 0], # Our AggType is a simple list
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> List[int]:
|
|
column_accessor = BlockColumnAccessor.for_column(block[self._target_col_name])
|
|
|
|
total_count = column_accessor.count(ignore_nulls=False)
|
|
|
|
null_count = pc.sum(
|
|
pc.is_null(
|
|
column_accessor._to_arrow_compatible_container(), nan_is_null=True
|
|
)
|
|
).as_py()
|
|
|
|
# Return our accumulator
|
|
return [null_count, total_count]
|
|
|
|
def combine(self, current_accumulator: List[int], new: List[int]) -> List[int]:
|
|
# Merge two accumulators by summing their components
|
|
assert len(current_accumulator) == len(new) == 2
|
|
return [
|
|
current_accumulator[0] + new[0], # Sum null counts
|
|
current_accumulator[1] + new[1], # Sum total counts
|
|
]
|
|
|
|
def finalize(self, accumulator: List[int]) -> Optional[float]:
|
|
# Calculate the final percentage
|
|
if accumulator[1] == 0:
|
|
return None
|
|
return (accumulator[0] / accumulator[1]) * 100.0
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class ZeroPercentage(AggregateFnV2[List[int], float]):
|
|
"""Calculates the percentage of zero values in a numeric column.
|
|
|
|
This aggregation computes the percentage of zero values in a numeric dataset column.
|
|
It can optionally ignore null values when calculating the percentage. The result is
|
|
a percentage value between 0.0 and 100.0, where 0.0 means no zero values and 100.0
|
|
means all non-null values are zero.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import ZeroPercentage
|
|
|
|
# Create a dataset with some zero values
|
|
ds = ray.data.from_items([
|
|
{"value": 0}, {"value": 1}, {"value": 0},
|
|
{"value": 3}, {"value": 0}
|
|
])
|
|
|
|
# Calculate zero value percentage
|
|
result = ds.aggregate(ZeroPercentage(on="value"))
|
|
# result: 60.0 (3 out of 5 values are zero)
|
|
|
|
# With null values and ignore_nulls=True (default)
|
|
ds = ray.data.from_items([
|
|
{"value": 0}, {"value": None}, {"value": 0},
|
|
{"value": 3}, {"value": 0}
|
|
])
|
|
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=True))
|
|
# result: 75.0 (3 out of 4 non-null values are zero)
|
|
|
|
# Using with groupby
|
|
ds = ray.data.from_items([
|
|
{"group": "A", "value": 0}, {"group": "A", "value": 1},
|
|
{"group": "B", "value": 0}, {"group": "B", "value": 0}
|
|
])
|
|
result = ds.groupby("group").aggregate(ZeroPercentage(on="value")).take_all()
|
|
# result: [{'group': 'A', 'zero_pct(value)': 50.0},
|
|
# {'group': 'B', 'zero_pct(value)': 100.0}]
|
|
|
|
Args:
|
|
on: The name of the column to calculate zero value percentage on.
|
|
Must be a numeric column.
|
|
ignore_nulls: Whether to ignore null values when calculating the percentage.
|
|
If True (default), null values are excluded from both numerator and denominator.
|
|
If False, null values are included in the denominator but not the numerator.
|
|
alias_name: Optional name for the resulting column. If not provided,
|
|
defaults to "zero_pct({column_name})".
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
ignore_nulls: bool = True,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
# Initialize with a list accumulator [zero_count, non_null_count]
|
|
super().__init__(
|
|
alias_name if alias_name else f"zero_pct({str(on)})",
|
|
on=on,
|
|
ignore_nulls=ignore_nulls,
|
|
zero_factory=lambda: [0, 0],
|
|
)
|
|
|
|
def aggregate_block(self, block: Block) -> List[int]:
|
|
column_accessor = BlockColumnAccessor.for_column(block[self._target_col_name])
|
|
|
|
count = column_accessor.count(ignore_nulls=self._ignore_nulls)
|
|
|
|
if count == 0:
|
|
return [0, 0]
|
|
|
|
arrow_compatible = column_accessor._to_arrow_compatible_container()
|
|
# Use PyArrow compute to count zeros
|
|
# First create a boolean mask for zero values
|
|
zero_mask = pc.equal(arrow_compatible, 0)
|
|
|
|
# Sum the boolean mask to get count of True values (zeros)
|
|
zero_count = pc.sum(zero_mask).as_py() or 0
|
|
|
|
return [zero_count, count]
|
|
|
|
def combine(self, current_accumulator: List[int], new: List[int]) -> List[int]:
|
|
return [
|
|
current_accumulator[0] + new[0], # Sum zero counts
|
|
current_accumulator[1] + new[1], # Sum non-null counts
|
|
]
|
|
|
|
def finalize(self, accumulator: List[int]) -> Optional[float]:
|
|
if accumulator[1] == 0:
|
|
return None
|
|
return (accumulator[0] / accumulator[1]) * 100.0
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class ApproximateQuantile(AggregateFnV2):
|
|
def _require_datasketches(self):
|
|
try:
|
|
from datasketches import kll_floats_sketch # type: ignore[import]
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"ApproximateQuantile requires the `datasketches` package. "
|
|
"Install it with `pip install datasketches`."
|
|
) from exc
|
|
return kll_floats_sketch
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
quantiles: List[float],
|
|
quantile_precision: int = 800,
|
|
alias_name: Optional[str] = None,
|
|
):
|
|
"""
|
|
Computes the approximate quantiles of a column by using a datasketches kll_floats_sketch.
|
|
https://datasketches.apache.org/docs/KLL/KLLSketch.html
|
|
|
|
The accuracy of the KLL quantile sketch is a function of the configured quantile precision, which also affects
|
|
the overall size of the sketch.
|
|
The KLL Sketch has absolute error. For example, a specified rank accuracy of 1% at the
|
|
median (rank = 0.50) means that the true quantile (if you could extract it from the set)
|
|
should be between getQuantile(0.49) and getQuantile(0.51). This same 1% error applied at a
|
|
rank of 0.95 means that the true quantile should be between getQuantile(0.94) and getQuantile(0.96).
|
|
In other words, the error is a fixed +/- epsilon for the entire range of ranks.
|
|
|
|
Typical single-sided rank error by quantile_precision (use for getQuantile/getRank):
|
|
- quantile_precision=100 → ~2.61%
|
|
- quantile_precision=200 → ~1.33%
|
|
- quantile_precision=400 → ~0.68%
|
|
- quantile_precision=800 → ~0.35%
|
|
|
|
See https://datasketches.apache.org/docs/KLL/KLLAccuracyAndSize.html for details on accuracy and size.
|
|
|
|
Null values in the target column are ignored when constructing the sketch.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import ApproximateQuantile
|
|
|
|
# Create a dataset with some values
|
|
ds = ray.data.from_items(
|
|
[{"value": 20.0}, {"value": 40.0}, {"value": 60.0},
|
|
{"value": 80.0}, {"value": 100.0}]
|
|
)
|
|
|
|
result = ds.aggregate(ApproximateQuantile(on="value", quantiles=[0.1, 0.5, 0.9]))
|
|
# Result: {'approx_quantile(value)': [20.0, 60.0, 100.0]}
|
|
|
|
|
|
Args:
|
|
on: The name of the column to calculate the quantile on. Must be a numeric column.
|
|
quantiles: The list of quantiles to compute. Must be between 0 and 1 inclusive. For example, quantiles=[0.5] computes the median. Null entries in the source column are skipped.
|
|
quantile_precision: Controls the accuracy and memory footprint of the sketch (K in KLL); higher values yield lower error but use more memory. Defaults to 800. See https://datasketches.apache.org/docs/KLL/KLLAccuracyAndSize.html for details on accuracy and size.
|
|
alias_name: Optional name for the resulting column. If not provided, defaults to "approx_quantile({column_name})".
|
|
"""
|
|
self._sketch_cls = self._require_datasketches()
|
|
self._quantiles = quantiles
|
|
self._quantile_precision = quantile_precision
|
|
super().__init__(
|
|
alias_name if alias_name else f"approx_quantile({str(on)})",
|
|
on=on,
|
|
ignore_nulls=True,
|
|
zero_factory=lambda: self.zero(quantile_precision).serialize(),
|
|
)
|
|
|
|
def zero(self, quantile_precision: int):
|
|
return self._sketch_cls(k=quantile_precision)
|
|
|
|
def aggregate_block(self, block: Block) -> bytes:
|
|
block_acc = BlockAccessor.for_block(block)
|
|
table = block_acc.to_arrow()
|
|
column = table.column(self.get_target_column())
|
|
sketch = self.zero(self._quantile_precision)
|
|
for value in column:
|
|
# we ignore nulls here
|
|
if value.as_py() is not None:
|
|
sketch.update(float(value.as_py()))
|
|
return sketch.serialize()
|
|
|
|
def combine(self, current_accumulator: bytes, new: bytes) -> bytes:
|
|
combined = self.zero(self._quantile_precision)
|
|
combined.merge(self._sketch_cls.deserialize(current_accumulator))
|
|
combined.merge(self._sketch_cls.deserialize(new))
|
|
return combined.serialize()
|
|
|
|
def finalize(self, accumulator: bytes) -> List[float]:
|
|
return self._sketch_cls.deserialize(accumulator).get_quantiles(self._quantiles)
|
|
|
|
|
|
@PublicAPI(stability="alpha")
|
|
class ApproximateTopK(AggregateFnV2):
|
|
def _require_datasketches(self):
|
|
try:
|
|
from datasketches import frequent_strings_sketch
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"ApproximateTopK requires the `datasketches` package. "
|
|
"Install it with `pip install datasketches`."
|
|
) from exc
|
|
return frequent_strings_sketch
|
|
|
|
def __init__(
|
|
self,
|
|
on: str,
|
|
k: int,
|
|
log_capacity: int = 15,
|
|
alias_name: Optional[str] = None,
|
|
encode_lists: bool = False,
|
|
):
|
|
"""
|
|
Computes the approximate top k items in a column by using a datasketches frequent_strings_sketch.
|
|
https://datasketches.apache.org/docs/Frequency/FrequentItemsOverview.html
|
|
|
|
Guarantees:
|
|
- Any item with true frequency > N / (2^log_capacity) is guaranteed to appear in the results
|
|
- Reported counts may have an error of at most ± N / (2^log_capacity).
|
|
|
|
|
|
If log_capacity is too small for your data:
|
|
- Low-frequency items may be evicted from the sketch, potentially causing the top-k
|
|
results to miss items that should appear in the output.
|
|
- The error bounds increase, reducing the accuracy of the reported counts.
|
|
|
|
Example:
|
|
|
|
.. testcode::
|
|
|
|
import ray
|
|
from ray.data.aggregate import ApproximateTopK
|
|
|
|
ds = ray.data.from_items([
|
|
{"word": "apple"}, {"word": "banana"}, {"word": "apple"},
|
|
{"word": "cherry"}, {"word": "apple"}
|
|
])
|
|
|
|
result = ds.aggregate(ApproximateTopK(on="word", k=2))
|
|
# Result: {'approx_topk(word)': [{'word': 'apple', 'count': 3}, {'word': 'banana', 'count': 1}]}
|
|
|
|
Args:
|
|
on: The name of the column to aggregate.
|
|
k: The number of top items to return.
|
|
log_capacity: Base 2 logarithm of the maximum size of the internal hash map.
|
|
Higher values increase accuracy but use more memory. Defaults to 15.
|
|
alias_name: The name of the aggregate. Defaults to None.
|
|
encode_lists: If `True`, encode list elements. If `False`, encode
|
|
whole lists (i.e., the entire list is considered as a single object).
|
|
`False` by default. Note that this is a top-level flatten (not a recursive
|
|
flatten) operation.
|
|
"""
|
|
|
|
self.k = k
|
|
self._log_capacity = log_capacity
|
|
self._frequent_strings_sketch = self._require_datasketches()
|
|
self._encode_lists = encode_lists
|
|
|
|
super().__init__(
|
|
alias_name if alias_name else f"approx_topk({str(on)})",
|
|
on=on,
|
|
ignore_nulls=True,
|
|
zero_factory=lambda: self.zero(log_capacity).serialize(),
|
|
)
|
|
|
|
def zero(self, log_capacity: int):
|
|
return self._frequent_strings_sketch(lg_max_k=log_capacity)
|
|
|
|
def aggregate_block(self, block: Block) -> bytes:
|
|
# Note: The datasketches Python bindings only expose frequent_strings_sketch
|
|
# (not type-specific variants like frequent_ints_sketch). We use pickle
|
|
# serialization as a workaround, which is less performant than native
|
|
# type-specific sketches. Revisit if type-specific bindings are added.
|
|
block_acc = BlockAccessor.for_block(block)
|
|
table = block_acc.to_arrow()
|
|
column = table.column(self.get_target_column())
|
|
sketch = self.zero(self._log_capacity)
|
|
for value in column:
|
|
py_value = value.as_py()
|
|
if self._encode_lists and isinstance(py_value, list):
|
|
for item in py_value:
|
|
if item is None:
|
|
continue
|
|
dump = pickle.dumps(item).hex()
|
|
sketch.update(dump)
|
|
elif py_value is not None:
|
|
dump = pickle.dumps(py_value).hex()
|
|
sketch.update(dump)
|
|
return sketch.serialize()
|
|
|
|
def combine(self, current_accumulator: bytes, new: bytes) -> bytes:
|
|
combined = self.zero(self._log_capacity)
|
|
combined.merge(self._frequent_strings_sketch.deserialize(current_accumulator))
|
|
combined.merge(self._frequent_strings_sketch.deserialize(new))
|
|
return combined.serialize()
|
|
|
|
def finalize(self, accumulator: bytes) -> List[Dict[str, Any]]:
|
|
from datasketches import frequent_items_error_type
|
|
|
|
column = self.get_target_column()
|
|
|
|
frequent_items = self._frequent_strings_sketch.deserialize(
|
|
accumulator
|
|
).get_frequent_items(frequent_items_error_type.NO_FALSE_NEGATIVES)
|
|
|
|
return [
|
|
{column: pickle.loads(bytes.fromhex(item[0])), "count": int(item[1])}
|
|
for item in frequent_items[: self.k]
|
|
]
|