perf(Segment Membership): Count all segments in one scan per environment (#8018)

Author: khvn26

api/segment_membership/services.py

@@ -110,7 +110,10 @@ def compute_segment_counts_for_project(
     project: Project, cursor: CursorWrapper
 ) -> list[SegmentMembershipCount]:
     """Count identity matches per (canonical-segment, environment) for
-    `project` in one `UNION ALL` query.
+    `project`, scanning each environment once.
+
+    A single `GROUP BY environment_id` over `IDENTITIES FINAL` counts every
+    segment in one pass via `countIf(<predicate>)` per segment.
 
     Returns unsaved `SegmentMembershipCount` instances with `count` and
     keys populated; the caller stamps `last_synced_at` consistently
@@ -128,7 +131,8 @@ def compute_segment_counts_for_project(
 
     dialect = ClickHouseDialect()
     binder = Binder(PyformatParamStyle())
-    select_clauses: list[str] = []
+    count_columns: list[str] = []
+    counted_segment_ids: list[int] = []
     for seg in segments:
         translate_ctx = TranslateContext(
             evaluation_context=EvaluationContext(
@@ -149,33 +153,35 @@ def compute_segment_counts_for_project(
                 reason="untranslatable",
             )
             continue
-        select_clauses.append(
-            f"SELECT {seg.id} AS segment_id, "
-            f"i.environment_id AS env_key, count() AS c "
-            f"FROM IDENTITIES AS i FINAL "
-            f"WHERE i.environment_id IN %(env_keys)s "
-            f"AND i.is_deleted = false AND ({predicate}) "
-            f"GROUP BY i.environment_id"
-        )
+        count_columns.append(f"countIf({predicate}) AS c{seg.id}")
+        counted_segment_ids.append(seg.id)
 
-    if not select_clauses:
+    if not count_columns:
         return []
 
-    sql = "\nUNION ALL\n".join(select_clauses)
+    sql = (
+        f"SELECT i.environment_id AS env_key, {', '.join(count_columns)} "
+        f"FROM IDENTITIES AS i FINAL "
+        f"WHERE i.environment_id IN %(env_keys)s AND i.is_deleted = false "
+        f"GROUP BY i.environment_id"
+    )
     cursor.execute(sql, {"env_keys": tuple(env_id_by_key), **binder.params})
     rows: list[tuple[Any, ...]] = cursor.fetchall()
     membership_counts: list[SegmentMembershipCount] = []
     for row in rows:
-        env_id = env_id_by_key.get(str(row[1]))
+        env_id = env_id_by_key.get(str(row[0]))
         if env_id is None:
             continue
-        membership_counts.append(
-            SegmentMembershipCount(
-                segment_id=int(row[0]),
-                environment_id=env_id,
-                count=int(row[2]),
-            )
-        )
+        # Columns line up with counted_segment_ids; a zero-match pair is absent.
+        for segment_id, count in zip(counted_segment_ids, row[1:]):
+            if count:
+                membership_counts.append(
+                    SegmentMembershipCount(
+                        segment_id=segment_id,
+                        environment_id=env_id,
+                        count=int(count),
+                    )
+                )
     return membership_counts
 
 

api/tests/unit/segment_membership/test_unit_segment_membership_services.py

@@ -6,7 +6,7 @@
 from django.db import connections
 from django.utils import timezone
 from flag_engine.segments.constants import EQUAL, REGEX
-from pytest_django.fixtures import SettingsWrapper
+from pytest_django.fixtures import DjangoAssertNumQueries, SettingsWrapper
 from pytest_mock import MockerFixture
 from task_processor.models import Task
 from task_processor.task_run_method import TaskRunMethod
@@ -131,7 +131,7 @@ def test_compute_segment_counts_for_project__unknown_env_key_in_row__skips(
         return_value="TRUE",
     )
     cursor = MagicMock()
-    cursor.fetchall.return_value = [(segment.id, "ghost-env", 99)]
+    cursor.fetchall.return_value = [("ghost-env", 99)]
 
     # When
     result = compute_segment_counts_for_project(project, cursor)
@@ -198,6 +198,27 @@ def percent_regex_segment(segment: Segment) -> Segment:
     return segment
 
 
+@pytest.fixture
+def three_segments(matching_segment: Segment, project: Project) -> dict[str, Segment]:
+    """`matching_segment` matches foo == bar (alice, bob); add one matching a
+    different value (carol) and one matching nobody."""
+    seg_baz = Segment.objects.create(name="baz", project=project)
+    Condition.objects.create(
+        rule=SegmentRule.objects.create(segment=seg_baz, type=SegmentRule.ALL_RULE),
+        property="foo",
+        operator=EQUAL,
+        value="baz",
+    )
+    seg_none = Segment.objects.create(name="none", project=project)
+    Condition.objects.create(
+        rule=SegmentRule.objects.create(segment=seg_none, type=SegmentRule.ALL_RULE),
+        property="foo",
+        operator=EQUAL,
+        value="zzz",
+    )
+    return {"bar": matching_segment, "baz": seg_baz, "none": seg_none}
+
+
 @pytest.mark.clickhouse
 def test_get_segment_members_page__deleted_identity__excluded(
     segment_membership_identities: None,
@@ -288,6 +309,43 @@ def test_compute_segment_counts_for_project__deleted_identity__excluded_from_cou
     assert counts[0].count == 2
 
 
+@pytest.mark.clickhouse
+def test_compute_segment_counts_for_project__multiple_segments__maps_each_count_drops_zeros(
+    segment_membership_identities: None,
+    three_segments: dict[str, Segment],
+    project: Project,
+) -> None:
+    # Given / When
+    with connections["clickhouse"].cursor() as cursor:
+        counts = compute_segment_counts_for_project(project, cursor)
+
+    # Then
+    by_segment = {c.segment_id: c.count for c in counts}
+    assert by_segment == {
+        three_segments["bar"].id: 2,
+        three_segments["baz"].id: 1,
+    }
+    assert three_segments["none"].id not in by_segment
+
+
+@pytest.mark.clickhouse
+def test_compute_segment_counts_for_project__multiple_segments__uses_single_query(
+    segment_membership_identities: None,
+    three_segments: dict[str, Segment],
+    project: Project,
+    django_assert_num_queries: DjangoAssertNumQueries,
+) -> None:
+    # Given
+    # three_segments + identities are set up by fixtures
+
+    # When / Then
+    with (
+        django_assert_num_queries(1, connection=connections["clickhouse"]),
+        connections["clickhouse"].cursor() as cursor,
+    ):
+        compute_segment_counts_for_project(project, cursor)
+
+
 @pytest.mark.clickhouse
 def test_get_segment_members_page__regex_with_percent__returns_matches(
     segment_membership_identities: None,

docs/docs/deployment-self-hosting/observability/_events-catalogue.md

@@ -369,7 +369,7 @@ Attributes:
 ### `segment_membership.compute.segment.skipped`
 
 Logged at `error` from:
- - `api/segment_membership/services.py:145`
+ - `api/segment_membership/services.py:149`
 
 Attributes:
  - `project.id`
@@ -379,7 +379,7 @@ Attributes:
 ### `segment_membership.members.segment.skipped`
 
 Logged at `error` from:
- - `api/segment_membership/services.py:209`
+ - `api/segment_membership/services.py:215`
 
 Attributes:
  - `reason`