Skip to content

Commit

Permalink
feat(hogql): cohort left join conjoined (#19725)
Browse files Browse the repository at this point in the history 
* WIP

* Added a new cohort modifier for a conjoined left query

* Missing changes

* Regenned mypy sync file

* Fixed tests

* Update query snapshots

* Seperate out dynamic and static cohorts

* Use version when filtering dynamic cohorts

* Fixed tests

* Update query snapshots

* Updated mypy baseline

* Removed group by

* Update query snapshots

* Update query snapshots

* Update query snapshots

---------

Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Thomas Obermüller <[email protected]>
  • Loading branch information
@Gilbert09 @github-actions @thmsobrmlr
3 people authored Jan 16, 2024
1 parent a6be5b9 commit 626e880
Show file tree
Hide file tree
Showing 18 changed files with 509 additions and 92 deletions.
2 changes: 1 addition & 1 deletion frontend/src/queries/schema.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1709,7 +1709,7 @@
"description": "HogQL Query Options are automatically set per team. However, they can be overriden in the query.",
"properties": {
"inCohortVia": {
"enum": ["leftjoin", "subquery"],
"enum": ["leftjoin", "subquery", "leftjoin_conjoined"],
"type": "string"
},
"materializationMode": {
Expand Down
2 changes: 1 addition & 1 deletion frontend/src/queries/schema.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ export interface DataNode extends Node {
export interface HogQLQueryModifiers {
personsOnEventsMode?: 'disabled' | 'v1_enabled' | 'v1_mixed' | 'v2_enabled'
personsArgMaxVersion?: 'auto' | 'v1' | 'v2'
inCohortVia?: 'leftjoin' | 'subquery'
inCohortVia?: 'leftjoin' | 'subquery' | 'leftjoin_conjoined'
materializationMode?: 'auto' | 'legacy_null_as_string' | 'legacy_null_as_null' | 'disabled'
}

Expand Down
29 changes: 1 addition & 28 deletions mypy-baseline.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,6 @@ posthog/hogql/transforms/property_types.py:0: error: Statement is unreachable [
posthog/hogql/transforms/property_types.py:0: error: Argument 2 to "_get_materialized_column" of "PropertySwapper" has incompatible type "str | int"; expected "str" [arg-type]
posthog/hogql/modifiers.py:0: error: Incompatible types in assignment (expression has type "PersonOnEventsMode", variable has type "PersonsOnEventsMode | None") [assignment]
posthog/hogql/modifiers.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "PersonsArgMaxVersion | None") [assignment]
posthog/hogql/modifiers.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "InCohortVia | None") [assignment]
posthog/hogql/functions/cohort.py:0: error: Argument 1 to "escape_clickhouse_string" has incompatible type "str | None"; expected "float | int | str | list[Any] | tuple[Any, ...] | date | datetime | UUID | UUIDT" [arg-type]
posthog/hogql/functions/cohort.py:0: error: Argument 1 to "escape_clickhouse_string" has incompatible type "str | None"; expected "float | int | str | list[Any] | tuple[Any, ...] | date | datetime | UUID | UUIDT" [arg-type]
posthog/hogql/functions/cohort.py:0: error: Incompatible types in assignment (expression has type "ValuesQuerySet[Cohort, tuple[int, bool | None]]", variable has type "ValuesQuerySet[Cohort, tuple[int, bool | None, str | None]]") [assignment]
Expand Down Expand Up @@ -366,11 +365,6 @@ posthog/hogql/query.py:0: error: Argument 1 to "get_default_limit_for_context" h
posthog/hogql/query.py:0: error: "SelectQuery" has no attribute "select_queries" [attr-defined]
posthog/hogql/query.py:0: error: Subclass of "SelectQuery" and "SelectUnionQuery" cannot exist: would have incompatible method signatures [unreachable]
posthog/hogql_queries/query_runner.py:0: error: Incompatible types in assignment (expression has type "HogQLQuery | TrendsQuery | LifecycleQuery | InsightActorsQuery | EventsQuery | ActorsQuery | RetentionQuery | SessionsTimelineQuery | WebOverviewQuery | WebTopClicksQuery | WebStatsTableQuery | StickinessQuery | BaseModel | dict[str, Any]", variable has type "HogQLQuery | TrendsQuery | LifecycleQuery | InsightActorsQuery | EventsQuery | ActorsQuery | RetentionQuery | SessionsTimelineQuery | WebOverviewQuery | WebTopClicksQuery | WebStatsTableQuery | StickinessQuery") [assignment]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Argument "chain" to "Field" has incompatible type "list[str]"; expected "list[str | int]" [arg-type]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance
posthog/hogql_queries/insights/trends/breakdown_values.py:0: note: Consider using "Sequence" instead, which is covariant
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Argument "breakdown_type" to "get_properties_chain" has incompatible type "str"; expected "Literal['person', 'session', 'group', 'event']" [arg-type]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Argument "breakdown_field" to "get_properties_chain" has incompatible type "str | float"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Incompatible types in assignment (expression has type "float | int", variable has type "int") [assignment]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Item "SelectUnionQuery" of "SelectQuery | SelectUnionQuery" has no attribute "select" [union-attr]
posthog/hogql_queries/insights/trends/breakdown_values.py:0: error: Value of type "list[Any] | None" is not indexable [index]
Expand All @@ -380,29 +374,18 @@ posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "Bre
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument 1 to "parse_expr" has incompatible type "str | float | list[str | float] | Any | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument 1 to "int" has incompatible type "str | float | list[str | float] | Any | None"; expected "str | Buffer | SupportsInt | SupportsIndex | SupportsTrunc" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument 1 to "parse_expr" has incompatible type "str | float | list[str | float] | Any | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Statement is unreachable [unreachable]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument 1 to "int" has incompatible type "str | float | list[str | float] | Any | None"; expected "str | Buffer | SupportsInt | SupportsIndex | SupportsTrunc" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument 1 to "parse_expr" has incompatible type "str | float | list[str | float] | Any | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Statement is unreachable [unreachable]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "exprs" to "Or" has incompatible type "list[CompareOperation]"; expected "list[Expr]" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance
posthog/hogql_queries/insights/trends/breakdown.py:0: note: Consider using "Sequence" instead, which is covariant
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Incompatible types in assignment (expression has type "float", variable has type "int") [assignment]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Incompatible types in assignment (expression has type "float", variable has type "int") [assignment]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Incompatible types in assignment (expression has type "str", variable has type "int") [assignment]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "event_name" to "BreakdownValues" has incompatible type "str | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "breakdown_field" to "BreakdownValues" has incompatible type "str | float | list[str | float] | Any | None"; expected "str | float" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "breakdown_type" to "BreakdownValues" has incompatible type "BreakdownType | Any | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_histogram_bin_count" [union-attr]
Expand All @@ -417,7 +400,6 @@ posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "breakdown
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Argument "breakdown_field" to "get_properties_chain" has incompatible type "str | float | list[str | float] | Any | None"; expected "str" [arg-type]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_group_type_index" [union-attr]
posthog/hogql_queries/insights/trends/breakdown.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
posthog/hogql_queries/hogql_query_runner.py:0: error: Statement is unreachable [unreachable]
posthog/hogql_queries/hogql_query_runner.py:0: error: Argument "placeholders" to "parse_select" has incompatible type "dict[str, Constant] | None"; expected "dict[str, Expr] | None" [arg-type]
posthog/hogql_queries/hogql_query_runner.py:0: error: Incompatible types in assignment (expression has type "Expr", variable has type "SelectQuery | SelectUnionQuery") [assignment]
Expand All @@ -439,13 +421,9 @@ posthog/api/person.py:0: error: Argument 1 to "loads" has incompatible type "str
posthog/api/person.py:0: error: Argument "user" to "log_activity" has incompatible type "User | AnonymousUser"; expected "User | None" [arg-type]
posthog/api/person.py:0: error: Argument "user" to "log_activity" has incompatible type "User | AnonymousUser"; expected "User | None" [arg-type]
posthog/hogql_queries/web_analytics/web_analytics_query_runner.py:0: error: Argument 1 to "append" of "list" has incompatible type "EventPropertyFilter"; expected "Expr" [arg-type]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Return type "list[SelectQuery]" of "to_query" incompatible with return type "SelectQuery | SelectUnionQuery" in supertype "QueryRunner" [override]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Incompatible return value type (got "list[SelectQuery | SelectUnionQuery]", expected "list[SelectQuery]") [return-value]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Need type annotation for "timings" (hint: "timings: List[<type>] = ...") [var-annotated]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Argument 1 to "extend" of "list" has incompatible type "list[QueryTiming] | None"; expected "Iterable[Any]" [arg-type]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Statement is unreachable [unreachable]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: List item 0 has incompatible type "str | float | None"; expected "str | float" [list-item]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown" [union-attr]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Argument 1 to "FormulaAST" has incompatible type "map[Any]"; expected "list[list[float]]" [arg-type]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Argument 1 to "FormulaAST" has incompatible type "map[Any]"; expected "list[list[float]]" [arg-type]
posthog/hogql_queries/insights/trends/trends_query_runner.py:0: error: Item "None" of "BreakdownFilter | None" has no attribute "breakdown_type" [union-attr]
Expand Down Expand Up @@ -634,11 +612,6 @@ posthog/hogql_queries/insights/test/test_events_query_runner.py:0: error: Item "
posthog/hogql_queries/insights/test/test_events_query_runner.py:0: error: Item "None" of "Expr | None" has no attribute "exprs" [union-attr]
posthog/hogql_queries/insights/test/test_events_query_runner.py:0: error: Item "Expr" of "Expr | None" has no attribute "exprs" [union-attr]
posthog/hogql_queries/insights/test/test_events_query_runner.py:0: error: Item "None" of "Expr | None" has no attribute "exprs" [union-attr]
posthog/hogql/transforms/test/test_in_cohort.py:0: error: "TestInCohort" has no attribute "snapshot" [attr-defined]
posthog/hogql/transforms/test/test_in_cohort.py:0: error: Argument 1 to "len" has incompatible type "list[Any] | None"; expected "Sized" [arg-type]
posthog/hogql/transforms/test/test_in_cohort.py:0: error: Value of type "list[Any] | None" is not indexable [index]
posthog/hogql/transforms/test/test_in_cohort.py:0: error: "TestInCohort" has no attribute "snapshot" [attr-defined]
posthog/hogql/transforms/test/test_in_cohort.py:0: error: "TestInCohort" has no attribute "snapshot" [attr-defined]
posthog/hogql/test/test_timings.py:0: error: No overload variant of "__setitem__" of "list" matches argument types "int", "float" [call-overload]
posthog/hogql/test/test_timings.py:0: note: Possible overload variants:
posthog/hogql/test/test_timings.py:0: note: def __setitem__(self, SupportsIndex, int, /) -> None
Expand Down
6 changes: 4 additions & 2 deletions posthog/hogql/functions/cohort.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,10 @@ def cohort(node: ast.Expr, args: List[ast.Expr], context: HogQLContext) -> ast.E

from posthog.models import Cohort

if isinstance(arg.value, int) and not isinstance(arg.value, bool):
cohorts = Cohort.objects.filter(id=arg.value, team_id=context.team_id).values_list("id", "is_static", "name")
if (isinstance(arg.value, int) or isinstance(arg.value, float)) and not isinstance(arg.value, bool):
cohorts = Cohort.objects.filter(id=int(arg.value), team_id=context.team_id).values_list(
"id", "is_static", "name"
)
if len(cohorts) == 1:
context.add_notice(
start=arg.start,
Expand Down
4 changes: 2 additions & 2 deletions posthog/hogql/modifiers.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from typing import Optional, TYPE_CHECKING

from posthog.schema import HogQLQueryModifiers, MaterializationMode
from posthog.schema import HogQLQueryModifiers, InCohortVia, MaterializationMode
from posthog.utils import PersonOnEventsMode

if TYPE_CHECKING:
Expand All @@ -22,7 +22,7 @@ def create_default_modifiers_for_team(
modifiers.personsArgMaxVersion = "auto"

if modifiers.inCohortVia is None:
modifiers.inCohortVia = "subquery"
modifiers.inCohortVia = InCohortVia.subquery

if modifiers.materializationMode is None or modifiers.materializationMode == MaterializationMode.auto:
modifiers.materializationMode = MaterializationMode.legacy_null_as_null
Expand Down
9 changes: 6 additions & 3 deletions posthog/hogql/printer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,15 +33,15 @@
from posthog.hogql.modifiers import create_default_modifiers_for_team
from posthog.hogql.resolver import ResolverException, resolve_types
from posthog.hogql.resolver_utils import lookup_field_by_name
from posthog.hogql.transforms.in_cohort import resolve_in_cohorts
from posthog.hogql.transforms.in_cohort import resolve_in_cohorts, resolve_in_cohorts_conjoined
from posthog.hogql.transforms.lazy_tables import resolve_lazy_tables
from posthog.hogql.transforms.property_types import resolve_property_types
from posthog.hogql.visitor import Visitor, clone_expr
from posthog.models.property import PropertyName, TableColumn
from posthog.models.team.team import WeekStartDay
from posthog.models.team import Team
from posthog.models.utils import UUIDT
from posthog.schema import MaterializationMode
from posthog.schema import InCohortVia, MaterializationMode
from posthog.utils import PersonOnEventsMode


Expand Down Expand Up @@ -99,9 +99,12 @@ def prepare_ast_for_printing(
with context.timings.measure("create_hogql_database"):
context.database = context.database or create_hogql_database(context.team_id, context.modifiers)

if context.modifiers.inCohortVia == InCohortVia.leftjoin_conjoined:
with context.timings.measure("resolve_in_cohorts_conjoined"):
resolve_in_cohorts_conjoined(node, dialect, context, stack)
with context.timings.measure("resolve_types"):
node = resolve_types(node, context, dialect=dialect, scopes=[node.type for node in stack] if stack else None)
if context.modifiers.inCohortVia == "leftjoin":
if context.modifiers.inCohortVia == InCohortVia.leftjoin:
with context.timings.measure("resolve_in_cohorts"):
resolve_in_cohorts(node, dialect, stack, context)
if dialect == "clickhouse":
Expand Down
Loading

0 comments on commit 626e880

Please sign in to comment.