feat(experiments): add metric breakdown injector for mean metrics

products/experiments/backend/hogql_queries/experiment_mean_query_builder.py

@@ -10,7 +10,6 @@
     is_session_property_metric,
     validate_session_property,
 )
-from products.experiments.backend.hogql_queries.breakdown_injector import BreakdownInjector
 from products.experiments.backend.hogql_queries.metric_source import MetricSourceInfo
 
 if TYPE_CHECKING:
@@ -393,9 +392,9 @@ def build_mean_query(self) -> ast.SelectQuery:
 
         assert isinstance(query, ast.SelectQuery)
 
-        # Inject breakdown columns into the query AST. The metric-event injector is funnel-only,
-        # so mean metrics always use the property-from-exposure BreakdownInjector.
-        if isinstance(self._b.breakdown_injector, BreakdownInjector):
+        # Inject breakdown columns into the query AST. Both injectors expose
+        # inject_mean_breakdown_columns (old: exposure source; new: metric-event source + limit).
+        if self._b.breakdown_injector:
             self._b.breakdown_injector.inject_mean_breakdown_columns(query, final_cte_name="entity_metrics")
 
         return query
@@ -492,9 +491,9 @@ def build_mean_query_with_winsorization(self) -> ast.SelectQuery:
 
         assert isinstance(query, ast.SelectQuery)
 
-        # Inject breakdown columns into the query AST. The metric-event injector is funnel-only,
-        # so mean metrics always use the property-from-exposure BreakdownInjector.
-        if isinstance(self._b.breakdown_injector, BreakdownInjector):
+        # Inject breakdown columns into the query AST. Both injectors expose
+        # inject_mean_breakdown_columns (old: exposure source; new: metric-event source + limit).
+        if self._b.breakdown_injector:
             self._b.breakdown_injector.inject_mean_breakdown_columns(query, final_cte_name="winsorized_entity_metrics")
 
         return query

products/experiments/backend/hogql_queries/experiment_query_runner.py

@@ -544,12 +544,15 @@ def _get_experiment_query(self) -> ast.SelectQuery:
 
         breakdowns = self._get_breakdowns_for_builder()
         breakdown_injector: MetricBreakdownInjector | None = None
-        # Data warehouse funnels route through the legacy UNION ALL builder, which replaces the
-        # metric_events CTE and drops the injected breakdown columns, so gate them out until the
-        # injector supports that path.
+        # Metric types migrated to the metric-event breakdown injector. Others fall back to the
+        # old BreakdownInjector until they migrate (then BreakdownInjector is deleted).
+        migrated_metric = isinstance(self.metric, ExperimentFunnelMetric | ExperimentMeanMetric)
+        # Data warehouse sources aren't supported by the injector yet: funnels route through the
+        # legacy builder that drops the injected columns, and DW mean metrics select from the
+        # warehouse table where event/person/session breakdown chains don't resolve.
         if (
             breakdowns
-            and isinstance(self.metric, ExperimentFunnelMetric)
+            and migrated_metric
             and not self.is_data_warehouse_query
             and self._metric_event_breakdowns_enabled()
         ):

products/experiments/backend/hogql_queries/metric_breakdown_injector.py

@@ -1,23 +1,26 @@
 """
-Metric-event breakdown injection for experiment funnel queries.
+Metric-event breakdown injection for experiment queries.
 
 Unlike ``BreakdownInjector`` (which attributes the breakdown value from the
 exposure event), this injector reads the breakdown property off the *metric*
-event and attributes it across funnel steps using the metric's configured
-``breakdownAttributionType`` — aligning experiment funnel breakdowns with how
-insights funnels behave.
+event. Funnels attribute across steps using the metric's configured
+``breakdownAttributionType``; mean metrics attribute first-touch. This aligns
+experiment breakdowns with how insights behave.
 
-Scope: funnel metrics only. This is a standalone class (it does not subclass
+Scope: funnel and mean metrics. This is a standalone class (it does not subclass
 ``BreakdownInjector``); the old injector is removed once all metric types
 migrate to metric-event breakdowns.
 """
 
