funnel udfs python

PostHog · Oct 2, 2024 · ff7d4e7 · ff7d4e7
1 parent cbf1ffe
commit ff7d4e7
Show file tree

Hide file tree

Showing 36 changed files with 24,670 additions and 5,492 deletions.
diff --git a/docker/clickhouse/user_defined_function.xml b/docker/clickhouse/user_defined_function.xml
@@ -2,7 +2,7 @@
     <function>
         <type>executable</type>
         <name>aggregate_funnel</name>
-        <return_type>Array(Tuple(Int8, Nullable(String), Array(Float64)))</return_type>
+        <return_type>Array(Tuple(Int8, Nullable(String), Array(Float64), Array(Array(UUID))))</return_type>
         <return_name>result</return_name>
         <argument>
             <type>UInt8</type>
@@ -25,7 +25,7 @@
             <name>prop_vals</name>
         </argument>
         <argument>
-            <type>Array(Tuple(Nullable(Float64), Nullable(String), Array(Int8)))</type>
+            <type>Array(Tuple(Nullable(Float64), UUID, Nullable(String), Array(Int8)))</type>
             <name>value</name>
         </argument>
         <format>JSONEachRow</format>
@@ -35,7 +35,7 @@
     <function>
         <type>executable</type>
         <name>aggregate_funnel_cohort</name>
-        <return_type>Array(Tuple(Int8, UInt64, Array(Float64)))</return_type>
+        <return_type>Array(Tuple(Int8, UInt64, Array(Float64), Array(Array(UUID))))</return_type>
         <return_name>result</return_name>
         <argument>
             <type>UInt8</type>
@@ -58,7 +58,7 @@
             <name>prop_vals</name>
         </argument>
         <argument>
-            <type>Array(Tuple(Nullable(Float64), UInt64, Array(Int8)))</type>
+            <type>Array(Tuple(Nullable(Float64), UUID, UInt64, Array(Int8)))</type>
             <name>value</name>
         </argument>
         <format>JSONEachRow</format>
@@ -68,7 +68,7 @@
     <function>
         <type>executable</type>
         <name>aggregate_funnel_array</name>
-        <return_type>Array(Tuple(Int8, Array(String), Array(Float64)))</return_type>
+        <return_type>Array(Tuple(Int8, Array(String), Array(Float64), Array(Array(UUID))))</return_type>
         <return_name>result</return_name>
         <argument>
             <type>UInt8</type>
@@ -91,7 +91,7 @@
             <name>prop_vals</name>
         </argument>
         <argument>
-            <type>Array(Tuple(Nullable(Float64), Array(String), Array(Int8)))</type>
+            <type>Array(Tuple(Nullable(Float64), UUID, Array(String), Array(Int8)))</type>
             <name>value</name>
         </argument>
         <format>JSONEachRow</format>
@@ -124,7 +124,7 @@
             <name>prop_vals</name>
         </argument>
         <argument>
-            <type>Array(Tuple(Nullable(Float64), Nullable(String), Array(Int8)))</type>
+            <type>Array(Tuple(Nullable(Float64), UUID, Nullable(String), Array(Int8)))</type>
             <name>value</name>
         </argument>
         <format>JSONEachRow</format>

