feat(experiments): Initial data warehouse Trend support #26356
+407
−33
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,6 +3,9 @@ | |
| from django.conf import settings | ||
| from posthog.constants import ExperimentNoResultsErrorKeys | ||
| from posthog.hogql import ast | ||
| from posthog.hogql.context import HogQLContext | ||
| from posthog.hogql.database.database import create_hogql_database | ||
| from posthog.hogql.database.models import LazyJoin | ||
| from posthog.hogql_queries.experiments import CONTROL_VARIANT_KEY | ||
| from posthog.hogql_queries.experiments.trends_statistics import ( | ||
| are_results_significant, | ||
|
|
@@ -19,6 +22,8 @@ | |
| BreakdownFilter, | ||
| CachedExperimentTrendsQueryResponse, | ||
| ChartDisplayType, | ||
| DataWarehouseNode, | ||
| DataWarehousePropertyFilter, | ||
| EventPropertyFilter, | ||
| EventsNode, | ||
| ExperimentSignificanceCode, | ||
|
|
@@ -27,11 +32,12 @@ | |
| ExperimentVariantTrendsBaseStats, | ||
| InsightDateRange, | ||
| PropertyMathType, | ||
| PropertyOperator, | ||
| TrendsFilter, | ||
| TrendsQuery, | ||
| TrendsQueryResponse, | ||
| ) | ||
| from typing import Any, Optional, cast | ||
| import threading | ||
|
|
||
|
|
||
|
|
@@ -89,12 +95,18 @@ def _get_insight_date_range(self) -> InsightDateRange: | |
| explicitDate=True, | ||
| ) | ||
|
|
||
| def _get_event_breakdown_filter(self) -> BreakdownFilter: | ||
| return BreakdownFilter( | ||
| breakdown=self.breakdown_key, | ||
| breakdown_type="event", | ||
| ) | ||
|
|
||
| def _get_data_warehouse_breakdown_filter(self) -> BreakdownFilter: | ||
| return BreakdownFilter( | ||
| breakdown=f"events.properties.{self.breakdown_key}", | ||
| breakdown_type="data_warehouse", | ||
| ) | ||
|
|
||
| def _prepare_count_query(self) -> TrendsQuery: | ||
| """ | ||
| This method takes the raw trend query and adapts it | ||
|
|
@@ -118,15 +130,32 @@ def _prepare_count_query(self) -> TrendsQuery: | |
|
|
||
| prepared_count_query.trendsFilter = TrendsFilter(display=ChartDisplayType.ACTIONS_LINE_GRAPH_CUMULATIVE) | ||
| prepared_count_query.dateRange = self._get_insight_date_range() | ||
| if self._is_data_warehouse_query(prepared_count_query): | ||
| prepared_count_query.breakdownFilter = self._get_data_warehouse_breakdown_filter() | ||
| prepared_count_query.properties = [ | ||
| DataWarehousePropertyFilter( | ||
| key="events.event", | ||
| value="$feature_flag_called", | ||
| operator=PropertyOperator.EXACT, | ||
| type="data_warehouse", | ||
| ), | ||
| DataWarehousePropertyFilter( | ||
| key=f"events.properties.{self.breakdown_key}", | ||
| value=self.variants, | ||
| operator=PropertyOperator.EXACT, | ||
| type="data_warehouse", | ||
| ), | ||
| ] | ||
| else: | ||
| prepared_count_query.breakdownFilter = self._get_event_breakdown_filter() | ||
| prepared_count_query.properties = [ | ||
| EventPropertyFilter( | ||
| key=self.breakdown_key, | ||
| value=self.variants, | ||
| operator=PropertyOperator.EXACT, | ||
| type="event", | ||
| ) | ||
| ] | ||
|
|
||
| return prepared_count_query | ||
|
|
||
|
|
@@ -152,7 +181,7 @@ def _prepare_exposure_query(self) -> TrendsQuery: | |
|
|
||
| if hasattr(count_event, "event"): | ||
| prepared_exposure_query.dateRange = self._get_insight_date_range() | ||
| prepared_exposure_query.breakdownFilter = self._get_event_breakdown_filter() | ||
| prepared_exposure_query.trendsFilter = TrendsFilter( | ||
| display=ChartDisplayType.ACTIONS_LINE_GRAPH_CUMULATIVE | ||
| ) | ||
|
|
@@ -166,7 +195,7 @@ def _prepare_exposure_query(self) -> TrendsQuery: | |
| EventPropertyFilter( | ||
| key=self.breakdown_key, | ||
| value=self.variants, | ||
| operator=PropertyOperator.EXACT, | ||
| type="event", | ||
| ) | ||
| ] | ||
|
|
@@ -177,13 +206,13 @@ def _prepare_exposure_query(self) -> TrendsQuery: | |
| elif self.query.exposure_query: | ||
| prepared_exposure_query = TrendsQuery(**self.query.exposure_query.model_dump()) | ||
| prepared_exposure_query.dateRange = self._get_insight_date_range() | ||
| prepared_exposure_query.trendsFilter = TrendsFilter(display=ChartDisplayType.ACTIONS_LINE_GRAPH_CUMULATIVE) | ||
| prepared_exposure_query.breakdownFilter = self._get_event_breakdown_filter() | ||
| prepared_exposure_query.properties = [ | ||
| EventPropertyFilter( | ||
| key=self.breakdown_key, | ||
| value=self.variants, | ||
| operator=PropertyOperator.EXACT, | ||
| type="event", | ||
| ) | ||
| ] | ||
|
|
@@ -206,13 +235,13 @@ def _prepare_exposure_query(self) -> TrendsQuery: | |
| EventPropertyFilter( | ||
| key="$feature_flag_response", | ||
| value=self.variants, | ||
| operator=PropertyOperator.EXACT, | ||
| type="event", | ||
| ), | ||
| EventPropertyFilter( | ||
| key="$feature_flag", | ||
| value=[self.feature_flag.key], | ||
| operator=PropertyOperator.EXACT, | ||
| type="event", | ||
| ), | ||
| ], | ||
|
|
@@ -226,7 +255,63 @@ def calculate(self) -> ExperimentTrendsQueryResponse: | |
|
|
||
| def run(query_runner: TrendsQueryRunner, result_key: str, is_parallel: bool): | ||
| try: | ||
| database = create_hogql_database(team_id=self.team.pk) | ||
| if self._is_data_warehouse_query(query_runner.query): | ||
| series_node = cast(DataWarehouseNode, query_runner.query.series[0]) | ||
| table = database.get_table(series_node.table_name) | ||
| table.fields["events"] = LazyJoin( | ||
| from_field=[series_node.distinct_id_field], | ||
| join_table=database.get_table("events"), | ||
| join_function=lambda join_to_add, context, node: ( | ||
| ast.JoinExpr( | ||
| table=ast.SelectQuery( | ||
|
There was a problem hiding this comment. This is currently joining over all events, but we should join only over the In the future we should also support custom exposure events here, but for now feel free to hardcode There was a problem hiding this comment. Good catch, handled with f131df9 |
||
| select=[ | ||
| ast.Alias(alias=name, expr=ast.Field(chain=["events", *chain])) | ||
| for name, chain in { | ||
| **join_to_add.fields_accessed, | ||
| "timestamp": ["timestamp"], | ||
| "distinct_id": ["distinct_id"], | ||
| "properties": ["properties"], | ||
| }.items() | ||
| ], | ||
| select_from=ast.JoinExpr(table=ast.Field(chain=["events"])), | ||
| ), | ||
| join_type="ASOF LEFT JOIN", | ||
|
There was a problem hiding this comment. I'd love to have a comment here explaining what we're doing, especially because some of how the ASOF JOIN works is only explicit in the Clickhouse docs: (feel free to adjust) There was a problem hiding this comment. Good suggestion, added the |
||
| alias=join_to_add.to_table, | ||
| constraint=ast.JoinConstraint( | ||
| expr=ast.And( | ||
| exprs=[ | ||
| ast.CompareOperation( | ||
| left=ast.Field( | ||
| chain=[ | ||
| join_to_add.from_table, | ||
| series_node.distinct_id_field, | ||
| ] | ||
| ), | ||
| op=ast.CompareOperationOp.Eq, | ||
| right=ast.Field(chain=[join_to_add.to_table, "distinct_id"]), | ||
| ), | ||
| ast.CompareOperation( | ||
| left=ast.Field( | ||
| chain=[ | ||
| join_to_add.from_table, | ||
| series_node.timestamp_field, | ||
| ] | ||
| ), | ||
| op=ast.CompareOperationOp.GtEq, | ||
| right=ast.Field(chain=[join_to_add.to_table, "timestamp"]), | ||
| ), | ||
| ] | ||
| ), | ||
| constraint_type="ON", | ||
| ), | ||
| ) | ||
| ), | ||
| ) | ||
|
|
||
| context = HogQLContext(team_id=self.team.pk, database=database) | ||
|
|
||
| result = query_runner.calculate(context=context) | ||
| shared_results[result_key] = result | ||
| except Exception as e: | ||
| errors.append(e) | ||
|
|
@@ -362,5 +447,8 @@ def _validate_event_variants(self, count_result: TrendsQueryResponse): | |
| if has_errors: | ||
| raise ValidationError(detail=json.dumps(errors)) | ||
|
|
||
| def _is_data_warehouse_query(self, query: TrendsQuery) -> bool: | ||
| return isinstance(query.series[0], DataWarehouseNode) | ||
|
|
||
| def to_query(self) -> ast.SelectQuery: | ||
| raise ValueError(f"Cannot convert source query of type {self.query.count_query.kind} to query") | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I just want to confirm my understanding: the purpose of this is to create a
Databaseinstance, which is an object that lets you run HogQL queries across various data sources, like person-events or data warehouse tables.We do this because we need to build a custom context with our own
Database, including our own table where we define the join in a custom way. Then, we pass this custom context to the query runner.Please check if the above is correct and leave a comment in the code :)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is correct :) Left a comment in a07bb2a