[core][feat] Allow aggregation function for time series data (#2031)

* [core][feat] Allow aggregation function for time series data

* fix default value

Author: aquamatthias

fixcore/fixcore/cli/command.py

@@ -6242,6 +6242,7 @@ def args_info(self) -> ArgsInfo:
                 ArgInfo("--group", expects_value=True, can_occur_multiple_times=True, help_text="<group>"),
                 ArgInfo("--filter", expects_value=True, can_occur_multiple_times=True, help_text="<var><op><value>"),
                 ArgInfo("--granularity", expects_value=True, help_text="<duration|integer>"),
+                ArgInfo("--aggregation", expects_value=True, possible_values=["avg", "sum", "min", "max"]),
             ],
             "downsample": [],
         }
@@ -6273,10 +6274,18 @@ def parse_duration_or_int(s: str) -> Union[int, timedelta]:
             parser.add_argument("--group", type=str, nargs="*", default=None)
             parser.add_argument("--filter", type=predicate_term.parse, nargs="*", default=None)
             parser.add_argument("--granularity", type=parse_duration_or_int, default=5)
+            parser.add_argument("--aggregation", choices=["avg", "sum", "min", "max"], default="avg")
             p = parser.parse_args(args_parts_unquoted_parser.parse(part))
             timeout = if_set(ctx.env.get("search_timeout"), duration)
             cursor = await self.dependencies.db_access.time_series_db.load_time_series(
-                p.name, p.start, p.end, group_by=p.group, filter_by=p.filter, granularity=p.granularity, timeout=timeout
+                p.name,
+                p.start,
+                p.end,
+                group_by=p.group,
+                filter_by=p.filter,
+                granularity=p.granularity,
+                timeout=timeout,
+                aggregation=p.aggregation,
             )
             return CLISourceContext(cursor.count(), cursor.full_count()), cursor
 

fixcore/fixcore/db/arango_query.py

@@ -875,6 +875,7 @@ def load_time_series(
     start: datetime,
     end: datetime,
     granularity: timedelta,
+    group_aggregation: Literal["avg", "sum", "min", "max"] = "avg",
     group_by: Optional[Collection[str]] = None,
     group_filter: Optional[List[Predicate]] = None,
 ) -> Tuple[str, Json]:
@@ -911,7 +912,7 @@ def load_time_series(
         group = f"group: {{ {', '.join(parts)} }},"
 
     query += f" COLLECT {', '.join(collect)} INTO group"
-    query += f" SORT group_slot RETURN {{at: group_slot, {group} v: AVG(group[*].d.v)}}"
+    query += f" SORT group_slot RETURN {{at: group_slot, {group} v: {group_aggregation}(group[*].d.v)}}"
     return query, ctx.bind_vars
 
 

fixcore/fixcore/db/timeseriesdb.py

@@ -6,7 +6,7 @@
 from datetime import timedelta, datetime, timezone
 from functools import partial
 from numbers import Number
-from typing import Optional, List, Set, Union, cast, Callable, Dict, AsyncIterator
+from typing import Optional, List, Set, Union, cast, Callable, Dict, AsyncIterator, Literal
 
 from attr import evolve, define
 from fixcore.core_config import CoreConfig
@@ -115,6 +115,7 @@ async def load_time_series(
         trafo: Optional[Callable[[Json], Json]] = None,
         batch_size: Optional[int] = None,
         timeout: Optional[timedelta] = None,
+        aggregation: Literal["avg", "sum", "min", "max"] = "avg",
     ) -> AsyncCursor:
         """
         Load time series data.
@@ -129,6 +130,7 @@ async def load_time_series(
         :param trafo: Optional transformation function to apply to each result.
         :param batch_size: Optional batch size for the query.
         :param timeout: Timeout for the query to run.
+        :param aggregation: The aggregation function to use.
         :return: A cursor to iterate over the time series data.
         """
         assert start < end, "start must be before end"
@@ -145,7 +147,9 @@ async def load_time_series(
 
         grl = max(grl, timedelta(hours=1))
 
-        qs, bv = arango_query.load_time_series(self.collection_name, name, start, end, grl, group_by, filter_by)
+        qs, bv = arango_query.load_time_series(
+            self.collection_name, name, start, end, grl, aggregation, group_by, filter_by
+        )
 
         def result_trafo(js: Json) -> Json:
             js["at"] = utc_str(datetime.fromtimestamp(js["at"], timezone.utc))
@@ -192,6 +196,7 @@ def ts_format(ts: str, js: Json) -> Json:
                             trafo=partial(ts_format, ts.name),
                             batch_size=100_000,  # The values are tiny. Use a large batch size.
                             timeout=timedelta(seconds=60),
+                            aggregation="avg",
                         )
                     ]:
                         log.info(f"Compact {ts.name} bucket {bucket} to {len(ts_data)} entries (last={ts_bucket_last})")

