Add a basic funnel test

posthog/hogql_queries/experiments/experiment_query_runner.py

@@ -20,28 +20,36 @@
     calculate_credible_intervals_v2_continuous,
     calculate_probabilities_v2_continuous,
 )
+from posthog.hogql_queries.experiments.funnels_statistics_v2 import (
+    calculate_probabilities_v2 as calculate_probabilities_v2_funnel,
+    are_results_significant_v2 as are_results_significant_v2_funnel,
+    calculate_credible_intervals_v2 as calculate_credible_intervals_v2_funnel,
+)
 from posthog.hogql_queries.query_runner import QueryRunner
 from posthog.models.experiment import Experiment
 from rest_framework.exceptions import ValidationError
 from posthog.schema import (
+    CachedExperimentFunnelsQueryResponse,
     CachedExperimentTrendsQueryResponse,
     ExperimentDataWarehouseMetricConfig,
     ExperimentEventMetricConfig,
+    ExperimentFunnelsQueryResponse,
     ExperimentMetricType,
     ExperimentSignificanceCode,
     ExperimentQuery,
     ExperimentTrendsQueryResponse,
+    ExperimentVariantFunnelsBaseStats,
     ExperimentVariantTrendsBaseStats,
     DateRange,
 )
-from typing import Optional
+from typing import Optional, Union
 from datetime import datetime, timedelta, UTC
 
 
 class ExperimentQueryRunner(QueryRunner):
     query: ExperimentQuery
-    response: ExperimentTrendsQueryResponse
-    cached_response: CachedExperimentTrendsQueryResponse
+    response: Union[ExperimentTrendsQueryResponse, ExperimentFunnelsQueryResponse]
+    cached_response: Union[CachedExperimentTrendsQueryResponse, CachedExperimentFunnelsQueryResponse]
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -278,15 +286,25 @@ def _evaluate_experiment_query(self) -> list[ExperimentVariantTrendsBaseStats]:
             modifiers=create_default_modifiers_for_team(self.team),
         )
 
-        variants: list[ExperimentVariantTrendsBaseStats] = [
-            ExperimentVariantTrendsBaseStats(
-                absolute_exposure=result[1],
-                count=result[2],
-                exposure=result[1],
-                key=result[0],
-            )
-            for result in response.results
-        ]
+        if self.metric.metric_type == ExperimentMetricType.FUNNEL:
+            variants: list[ExperimentVariantFunnelsBaseStats] = [
+                ExperimentVariantFunnelsBaseStats(
+                    failure_count=result[1] - result[2],
+                    key=result[0],
+                    success_count=result[2],
+                )
+                for result in response.results
+            ]
+        else:
+            variants: list[ExperimentVariantTrendsBaseStats] = [
+                ExperimentVariantTrendsBaseStats(
+                    absolute_exposure=result[1],
+                    count=result[2],
+                    exposure=result[1],
+                    key=result[0],
+                )
+                for result in response.results
+            ]
 
         return variants
 
@@ -314,29 +332,52 @@ def calculate(self) -> ExperimentTrendsQueryResponse:
                         control_variant, test_variants, probabilities
                     )
                     credible_intervals = calculate_credible_intervals_v2_count([control_variant, *test_variants])
+                case ExperimentMetricType.FUNNEL:
+                    probabilities = calculate_probabilities_v2_funnel(control_variant, test_variants)
+                    significance_code, p_value = are_results_significant_v2_funnel(
+                        control_variant, test_variants, probabilities
+                    )
+                    credible_intervals = calculate_credible_intervals_v2_funnel([control_variant, *test_variants])
                 case _:
                     raise ValueError(f"Unsupported metric type: {self.metric.metric_type}")
         else:
             probabilities = calculate_probabilities(control_variant, test_variants)
             significance_code, p_value = are_results_significant(control_variant, test_variants, probabilities)
             credible_intervals = calculate_credible_intervals([control_variant, *test_variants])
 
