Repository: PostHog/posthog
Pull request: perf: Parallel execution of trends graph series
## Problem

We can massively speed up queries with multiple graph series by executing the clickhouse queries in parallel

closes #8834

## Changes

- Run multiple graph series in parallel

## How did you test this code?

<!-- If the answer is manually, please include a quick step-by-step on how to test this PR. -->

Unit tests
Create trends query with multiple graph series.

Hard to test the speed up locally, but specific queries for client should execute a lot faster (currently around 2 minutes if it finishes at all, with 3 graph series)

Problem statement:
Calculate multi-line trends in parallel
## Is your feature request related to a problem?

If you have a trend with multiple lines, clickhouse queries for each are run one-after-another. If each takes 10s to calculate, then user will need to wait >30s in total.

## Describe the solution you'd like

Run queries in parallel.

## Describe alternatives you've considered

## Additional context

cc @EDsCODE @neilkakkar This seems like a quick win query-api wise though not 100% sure on how best implement this. Would you be up for prioritizing this _next_ sprint?

#### *Thank you* for your feature request – we love each and every one!

Hints from issue discussion:
Naive question: Is there a risk with potentially running lots of expensive queries running in parallel they end up causing each-other to fail (due to peak resource consumption)?
I don't think so - the load should be comparable to normal parallelism from users doing queries or background cache refreshes. 

But if not too hard to hard, having _some_ limitation on parallel execution on this API call (e.g. no more than 5 queries at once) wouldn't hurt and still solve the 99.9% usecase.

Commit message: remove unused lines

Patch to review:
diff --git a/ee/clickhouse/queries/trends/clickhouse_trends.py b/ee/clickhouse/queries/trends/clickhouse_trends.py
index 9053bf4fe3d8f..f6b0d3f90f1c3 100644
--- a/ee/clickhouse/queries/trends/clickhouse_trends.py
+++ b/ee/clickhouse/queries/trends/clickhouse_trends.py
@@ -1,6 +1,7 @@
 import copy
+import threading
 from itertools import accumulate
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Tuple, Union
 
 from django.db.models.query import Prefetch
 from django.utils import timezone
