1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
|
from typing import List, Callable
from sortedcontainers import SortedList
from fakeredis import _msgs as msgs
from fakeredis._command_args_parsing import extract_args
from fakeredis._commands import command, CommandItem, Int, Key, Float
from fakeredis._helpers import SimpleString, SimpleError, OK, Database
class TDigest(SortedList):
def __init__(self, compression: int = 100):
super().__init__()
self.compression = compression
class TDigestCommandsMixin:
_encodefloat: Callable[[float, bool], bytes]
def __init__(self, *args, **kwargs):
self._db: Database
@command(
name="TDIGEST.CREATE",
fixed=(Key(TDigest),),
repeat=(bytes,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_create(self, key: CommandItem, *args: bytes) -> SimpleString:
if key.value is not None:
raise SimpleError(msgs.TDIGEST_KEY_EXISTS)
(compression,), left_args = extract_args(args, ("+compression",))
if compression is None:
compression = 100
key.update(TDigest(compression))
return OK
@command(
name="TDIGEST.RESET",
fixed=(Key(TDigest),),
repeat=(),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_reset(self, key: CommandItem) -> SimpleString:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
key.value.clear()
return OK
@command(
name="TDIGEST.ADD",
fixed=(Key(TDigest), Float),
repeat=(Float,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_add(self, key: CommandItem, *values: float) -> SimpleString:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
# parsing
try:
values_to_add = [float(val) for val in values]
except ValueError:
raise SimpleError(msgs.TDIGEST_ERROR_PARSING_VALUE)
# adding
key.value.update(values_to_add)
return OK
@command(
name="TDIGEST.MERGE",
fixed=(Key(TDigest), Int, bytes),
repeat=(bytes,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_merge(self, dest: CommandItem, numkeys: int, *args: bytes) -> SimpleString:
if len(args) < numkeys:
raise SimpleError(msgs.WRONG_ARGS_MSG6.format("tdigest.merge"))
sources_names = args[:numkeys]
(compression, override), _ = extract_args(args[numkeys:], ("+compression", "override"))
sources = [self._db.get(name).value for name in sources_names if name in self._db]
if len(sources) != len(sources_names):
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if override:
if dest.value is None:
compression = compression or max([source.compression for source in sources])
dest.value = TDigest(compression)
else:
dest.value.clear()
if dest.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
for source in sources:
dest.value.update(source)
dest.updated()
return OK
@command(
name="TDIGEST.MAX", fixed=(Key(TDigest),), repeat=(), flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL
)
def tdigest_max(self, key: CommandItem) -> bytes:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return b"nan"
return str(key.value[-1]).encode()
@command(
name="TDIGEST.MIN", fixed=(Key(TDigest),), repeat=(), flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL
)
def tdigest_min(self, key: CommandItem) -> bytes:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return b"nan"
return str(key.value[0]).encode()
@command(
name="TDIGEST.RANK",
fixed=(Key(TDigest), Float),
repeat=(Float,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_rank(self, key: CommandItem, *values: float) -> List[int]:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return [
-2,
]
res = []
for v in values:
if v > key.value[-1]:
res.append(len(key.value))
else:
res.append(key.value.bisect_right(v) - 1)
return res
@command(
name="TDIGEST.REVRANK",
fixed=(Key(TDigest), Float),
repeat=(Float,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_revrank(self, key: CommandItem, *values: float) -> List[int]:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return [
-2,
]
res = []
length = len(key.value)
for v in values:
loc = key.value.bisect_right(v)
if loc == length:
loc += 1
res.append(length - loc)
return res
@command(
name="TDIGEST.QUANTILE",
fixed=(Key(TDigest), Float),
repeat=(Float,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_quantile(self, key: CommandItem, *quantiles: float) -> List[bytes]:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) <= 1:
return [
b"nan",
]
res: List[bytes] = []
for q in quantiles:
if q < 0 or q > 1:
raise SimpleError(msgs.TDIGEST_BAD_QUANTILE)
ind = int(q * len(key.value))
if ind == len(key.value):
ind -= 1
res.append(self._encodefloat(key.value[ind], True))
return res
@command(
name="TDIGEST.CDF",
fixed=(Key(TDigest), Float),
repeat=(Float,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_cdf(self, key: CommandItem, *values: float) -> List[bytes]: # Cumulative Distribution Function
"""Returns, for each input value, an estimation of the fraction (floating-point) of
(observations smaller than the given value + half the observations equal to the given value).
"""
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
res: List[bytes] = []
for v in values:
left = key.value.bisect_left(v)
right = key.value.bisect_right(v)
if right == 0:
res.append(b"0")
elif left == len(key.value):
res.append(b"1")
else:
res.append(self._encodefloat(float((left + right) / 2) / len(key.value), True))
return res
@command(
name="TDIGEST.INFO", fixed=(Key(TDigest),), repeat=(), flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL
)
def tdigest_info(self, key: CommandItem) -> List[bytes]:
return [
b"Compression",
key.value.compression,
b"Capacity",
len(key.value),
b"Merged nodes",
len(key.value),
b"Unmerged nodes",
0,
b"Merged weight",
len(key.value),
b"Unmerged weight",
0,
b"Observations",
len(key.value),
b"Total compressions",
len(key.value),
b"Memory usage",
len(key.value),
]
@command(
name="TDIGEST.TRIMMED_MEAN",
fixed=(Key(TDigest), Float, Float),
repeat=(),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_trimmed_mean(self, key: CommandItem, lower: float, upper: float) -> bytes:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if lower < 0 or upper > 1 or lower > upper:
raise SimpleError(msgs.TDIGEST_BAD_QUANTILE)
if len(key.value) == 0:
return b"nan"
left = int(lower * len(key.value))
right = int(upper * len(key.value))
res = key.value[(left + right) // 2]
if right == left + 1:
res = (res + key.value[right]) / 2
return self._encodefloat(res, True)
@command(
name="TDIGEST.BYRANK",
fixed=(Key(TDigest), Int),
repeat=(Int,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_byrank(self, key: CommandItem, *ranks: int) -> List[bytes]:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return [
b"nan",
]
res: List[bytes] = []
for rank in ranks:
if rank < 0:
raise SimpleError(msgs.TDIGEST_BAD_RANK)
if rank >= len(key.value):
res.append(b"inf")
else:
res.append(self._encodefloat(key.value[rank], True))
return res
@command(
name="TDIGEST.BYREVRANK",
fixed=(Key(TDigest), Int),
repeat=(Int,),
flags=msgs.FLAG_DO_NOT_CREATE + msgs.FLAG_LEAVE_EMPTY_VAL,
)
def tdigest_byrevrank(self, key: CommandItem, *ranks: int) -> List[bytes]:
if key.value is None:
raise SimpleError(msgs.TDIGEST_KEY_NOT_EXISTS)
if len(key.value) == 0:
return [
b"nan",
]
res: List[bytes] = []
for rank in ranks:
if rank < 0:
raise SimpleError(msgs.TDIGEST_BAD_RANK)
if rank >= len(key.value):
res.append(b"-inf")
else:
res.append(self._encodefloat(key.value[-rank - 1], True))
return res
|