-        return ExperimentTrendsQueryResponse(
-            kind="ExperimentTrendsQuery",
-            insight=[],
-            count_query=None,
-            exposure_query=None,
-            variants=[variant.model_dump() for variant in [control_variant, *test_variants]],
-            probability={
-                variant.key: probability
-                for variant, probability in zip([control_variant, *test_variants], probabilities)
-            },
-            significant=significance_code == ExperimentSignificanceCode.SIGNIFICANT,
-            significance_code=significance_code,
-            stats_version=self.stats_version,
-            p_value=p_value,
-            credible_intervals=credible_intervals,
-        )
+        if self.metric.metric_type == ExperimentMetricType.FUNNEL:
+            return ExperimentFunnelsQueryResponse(
+                kind="ExperimentFunnelsQuery",
+                funnels_query=None,
+                insight=[[{"name": "foo"}], [{"name": "bar"}]],
+                variants=[variant.model_dump() for variant in [control_variant, *test_variants]],
+                probability={
+                    variant.key: probability
+                    for variant, probability in zip([control_variant, *test_variants], probabilities)
+                },
+                significant=significance_code == ExperimentSignificanceCode.SIGNIFICANT,
+                significance_code=significance_code,
+                stats_version=self.stats_version,
+                expected_loss=0,
+                credible_intervals=credible_intervals,
+            )
+        else:
+            return ExperimentTrendsQueryResponse(
+                kind="ExperimentTrendsQuery",
+                insight=[],
+                count_query=None,
+                exposure_query=None,
+                variants=[variant.model_dump() for variant in [control_variant, *test_variants]],
+                probability={
+                    variant.key: probability
+                    for variant, probability in zip([control_variant, *test_variants], probabilities)
+                },
+                significant=significance_code == ExperimentSignificanceCode.SIGNIFICANT,
+                significance_code=significance_code,
+                stats_version=self.stats_version,
+                p_value=p_value,
+                credible_intervals=credible_intervals,
+            )
 
     def to_query(self) -> ast.SelectQuery:
         raise ValueError(f"Cannot convert source query of type {self.query.metric.kind} to query")

posthog/hogql_queries/experiments/test/__snapshots__/test_experiment_query_runner.ambr

@@ -1,4 +1,50 @@
 # serializer version: 1
