Coverage for adrs/data/processor.py: 93%

1import logging 

2import polars as pl 

3import polars.selectors as cs 

4 

5from functools import reduce 

6from datetime import datetime, timedelta 

7from typing import cast 

8from abc import abstractmethod 

9 

10from adrs.types import Topic 

11 

12from .types import DataInfo 

13from .datamap import Datamap 

14 

15 

16logger = logging.getLogger(__name__) 

17 

18 

19class DataProcessor: 

20 def __init__(self, data_infos: list[DataInfo] = []): 

21 self.data_infos = data_infos 

22 

23 @abstractmethod 

24 def process( 

25 self, 

26 datamap: Datamap, 

27 last_closed_time: datetime | None = None, 

28 ) -> pl.DataFrame | None: 

29 intervals = list( 

30 Topic.from_str(info.topic).interval() for info in self.data_infos 

31 ) 

32 if None in intervals: 

33 raise Exception( 

34 "Requires custom implementation of DataProcessor since there are topics that does not have an interval." 

35 ) 

36 intervals = cast(list[timedelta], intervals) 

37 if not all(interval == intervals[0] for interval in intervals): 

38 raise Exception( 

39 "Requires custom implementation of DataProcessor since there are topics that does not have the same interval." 

40 ) 

41 

42 lookback_sizes = list(info.lookback_size for info in self.data_infos) 

43 if not all( 

44 lookback_size == lookback_sizes[0] for lookback_size in lookback_sizes 

45 ): 

46 raise Exception( 

47 "Requires custom implementation of DataProcessor since there are topics that does not have the same lookback size." 

48 ) 

49 

50 # Make sure we have collected all topics for current time 

51 topic_map = {topic: sdl for topic, sdl in datamap.items()} 

52 for info in self.data_infos: 

53 if Topic.from_str(info.topic) not in topic_map: 

54 raise Exception(f"Topic '{info.topic}' not found in datamap.") 

55 dfs = [ 

56 topic_map[Topic.from_str(info.topic)] 

57 .to_df() 

58 .select( 

59 [ 

60 "start_time", 

61 *(pl.col(col.src).alias(col.dst) for col in info.columns), 

62 ] 

63 ) 

64 for info in self.data_infos 

65 ] 

66 

67 df = reduce(lambda acc, x: acc.join(x, on="start_time"), dfs) 

68 last_joined_data_time = df["start_time"][-1] 

69 if last_closed_time is not None and last_joined_data_time < last_closed_time: 

70 logging.warning( 

71 f"[WAIT_DATA] last_closed_time is {last_closed_time} but currently only have all data up to {last_joined_data_time}" 

72 ) 

73 return None 

74 

75 missing_data_df = df.upsample( 

76 time_column="start_time", every=intervals[0] 

77 ).join(df, on="start_time", how="anti") 

78 if missing_data_df.height > 0: 

79 logging.warning(f"Missing data in datas_df: {missing_data_df}") 

80 

81 return ( 

82 df.drop_nulls() 

83 .drop_nans() 

84 .filter(~pl.any_horizontal(cs.numeric().is_infinite())) 

85 )