basic breakdown working

posthog/hogql_queries/insights/trends/breakdown.py

@@ -12,7 +12,6 @@
 from posthog.hogql_queries.insights.trends.display import TrendsDisplay
 from posthog.hogql_queries.insights.trends.utils import (
     get_properties_chain,
-    series_event_name,
 )
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
 from posthog.models.filters.mixins.utils import cached_property
@@ -222,7 +221,7 @@ def _get_breakdown_values(self) -> List[str | int]:
         with self.timings.measure("breakdown_values_query"):
             breakdown = BreakdownValues(
                 team=self.team,
-                event_name=series_event_name(self.series) or "",
+                series=self.series,
                 events_filter=self.events_filter,
                 chart_display_type=self._trends_display().display_type,
                 breakdown_filter=self.query.breakdownFilter,

posthog/hogql_queries/insights/trends/breakdown_values.py

@@ -5,7 +5,8 @@
 from posthog.hogql.query import execute_hogql_query
 from posthog.hogql_queries.insights.trends.utils import get_properties_chain
 from posthog.models.team.team import Team
-from posthog.schema import BreakdownFilter, BreakdownType, ChartDisplayType
+from posthog.schema import BreakdownFilter, BreakdownType, ChartDisplayType, SeriesType, DataWarehouseNode
+from functools import cached_property
 
 BREAKDOWN_OTHER_STRING_LABEL = "$$_posthog_breakdown_other_$$"
 BREAKDOWN_OTHER_NUMERIC_LABEL = 9007199254740991  # pow(2, 53) - 1, for JS compatibility
@@ -15,7 +16,7 @@
 
 class BreakdownValues:
     team: Team
-    event_name: str
+    series: SeriesType
     breakdown_field: Union[str, float, List[Union[str, float]]]
     breakdown_type: BreakdownType
     events_filter: ast.Expr
@@ -28,13 +29,13 @@ class BreakdownValues:
     def __init__(
         self,
         team: Team,
-        event_name: str,
+        series: SeriesType,
         events_filter: ast.Expr,
         chart_display_type: ChartDisplayType,
         breakdown_filter: BreakdownFilter,
     ):
         self.team = team
-        self.event_name = event_name
+        self.series = series
         self.breakdown_field = breakdown_filter.breakdown  # type: ignore
         self.breakdown_type = breakdown_filter.breakdown_type  # type: ignore
         self.events_filter = events_filter
@@ -88,9 +89,8 @@ def get_breakdown_values(self) -> List[str | int]:
             """
                 SELECT
                     {select_field},
-                    count(e.uuid) as count
-                FROM
-                    events e
+                    count({id_field}) as count
+                FROM {table} e
                 WHERE
                     {events_where}
                 GROUP BY
@@ -104,6 +104,8 @@ def get_breakdown_values(self) -> List[str | int]:
                 "events_where": self.events_filter,
                 "select_field": select_field,
                 "breakdown_limit": ast.Constant(value=breakdown_limit),
+                "table": self._table,
+                "id_field": self._id_field,
             },
         )
 
@@ -166,3 +168,17 @@ def _to_bucketing_expression(self) -> ast.Expr:
             qunatile_expression = f"quantiles({','.join([f'{quantile:.2f}' for quantile in quantiles])})(value)"
 
         return parse_expr(f"arrayCompact(arrayMap(x -> floor(x, 2), {qunatile_expression}))")
+
+    @cached_property
+    def _id_field(self) -> ast.Field:
+        if isinstance(self.series, DataWarehouseNode):
+            return ast.Field(chain=["e", self.series.id_field])
+
+        return ast.Field(chain=["e", "uuid"])
+
+    @cached_property
+    def _table(self) -> ast.Field:
+        if isinstance(self.series, DataWarehouseNode):
+            return ast.Field(chain=[self.series.table_name])
+
+        return ast.Field(chain=["events"])

posthog/hogql_queries/insights/trends/test/test_data_warehouse_query_builder.py

@@ -7,6 +7,8 @@
 from posthog.hogql_queries.insights.trends.data_warehouse_trends_query_builder import DataWarehouseTrendsQueryBuilder
 from posthog.hogql_queries.utils.query_date_range import QueryDateRange
 from posthog.schema import (
+    BreakdownFilter,
+    BreakdownType,
     DateRange,
     DataWarehouseNode,
     TrendsQuery,
@@ -187,3 +189,33 @@ def test_trends_property(self):
         assert response.columns is not None
         assert set(response.columns).issubset({"date", "total"})
         assert response.results[0][1] == [1, 0, 0, 0, 0, 0, 0]
+
+    def test_trends_breakdown(self):
+        table_name = self.create_parquet_file()
+
+        trends_query = TrendsQuery(
+            kind="TrendsQuery",
+            dateRange=DateRange(date_from="2023-01-01"),
+            series=[DataWarehouseNode(table_name=table_name, id_field="id", timestamp_field="created")],
+            breakdownFilter=BreakdownFilter(breakdown_type=BreakdownType.data_warehouse, breakdown="prop_1"),
+        )
+
+        with freeze_time("2023-01-07"):
+            response = self.get_response(trends_query=trends_query)
+
+        assert response.columns is not None
+        assert set(response.columns).issubset({"date", "total", "breakdown_value"})
+        assert response.results[0][1] == [1, 0, 0, 0, 0, 0, 0]
+        assert response.results[0][2] == "a"
+
+        assert response.results[1][1] == [0, 1, 0, 0, 0, 0, 0]
+        assert response.results[1][2] == "b"
+
+        assert response.results[2][1] == [0, 0, 1, 0, 0, 0, 0]
+        assert response.results[2][2] == "c"
+
+        assert response.results[3][1] == [0, 0, 0, 1, 0, 0, 0]
+        assert response.results[3][2] == "d"
+
+        assert response.results[4][1] == [0, 0, 0, 0, 0, 0, 0]
+        assert response.results[4][2] == "$$_posthog_breakdown_other_$$"

posthog/hogql_queries/insights/trends/utils.py

@@ -9,7 +9,9 @@ def series_event_name(series: SeriesType) -> str | None:
 
 
 def get_properties_chain(
-    breakdown_type: Union[Literal["person"], Literal["session"], Literal["group"], Literal["event"]],
+    breakdown_type: Union[
+        Literal["person"], Literal["session"], Literal["group"], Literal["event"], Literal["data_warehouse"]
+    ],
     breakdown_field: str,
     group_type_index: Optional[float | int],
 ) -> List[str | int]:
@@ -25,4 +27,7 @@ def get_properties_chain(
     elif breakdown_type == "group" and group_type_index is None:
         raise Exception("group_type_index missing from params")
 
+    if breakdown_type == "data_warehouse":
+        return [breakdown_field]
+
     return ["properties", breakdown_field]

posthog/schema.py

@@ -97,6 +97,7 @@ class BreakdownAttributionType(str, Enum):
 class BreakdownType(str, Enum):
     cohort = "cohort"
     person = "person"
+    data_warehouse = "data_warehouse"
     event = "event"
     group = "group"
     session = "session"