-from typing import cast
+from typing import Union, cast
 
 from posthog.schema import (
     Breakdown,
     BreakdownAttributionType,
     ExperimentFunnelMetric,
+    ExperimentMeanMetric,
+    ExperimentRatioMetric,
+    ExperimentRetentionMetric,
     MultipleBreakdownType,
     StepOrderValue,
 )
@@ -29,9 +32,11 @@
 from posthog.hogql_queries.insights.trends.utils import get_properties_chain
 from posthog.hogql_queries.insights.utils.breakdowns import BREAKDOWN_NULL_STRING_LABEL, BREAKDOWN_OTHER_STRING_LABEL
 
+MetricType = Union[ExperimentFunnelMetric, ExperimentMeanMetric, ExperimentRatioMetric, ExperimentRetentionMetric]
+
 
 class MetricBreakdownInjector:
-    def __init__(self, breakdowns: list[Breakdown], metric: ExperimentFunnelMetric):
+    def __init__(self, breakdowns: list[Breakdown], metric: MetricType):
         self.breakdowns = breakdowns
         self.metric = metric
 
@@ -94,6 +99,7 @@ def _attribution_agg_and_steps(self) -> tuple[str, list[int]]:
         - step (unordered, "Any step"): argMinIf across all metric step columns, since an
           unordered funnel has no fixed step position and any matching step attributes.
         """
+        assert isinstance(self.metric, ExperimentFunnelMetric)
         attribution = self.metric.breakdownAttributionType or BreakdownAttributionType.FIRST_TOUCH
         num_metric_steps = len(self.metric.series)
         is_unordered = self.metric.funnel_order_type == StepOrderValue.UNORDERED
@@ -180,20 +186,27 @@ def _top_breakdowns_subquery(self, aliases: list[str]) -> ast.SelectQuery:
         )
         return subquery
 
-    def _inject_final_breakdown_columns(self, query: ast.SelectQuery, aliases: list[str]) -> None:
+    def _inject_final_breakdown_columns(
+        self, query: ast.SelectQuery, aliases: list[str], final_cte_name: str = "entity_metrics"
+    ) -> None:
         """Surface breakdown columns in the outer SELECT (after variant) and GROUP BY.
 
         Applies the top-N + "Other" limit: breakdown tuples beyond the limit (ranked by
         user count) are relabeled to BREAKDOWN_OTHER_STRING_LABEL before the final GROUP BY,
         capping output cardinality. The relabel collapses the whole tuple together so all
         breakdown columns of an "Other" row carry the Other label.
