Added support for formulas

posthog/hogql_queries/trends_query_runner.py

@@ -1,3 +1,5 @@
+from itertools import groupby
+from operator import itemgetter
 from typing import List, Optional, Any, Dict
 
 from django.utils.timezone import datetime
@@ -8,6 +10,7 @@
 from posthog.hogql.query import execute_hogql_query
 from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.query_runner import QueryRunner
+from posthog.hogql_queries.utils.formula_ast import FormulaAST
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
 from posthog.hogql_queries.utils.query_previous_period_date_range import QueryPreviousPeriodDateRange
 from posthog.models import Team
@@ -114,6 +117,9 @@ def calculate(self):
 
             res.extend(self.build_series_response(response, series_with_extra))
 
+        if self.query.trendsFilter.formula is not None:
+            res = self.apply_formula(self.query.trendsFilter.formula, res)
+
         return TrendsQueryResponse(result=res, timings=timings)
 
     def build_series_response(self, response: HogQLQueryResponse, series: SeriesWithExtras):
@@ -246,3 +252,31 @@ def setup_series(self) -> List[SeriesWithExtras]:
             return updated_series
 
         return [SeriesWithExtras(series, is_previous_period_series=False) for series in self.query.series]
+
+    def apply_formula(self, formula: str, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        if self.query.trendsFilter.compare:
+            sorted_results = sorted(results, key=itemgetter("compare_label"))
+            res = []
+            for _, group in groupby(sorted_results, key=itemgetter("compare_label")):
+                group_list = list(group)
+
+                series_data = map(lambda s: s["data"], group_list)
+                new_series_data = FormulaAST(series_data).call(formula)
+
+                new_result = group_list[0]
+                new_result["data"] = new_series_data
+                new_result["count"] = float(sum(new_series_data))
+                new_result["label"] = f"Formula ({formula})"
+
+                res.append(new_result)
+            return res
+
+        series_data = map(lambda s: s["data"], results)
+        new_series_data = FormulaAST(series_data).call(formula)
+        new_result = results[0]
+
+        new_result["data"] = new_series_data
+        new_result["count"] = float(sum(new_series_data))
+        new_result["label"] = f"Formula ({formula})"
+
+        return [new_result]

posthog/hogql_queries/utils/formula_ast.py

@@ -0,0 +1,67 @@
+import ast
+import operator
+from typing import Any, Dict, List
+
+
+class FormulaAST:
+    op_map = {
+        ast.Add: operator.add,
+        ast.Sub: operator.sub,
+        ast.Mult: operator.mul,
+        ast.Div: operator.truediv,
+        ast.Mod: operator.mod,
+        ast.Pow: operator.pow,
+    }
+    zipped_data: List[tuple[float]]
+
+    def __init__(self, data: List[List[float]]):
+        self.zipped_data = list(zip(*data))
+
+    def call(self, node: str):
+        res = []
+        for consts in self.zipped_data:
+            map = {}
+            for index, value in enumerate(consts):
+                map[chr(ord("`") + index + 1)] = value
+            result = self._evaluate(node.lower(), map)
+            res.append(result)
+        return res
+
+    def _evaluate(self, node, const_map: Dict[str, Any]):
+        if isinstance(node, (list, tuple)):
+            return [self._evaluate(sub_node, const_map) for sub_node in node]
+
+        elif isinstance(node, str):
+            return self._evaluate(ast.parse(node), const_map)
+
+        elif isinstance(node, ast.Module):
+            values = []
+            for body in node.body:
+                values.append(self._evaluate(body, const_map))
+            if len(values) == 1:
+                values = values[0]
+            return values
+
+        elif isinstance(node, ast.Expr):
+            return self._evaluate(node.value, const_map)
+
+        elif isinstance(node, ast.BinOp):
+            left = self._evaluate(node.left, const_map)
+            op = node.op
+            right = self._evaluate(node.right, const_map)
+
+            try:
+                return self.op_map[type(op)](left, right)
+            except KeyError:
+                raise ValueError(f"Operator {op.__class__.__name__} not supported")
+
+        elif isinstance(node, ast.Num):
+            return node.n
+
+        elif isinstance(node, ast.Name):
+            try:
+                return const_map[node.id]
+            except KeyError:
+                raise ValueError(f"Constant {node.id} not supported")
+
+        raise TypeError(f"Unsupported operation: {node.__class__.__name__}")