feat(performance): query improvements for trends (load less people)

frontend/src/lib/components/CommandBar/commandBarLogic.ts

@@ -54,7 +54,7 @@ export const commandBarLogic = kea<commandBarLogicType>([
             if (shouldIgnoreInput(event)) {
                 return
             }
-            if ((event.ctrlKey || event.metaKey) && event.key === 'k') {
+            if ((event.ctrlKey || event.metaKey) && (event.key === 'k' || event.key === 'K')) {
                 event.preventDefault()
                 if (event.shiftKey) {
                     // cmd+shift+k opens actions

posthog/api/test/__snapshots__/test_query.ambr

@@ -486,7 +486,7 @@
      FROM person_distinct_id2
      WHERE equals(person_distinct_id2.team_id, 2)
      GROUP BY person_distinct_id2.distinct_id
-     HAVING ifNull(equals(argMax(person_distinct_id2.is_deleted, person_distinct_id2.version), 0), 0)) AS events__pdi ON equals(events.distinct_id, events__pdi.distinct_id)
+     HAVING ifNull(equals(argMax(person_distinct_id2.is_deleted, person_distinct_id2.version), 0), 0) SETTINGS optimize_aggregation_in_order=1) AS events__pdi ON equals(events.distinct_id, events__pdi.distinct_id)
   LEFT JOIN
     (SELECT person.id AS id,
             replaceRegexpAll(nullIf(nullIf(JSONExtractRaw(person.properties, 'email'), ''), 'null'), '^"|"$', '') AS properties___email
@@ -496,7 +496,7 @@
                                                        FROM person
                                                        WHERE equals(person.team_id, 2)
                                                        GROUP BY person.id
-                                                       HAVING and(ifNull(equals(argMax(person.is_deleted, person.version), 0), 0), ifNull(less(argMax(person.created_at, person.version), plus(now64(6, 'UTC'), toIntervalDay(1))), 0)))), 0)) SETTINGS optimize_aggregation_in_order=1) AS events__pdi__person ON equals(events__pdi.events__pdi___person_id, events__pdi__person.id)
+                                                       HAVING and(ifNull(equals(argMax(person.is_deleted, person.version), 0), 0), ifNull(less(argMax(toTimeZone(person.created_at, 'UTC'), person.version), plus(now64(6, 'UTC'), toIntervalDay(1))), 0)))), 0)) SETTINGS optimize_aggregation_in_order=1) AS events__pdi__person ON equals(events__pdi.events__pdi___person_id, events__pdi__person.id)
   WHERE and(equals(events.team_id, 2), ifNull(equals(events__pdi__person.properties___email, '[email protected]'), 0), less(toTimeZone(events.timestamp, 'UTC'), toDateTime64('2020-01-10 12:14:05.000000', 6, 'UTC')), greater(toTimeZone(events.timestamp, 'UTC'), toDateTime64('2020-01-09 12:14:00.000000', 6, 'UTC')))
   ORDER BY events.event ASC
   LIMIT 101
@@ -526,7 +526,7 @@
      FROM person_distinct_id2
      WHERE equals(person_distinct_id2.team_id, 2)
      GROUP BY person_distinct_id2.distinct_id
-     HAVING ifNull(equals(argMax(person_distinct_id2.is_deleted, person_distinct_id2.version), 0), 0)) AS events__pdi ON equals(events.distinct_id, events__pdi.distinct_id)
+     HAVING ifNull(equals(argMax(person_distinct_id2.is_deleted, person_distinct_id2.version), 0), 0) SETTINGS optimize_aggregation_in_order=1) AS events__pdi ON equals(events.distinct_id, events__pdi.distinct_id)
   LEFT JOIN
     (SELECT person.id AS id,
             nullIf(nullIf(person.pmat_email, ''), 'null') AS properties___email
@@ -536,7 +536,7 @@
                                                        FROM person
                                                        WHERE equals(person.team_id, 2)
                                                        GROUP BY person.id
-                                                       HAVING and(ifNull(equals(argMax(person.is_deleted, person.version), 0), 0), ifNull(less(argMax(person.created_at, person.version), plus(now64(6, 'UTC'), toIntervalDay(1))), 0)))), 0)) SETTINGS optimize_aggregation_in_order=1) AS events__pdi__person ON equals(events__pdi.events__pdi___person_id, events__pdi__person.id)
+                                                       HAVING and(ifNull(equals(argMax(person.is_deleted, person.version), 0), 0), ifNull(less(argMax(toTimeZone(person.created_at, 'UTC'), person.version), plus(now64(6, 'UTC'), toIntervalDay(1))), 0)))), 0)) SETTINGS optimize_aggregation_in_order=1) AS events__pdi__person ON equals(events__pdi.events__pdi___person_id, events__pdi__person.id)
   WHERE and(equals(events.team_id, 2), ifNull(equals(events__pdi__person.properties___email, '[email protected]'), 0), less(toTimeZone(events.timestamp, 'UTC'), toDateTime64('2020-01-10 12:14:05.000000', 6, 'UTC')), greater(toTimeZone(events.timestamp, 'UTC'), toDateTime64('2020-01-09 12:14:00.000000', 6, 'UTC')))
   ORDER BY events.event ASC
   LIMIT 101

