262 lines
7.1 KiB
Python
262 lines
7.1 KiB
Python
from typing import Iterable, Union
|
|
|
|
from .aggregation import Asc, Desc, Reducer, SortDirection
|
|
|
|
|
|
def _ensure_at_prefix(name: str) -> str:
|
|
"""Return ``name`` with a single leading ``@`` prefix."""
|
|
return name if name.startswith("@") else "@" + name
|
|
|
|
|
|
class FieldOnlyReducer(Reducer):
|
|
"""See https://redis.io/docs/interact/search-and-query/search/aggregations/"""
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
self._field = field
|
|
|
|
|
|
class count(Reducer):
|
|
"""
|
|
Counts the number of results in the group
|
|
"""
|
|
|
|
NAME = "COUNT"
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
|
|
|
|
class sum(FieldOnlyReducer):
|
|
"""
|
|
Calculates the sum of all the values in the given fields within the group
|
|
"""
|
|
|
|
NAME = "SUM"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class min(FieldOnlyReducer):
|
|
"""
|
|
Calculates the smallest value in the given field within the group
|
|
"""
|
|
|
|
NAME = "MIN"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class max(FieldOnlyReducer):
|
|
"""
|
|
Calculates the largest value in the given field within the group
|
|
"""
|
|
|
|
NAME = "MAX"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class avg(FieldOnlyReducer):
|
|
"""
|
|
Calculates the mean value in the given field within the group
|
|
"""
|
|
|
|
NAME = "AVG"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class tolist(FieldOnlyReducer):
|
|
"""
|
|
Returns all the matched properties in a list
|
|
"""
|
|
|
|
NAME = "TOLIST"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class count_distinct(FieldOnlyReducer):
|
|
"""
|
|
Calculate the number of distinct values contained in all the results in
|
|
the group for the given field
|
|
"""
|
|
|
|
NAME = "COUNT_DISTINCT"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class count_distinctish(FieldOnlyReducer):
|
|
"""
|
|
Calculate the number of distinct values contained in all the results in the
|
|
group for the given field. This uses a faster algorithm than
|
|
`count_distinct` but is less accurate
|
|
"""
|
|
|
|
NAME = "COUNT_DISTINCTISH"
|
|
|
|
|
|
class quantile(Reducer):
|
|
"""
|
|
Return the value for the nth percentile within the range of values for the
|
|
field within the group.
|
|
"""
|
|
|
|
NAME = "QUANTILE"
|
|
|
|
def __init__(self, field: str, pct: float) -> None:
|
|
super().__init__(field, str(pct))
|
|
self._field = field
|
|
|
|
|
|
class stddev(FieldOnlyReducer):
|
|
"""
|
|
Return the standard deviation for the values within the group
|
|
"""
|
|
|
|
NAME = "STDDEV"
|
|
|
|
def __init__(self, field: str) -> None:
|
|
super().__init__(field)
|
|
|
|
|
|
class first_value(Reducer):
|
|
"""
|
|
Selects the first value within the group according to sorting parameters
|
|
"""
|
|
|
|
NAME = "FIRST_VALUE"
|
|
|
|
def __init__(self, field: str, *byfields: Union[Asc, Desc]) -> None:
|
|
"""
|
|
Selects the first value of the given field within the group.
|
|
|
|
### Parameter
|
|
|
|
- **field**: Source field used for the value
|
|
- **byfields**: How to sort the results. This can be either the
|
|
*class* of `aggregation.Asc` or `aggregation.Desc` in which
|
|
case the field `field` is also used as the sort input.
|
|
|
|
`byfields` can also be one or more *instances* of `Asc` or `Desc`
|
|
indicating the sort order for these fields
|
|
"""
|
|
|
|
fieldstrs = []
|
|
if (
|
|
len(byfields) == 1
|
|
and isinstance(byfields[0], type)
|
|
and issubclass(byfields[0], SortDirection)
|
|
):
|
|
byfields = [byfields[0](field)]
|
|
|
|
for f in byfields:
|
|
fieldstrs += [f.field, f.DIRSTRING]
|
|
|
|
args = [field]
|
|
if fieldstrs:
|
|
args += ["BY"] + fieldstrs
|
|
super().__init__(*args)
|
|
self._field = field
|
|
|
|
|
|
class random_sample(Reducer):
|
|
"""
|
|
Returns a random sample of items from the dataset, from the given property
|
|
"""
|
|
|
|
NAME = "RANDOM_SAMPLE"
|
|
|
|
def __init__(self, field: str, size: int) -> None:
|
|
"""
|
|
### Parameter
|
|
|
|
**field**: Field to sample from
|
|
**size**: Return this many items (can be less)
|
|
"""
|
|
args = [field, str(size)]
|
|
super().__init__(*args)
|
|
self._field = field
|
|
|
|
|
|
class collect(Reducer):
|
|
"""
|
|
Gathers the rows of each ``GROUPBY`` group, projects a chosen set of
|
|
fields from every row, optionally deduplicates, sorts, and limits them,
|
|
and returns them as an array of per-entry maps under the reducer alias.
|
|
|
|
``COLLECT`` is a preview feature gated behind
|
|
``search-enable-unstable-features``. See
|
|
`FT.AGGREGATE <https://redis.io/commands/ft.aggregate>`_.
|
|
"""
|
|
|
|
NAME = "COLLECT"
|
|
|
|
def __init__(
|
|
self,
|
|
fields: str | Iterable[str] = "*",
|
|
distinct: bool = False,
|
|
sort_by: Asc | Desc | Iterable[Asc | Desc] | None = None,
|
|
limit: tuple[int, int] | None = None,
|
|
) -> None:
|
|
"""
|
|
### Parameters
|
|
|
|
- **fields**: The fields to project from each collected row. Either the
|
|
literal ``"*"`` (project every field present in the pipeline at this
|
|
stage) or a field name / iterable of field names. Names are
|
|
normalized to a single leading ``@`` on the wire; output map keys
|
|
are the bare names.
|
|
- **distinct**: When ``True``, emit ``DISTINCT`` to deduplicate entries
|
|
with identical projected fields. Forward-compatible: this option is
|
|
not yet implemented by the server and currently produces a server
|
|
error when sent.
|
|
- **sort_by**: An ``Asc``/``Desc`` instance or an iterable of them, used
|
|
to order the collected entries within each group.
|
|
- **limit**: An ``(offset, count)`` pair. Returns at most ``count``
|
|
entries per group after skipping ``offset``. With ``sort_by`` this
|
|
acts as a top-N selection.
|
|
"""
|
|
args: list[str] = []
|
|
|
|
# FIELDS (required)
|
|
if fields == "*":
|
|
args += ["FIELDS", "*"]
|
|
else:
|
|
names = [fields] if isinstance(fields, str) else list(fields)
|
|
if not names or any(not n.strip() for n in names):
|
|
raise ValueError(
|
|
"collect fields must be '*' or a non-empty list of names"
|
|
)
|
|
names = [_ensure_at_prefix(n) for n in names]
|
|
args += ["FIELDS", str(len(names))] + names
|
|
|
|
# DISTINCT (optional)
|
|
if distinct:
|
|
args += ["DISTINCT"]
|
|
|
|
# SORTBY (optional)
|
|
if sort_by is not None:
|
|
sort_fields = [sort_by] if isinstance(sort_by, (Asc, Desc)) else sort_by
|
|
sort_args: list[str] = []
|
|
for f in sort_fields:
|
|
sort_args += [_ensure_at_prefix(f.field), f.DIRSTRING]
|
|
if not sort_args:
|
|
raise ValueError("collect sort_by must contain at least one field")
|
|
args += ["SORTBY", str(len(sort_args))] + sort_args
|
|
|
|
# LIMIT (optional)
|
|
if limit is not None:
|
|
offset, count = limit
|
|
args += ["LIMIT", str(offset), str(count)]
|
|
|
|
super().__init__(*args)
|