fixcore/fixcore/web/api.py

@@ -772,12 +772,14 @@ def parse_duration_or_int(s: str) -> Union[int, timedelta]:
         start = if_set(body.get("start"), parse_utc, utc() - timedelta(days=7))
         end = if_set(body.get("end"), parse_utc, utc())
         group_by: Optional[Set[str]] = if_set(body.get("group"), set)
+        aggregation: Literal["avg", "sum", "min", "max"] = body.get("aggregation", "avg")
+        assert aggregation in ["avg", "sum", "min", "max"], f"Invalid aggregation {aggregation}"
         filter_by: Optional[List[Predicate]] = if_set(
             body.get("filter"), lambda x: [predicate_term.parse(y) for y in x]  # type: ignore
         )
         granularity: Optional[Union[int, timedelta]] = if_set(body.get("granularity"), parse_duration_or_int)
         cursor = await deps.db_access.time_series_db.load_time_series(
-            name, start, end, group_by=group_by, filter_by=filter_by, granularity=granularity
+            name, start, end, group_by=group_by, filter_by=filter_by, granularity=granularity, aggregation=aggregation
         )
         return await self.stream_response_from_gen(
             request, cursor, count=cursor.count(), total_count=cursor.full_count()

fixcore/tests/fixcore/db/arango_query_test.py

@@ -324,7 +324,7 @@ def test_load_time_series() -> None:
         "LET m0 = (FLOOR(d.at / @b3) * @b3) + @b4 "
         "COLLECT group_slot=m0 INTO group "
         "SORT group_slot "
-        "RETURN {at: group_slot,  v: AVG(group[*].d.v)}"
+        "RETURN {at: group_slot,  v: avg(group[*].d.v)}"
     )
     assert bv == {"b0": "foo", "b1": 1699913600, "b2": 1700000000, "b3": 3600, "b4": 800}
     # no group by defined --> group by all values
@@ -334,7 +334,7 @@ def test_load_time_series() -> None:
         "LET m0 = (FLOOR(d.at / @b3) * @b3) + @b4 "
         "COLLECT group_slot=m0, complete_group=d.group INTO group "
         "SORT group_slot "
-        "RETURN {at: group_slot, group: complete_group, v: AVG(group[*].d.v)}"
+        "RETURN {at: group_slot, group: complete_group, v: avg(group[*].d.v)}"
     )
     assert bv == {"b0": "foo", "b1": 1699913600, "b2": 1700000000, "b3": 3600, "b4": 800}
     # group by specific group variables
@@ -344,7 +344,7 @@ def test_load_time_series() -> None:
         "LET m0 = (FLOOR(d.at / @b3) * @b3) + @b4 "
         "COLLECT group_slot=m0, group_a=d.group.a, group_b=d.group.b INTO group "
         "SORT group_slot "
-        "RETURN {at: group_slot, group: { a: group_a, b: group_b }, v: AVG(group[*].d.v)}"
+        "RETURN {at: group_slot, group: { a: group_a, b: group_b }, v: avg(group[*].d.v)}"
     )
     assert bv == {"b0": "foo", "b1": 1699913600, "b2": 1700000000, "b3": 3600, "b4": 800}
     # group by specific group variables and filter by group variables
@@ -357,6 +357,6 @@ def test_load_time_series() -> None:
         "LET m0 = (FLOOR(d.at / @b4) * @b4) + @b5 "
         "COLLECT group_slot=m0, group_a=d.group.a, group_b=d.group.b INTO group "
         "SORT group_slot "
-        "RETURN {at: group_slot, group: { a: group_a, b: group_b }, v: AVG(group[*].d.v)}"
+        "RETURN {at: group_slot, group: { a: group_a, b: group_b }, v: avg(group[*].d.v)}"
     )
     assert bv == {"b0": "foo", "b1": 1699913600, "b2": 1700000000, "b3": "a", "b4": 3600, "b5": 800}