posthog/api/test/test_cohort.py

@@ -33,7 +33,7 @@
 class TestCohort(TestExportMixin, ClickhouseTestMixin, APIBaseTest, QueryMatchingTest):
     # select all queries for snapshots
     def capture_select_queries(self):
-        return self.capture_queries(("INSERT INTO cohortpeople", "SELECT", "ALTER", "select", "DELETE"))
+        return self.capture_queries_startswith(("INSERT INTO cohortpeople", "SELECT", "ALTER", "select", "DELETE"))
 
     def _get_cohort_activity(
         self,
@@ -101,7 +101,7 @@ def test_creating_update_and_calculating(self, patch_sync_execute, patch_calcula
             },
         )
 
-        with self.capture_queries("INSERT INTO cohortpeople") as insert_statements:
+        with self.capture_queries_startswith("INSERT INTO cohortpeople") as insert_statements:
             response = self.client.patch(
                 f"/api/projects/{self.team.id}/cohorts/{response.json()['id']}",
                 data={

posthog/hogql/database/database.py

@@ -47,7 +47,11 @@
     PersonDistinctIdsTable,
     RawPersonDistinctIdsTable,
 )
-from posthog.hogql.database.schema.persons import PersonsTable, RawPersonsTable, join_with_persons_table
+from posthog.hogql.database.schema.persons import (
+    PersonsTable,
+    RawPersonsTable,
+    join_with_persons_table,
+)
 from posthog.hogql.database.schema.session_replay_events import (
     RawSessionReplayEventsTable,
     SessionReplayEventsTable,

posthog/hogql/database/schema/person_distinct_id_overrides.py

@@ -1,4 +1,5 @@
 from posthog.hogql.ast import SelectQuery
+from posthog.hogql.constants import HogQLQuerySettings
 from posthog.hogql.context import HogQLContext
 
 from posthog.hogql.database.argmax import argmax_select
@@ -32,13 +33,15 @@ def select_from_person_distinct_id_overrides_table(requested_fields: dict[str, l
     # Always include "person_id", as it's the key we use to make further joins, and it'd be great if it's available
     if "person_id" not in requested_fields:
         requested_fields = {**requested_fields, "person_id": ["person_id"]}
-    return argmax_select(
+    select = argmax_select(
         table_name="raw_person_distinct_id_overrides",
         select_fields=requested_fields,
         group_fields=["distinct_id"],
         argmax_field="version",
         deleted_field="is_deleted",
     )
+    select.settings = HogQLQuerySettings(optimize_aggregation_in_order=True)
+    return select
 
 
 def join_with_person_distinct_id_overrides_table(

posthog/hogql/database/schema/person_distinct_ids.py

@@ -1,4 +1,5 @@
 from posthog.hogql.ast import SelectQuery
+from posthog.hogql.constants import HogQLQuerySettings
 from posthog.hogql.context import HogQLContext
 
 from posthog.hogql.database.argmax import argmax_select
@@ -32,13 +33,15 @@ def select_from_person_distinct_ids_table(requested_fields: dict[str, list[str |
     # Always include "person_id", as it's the key we use to make further joins, and it'd be great if it's available
     if "person_id" not in requested_fields:
         requested_fields = {**requested_fields, "person_id": ["person_id"]}
-    return argmax_select(
+    select = argmax_select(
         table_name="raw_person_distinct_ids",
         select_fields=requested_fields,
         group_fields=["distinct_id"],
         argmax_field="version",
         deleted_field="is_deleted",
     )
+    select.settings = HogQLQuerySettings(optimize_aggregation_in_order=True)
+    return select
 
 
 def join_with_person_distinct_ids_table(

posthog/hogql/database/schema/persons.py

@@ -1,7 +1,9 @@
-from typing import cast
+from typing import cast, Optional
+from typing_extensions import Self
 import posthoganalytics
 
-from posthog.hogql.ast import SelectQuery, And
+from posthog.hogql.ast import SelectQuery, And, CompareOperation, CompareOperationOp, Field, JoinExpr
+from posthog.hogql.base import Expr
 from posthog.hogql.constants import HogQLQuerySettings
 from posthog.hogql.context import HogQLContext
 from posthog.hogql.database.argmax import argmax_select
@@ -21,6 +23,7 @@
 from posthog.hogql.database.schema.util.where_clause_extractor import WhereClauseExtractor
 from posthog.hogql.database.schema.persons_pdi import PersonsPDITable, persons_pdi_join
 from posthog.hogql.errors import ResolutionError
+from posthog.hogql.visitor import clone_expr
 from posthog.models.organization import Organization
 from posthog.schema import PersonsArgMaxVersion
 
@@ -38,7 +41,13 @@
 }
 
 
-def select_from_persons_table(join_or_table: LazyJoinToAdd | LazyTableToAdd, context: HogQLContext, node: SelectQuery):
+def select_from_persons_table(
+    join_or_table: LazyJoinToAdd | LazyTableToAdd,
+    context: HogQLContext,
+    node: SelectQuery,
+    *,
+    filter: Optional[Expr] = None,
+):
     version = context.modifiers.personsArgMaxVersion
     if version == PersonsArgMaxVersion.AUTO:
         version = PersonsArgMaxVersion.V1
@@ -67,6 +76,8 @@ def select_from_persons_table(join_or_table: LazyJoinToAdd | LazyTableToAdd, con
             ),
         )
         select.settings = HogQLQuerySettings(optimize_aggregation_in_order=True)
+        if filter is not None:
+            cast(ast.SelectQuery, cast(ast.CompareOperation, select.where).right).where = filter
 
         for field_name, field_chain in join_or_table.fields_accessed.items():
             # We need to always select the 'id' field for the join constraint. The field name here is likely to
@@ -88,6 +99,11 @@ def select_from_persons_table(join_or_table: LazyJoinToAdd | LazyTableToAdd, con
             timestamp_field_to_clamp="created_at",
         )
         select.settings = HogQLQuerySettings(optimize_aggregation_in_order=True)
+        if filter is not None:
+            if select.where:
+                select.where = And(exprs=[select.where, filter])
+            else:
+                select.where = filter
 
     if context.modifiers.optimizeJoinedFilters:
         extractor = WhereClauseExtractor(context)
@@ -162,10 +178,55 @@ def to_printed_hogql(self):
         return "raw_persons"
 
 
+# Persons is a lazy table that allows you to insert a where statement inside of the person subselect
+# It pulls any "persons.id in ()" statement inside of the argmax subselect
+# This is useful when executing a query for a large team.
 class PersonsTable(LazyTable):
     fields: dict[str, FieldOrTable] = PERSONS_FIELDS
+    filter: Optional[Expr] = None
+
+    @staticmethod
+    def _is_promotable_expr(expr, alias: Optional[str] = None):
+        return (
+            isinstance(expr, CompareOperation)
+            and expr.op == CompareOperationOp.In
+            and isinstance(expr.left, Field)
+            and expr.left.chain == [alias or "persons", "id"]
+        )
+
+    @staticmethod
+    def _partition_exprs(exprs, alias: Optional[str] = None):
+        not_promotable = []
+        promotable = []
+        for expr in exprs:
+            if PersonsTable._is_promotable_expr(expr, alias):
+                # Erase "persons" from the chain before bringing inside
+                expr.left = Field(chain=["id"])
+                promotable.append(expr)
+            else:
+                not_promotable.append(expr)
+
+        return promotable, not_promotable
+
+    # If the join has a clause we can bring inside the subselect, create a new table that represents that
+    def create_new_table_with_filter(self, join: JoinExpr) -> Self:
+        if join.constraint is not None and isinstance(join.constraint.expr, And):
+            exprs = cast(And, join.constraint.expr).exprs
+            promotable, not_promotable = PersonsTable._partition_exprs(exprs, join.alias)
+            if len(promotable) == 0:
+                return self
+            join.constraint.expr.exprs = not_promotable
+            p = self.model_copy()
+            if len(promotable) == 1:
+                p.filter = promotable[0]
+            elif len(promotable) > 1:
+                p.filter = And(exprs=promotable)
+            return p
+        return self
 
     def lazy_select(self, table_to_add: LazyTableToAdd, context, node):
+        if self.filter is not None:
+            return select_from_persons_table(table_to_add, context, node, filter=clone_expr(self.filter, True, True))
         return select_from_persons_table(table_to_add, context, node)
 
     def to_printed_clickhouse(self, context):

posthog/hogql/database/schema/persons_pdi.py

@@ -1,6 +1,7 @@
 import posthoganalytics
 
 from posthog.hogql.ast import SelectQuery
+from posthog.hogql.constants import HogQLQuerySettings
 from posthog.hogql.context import HogQLContext
 from posthog.hogql.database.argmax import argmax_select
 from posthog.hogql.database.models import (
@@ -21,13 +22,15 @@ def persons_pdi_select(requested_fields: dict[str, list[str | int]]):
     # Always include "person_id", as it's the key we use to make further joins, and it'd be great if it's available
     if "person_id" not in requested_fields:
         requested_fields = {**requested_fields, "person_id": ["person_id"]}
-    return argmax_select(
+    select = argmax_select(
         table_name="raw_person_distinct_ids",
         select_fields=requested_fields,
         group_fields=["distinct_id"],
         argmax_field="version",
         deleted_field="is_deleted",
     )
+    select.settings = HogQLQuerySettings(optimize_aggregation_in_order=True)
+    return select
 
 
 # :NOTE: We already have person_distinct_ids.py, which most tables link to. This persons_pdi.py is a hack to