@@ -52,6 +53,45 @@ def _run_query(self, filter: Filter, entity: Entity, team_id: int) -> List[Dict[
             serialized_data = self._handle_cumulative(serialized_data)
         return serialized_data
 
+    def _run_query_for_threading(self, result: List, index: int, sql, params):
+        result[index] = sync_execute(sql, params)
+
+    def _run_parallel(self, filter: Filter, team: Team) -> List[Dict[str, Any]]:
+        result: List[Union[None, List[Dict[str, Any]]]] = [None] * len(filter.entities)
+        jobs = []
+
+        for entity in filter.entities:
+            sql, params, _ = self._get_sql_for_entity(filter, entity, team.pk)
+            thread = threading.Thread(target=self._run_query_for_threading, args=(result, entity.index, sql, params),)
+            jobs.append(thread)
+
+        # Start the threads (i.e. calculate the random number lists)
+        for j in jobs:
+            j.start()
+
+        # Ensure all of the threads have finished
+        for j in jobs:
+            j.join()
+
+        # Parse results for each thread
+        for entity in filter.entities:
+            _, _, parse_function = self._get_sql_for_entity(filter, entity, team.pk)
+            serialized_data = parse_function(result[entity.index])
+            serialized_data = self._format_serialized(entity, serialized_data)
+
+            if filter.display == TRENDS_CUMULATIVE:
+                serialized_data = self._handle_cumulative(serialized_data)
+            result[entity.index] = serialized_data
+
+        # flatten results
+        flat_results: List[Dict[str, Any]] = []
+        for item in result:
+            if item is not None:  # Mypy thinks results might still be None
+                for flat in item:
+                    flat_results.append(flat)
+
+        return flat_results
+
     def run(self, filter: Filter, team: Team, *args, **kwargs) -> List[Dict[str, Any]]:
         actions = Action.objects.filter(team_id=team.pk).order_by("-id")
         if len(filter.actions) > 0:
@@ -63,15 +103,19 @@ def run(self, filter: Filter, team: Team, *args, **kwargs) -> List[Dict[str, Any
         if filter.formula:
             return handle_compare(filter, self._run_formula_query, team)
 
-        result = []
         for entity in filter.entities:
             if entity.type == TREND_FILTER_TYPE_ACTIONS:
                 try:
                     entity.name = actions.get(id=entity.id).name
                 except Action.DoesNotExist:
-                    continue
-            entities_list = handle_compare(filter, self._run_query, team, entity=entity)
-            result.extend(entities_list)
+                    return []
+
+        if len(filter.entities) == 1 or filter.compare:
+            result = []
+            for entity in filter.entities:
+                result.extend(handle_compare(filter, self._run_query, team, entity=entity))
+        else:
+            result = self._run_parallel(filter, team)
 
         return result
 


Code review comments:
diff_hunk: @@ -52,6 +53,45 @@ def _run_query(self, filter: Filter, entity: Entity, team_id: int) -> List[Dict[
             serialized_data = self._handle_cumulative(serialized_data)
         return serialized_data
 
+    def _run_query_for_threading(self, result: List, index: int, sql, params):
+        result[index] = sync_execute(sql, params)
+
+    def _run_parallel(self, filter: Filter, team: Team) -> List[Dict[str, Any]]:
+        result: List[Union[None, List[Dict[str, Any]]]] = [None] * len(filter.entities)
+        jobs = []
+
+        for entity in filter.entities:
+            sql, params, _ = self._get_sql_for_entity(filter, entity, team.pk)
path: ee/clickhouse/queries/trends/clickhouse_trends.py
text: @user1:
Q: Why not generate the sql/params in the thread as well? We could leverage `_run_query` (indirectly) then.

@author:
See above, but connecting to the database doesn't work in thredas.

diff_hunk: @@ -52,6 +53,45 @@ def _run_query(self, filter: Filter, entity: Entity, team_id: int) -> List[Dict[
             serialized_data = self._handle_cumulative(serialized_data)
         return serialized_data
 
+    def _run_query_for_threading(self, result: List, index: int, sql, params):
+        result[index] = sync_execute(sql, params)
+
+    def _run_parallel(self, filter: Filter, team: Team) -> List[Dict[str, Any]]:
+        result: List[Union[None, List[Dict[str, Any]]]] = [None] * len(filter.entities)
+        jobs = []
+
+        for entity in filter.entities:
+            sql, params, _ = self._get_sql_for_entity(filter, entity, team.pk)
+            thread = threading.Thread(target=self._run_query_for_threading, args=(result, entity.index, sql, params),)
+            jobs.append(thread)
+
+        # Start the threads (i.e. calculate the random number lists)
+        for j in jobs:
+            j.start()
+
+        # Ensure all of the threads have finished
+        for j in jobs:
+            j.join()
+
+        # Parse results for each thread
+        for entity in filter.entities:
+            _, _, parse_function = self._get_sql_for_entity(filter, entity, team.pk)
path: ee/clickhouse/queries/trends/clickhouse_trends.py
text: @user1:
Constructing this twice might not be ideal. What if each thread was also responsible for also serializing the data? This might avoid the parsing-related duplication that's occurring here and in`_run_query`

@author:
Connecting to the database doesn't work in threads, hence doing this outside.

@user2:
Ideally, we just pull the parse function needed directly. Otherwise, there's a lot of unnecessary logic being done to build the query

@author:
How about this? Cleaner than repeating all of the logic

@user2:
oh yep, nice

diff_hunk: @@ -52,6 +53,45 @@ def _run_query(self, filter: Filter, entity: Entity, team_id: int) -> List[Dict[
             serialized_data = self._handle_cumulative(serialized_data)
         return serialized_data
 
+    def _run_query_for_threading(self, result: List, index: int, sql, params):
+        result[index] = sync_execute(sql, params)
+
+    def _run_parallel(self, filter: Filter, team: Team) -> List[Dict[str, Any]]:
+        result: List[Union[None, List[Dict[str, Any]]]] = [None] * len(filter.entities)
+        jobs = []
+
+        for entity in filter.entities:
+            sql, params, _ = self._get_sql_for_entity(filter, entity, team.pk)
+            thread = threading.Thread(target=self._run_query_for_threading, args=(result, entity.index, sql, params),)
+            jobs.append(thread)
+
+        # Start the threads (i.e. calculate the random number lists)
+        for j in jobs:
+            j.start()
+
+        # Ensure all of the threads have finished
+        for j in jobs:
+            j.join()
+
+        # Parse results for each thread
+        for entity in filter.entities:
+            _, _, parse_function = self._get_sql_for_entity(filter, entity, team.pk)
+            serialized_data = parse_function(result[entity.index])
+            serialized_data = self._format_serialized(entity, serialized_data)
+
+            if filter.display == TRENDS_CUMULATIVE:
+                serialized_data = self._handle_cumulative(serialized_data)
+            result[entity.index] = serialized_data
+
+        # flatten results
+        flat_results: List[Dict[str, Any]] = []
+        for item in result:
+            if item is not None:  # Mypy thinks results might still be None
path: ee/clickhouse/queries/trends/clickhouse_trends.py
text: @user1:
You can use `cast()` function instead