Coverage for adrs/data/datamap.py: 56%
121 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-29 05:07 +0800
1import io
2import asyncio
3import pickle
4import logging
5import polars as pl
6from typing import Self
7from datetime import datetime, timezone, timedelta
9from adrs.types import Topic, Data, SortedDataList
10from adrs.data.dataloader import DataLoader
11from adrs.data.progress import progress_context
13from .types import DataInfo
15logger = logging.getLogger(__name__)
18def dedup_data_infos_by_max_lookback_size(
19 data_infos: list[DataInfo],
20) -> list[DataInfo]:
21 """Deduplicate DataInfo by keeping the one with the maximum lookback_size for each topic."""
22 info_map = {}
23 for info in data_infos:
24 if info not in info_map:
25 info_map[info] = info
26 else:
27 if info.lookback_size > info_map[info].lookback_size:
28 info_map[info] = info
29 return list(info_map.values())
32class Datamap:
33 map: dict[Topic, SortedDataList]
34 data_infos: list[DataInfo] # Dedupped, only holds info with greatest lookback
35 topics: set[Topic]
37 def __init__(self, data_infos: list[DataInfo] = []):
38 self.map = {}
39 self.data_infos = data_infos
41 def get_lookback_size(self, topic: Topic) -> int:
42 return max(
43 map(
44 lambda di: di.lookback_size,
45 filter(lambda di: Topic.from_str(di.topic) == topic, self.data_infos),
46 )
47 )
49 def update(self, topic: Topic, data: Data):
50 lookback_size = self.get_lookback_size(topic)
52 # check for race condition: duplicate data
53 if (
54 topic in self.map
55 and self.map[topic].last_start_time() == data["start_time"]
56 ):
57 logging.warning(f"Duplicate data for topic {topic} at {data['start_time']}")
58 self.map[topic].replace_last(data)
59 return
61 # maintain the datamap
62 if topic not in self.map:
63 self.map[topic] = SortedDataList([data])
64 else:
65 self.map[topic].append(data)
66 self.map[topic].tail(lookback_size)
68 def is_ready(self) -> bool:
69 has_init_all_df = len(self.topics) == len(self.map.keys())
70 has_enough_data = all(
71 len(self.map[Topic.from_str(info.topic)]) >= info.lookback_size
72 for info in self.data_infos
73 )
74 return has_init_all_df and has_enough_data
76 def write_ipc(self) -> bytes:
77 """Serialize the Datamap into bytes (IPC for DataFrames + pickle for metadata)."""
78 payload = {"map": {}, "data_infos": self.data_infos}
80 for topic, data_list in self.map.items():
81 df = data_list.to_df()
82 buf = io.BytesIO()
83 df.write_ipc(buf)
84 payload["map"][topic] = buf.getvalue() # raw bytes # type: ignore
86 return pickle.dumps(payload)
88 @classmethod
89 def read_ipc(cls, raw: bytes) -> Self:
90 """Deserialize from bytes back into a Datamap."""
91 payload = pickle.loads(raw)
92 obj = cls()
94 # rebuild map
95 for topic, ipc_bytes in payload["map"].items():
96 buf = io.BytesIO(ipc_bytes)
97 obj.map[topic] = SortedDataList.from_df(pl.read_ipc(buf))
99 return obj
101 async def _init(
102 self,
103 dataloader: DataLoader,
104 topic: Topic,
105 start_time: datetime,
106 end_time: datetime,
107 should_lookback: bool = True,
108 ):
109 interval = topic.interval()
110 if interval is None:
111 raise Exception(f"Topic {topic} does not have an interval")
112 lookback_size = self.get_lookback_size(topic)
114 start_time = (
115 start_time - interval * lookback_size if should_lookback else start_time
116 ).replace(hour=0, minute=0, second=0, microsecond=0)
118 last_closed = topic.last_closed_time_relative(end_time, is_collect=False)
119 if last_closed is None:
120 raise Exception(f"Topic {topic} does not have a last_closed_time_relative")
122 # Skip if have already loaded before
123 loaded = self.map.get(topic)
124 first_loaded = loaded.first_start_time() if loaded is not None else None
125 last_loaded = loaded.last_start_time() if loaded is not None else None
126 if (
127 first_loaded is not None
128 and last_loaded is not None
129 and first_loaded.replace(hour=0, minute=0, second=0, microsecond=0)
130 <= start_time
131 and last_loaded >= last_closed
132 ):
133 return
135 # Explicitly override today's data because it might be incomplete
136 if end_time.date() == datetime.now(tz=timezone.utc).date():
137 end_time = end_time.replace(hour=0, minute=0, second=0, microsecond=0)
138 logger.info(
139 f"Loading data for topic {topic} from {start_time} to {end_time}"
140 )
141 df = await dataloader.load(
142 topic=str(topic),
143 start_time=start_time,
144 end_time=end_time,
145 )
146 # Today's bar may not exist yet — slow topics (e.g. day/24h) have no
147 # closed bar until the period ends, and the datasource answers the
148 # speculative [today, today+1d] request with HTTP 400. That's not
149 # fatal: fall back to the historical data we already have.
150 sdl = SortedDataList.from_df(df)
151 try:
152 today_df = await dataloader.load(
153 topic=str(topic),
154 start_time=end_time,
155 end_time=end_time + timedelta(days=1),
156 override_existing=True,
157 )
158 sdl.merge_df(today_df)
159 except Exception as e:
160 logger.warning(
161 f"Today's bar for topic {topic} not available yet "
162 f"({type(e).__name__}: {e}); using data up to {end_time}."
163 )
164 df = sdl.to_df()
165 else:
166 logger.info(
167 f"Loading data for topic {topic} from {start_time} to {end_time}"
168 )
169 df = await dataloader.load(
170 topic=str(topic),
171 start_time=start_time,
172 end_time=end_time,
173 )
174 df = df.with_columns(pl.col("start_time").dt.replace_time_zone("UTC"))
175 self.map[topic] = SortedDataList.from_df(df)
176 logger.info(f"Loaded {len(df)} datapoints for topic {topic}")
178 async def init(
179 self,
180 dataloader: DataLoader,
181 infos: list[DataInfo],
182 start_time: datetime,
183 end_time: datetime,
184 should_lookback: bool = True,
185 ):
186 self.data_infos = dedup_data_infos_by_max_lookback_size(infos + self.data_infos)
187 self.topics = {Topic.from_str(data_info.topic) for data_info in self.data_infos}
189 async with progress_context(total_topics=len(self.topics)) as outer:
191 async def _init_and_tick(topic: Topic) -> None:
192 await self._init(
193 dataloader, topic, start_time, end_time, should_lookback
194 )
195 outer.advance(1)
197 async with asyncio.TaskGroup() as tg:
198 for topic in self.topics:
199 tg.create_task(_init_and_tick(topic))
201 async def resync(
202 self,
203 topic: Topic,
204 dataloader: DataLoader,
205 ):
206 current_time = datetime.now(tz=timezone.utc)
207 interval = topic.interval()
208 if interval is None:
209 logger.error(f"Failed to resync as topic {topic} has no interval")
210 raise ValueError(f"No interval from topic {topic}")
211 limit = self.get_lookback_size(topic)
213 datas = SortedDataList.from_df(
214 await dataloader.load(
215 topic=str(topic),
216 start_time=current_time - interval * limit,
217 end_time=current_time,
218 override_existing=True,
219 )
220 )
221 if topic not in self.map:
222 self.map[topic] = datas
223 else:
224 self.map[topic].merge_df(datas.to_df())
225 self.map[topic].tail(limit)
227 logger.info(
228 f"[resync] [{topic}] successfully resynced latest {len(datas)} datapoints"
229 )
231 def get(self, info: DataInfo) -> pl.DataFrame:
232 return (
233 self.map[Topic.from_str(info.topic)]
234 .to_df()
235 .select(
236 [
237 "start_time",
238 *(pl.col(col.src).alias(col.dst) for col in info.columns),
239 ]
240 )
241 )
243 def __getitem__(self, info: DataInfo) -> pl.DataFrame:
244 return self.get(info)
246 def keys(self):
247 return self.map.keys()
249 def values(self):
250 return self.map.values()
252 def items(self):
253 return self.map.items()
255 def __len__(self):
256 return len(self.map)