diff --git a/posthog/hogql_queries/insights/funnels/base.py b/posthog/hogql_queries/insights/funnels/base.py
@@ -435,31 +435,26 @@ def _get_inner_event_query_for_udf(
             extra_fields.append(prop)
 
         funnel_events_query = FunnelEventQuery(
-            context=self.context,
-            extra_fields=[*self._extra_event_fields, *extra_fields],
-            extra_event_properties=self._extra_event_properties,
+            context=self.context, extra_fields=[*self.extra_event_fields_and_properties, *extra_fields]
         ).to_query(
             skip_entity_filter=skip_entity_filter,
         )
-        # funnel_events_query, params = FunnelEventQuery(
-        #     extra_fields=[*self._extra_event_fields, *extra_fields],
-        #     extra_event_properties=self._extra_event_properties,
-        # ).get_query(entities_to_use, entity_name, skip_entity_filter=skip_entity_filter)
 
         all_step_cols: list[ast.Expr] = []
         all_exclusions: list[list[FunnelExclusionEventsNode | FunnelExclusionActionsNode]] = []
         for index, entity in enumerate(entities_to_use):
-            step_cols = self._get_step_col(entity, index, entity_name)
+            step_cols = self._get_step_col(entity, index, entity_name, for_udf=True)
             all_step_cols.extend(step_cols)
             all_exclusions.append([])
 
-        for excluded_entity in funnelsFilter.exclusions or []:
-            for i in range(excluded_entity.funnelFromStep + 1, excluded_entity.funnelToStep + 1):
-                all_exclusions[i].append(excluded_entity)
+        if funnelsFilter.exclusions:
+            for excluded_entity in funnelsFilter.exclusions:
+                for i in range(excluded_entity.funnelFromStep + 1, excluded_entity.funnelToStep + 1):
+                    all_exclusions[i].append(excluded_entity)
 
-        for index, exclusions in enumerate(all_exclusions):
-            exclusion_col_expr = self._get_exclusions_col(exclusions, index, entity_name)
-            all_step_cols.append(exclusion_col_expr)
+            for index, exclusions in enumerate(all_exclusions):
+                exclusion_col_expr = self._get_exclusions_col(exclusions, index, entity_name)
+                all_step_cols.append(exclusion_col_expr)
 
         breakdown_select_prop = self._get_breakdown_select_prop()
 
@@ -491,7 +486,6 @@ def _get_exclusions_col(
     ) -> ast.Expr:
         if not exclusions:
             return parse_expr(f"0 as exclusion_{index}")
-
         conditions = [self._build_step_query(exclusion, index, entity_name, "") for exclusion in exclusions]
         return parse_expr(
             f"if({{condition}}, 1, 0) as exclusion_{index}", placeholders={"condition": ast.Or(exprs=conditions)}
@@ -664,10 +658,10 @@ def _get_step_col(
                 parse_expr(f"if({step_prefix}step_{index} = 1, timestamp, null) as {step_prefix}latest_{index}")
             )
 
-        for field in self.extra_event_fields_and_properties:
-            step_cols.append(
-                parse_expr(f'if({step_prefix}step_{index} = 1, "{field}", null) as "{step_prefix}{field}_{index}"')
-            )
+            for field in self.extra_event_fields_and_properties:
+                step_cols.append(
+                    parse_expr(f'if({step_prefix}step_{index} = 1, "{field}", null) as "{step_prefix}{field}_{index}"')
+                )
 
         return step_cols
 
@@ -758,12 +752,15 @@ def _get_funnel_person_step_condition(self) -> ast.Expr:
 
         return ast.And(exprs=conditions)
 
-    def _get_funnel_person_step_events(self) -> list[ast.Expr]:
-        if (
+    def _include_matched_events(self):
+        return (
             hasattr(self.context, "actorsQuery")
             and self.context.actorsQuery is not None
             and self.context.actorsQuery.includeRecordings
-        ):
+        )
+
+    def _get_funnel_person_step_events(self) -> list[ast.Expr]:
+        if self._include_matched_events():
             if self.context.includeFinalMatchingEvents:
                 # Always returns the user's final step of the funnel
                 return [parse_expr("final_matching_events as matching_events")]

diff --git a/posthog/hogql_queries/insights/funnels/funnel_correlation_query_runner.py b/posthog/hogql_queries/insights/funnels/funnel_correlation_query_runner.py
@@ -4,6 +4,7 @@
 from posthog.constants import AUTOCAPTURE_EVENT
 from posthog.hogql.parser import parse_select
 from posthog.hogql.property import property_to_expr
+from posthog.hogql_queries.insights.funnels import FunnelUDF
 from posthog.hogql_queries.insights.funnels.funnel_event_query import FunnelEventQuery
 from posthog.hogql_queries.insights.funnels.funnel_persons import FunnelActors
 from posthog.hogql_queries.insights.funnels.funnel_strict_persons import FunnelStrictActors
@@ -19,7 +20,11 @@
 from posthog.hogql.query import execute_hogql_query
 from posthog.hogql.timings import HogQLTimings
 from posthog.hogql_queries.insights.funnels.funnel_query_context import FunnelQueryContext
-from posthog.hogql_queries.insights.funnels.utils import funnel_window_interval_unit_to_sql, get_funnel_actor_class
+from posthog.hogql_queries.insights.funnels.utils import (
+    funnel_window_interval_unit_to_sql,
+    get_funnel_actor_class,
+    use_udf,
+)
 from posthog.hogql_queries.query_runner import QueryRunner
 from posthog.models import Team
 from posthog.models.property.util import get_property_string_expr
@@ -93,7 +98,7 @@ class FunnelCorrelationQueryRunner(QueryRunner):
     actors_query: FunnelsActorsQuery
     correlation_actors_query: Optional[FunnelCorrelationActorsQuery]
 
-    _funnel_actors_generator: FunnelActors | FunnelStrictActors | FunnelUnorderedActors
+    _funnel_actors_generator: FunnelActors | FunnelStrictActors | FunnelUnorderedActors | FunnelUDF
 
     def __init__(
         self,
@@ -132,9 +137,11 @@ def __init__(
         self.context.actorsQuery = self.actors_query
 
         # Used for generating the funnel persons cte
-        funnel_order_actor_class = get_funnel_actor_class(self.context.funnelsFilter)(context=self.context)
+        funnel_order_actor_class = get_funnel_actor_class(
+            self.context.funnelsFilter, use_udf(self.context.funnelsFilter, self.team)
+        )(context=self.context)
         assert isinstance(
-            funnel_order_actor_class, FunnelActors | FunnelStrictActors | FunnelUnorderedActors
+            funnel_order_actor_class, FunnelActors | FunnelStrictActors | FunnelUnorderedActors | FunnelUDF
         )  # for typings
         self._funnel_actors_generator = funnel_order_actor_class