+
+        ``final_cte_name`` is the CTE the outer SELECT reads from (``entity_metrics`` for funnels
+        and non-winsorized mean, ``winsorized_entity_metrics`` for winsorized mean). Ranking is
+        always computed from ``entity_metrics`` since winsorization doesn't change the set of
+        (entity, variant, breakdown) tuples.
         """
-        # Membership test against the top-N set. Single breakdown compares the scalar value
-        # directly; multiple breakdowns compare the value tuple element-wise.
+        # Membership test against the top-N set, always ranked from entity_metrics (pre-winsorization).
+        # Single breakdown compares the scalar value directly; multiple breakdowns compare the tuple.
         if len(aliases) == 1:
-            left: ast.Expr = ast.Field(chain=["entity_metrics", aliases[0]])
+            left: ast.Expr = ast.Field(chain=[final_cte_name, aliases[0]])
         else:
-            left = ast.Tuple(exprs=[ast.Field(chain=["entity_metrics", alias]) for alias in aliases])
+            left = ast.Tuple(exprs=[ast.Field(chain=[final_cte_name, alias]) for alias in aliases])
         in_top = ast.CompareOperation(
             op=ast.CompareOperationOp.In,
             left=left,
@@ -205,7 +218,7 @@ def _inject_final_breakdown_columns(self, query: ast.SelectQuery, aliases: list[
                 "if({in_top}, {value}, {other})",
                 placeholders={
                     "in_top": in_top,
-                    "value": ast.Field(chain=["entity_metrics", alias]),
+                    "value": ast.Field(chain=[final_cte_name, alias]),
                     "other": ast.Constant(value=BREAKDOWN_OTHER_STRING_LABEL),
                 },
             )
@@ -267,3 +280,152 @@ def inject_funnel_breakdown_columns_optimized(self, query: ast.SelectQuery) -> N
                     )
 
         self._inject_final_breakdown_columns(query, aliases)
+
+    def inject_mean_breakdown_columns(self, query: ast.SelectQuery, final_cte_name: str = "entity_metrics") -> None:
+        """Inject breakdown columns into a mean query.
+
+        Reads the breakdown property off the *metric event* (in ``metric_events``) and
+        attributes each user to a single bucket via **first-touch** — ``argMin`` over the
+        metric event timestamp. Mean has no attribution modes (unlike funnels), so the
+        breakdown is always taken from the user's first qualifying metric event.
+
+        Known limitation: with CUPED the ``entity_metrics`` join widens to (conversion window
+        OR pre-exposure lookback), so ``argMin`` can attribute the breakdown from an event that
+        happened before exposure. Attribution here does not add its own exposure-time guard; a
+        follow-up should exclude pre-window events from the attribution when CUPED is enabled.
+
+        ``final_cte_name`` is ``entity_metrics`` for the plain query or
+        ``winsorized_entity_metrics`` for the winsorized query; winsorization needs the
+        breakdown carried through its percentiles/join CTEs.
+
+        Session-property metrics have an extra ``metric_events_by_session`` layer; their
+        ``metric_events`` CTE joins exposures and can't see raw event properties, so the
+        breakdown is read (and deduped) in the by-session layer instead. See
+        ``_read_breakdown_into_metric_events``.
+        """
+        if not self._has_breakdown():
+            return
+
+        aliases = self._get_breakdown_aliases()
+
+        # The session-property shape reads the breakdown a layer earlier; the attribution
+        # timestamp is the session's first event time rather than the metric event timestamp.
+        is_session_property = bool(query.ctes and "metric_events_by_session" in query.ctes)
+        timestamp_field = "first_event_timestamp" if is_session_property else "timestamp"
+
+        self._read_breakdown_into_metric_events(query, aliases, is_session_property)
+
+        # Attribute first-touch per user: argMin over the (session-)first event timestamp.
+        if query.ctes and "entity_metrics" in query.ctes:
+            entity_metrics_cte = query.ctes["entity_metrics"]
+            if isinstance(entity_metrics_cte, ast.CTE) and isinstance(entity_metrics_cte.expr, ast.SelectQuery):
+                for alias in aliases:
+                    entity_metrics_cte.expr.select.append(
+                        ast.Alias(
+                            alias=alias,
+                            expr=ast.Call(
+                                name="argMin",
+                                args=[
+                                    ast.Field(chain=["metric_events", alias]),
+                                    ast.Field(chain=["metric_events", timestamp_field]),
+                                ],
+                            ),
+                        )
+                    )
+
+        # Carry the attributed breakdown through the winsorization CTEs.
+        self._inject_winsorization_breakdown_columns(query, aliases, final_cte_name)
+
+        self._inject_final_breakdown_columns(query, aliases, final_cte_name=final_cte_name)
+
+    def _read_breakdown_into_metric_events(
+        self, query: ast.SelectQuery, aliases: list[str], is_session_property: bool
+    ) -> None:
+        """Surface the breakdown value (and an attribution timestamp) in the ``metric_events`` CTE.
+
+        Standard path: read the breakdown directly off the metric event in ``metric_events``.
+
+        Session-property path: ``metric_events`` joins exposures and can't see raw event
+        properties, so read the breakdown off raw events in ``metric_events_by_session``
+        (deduped per session via ``any()``, mirroring ``session_value``), then carry both the
+        breakdown and ``first_event_timestamp`` up into ``metric_events``.
+        """
+        if not query.ctes:
+            return
+
+        breakdown_exprs = self.build_breakdown_exprs(table_alias="")
+
+        if not is_session_property:
+            metric_events_cte = query.ctes.get("metric_events")
+            if isinstance(metric_events_cte, ast.CTE) and isinstance(metric_events_cte.expr, ast.SelectQuery):
+                for alias, expr in breakdown_exprs:
+                    metric_events_cte.expr.select.append(ast.Alias(alias=alias, expr=expr))
+            return
+
+        # Read + dedupe the breakdown per session in the by-session layer.
+        by_session_cte = query.ctes.get("metric_events_by_session")
+        if isinstance(by_session_cte, ast.CTE) and isinstance(by_session_cte.expr, ast.SelectQuery):
+            for alias, expr in breakdown_exprs:
+                by_session_cte.expr.select.append(ast.Alias(alias=alias, expr=ast.Call(name="any", args=[expr])))
+
+        # Carry the breakdown and first_event_timestamp up into metric_events for attribution.
+        metric_events_cte = query.ctes.get("metric_events")
+        if isinstance(metric_events_cte, ast.CTE) and isinstance(metric_events_cte.expr, ast.SelectQuery):
+            for alias in aliases:
+                metric_events_cte.expr.select.append(
+                    ast.Alias(alias=alias, expr=ast.Field(chain=["metric_events_by_session", alias]))
+                )
+            metric_events_cte.expr.select.append(
+                ast.Alias(
+                    alias="first_event_timestamp",
+                    expr=ast.Field(chain=["metric_events_by_session", "first_event_timestamp"]),
+                )
+            )
+
+    def _inject_winsorization_breakdown_columns(
+        self, query: ast.SelectQuery, aliases: list[str], final_cte_name: str
+    ) -> None:
+        """Carry breakdown columns through the percentiles and winsorized_entity_metrics CTEs.
+
+        Winsorization computes per-breakdown-group percentile thresholds (preserved here, not
+        changed), so the breakdown must be grouped in ``percentiles`` and joined into
+        ``winsorized_entity_metrics`` on the breakdown columns.
+        """
+        if not (query.ctes and "percentiles" in query.ctes):
+            return
+
+        percentiles_cte = query.ctes["percentiles"]
+        if isinstance(percentiles_cte, ast.CTE) and isinstance(percentiles_cte.expr, ast.SelectQuery):
+            for alias in aliases:
+                percentiles_cte.expr.select.append(
+                    ast.Alias(alias=alias, expr=ast.Field(chain=["entity_metrics", alias]))
+                )
+            if percentiles_cte.expr.group_by is None:
+                percentiles_cte.expr.group_by = []
+            for alias in aliases:
+                percentiles_cte.expr.group_by.append(ast.Field(chain=["entity_metrics", alias]))
+
+        if final_cte_name == "winsorized_entity_metrics" and "winsorized_entity_metrics" in query.ctes:
+            winsorized_cte = query.ctes["winsorized_entity_metrics"]
+            if isinstance(winsorized_cte, ast.CTE) and isinstance(winsorized_cte.expr, ast.SelectQuery):
+                for alias in aliases:
+                    winsorized_cte.expr.select.append(
+                        ast.Alias(alias=alias, expr=ast.Field(chain=["entity_metrics", alias]))
+                    )
+                # Convert the CROSS JOIN with percentiles into a per-breakdown-group INNER JOIN.
+                if winsorized_cte.expr.select_from:
+                    join_expr = winsorized_cte.expr.select_from.next_join
+                    if join_expr and isinstance(join_expr, ast.JoinExpr):
+                        join_expr.join_type = "JOIN"
+                        join_conditions: list[ast.Expr] = [
+                            ast.CompareOperation(
+                                op=ast.CompareOperationOp.Eq,
+                                left=ast.Field(chain=["percentiles", alias]),
+                                right=ast.Field(chain=["entity_metrics", alias]),
+                            )
+                            for alias in aliases
+                        ]
+                        condition_expr = join_conditions[0]
+                        for condition in join_conditions[1:]:
+                            condition_expr = ast.And(exprs=[condition_expr, condition])
+                        join_expr.constraint = ast.JoinConstraint(expr=condition_expr, constraint_type="ON")