+# name: TestExperimentQueryRunner.test_query_runner_funnel_metric
+  '''
+  SELECT maq.variant AS variant,
+         count(maq.distinct_id) AS num_users,
+         sum(maq.value) AS total_sum,
+         sum(power(maq.value, 2)) AS total_sum_of_squares
+  FROM
+    (SELECT base.variant AS variant,
+            base.distinct_id AS distinct_id,
+            sum(coalesce(eae.value, 0)) AS value
+     FROM
+       (SELECT events.distinct_id AS distinct_id,
+               replaceAll(JSONExtractRaw(events.properties, '$feature_flag_response'), '"', '') AS variant,
+               min(toTimeZone(events.timestamp, 'UTC')) AS first_exposure_time
+        FROM events
+        WHERE and(equals(events.team_id, 99999), and(equals(events.event, '$feature_flag_called'), ifNull(equals(replaceAll(JSONExtractRaw(events.properties, '$feature_flag'), '"', ''), 'test-experiment'), 0)))
+        GROUP BY variant,
+                 events.distinct_id) AS base
+     LEFT JOIN
+       (SELECT toTimeZone(events.timestamp, 'UTC') AS timestamp,
+               events.distinct_id AS distinct_id,
+               exposure.variant AS variant,
+               events.event AS event,
+               1 AS value
+        FROM events
+        INNER JOIN
+          (SELECT events.distinct_id AS distinct_id,
+                  replaceAll(JSONExtractRaw(events.properties, '$feature_flag_response'), '"', '') AS variant,
+                  min(toTimeZone(events.timestamp, 'UTC')) AS first_exposure_time
+           FROM events
+           WHERE and(equals(events.team_id, 99999), and(equals(events.event, '$feature_flag_called'), ifNull(equals(replaceAll(JSONExtractRaw(events.properties, '$feature_flag'), '"', ''), 'test-experiment'), 0)))
+           GROUP BY variant,
+                    events.distinct_id) AS exposure ON equals(events.distinct_id, exposure.distinct_id)
+        WHERE and(equals(events.team_id, 99999), greaterOrEquals(toTimeZone(events.timestamp, 'UTC'), exposure.first_exposure_time), equals(events.event, 'purchase'))) AS eae ON and(equals(base.distinct_id, eae.distinct_id), equals(base.variant, eae.variant))
+     GROUP BY base.variant,
+              base.distinct_id) AS maq
+  GROUP BY maq.variant
+  LIMIT 100 SETTINGS readonly=2,
+                     max_execution_time=60,
+                     allow_experimental_object_type=1,
+                     format_csv_allow_double_quotes=0,
+                     max_ast_elements=4000000,
+                     max_expanded_ast_elements=4000000,
+                     max_bytes_before_external_group_by=0
+  '''
+# ---
 # name: TestExperimentQueryRunner.test_query_runner_standard_flow_v2_stats
   '''
   SELECT maq.variant AS variant,

posthog/hogql_queries/experiments/test/test_experiment_query_runner.py

@@ -333,6 +333,68 @@ def create_data_warehouse_table_with_subscriptions(self):
 
         return subscription_table_name
 
+    @freeze_time("2020-01-01T12:00:00Z")
+    @snapshot_clickhouse_queries
+    def test_query_runner_funnel_metric(self):
+        feature_flag = self.create_feature_flag()
+        experiment = self.create_experiment(feature_flag=feature_flag)
+        experiment.stats_config = {"version": 2}
+        experiment.save()
+
+        feature_flag_property = f"$feature/{feature_flag.key}"
+
+        metric = ExperimentMetric(
+            metric_type=ExperimentMetricType.FUNNEL,
+            metric_config=ExperimentEventMetricConfig(event="purchase"),
+        )
+
+        experiment_query = ExperimentQuery(
+            experiment_id=experiment.id,
+            kind="ExperimentQuery",
+            metric=metric,
+        )
+
+        experiment.metrics = [metric.model_dump(mode="json")]
+        experiment.save()
+
+        for variant, purchase_count in [("control", 6), ("test", 8)]:
+            for i in range(10):
+                _create_person(distinct_ids=[f"user_{variant}_{i}"], team_id=self.team.pk)
+                _create_event(
+                    team=self.team,
+                    event="$feature_flag_called",
+                    distinct_id=f"user_{variant}_{i}",
+                    timestamp="2020-01-02T12:00:00Z",
+                    properties={
+                        feature_flag_property: variant,
+                        "$feature_flag_response": variant,
+                        "$feature_flag": feature_flag.key,
+                    },
+                )
+                if i < purchase_count:
+                    _create_event(
+                        team=self.team,
+                        event="purchase",
+                        distinct_id=f"user_{variant}_{i}",
+                        timestamp="2020-01-02T12:01:00Z",
+                        properties={feature_flag_property: variant},
+                    )
+
+        flush_persons_and_events()
+
+        query_runner = ExperimentQueryRunner(query=experiment_query, team=self.team)
+        result = query_runner.calculate()
+
+        self.assertEqual(len(result.variants), 2)
+
+        control_variant = next(variant for variant in result.variants if variant.key == "control")
+        test_variant = next(variant for variant in result.variants if variant.key == "test")
+
+        self.assertEqual(control_variant.success_count, 6)
+        self.assertEqual(control_variant.failure_count, 4)
+        self.assertEqual(test_variant.success_count, 8)
+        self.assertEqual(test_variant.failure_count, 2)
+
     @freeze_time("2020-01-01T12:00:00Z")
     @snapshot_clickhouse_queries
     def test_query_runner_standard_flow_v2_stats(self):