-
Notifications
You must be signed in to change notification settings - Fork 37
Expand file tree
/
Copy pathdatapoints.py
More file actions
2688 lines (2368 loc) · 143 KB
/
Copy pathdatapoints.py
File metadata and controls
2688 lines (2368 loc) · 143 KB
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
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
from __future__ import annotations
import asyncio
import datetime
import functools
import heapq
import itertools
import math
from abc import ABC, abstractmethod
from collections import Counter, defaultdict
from collections.abc import AsyncIterator, Iterable, Iterator, MutableSequence, Sequence
from itertools import chain
from operator import itemgetter
from typing import (
TYPE_CHECKING,
Any,
Literal,
NamedTuple,
TypeGuard,
TypeVar,
cast,
overload,
)
from zoneinfo import ZoneInfo
from typing_extensions import Self
from cognite.client._api.datapoint_tasks import (
BaseDpsFetchSubtask,
BaseTaskOrchestrator,
_DpsQueryValidator,
_FullDatapointsQuery,
)
from cognite.client._api.synthetic_time_series import SyntheticDatapointsAPI
from cognite.client._api_client import APIClient
from cognite.client._proto.data_point_list_response_pb2 import DataPointListItem, DataPointListResponse
from cognite.client.constants import DEFAULT_DATAPOINTS_CHUNK_SIZE
from cognite.client.data_classes import (
Datapoints,
DatapointsArray,
DatapointsArrayList,
DatapointsList,
DatapointsQuery,
LatestDatapoint,
LatestDatapointList,
LatestDatapointQuery,
)
from cognite.client.data_classes.data_modeling.ids import NodeId
from cognite.client.data_classes.datapoint_aggregates import Aggregate
from cognite.client.exceptions import CogniteAPIError, CogniteNotFoundError
from cognite.client.utils import _json_extended as _json
from cognite.client.utils._auxiliary import (
exactly_one_is_not_none,
find_duplicates,
is_finite,
is_positive,
split_into_chunks,
split_into_n_parts,
unpack_items,
unpack_items_in_payload,
)
from cognite.client.utils._concurrency import AsyncSDKTask, execute_async_tasks
from cognite.client.utils._identifier import Identifier, IdentifierSequence, IdentifierSequenceCore
from cognite.client.utils._importing import local_import
from cognite.client.utils._time import (
timestamp_to_ms,
)
from cognite.client.utils._validation import validate_user_input_dict_with_identifier
from cognite.client.utils.useful_types import SequenceNotStr
if TYPE_CHECKING:
import pandas as pd
from cognite.client import AsyncCogniteClient
from cognite.client.config import ClientConfig
PoolSubtaskType = tuple[float, int, BaseDpsFetchSubtask]
_T = TypeVar("_T")
class DpsFetchStrategy(ABC):
def __init__(self, dps_client: DatapointsAPI, all_queries: list[DatapointsQuery]) -> None:
from cognite.client import global_config
self.dps_client = dps_client
self.all_queries = all_queries
self.agg_queries, self.raw_queries = self.split_queries(all_queries)
self.concurrency_limit = global_config.concurrency_settings.datapoints.read
self.n_queries = len(all_queries)
self.semaphore = dps_client._get_semaphore("read")
@staticmethod
def split_queries(all_queries: list[DatapointsQuery]) -> tuple[list[DatapointsQuery], list[DatapointsQuery]]:
split_qs: tuple[list[DatapointsQuery], list[DatapointsQuery]] = [], []
for query in all_queries:
split_qs[query.is_raw_query].append(query)
return split_qs
async def fetch_all_datapoints(self) -> DatapointsList:
return DatapointsList(
[ts_task.get_result(use_numpy=False) async for ts_task in self._fetch_all(use_numpy=False)],
).set_client_ref(self.dps_client._cognite_client)
async def fetch_all_datapoints_numpy(self) -> DatapointsArrayList:
return DatapointsArrayList(
[ts_task.get_result(use_numpy=True) async for ts_task in self._fetch_all(use_numpy=True)],
).set_client_ref(self.dps_client._cognite_client)
async def _request_datapoints(self, payload: dict[str, Any]) -> Sequence[DataPointListItem]:
(res := DataPointListResponse()).MergeFromString(
(
await self.dps_client._post(
f"{self.dps_client._RESOURCE_PATH}/list",
json=payload,
headers={"accept": "application/protobuf"},
semaphore=self.semaphore,
)
).content
)
return res.items
def _raise_if_missing(self, to_raise: set[DatapointsQuery]) -> None:
if to_raise:
raise CogniteNotFoundError(
"Time series not found",
code=400,
missing=[q.identifier.as_dict(camel_case=False) for q in to_raise],
x_request_id="<no failing request was made>",
cluster=self.dps_client._config.cdf_cluster,
project=self.dps_client._config.project,
)
def _raise_if_unknown_failed(self, unknown_failed: list[BaseException]) -> None:
if unknown_failed:
# Typically this happens when asking for aggregates for a string time series.
# We don't do anything fancy like ExceptionGroup (Python 3.11+), just raise the first:
raise unknown_failed[0]
@abstractmethod
def _fetch_all(self, use_numpy: bool) -> AsyncIterator[BaseTaskOrchestrator]:
raise NotImplementedError
class EagerDpsFetcher(DpsFetchStrategy):
"""A datapoints fetching strategy to make small queries as fast as possible.
Is used when the number of time series to fetch is smaller than or equal to the number of `concurrency_limit`, so
that each worker only fetches datapoints for a single time series per request (this maximises throughput
according to the API docs). This does -not- mean that we assign a time series to each worker! All available
workers will fetch data for the same time series to speed up fetching. To make this work, the time domain is
split based on the density of datapoints returned and other heuristics like granularity (e.g. given '1h', at
most 168 datapoints exist per week).
"""
async def _fetch_all(self, use_numpy: bool) -> AsyncIterator[BaseTaskOrchestrator]:
missing_to_raise: set[DatapointsQuery] = set()
futures_dct, ts_task_lookup = self._create_initial_tasks(use_numpy)
# Run until all top level tasks are complete:
while futures_dct:
done, _ = await asyncio.wait(futures_dct, return_when=asyncio.FIRST_COMPLETED)
for future in done: # not guaranteed to be just one done task
ts_task = (subtask := futures_dct.pop(future)).parent
res = self._get_result_with_exception_handling(future, ts_task, ts_task_lookup, missing_to_raise)
if res is None:
continue
elif missing_to_raise:
# We are going to raise anyway, kill task:
ts_task.is_done = True # The 'missing' may come from a different task, so we need this
continue
# We may dynamically split subtasks based on what % of time range was returned:
if new_subtasks := subtask.store_partial_result(res):
self._queue_new_subtasks(futures_dct, new_subtasks)
if ts_task.is_done:
# Reduce peak memory consumption by finalizing as soon as tasks finish:
ts_task.finalize_datapoints()
continue
elif subtask.is_done:
continue
# Put the subtask back into the pool:
self._queue_new_subtasks(futures_dct, [subtask])
self._raise_if_missing(missing_to_raise)
# Return only non-missing time series tasks in correct order given by `all_queries`:
for task in filter(None, map(ts_task_lookup.get, self.all_queries)):
yield task
def _create_initial_tasks(
self,
use_numpy: bool,
) -> tuple[dict[asyncio.Task, BaseDpsFetchSubtask], dict[DatapointsQuery, BaseTaskOrchestrator]]:
futures_dct: dict[asyncio.Task, BaseDpsFetchSubtask] = {}
ts_task_lookup = {}
for query in self.all_queries:
ts_task = ts_task_lookup[query] = query.task_orchestrator(query=query, eager_mode=True, use_numpy=use_numpy)
for subtask in ts_task.split_into_subtasks(self.concurrency_limit, self.n_queries):
payload = {"items": [subtask.get_next_payload_item()], "ignoreUnknownIds": False}
future = asyncio.create_task(self._request_datapoints(payload))
futures_dct[future] = subtask
return futures_dct, ts_task_lookup
def _queue_new_subtasks(
self,
futures_dct: dict[asyncio.Task, BaseDpsFetchSubtask],
new_subtasks: Sequence[BaseDpsFetchSubtask],
) -> None:
for subtask in new_subtasks:
payload = {"items": [subtask.get_next_payload_item()]}
future = asyncio.create_task(self._request_datapoints(payload))
futures_dct[future] = subtask
def _get_result_with_exception_handling(
self,
future: asyncio.Task,
ts_task: BaseTaskOrchestrator,
ts_task_lookup: dict[DatapointsQuery, BaseTaskOrchestrator],
missing_to_raise: set[DatapointsQuery],
) -> DataPointListItem | None:
try:
return future.result()[0]
except CogniteAPIError as err:
# If the error is not "missing ts", we immediately reraise:
if not err.missing or err.code != 400:
raise
# The query decides if we can ignore it. If not, we store it so that we later can
# raise one exception with -all- missing-non-ignorable time series:
if not ts_task.query.ignore_unknown_ids:
missing_to_raise.add(ts_task.query)
ts_task.is_done = True
ts_task_lookup.pop(ts_task.query, None)
return None
class ChunkingDpsFetcher(DpsFetchStrategy):
"""A datapoints fetching strategy to make large queries faster through the grouping of more than one
time series per request.
The main underlying assumption is that "the more time series are queried, the lower the average density".
Is used when the number of time series to fetch is larger than the number of `concurrency_limit`. How many
time series are chunked per request is dynamic and is decided by the overall number to fetch, their
individual number of datapoints and whether raw- or aggregate datapoints are asked for since
they are independent in requests - as long as the total number of time series does not exceed `_FETCH_TS_LIMIT`.
"""
def __init__(self, *args: Any) -> None:
super().__init__(*args)
self._counter = itertools.count().__next__
# To chunk efficiently, we have subtask pools (heap queues) that we use to prioritise subtasks
# when building/combining subtasks into a full query:
self.raw_subtask_pool: list[PoolSubtaskType] = []
self.agg_subtask_pool: list[PoolSubtaskType] = []
self.subtask_pools = (self.agg_subtask_pool, self.raw_subtask_pool)
async def _fetch_all(self, use_numpy: bool) -> AsyncIterator[BaseTaskOrchestrator]:
# The initial tasks are important - as they tell us which time series are missing, which
# are string, which are sparse... We use this info when we choose the best fetch-strategy.
ts_task_lookup, missing_to_raise = {}, set()
initial_query_limits, initial_futures_dct = self._create_initial_tasks()
# Note to future dev: As of 3.13, we can do 'async for future in asyncio.as_completed(...)' and it
# will return the Tasks - not coroutines. Thus, we can't do this micro-optimization yet and need to
# use asyncio.wait instead:
unknown_failed = []
done, _ = await asyncio.wait(initial_futures_dct, return_when=asyncio.ALL_COMPLETED)
for future in done:
if exc := future.exception():
# We don't immediately reraise here to avoid 'Task exception was never retrieved' (when multiple fail):
unknown_failed.append(exc)
continue
res_lst = future.result()
new_ts_tasks, chunk_missing = self._create_ts_tasks_and_handle_missing(
res_lst, initial_futures_dct.pop(future), initial_query_limits, use_numpy
)
missing_to_raise.update(chunk_missing)
ts_task_lookup.update(new_ts_tasks)
self._raise_if_unknown_failed(unknown_failed)
self._raise_if_missing(missing_to_raise)
if ts_tasks_left := self._update_queries_with_new_chunking_limit(ts_task_lookup):
self._add_to_subtask_pools(
chain.from_iterable(
task.split_into_subtasks(concurrency_limit=self.concurrency_limit, n_tot_queries=len(ts_tasks_left))
for task in ts_tasks_left
)
)
futures_dct: dict[asyncio.Task, list[BaseDpsFetchSubtask]] = {}
await self._queue_new_subtasks(futures_dct)
await self._fetch_until_complete(futures_dct)
# Return only non-missing time series tasks in correct order given by `all_queries`:
for task in filter(None, map(ts_task_lookup.get, self.all_queries)):
yield task
async def _fetch_until_complete(self, futures_dct: dict[asyncio.Task, list[BaseDpsFetchSubtask]]) -> None:
while futures_dct:
done, _ = await asyncio.wait(futures_dct, return_when=asyncio.FIRST_COMPLETED)
res_lst_all, subtask_lst_all = [], []
for future in done: # not guaranteed to be just one done task
res_lst, subtask_lst = future.result(), futures_dct.pop(future)
subtask_lst_all.extend(subtask_lst)
res_lst_all.extend(res_lst)
for subtask, res in zip(subtask_lst_all, res_lst_all):
# We may dynamically split subtasks based on what % of time range was returned:
if new_subtasks := subtask.store_partial_result(res):
self._add_to_subtask_pools(new_subtasks)
if not subtask.is_done:
self._add_to_subtask_pools([subtask])
# Check each ts task in current batch once if finished:
for ts_task in {sub.parent for sub in subtask_lst_all}:
if ts_task.is_done:
ts_task.finalize_datapoints()
await self._queue_new_subtasks(futures_dct)
def _create_initial_tasks(
self,
) -> tuple[dict[DatapointsQuery, int], dict[asyncio.Task, tuple[list[DatapointsQuery], list[DatapointsQuery]]]]:
initial_query_limits: dict[DatapointsQuery, int] = {}
initial_futures_dct: dict[asyncio.Task, tuple[list[DatapointsQuery], list[DatapointsQuery]]] = {}
# Optimal queries uses the entire worker pool. We may be forced to use more (queue) when we
# can't fit all individual time series (maxes out at `_FETCH_TS_LIMIT * concurrency_limit`):
n_queries = max(self.concurrency_limit, math.ceil(self.n_queries / self.dps_client._FETCH_TS_LIMIT))
for query_chunks in zip(
split_into_n_parts(self.agg_queries, n=n_queries),
split_into_n_parts(self.raw_queries, n=n_queries),
strict=True,
):
if not any(query_chunks):
break # Not all workers needed (at least for now)
# Agg and raw limits are independent in the query, so we max out on both:
items = []
for queries, max_lim in zip(query_chunks, (self.dps_client._DPS_LIMIT_AGG, self.dps_client._DPS_LIMIT_RAW)):
maxed_limits = self._find_initial_query_limits([q.capped_limit for q in queries], max_lim)
chunk_query_limits = dict(zip(queries, maxed_limits))
initial_query_limits.update(chunk_query_limits)
for query, limit in chunk_query_limits.items():
(item := query.to_payload_item())["limit"] = limit
items.append(item)
payload = {"items": items, "ignoreUnknownIds": True}
future = asyncio.create_task(self._request_datapoints(payload))
initial_futures_dct[future] = query_chunks
return initial_query_limits, initial_futures_dct
def _create_ts_tasks_and_handle_missing(
self,
res: Sequence[DataPointListItem],
chunk_queues: tuple[list[DatapointsQuery], list[DatapointsQuery]],
initial_query_limits: dict[DatapointsQuery, int],
use_numpy: bool,
) -> tuple[dict[DatapointsQuery, BaseTaskOrchestrator], set[DatapointsQuery]]:
if len(res) == sum(map(len, chunk_queues)):
to_raise: set[DatapointsQuery] = set()
else:
# We have at least 1 missing time series:
chunk_queues, to_raise = self._handle_missing_ts(res, *chunk_queues)
# Align initial res with corresponding queries and create tasks:
ts_tasks = {
query: query.task_orchestrator(
query=query,
eager_mode=False,
use_numpy=use_numpy,
first_dps_batch=res,
first_limit=initial_query_limits[query],
)
for res, query in zip(res, chain(*chunk_queues))
}
return ts_tasks, to_raise
def _add_to_subtask_pools(self, new_subtasks: Iterable[BaseDpsFetchSubtask]) -> None:
for task in new_subtasks:
# We leverage how tuples are compared to prioritise items. First `payload limit` (to easily group
# smaller queries), then counter to always break ties, but keep order (never use tasks themselves):
limit = min(task.parent.get_remaining_limit(), task.max_query_limit)
new_subtask: PoolSubtaskType = (limit, self._counter(), task)
heapq.heappush(self.subtask_pools[task.parent.query.is_raw_query], new_subtask)
async def _queue_new_subtasks(
self,
futures_dct: dict[asyncio.Task, list[BaseDpsFetchSubtask]],
) -> None:
# This may seem silly (to ask the semaphore for unused capacity), but the logic is sound:
# we want to combine subtasks into requests *as late as possible* for optimal chunking.
# So, simply put, we wait to create a new request until we know we can schedule it immediately.
# That leaves room for as many done tasks as possible to have been added to the subtask pools
# in the meantime (we ofc also check that we have subtasks to schedule).
while self.semaphore_has_bandwidth(self.semaphore) and any(self.subtask_pools):
payload, subtask_lst = self._combine_subtasks_into_new_request()
future = asyncio.create_task(self._request_datapoints(payload))
futures_dct[future] = subtask_lst
await asyncio.sleep(0) # Yield control
@staticmethod
def semaphore_has_bandwidth(sem: asyncio.BoundedSemaphore) -> bool:
# We should ideally not use private attributes, but the semaphore is simple enough and we don't need
# guarantees - just a hint that it is time to combine time series task into a new "request task" that
# we can then queue.
return sem._value > 0
def _combine_subtasks_into_new_request(self) -> tuple[dict[str, list], list[BaseDpsFetchSubtask]]:
next_items: list[dict[str, Any]] = []
next_subtasks: list[BaseDpsFetchSubtask] = []
fetch_limits = (self.dps_client._DPS_LIMIT_AGG, self.dps_client._DPS_LIMIT_RAW)
for task_pool, request_max_limit, is_raw in zip(self.subtask_pools, fetch_limits, (False, True)):
if not task_pool:
continue
limit_used = 0 # Dps limit for raw and agg is independent (in the same query)
while task_pool:
if len(next_items) + 1 > self.dps_client._FETCH_TS_LIMIT:
# Hard limit on N ts, quit immediately (even if below dps limit):
payload = {"items": next_items}
return payload, next_subtasks
# Highest priority task i.e. the smallest limit, is always at index 0 (heap magic):
*_, next_task = task_pool[0]
next_payload_item = next_task.get_next_payload_item()
next_limit = next_payload_item["limit"]
if limit_used + next_limit <= request_max_limit:
next_items.append(next_payload_item)
next_subtasks.append(next_task)
limit_used += next_limit
heapq.heappop(task_pool)
else:
break
return {"items": next_items}, next_subtasks
@staticmethod
def _decide_individual_query_limit(query: DatapointsQuery, ts_task: BaseTaskOrchestrator, n_ts_limit: int) -> int:
# For a better estimate, we use first ts of first batch instead of `query.start`:
batch_start, batch_end = ts_task.start_ts_first_batch, ts_task.end_ts_first_batch
est_remaining_dps = ts_task.n_dps_first_batch * (query.end_ms - batch_end) / (batch_end - batch_start)
# To use the full request limit on a single ts, the estimate must be >> max_limit (raw/agg dependent):
if est_remaining_dps > 5 * (max_limit := query.max_query_limit):
return max_limit
# To build full queries, we want dps chunk sizes that easily combines to 100 %, e.g. we don't want
# 1/3, because 1/3 + 1/2 -> 83 % full, but 1/8 is fine as 4 x 1/8 + 1/2 = 100 %:
for chunk_size in (1, 2, 4, 8, 16, 32):
if est_remaining_dps > max_limit // chunk_size:
return max_limit // (2 * chunk_size)
return max_limit // n_ts_limit
def _update_queries_with_new_chunking_limit(
self, ts_task_lookup: dict[DatapointsQuery, BaseTaskOrchestrator]
) -> list[BaseTaskOrchestrator]:
remaining_tasks = {}
for query, ts_task in ts_task_lookup.items():
if ts_task.is_done:
ts_task.finalize_datapoints()
else:
remaining_tasks[query] = ts_task
tot_raw = sum(q.is_raw_query for q in remaining_tasks)
if tot_raw <= self.concurrency_limit >= len(remaining_tasks) - tot_raw:
# Number of raw and agg tasks independently <= concurrency_limit, so we're basically doing "eager fetching",
# but it's worth noting that we'll still chunk 1 raw + 1 agg per query (if they exist).
return list(remaining_tasks.values())
# Many tasks left, decide how we'll chunk'em by estimating which are dense (and need little to
# no chunking), and which are not (...and may be grouped - and how "tightly"):
for query, ts_task in remaining_tasks.items():
est_limit = self._decide_individual_query_limit(query, ts_task, self.dps_client._FETCH_TS_LIMIT)
query.max_query_limit = est_limit
return list(remaining_tasks.values())
@staticmethod
def _find_initial_query_limits(limits: list[int], max_limit: int) -> list[int]:
actual_lims = [0] * len(limits)
not_done = set(range(len(limits)))
while not_done:
part = max_limit // len(not_done)
if not part:
# We still might not have not reached max_limit, but we can no longer distribute evenly
break
rm_idx = set()
for i in not_done:
i_part = min(part, limits[i]) # A query of limit=10 does not need more of max_limit than 10
actual_lims[i] += i_part
max_limit -= i_part
if i_part == limits[i]:
rm_idx.add(i)
else:
limits[i] -= i_part
not_done -= rm_idx
return actual_lims
@staticmethod
def _handle_missing_ts(
res: Sequence[DataPointListItem],
agg_queries: list[DatapointsQuery],
raw_queries: list[DatapointsQuery],
) -> tuple[tuple[list[DatapointsQuery], list[DatapointsQuery]], set[DatapointsQuery]]:
to_raise = set()
not_missing = (
{("id", r.id) for r in res}
| {("externalId", r.externalId) for r in res}
| {("instanceId", NodeId(r.instanceId.space, r.instanceId.externalId)) for r in res}
)
for query in chain(agg_queries, raw_queries):
query.is_missing = query.identifier.as_tuple() not in not_missing
# Only raise for those time series that can't be missing (individually customisable parameter):
if query.is_missing and not query.ignore_unknown_ids:
to_raise.add(query)
agg_queries = [q for q in agg_queries if not q.is_missing]
raw_queries = [q for q in raw_queries if not q.is_missing]
return (agg_queries, raw_queries), to_raise
class DatapointsAPI(APIClient):
_RESOURCE_PATH = "/timeseries/data"
def __init__(self, config: ClientConfig, api_version: str | None, cognite_client: AsyncCogniteClient) -> None:
super().__init__(config, api_version, cognite_client)
self.synthetic = SyntheticDatapointsAPI(config, api_version, cognite_client)
self._FETCH_TS_LIMIT = 100
self._DPS_LIMIT_AGG = 10_000
self._DPS_LIMIT_RAW = 100_000
self._DPS_INSERT_LIMIT = 100_000
self._RETRIEVE_LATEST_LIMIT = 100
self._POST_DPS_OBJECTS_LIMIT = 10_000
self.query_validator = _DpsQueryValidator(dps_limit_raw=self._DPS_LIMIT_RAW, dps_limit_agg=self._DPS_LIMIT_AGG)
def _get_semaphore(self, operation: Literal["read", "write", "delete"]) -> asyncio.BoundedSemaphore:
from cognite.client import global_config
return global_config.concurrency_settings.datapoints._semaphore_factory(
operation, project=self._cognite_client.config.project
)
@overload
def __call__(
self,
queries: DatapointsQuery,
*,
return_arrays: Literal[True] = True,
chunk_size_datapoints: int = DEFAULT_DATAPOINTS_CHUNK_SIZE,
chunk_size_time_series: int | None = None,
) -> AsyncIterator[DatapointsArray]: ...
@overload
def __call__(
self,
queries: Sequence[DatapointsQuery],
*,
return_arrays: Literal[True] = True,
chunk_size_datapoints: int = DEFAULT_DATAPOINTS_CHUNK_SIZE,
chunk_size_time_series: int | None = None,
) -> AsyncIterator[DatapointsArrayList]: ...
@overload
def __call__(
self,
queries: DatapointsQuery,
*,
return_arrays: Literal[False],
chunk_size_datapoints: int = DEFAULT_DATAPOINTS_CHUNK_SIZE,
chunk_size_time_series: int | None = None,
) -> AsyncIterator[Datapoints]: ...
@overload
def __call__(
self,
queries: Sequence[DatapointsQuery],
*,
return_arrays: Literal[False],
chunk_size_datapoints: int = DEFAULT_DATAPOINTS_CHUNK_SIZE,
chunk_size_time_series: int | None = None,
) -> AsyncIterator[DatapointsList]: ...
async def __call__(
self,
queries: DatapointsQuery | Sequence[DatapointsQuery],
*,
chunk_size_datapoints: int = DEFAULT_DATAPOINTS_CHUNK_SIZE,
chunk_size_time_series: int | None = None,
return_arrays: bool = True,
) -> AsyncIterator[DatapointsArray | DatapointsArrayList | Datapoints | DatapointsList]:
"""`Iterate through datapoints in chunks, for one or more time series. <https://api-docs.cognite.com/20230101/tag/Time-series/operation/getMultiTimeSeriesDatapoints>`_
Note:
Control memory usage by specifying ``chunk_size_time_series``, how many time series to iterate simultaneously and ``chunk_size_datapoints``,
how many datapoints to yield per iteration (per individual time series). See full example in examples. Note that in order to make efficient
use of the API request limits, this method will never hold less than 100k datapoints in memory at a time, per time series.
If you run with memory constraints, use ``return_arrays=True`` (the default).
No empty chunk is ever returned.
Args:
queries (DatapointsQuery | Sequence[DatapointsQuery]): Query, or queries, using id, external_id or instance_id for the time series to fetch data for, with individual settings specified. The options 'limit' and 'include_outside_points' are not supported when iterating.
chunk_size_datapoints (int): The number of datapoints per time series to yield per iteration. Must evenly divide 100k OR be an integer multiple of 100k. Default: 100_000.
chunk_size_time_series (int | None): The max number of time series to yield per iteration (varies as time series get exhausted, but is never empty). Default: None (all given queries are iterated at the same time).
return_arrays (bool): Whether to return the datapoints as numpy arrays. Default: True.
Yields:
DatapointsArray | DatapointsArrayList | Datapoints | DatapointsList: If return_arrays=True, a ``DatapointsArray`` object containing the datapoints chunk, or a ``DatapointsArrayList`` if multiple time series were asked for. When False, a ``Datapoints`` object containing the datapoints chunk, or a ``DatapointsList`` if multiple time series were asked for.
Examples:
Iterate through the datapoints of a single time series with external_id="foo", in chunks of 25k:
>>> from cognite.client import CogniteClient
>>> from cognite.client.data_classes import DatapointsQuery
>>> client = CogniteClient()
>>> # async_client = AsyncCogniteClient() # another option
>>> query = DatapointsQuery(external_id="foo", start="2w-ago")
>>> for chunk in client.time_series.data(query, chunk_size_datapoints=25_000):
... pass # do something with the datapoints chunk
Iterate through datapoints from multiple time series, and do not return them as memory-efficient numpy arrays.
As one or more time series get exhausted (no more data), they are no longer part of the returned "chunk list".
Note that the order is still preserved (for the remaining).
If you run with ``chunk_size_time_series=None``, an easy way to check when a time series is exhausted is to
use the ``.get`` method, as illustrated below:
>>> from cognite.client.data_classes.data_modeling import NodeId
>>> queries = [
... DatapointsQuery(id=123),
... DatapointsQuery(external_id="foo"),
... DatapointsQuery(instance_id=NodeId("my-space", "my-ts-xid")),
... ]
>>> for chunk_lst in client.time_series.data(query, return_arrays=False):
... if chunk_lst.get(id=123) is None:
... print("Time series with id=123 has no more datapoints!")
A likely use case for iterating datapoints is to clone data from one project to another, while keeping a low memory
footprint and without having to write very custom logic involving count aggregates (which won't work for string data)
or do time-domain splitting yourself.
Here's an example of how to do so efficiently, while including bad- and uncertain data (``ignore_bad_datapoints=False``) and
copying status codes (``include_status=True``). This is automatically taken care of when the Datapoints(-Array) objects are passed
directly to an insert method. The only assumption below is that the time series have already been created in the target project.
>>> from cognite.client.utils import MIN_TIMESTAMP_MS, MAX_TIMESTAMP_MS
>>> target_client = CogniteClient()
>>> ts_to_copy = client.time_series.list(data_set_external_ids="my-use-case")
>>> queries = [
... DatapointsQuery(
... external_id=ts.external_id,
... include_status=True,
... ignore_bad_datapoints=False,
... start=MIN_TIMESTAMP_MS,
... end=MAX_TIMESTAMP_MS + 1, # end is exclusive
... )
... for ts in ts_to_copy
... ]
>>> for dps_chunk in client.time_series.data(
... queries, # may be several thousand time series...
... chunk_size_time_series=20, # control memory usage by specifying how many to iterate at a time
... chunk_size_datapoints=100_000,
... ):
... target_client.time_series.data.insert_multiple(
... [{"external_id": dps.external_id, "datapoints": dps} for dps in dps_chunk]
... )
"""
# To make efficient usage of the API, we don't want a chunk size like 10 to send a million API requests when we can
# get 10k/100k datapoints per request. Thus, we round up the given chunk size to the nearest integer multiple of 100k,
# then subdivide and yield client-side (we use the raw limit also when dealing with aggregates):
request_limit = self._DPS_LIMIT_RAW * math.ceil(chunk_size_datapoints / self._DPS_LIMIT_RAW)
if not is_finite(chunk_size_datapoints) or (
chunk_size_datapoints != request_limit and request_limit % chunk_size_datapoints
):
raise ValueError(
"The 'chunk_size_datapoints' must be a positive integer that evenly divides 100k OR an integer multiple of 100k "
f"(to ensure efficient API usage), not {chunk_size_datapoints}."
)
if not (chunk_size_time_series is None or is_positive(chunk_size_time_series)):
raise ValueError(
f"'chunk_size_time_series' must be a positive integer or None, not {chunk_size_time_series}"
)
user_queries = [queries] if (is_single := isinstance(queries, DatapointsQuery)) else queries
dps_lst_cls: type[DatapointsArrayList | DatapointsList] = (
DatapointsArrayList if return_arrays else DatapointsList
)
for uq in user_queries:
if uq.include_outside_points is True or uq.limit is not DatapointsQuery._NOT_SET:
raise ValueError(
"When iterating datapoints, the options 'include_outside_points' and 'limit' are not supported."
)
if dupes := find_duplicates(uq.identifier.as_primitive() for uq in user_queries):
raise ValueError(f"When iterating datapoints, identifiers must be unique! Duplicates found for: {dupes}")
alive_queries = {
uq.identifier: DatapointsQuery.valid_from_user_query(uq, limit=request_limit, include_outside_points=False)
for uq in user_queries
}
self.query_validator(alive_queries.values())
dps_lst: DatapointsArrayList | DatapointsList
chunk_fn = functools.partial(split_into_chunks, chunk_size=chunk_size_datapoints)
while alive_queries:
to_fetch_queries = list(itertools.islice(alive_queries.values(), chunk_size_time_series))
fetcher = self._select_dps_fetch_strategy(to_fetch_queries)(self, to_fetch_queries)
if return_arrays:
dps_lst = await fetcher.fetch_all_datapoints_numpy()
else:
dps_lst = await fetcher.fetch_all_datapoints()
self._update_alive_queries_and_do_manual_cursoring(alive_queries, dps_lst, to_fetch_queries, request_limit)
# We should never yield an empty chunk, so we filter out empty or exhausted time series from result
# (need to rebuild to not keep references to those empty in various private "id lookups")
dps_lst = dps_lst_cls(list(filter(None, dps_lst)))
if not any(dps_lst):
if alive_queries:
continue
break
if chunk_size_datapoints == request_limit:
yield dps_lst[0] if is_single else dps_lst
elif is_single:
for chunk in chunk_fn(dps_lst[0]):
yield chunk # type: ignore [misc]
else:
for all_chunks in itertools.zip_longest(*map(chunk_fn, dps_lst)):
# Filter out dps as ts get exhausted, then rebuild the Dps(Array)List container and yield chunk:
yield dps_lst_cls(list(filter(None, all_chunks))) # type: ignore [arg-type]
@staticmethod
def _update_alive_queries_and_do_manual_cursoring(
alive_queries: dict[Identifier, DatapointsQuery],
dps_lst: DatapointsArrayList | DatapointsList,
to_fetch_queries: list[DatapointsQuery],
request_limit: int,
) -> None:
for query in to_fetch_queries:
ident = query.identifier
dps = dps_lst.get(**{ident.name(camel_case=False): ident.as_primitive()})
if isinstance(dps, list):
raise RuntimeError(
"When iterating datapoints, identifiers must be unique! You cannot get around this by passing "
"several of [id, external_id, instance_id] for the same underlying time series."
)
# Update query.start for next iteration if ts is not yet exhausted:
if dps and len(dps) == request_limit:
new_start = dps[-1].timestamp + 1
if query.end_ms > new_start:
query.start = new_start # manual cursoring ftw
continue
alive_queries.pop(ident)
@overload
async def retrieve(
self,
*,
id: int | DatapointsQuery,
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> Datapoints | None: ...
@overload
async def retrieve(
self,
*,
id: Sequence[int | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
external_id: str | DatapointsQuery,
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> Datapoints | None: ...
@overload
async def retrieve(
self,
*,
external_id: SequenceNotStr[str | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
instance_id: NodeId | DatapointsQuery,
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> Datapoints | None: ...
@overload
async def retrieve(
self,
*,
instance_id: Sequence[NodeId | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
id: None | int | DatapointsQuery | Sequence[int | DatapointsQuery],
external_id: None | str | DatapointsQuery | SequenceNotStr[str | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
id: None | int | DatapointsQuery | Sequence[int | DatapointsQuery],
instance_id: None | NodeId | DatapointsQuery | Sequence[NodeId | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
external_id: None | str | DatapointsQuery | SequenceNotStr[str | DatapointsQuery],
instance_id: None | NodeId | DatapointsQuery | Sequence[NodeId | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
@overload
async def retrieve(
self,
*,
id: None | int | DatapointsQuery | Sequence[int | DatapointsQuery],
external_id: None | str | DatapointsQuery | SequenceNotStr[str | DatapointsQuery],
instance_id: None | NodeId | DatapointsQuery | Sequence[NodeId | DatapointsQuery],
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> DatapointsList: ...
async def retrieve(
self,
*,
id: None | int | DatapointsQuery | Sequence[int | DatapointsQuery] = None,
external_id: None | str | DatapointsQuery | SequenceNotStr[str | DatapointsQuery] = None,
instance_id: None | NodeId | DatapointsQuery | Sequence[NodeId | DatapointsQuery] = None,
start: int | str | datetime.datetime | None = None,
end: int | str | datetime.datetime | None = None,
aggregates: Aggregate | str | list[Aggregate | str] | None = None,
granularity: str | None = None,
timezone: str | datetime.timezone | ZoneInfo | None = None,
target_unit: str | None = None,
target_unit_system: str | None = None,
limit: int | None = None,
include_outside_points: bool = False,
ignore_unknown_ids: bool = False,
include_status: bool = False,
ignore_bad_datapoints: bool = True,
treat_uncertain_as_bad: bool = True,
) -> Datapoints | DatapointsList | None:
"""`Retrieve datapoints for one or more time series <https://api-docs.cognite.com/20230101/tag/Time-series/operation/getMultiTimeSeriesDatapoints>`_.
**Performance guide**:
In order to retrieve millions of datapoints as efficiently as possible, here are a few guidelines:
1. Make *one* call to retrieve and fetch all time series in go, rather than making multiple calls (if your memory allows it). The SDK will optimize retrieval strategy for you!
2. For best speed, and significantly lower memory usage, consider using ``retrieve_arrays(...)`` which uses ``numpy.ndarrays`` for data storage.
3. Unlimited queries (``limit=None``) are most performant as they are always fetched in parallel, for any number of requested time series, even one.
4. Limited queries, (e.g. ``limit=500_000``) are much less performant, at least for large limits, as each individual time series is fetched serially (we can't predict where on the timeline the datapoints are). Thus parallelisation is only used when asking for multiple "limited" time series.
5. Try to avoid specifying `start` and `end` to be very far from the actual data: If you have data from 2000 to 2015, don't use start=0 (1970).
6. Using ``timezone`` and/or calendar granularities like month/quarter/year in aggregate queries comes at a penalty as they are expensive for the API to compute.
Warning:
When using the AsyncCogniteClient, always ``await`` the result of this method and never run multiple calls concurrently (e.g. using asyncio.gather).
You can pass as many queries as you like to a single call, and the SDK will optimize the retrieval strategy for you intelligently.