aggregation.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372
  1. from typing import List, Union
  2. FIELDNAME = object()
  3. class Limit:
  4. def __init__(self, offset: int = 0, count: int = 0) -> None:
  5. self.offset = offset
  6. self.count = count
  7. def build_args(self):
  8. if self.count:
  9. return ["LIMIT", str(self.offset), str(self.count)]
  10. else:
  11. return []
  12. class Reducer:
  13. """
  14. Base reducer object for all reducers.
  15. See the `redisearch.reducers` module for the actual reducers.
  16. """
  17. NAME = None
  18. def __init__(self, *args: List[str]) -> None:
  19. self._args = args
  20. self._field = None
  21. self._alias = None
  22. def alias(self, alias: str) -> "Reducer":
  23. """
  24. Set the alias for this reducer.
  25. ### Parameters
  26. - **alias**: The value of the alias for this reducer. If this is the
  27. special value `aggregation.FIELDNAME` then this reducer will be
  28. aliased using the same name as the field upon which it operates.
  29. Note that using `FIELDNAME` is only possible on reducers which
  30. operate on a single field value.
  31. This method returns the `Reducer` object making it suitable for
  32. chaining.
  33. """
  34. if alias is FIELDNAME:
  35. if not self._field:
  36. raise ValueError("Cannot use FIELDNAME alias with no field")
  37. # Chop off initial '@'
  38. alias = self._field[1:]
  39. self._alias = alias
  40. return self
  41. @property
  42. def args(self) -> List[str]:
  43. return self._args
  44. class SortDirection:
  45. """
  46. This special class is used to indicate sort direction.
  47. """
  48. DIRSTRING = None
  49. def __init__(self, field: str) -> None:
  50. self.field = field
  51. class Asc(SortDirection):
  52. """
  53. Indicate that the given field should be sorted in ascending order
  54. """
  55. DIRSTRING = "ASC"
  56. class Desc(SortDirection):
  57. """
  58. Indicate that the given field should be sorted in descending order
  59. """
  60. DIRSTRING = "DESC"
  61. class AggregateRequest:
  62. """
  63. Aggregation request which can be passed to `Client.aggregate`.
  64. """
  65. def __init__(self, query: str = "*") -> None:
  66. """
  67. Create an aggregation request. This request may then be passed to
  68. `client.aggregate()`.
  69. In order for the request to be usable, it must contain at least one
  70. group.
  71. - **query** Query string for filtering records.
  72. All member methods (except `build_args()`)
  73. return the object itself, making them useful for chaining.
  74. """
  75. self._query = query
  76. self._aggregateplan = []
  77. self._loadfields = []
  78. self._loadall = False
  79. self._max = 0
  80. self._with_schema = False
  81. self._verbatim = False
  82. self._cursor = []
  83. self._dialect = None
  84. def load(self, *fields: List[str]) -> "AggregateRequest":
  85. """
  86. Indicate the fields to be returned in the response. These fields are
  87. returned in addition to any others implicitly specified.
  88. ### Parameters
  89. - **fields**: If fields not specified, all the fields will be loaded.
  90. Otherwise, fields should be given in the format of `@field`.
  91. """
  92. if fields:
  93. self._loadfields.extend(fields)
  94. else:
  95. self._loadall = True
  96. return self
  97. def group_by(
  98. self, fields: List[str], *reducers: Union[Reducer, List[Reducer]]
  99. ) -> "AggregateRequest":
  100. """
  101. Specify by which fields to group the aggregation.
  102. ### Parameters
  103. - **fields**: Fields to group by. This can either be a single string,
  104. or a list of strings. both cases, the field should be specified as
  105. `@field`.
  106. - **reducers**: One or more reducers. Reducers may be found in the
  107. `aggregation` module.
  108. """
  109. fields = [fields] if isinstance(fields, str) else fields
  110. reducers = [reducers] if isinstance(reducers, Reducer) else reducers
  111. ret = ["GROUPBY", str(len(fields)), *fields]
  112. for reducer in reducers:
  113. ret += ["REDUCE", reducer.NAME, str(len(reducer.args))]
  114. ret.extend(reducer.args)
  115. if reducer._alias is not None:
  116. ret += ["AS", reducer._alias]
  117. self._aggregateplan.extend(ret)
  118. return self
  119. def apply(self, **kwexpr) -> "AggregateRequest":
  120. """
  121. Specify one or more projection expressions to add to each result
  122. ### Parameters
  123. - **kwexpr**: One or more key-value pairs for a projection. The key is
  124. the alias for the projection, and the value is the projection
  125. expression itself, for example `apply(square_root="sqrt(@foo)")`
  126. """
  127. for alias, expr in kwexpr.items():
  128. ret = ["APPLY", expr]
  129. if alias is not None:
  130. ret += ["AS", alias]
  131. self._aggregateplan.extend(ret)
  132. return self
  133. def limit(self, offset: int, num: int) -> "AggregateRequest":
  134. """
  135. Sets the limit for the most recent group or query.
  136. If no group has been defined yet (via `group_by()`) then this sets
  137. the limit for the initial pool of results from the query. Otherwise,
  138. this limits the number of items operated on from the previous group.
  139. Setting a limit on the initial search results may be useful when
  140. attempting to execute an aggregation on a sample of a large data set.
  141. ### Parameters
  142. - **offset**: Result offset from which to begin paging
  143. - **num**: Number of results to return
  144. Example of sorting the initial results:
  145. ```
  146. AggregateRequest("@sale_amount:[10000, inf]")\
  147. .limit(0, 10)\
  148. .group_by("@state", r.count())
  149. ```
  150. Will only group by the states found in the first 10 results of the
  151. query `@sale_amount:[10000, inf]`. On the other hand,
  152. ```
  153. AggregateRequest("@sale_amount:[10000, inf]")\
  154. .limit(0, 1000)\
  155. .group_by("@state", r.count()\
  156. .limit(0, 10)
  157. ```
  158. Will group all the results matching the query, but only return the
  159. first 10 groups.
  160. If you only wish to return a *top-N* style query, consider using
  161. `sort_by()` instead.
  162. """
  163. _limit = Limit(offset, num)
  164. self._aggregateplan.extend(_limit.build_args())
  165. return self
  166. def sort_by(self, *fields: List[str], **kwargs) -> "AggregateRequest":
  167. """
  168. Indicate how the results should be sorted. This can also be used for
  169. *top-N* style queries
  170. ### Parameters
  171. - **fields**: The fields by which to sort. This can be either a single
  172. field or a list of fields. If you wish to specify order, you can
  173. use the `Asc` or `Desc` wrapper classes.
  174. - **max**: Maximum number of results to return. This can be
  175. used instead of `LIMIT` and is also faster.
  176. Example of sorting by `foo` ascending and `bar` descending:
  177. ```
  178. sort_by(Asc("@foo"), Desc("@bar"))
  179. ```
  180. Return the top 10 customers:
  181. ```
  182. AggregateRequest()\
  183. .group_by("@customer", r.sum("@paid").alias(FIELDNAME))\
  184. .sort_by(Desc("@paid"), max=10)
  185. ```
  186. """
  187. if isinstance(fields, (str, SortDirection)):
  188. fields = [fields]
  189. fields_args = []
  190. for f in fields:
  191. if isinstance(f, SortDirection):
  192. fields_args += [f.field, f.DIRSTRING]
  193. else:
  194. fields_args += [f]
  195. ret = ["SORTBY", str(len(fields_args))]
  196. ret.extend(fields_args)
  197. max = kwargs.get("max", 0)
  198. if max > 0:
  199. ret += ["MAX", str(max)]
  200. self._aggregateplan.extend(ret)
  201. return self
  202. def filter(self, expressions: Union[str, List[str]]) -> "AggregateRequest":
  203. """
  204. Specify filter for post-query results using predicates relating to
  205. values in the result set.
  206. ### Parameters
  207. - **fields**: Fields to group by. This can either be a single string,
  208. or a list of strings.
  209. """
  210. if isinstance(expressions, str):
  211. expressions = [expressions]
  212. for expression in expressions:
  213. self._aggregateplan.extend(["FILTER", expression])
  214. return self
  215. def with_schema(self) -> "AggregateRequest":
  216. """
  217. If set, the `schema` property will contain a list of `[field, type]`
  218. entries in the result object.
  219. """
  220. self._with_schema = True
  221. return self
  222. def verbatim(self) -> "AggregateRequest":
  223. self._verbatim = True
  224. return self
  225. def cursor(self, count: int = 0, max_idle: float = 0.0) -> "AggregateRequest":
  226. args = ["WITHCURSOR"]
  227. if count:
  228. args += ["COUNT", str(count)]
  229. if max_idle:
  230. args += ["MAXIDLE", str(max_idle * 1000)]
  231. self._cursor = args
  232. return self
  233. def build_args(self) -> List[str]:
  234. # @foo:bar ...
  235. ret = [self._query]
  236. if self._with_schema:
  237. ret.append("WITHSCHEMA")
  238. if self._verbatim:
  239. ret.append("VERBATIM")
  240. if self._cursor:
  241. ret += self._cursor
  242. if self._loadall:
  243. ret.append("LOAD")
  244. ret.append("*")
  245. elif self._loadfields:
  246. ret.append("LOAD")
  247. ret.append(str(len(self._loadfields)))
  248. ret.extend(self._loadfields)
  249. if self._dialect:
  250. ret.extend(["DIALECT", self._dialect])
  251. ret.extend(self._aggregateplan)
  252. return ret
  253. def dialect(self, dialect: int) -> "AggregateRequest":
  254. """
  255. Add a dialect field to the aggregate command.
  256. - **dialect** - dialect version to execute the query under
  257. """
  258. self._dialect = dialect
  259. return self
  260. class Cursor:
  261. def __init__(self, cid: int) -> None:
  262. self.cid = cid
  263. self.max_idle = 0
  264. self.count = 0
  265. def build_args(self):
  266. args = [str(self.cid)]
  267. if self.max_idle:
  268. args += ["MAXIDLE", str(self.max_idle)]
  269. if self.count:
  270. args += ["COUNT", str(self.count)]
  271. return args
  272. class AggregateResult:
  273. def __init__(self, rows, cursor: Cursor, schema) -> None:
  274. self.rows = rows
  275. self.cursor = cursor
  276. self.schema = schema
  277. def __repr__(self) -> (str, str):
  278. cid = self.cursor.cid if self.cursor else -1
  279. return (
  280. f"<{self.__class__.__name__} at 0x{id(self):x} "
  281. f"Rows={len(self.rows)}, Cursor={cid}>"
  282. )