fix(flags): bound the release condition size query's memory

products/feature_flags/backend/api/feature_flag.py

@@ -132,7 +132,11 @@
     teams_gating_replay_on_flag,
 )
 from products.feature_flags.backend.types import PropertyFilterType
-from products.feature_flags.backend.user_blast_radius import get_user_blast_radius
+from products.feature_flags.backend.user_blast_radius import (
+    get_person_blast_radius_v2,
+    get_user_blast_radius,
+    use_blast_radius_query_v2,
+)
 from products.feature_flags.backend.version_history import (
     VersionHistoryIncomplete,
     VersionNotFound,
@@ -249,7 +253,7 @@
 ENCRYPTED_VERSION_HISTORY_UNAVAILABLE = "Version history is not available for flags with encrypted payloads."
 
 
 def parse_created_by_ids(value: Any) -> list[int]:
     """Parse a `created_by_id` filter value into a list of user IDs.
 
     Accepts a single int/str, a comma-separated string, or a JSON-encoded list
@@ -870,7 +874,7 @@
     return "*" in settings.FEATURE_FLAG_FILTERS_ENFORCED_RULES
 
 
 def _reject_serde_unsafe_filters(filters: Any) -> None:
     """Serde-fidelity guard on incoming filters, independent of which rules are enforced.
 
     These are the pre-#50084 procedural type/bounds checks: values that fail serde in the
@@ -1604,7 +1608,7 @@
                 code="unsupported_config_version",
             )
 
     def _validate_filters_inner(self, filters, operation: str):
         # Unknown keys survive normalization during the validation rollout. Reserve the
         # config discriminator before that path can store an unsupported format.
         if "version" in filters:
@@ -2160,7 +2164,7 @@
         return instance
 
     @approval_gate(["feature_flag.enable", "feature_flag.disable", "feature_flag.update"])
     def update(self, instance: FeatureFlag, validated_data: dict, *args: Any, **kwargs: Any) -> FeatureFlag:
         request = self.context["request"]
         # Service-layer callers may carry no body, and an endpoint that declares no request schema
         # can hand us a non-dict one (e.g. a JSON array posted to /experiments/:id/launch/).
@@ -2487,7 +2491,7 @@
         flag_status, _ = checker.get_status()
         return flag_status.name
 
     def to_representation(self, instance):
         representation = super().to_representation(instance)
         filters = representation.get("filters", {})
 
@@ -4729,7 +4733,10 @@
         condition = request.data.get("condition") or {}
         group_type_index = request.data.get("group_type_index", None)
 
-        result = get_user_blast_radius(self.team, condition, group_type_index)
+        if group_type_index is None and use_blast_radius_query_v2(self.team):
+            result = get_person_blast_radius_v2(self.team, condition)
+        else:
+            result = get_user_blast_radius(self.team, condition, group_type_index)
 
         return Response({"affected": result.affected, "total": result.total})
 

products/feature_flags/backend/person_sampling.py

@@ -0,0 +1,158 @@
+from collections.abc import Callable
+from typing import Optional
+
+from posthog.schema import HogQLQueryModifiers
+
+from posthog.hogql import ast
+from posthog.hogql.constants import HogQLGlobalSettings
+from posthog.hogql.context import HogQLContext
+from posthog.hogql.database.database import Database
+from posthog.hogql.property import property_to_expr
+from posthog.hogql.query import execute_hogql_query
+
+from posthog.models.filters import Filter
+from posthog.models.team.team import Team
+from posthog.schema_enums import PersonsArgMaxVersion
+
+# 1-in-64 sample of the person keyspace. Sampling keys on person id (the dedup key), so every
+# version row of a sampled person lands in the sample and the argMax dedup stays exact within it.
+SAMPLE_MODULUS = 64
+
+# Below this many sampled matches the extrapolation is too noisy, so rerun exact. The relative
+# error of the estimate is about sqrt(63 / matched_persons), so 10,000 sampled matches
+# (~640k matched persons) keeps the worst case near 3%.
+#
+# The threshold reads the match count, not the table size, so a large team whose condition
+# matches few persons reruns exact. That rerun reads the whole table, but it does not rebuild
+# the shape this sampling exists to avoid: the id prefilter holds the dedup to the matched
+# candidate set, and count_settings adds in-order aggregation and a spill threshold. The cost
+# is scan time instead of a memory-limit error.
+# Caveat: the prefilter matches any historical row version, so a churny property (many persons
+# matched once, few match now) can make the exact rerun carry a candidate set far larger than
+# the current match count suggests.
+MIN_SAMPLED_MATCHES = 10_000
+
+
+def bounded_memory_settings() -> HogQLGlobalSettings:
+    """
+    Execution settings that keep person-dedup aggregations memory-bounded on large teams.
+
+    The person table sorts by (team_id, id), so with in-order aggregation the dedup
+    GROUP BY id streams instead of holding every person in a hash table. The spill
+    threshold is the backstop for aggregations where in-order cannot apply (for
+    example cohort subqueries): degrade to disk instead of a memory-limit error.
+    """
+    return HogQLGlobalSettings(
+        # A partial result is worse than an error here: the counts get extrapolated and the
+        # person pages get sent to, so a timed-out query must fail rather than come up short.
+        # Left unset it inherits the cluster profile, which may accept partial results.
+        timeout_overflow_mode="throw",
+        optimize_aggregation_in_order=True,
+        max_bytes_before_external_group_by=4 * 1024**3,
+    )
+
+
+def count_settings(sample_modulus: Optional[int]) -> HogQLGlobalSettings:
+    # A sampled count keeps the fast parallel hash aggregation: the sample already bounds
+    # the hash table, and in-order aggregation only makes it slower. The exact runs need
+    # the streaming mode to stay memory-bounded on large teams. Both throw on timeout,
+    # because the sampled count multiplies whatever it read by the modulus.
+    if sample_modulus is not None:
+        return HogQLGlobalSettings(
+            timeout_overflow_mode="throw",
+            max_bytes_before_external_group_by=4 * 1024**3,
+        )
+    return bounded_memory_settings()
+
+
+def sampled_or_exact_count(run_count: Callable[[Optional[int]], int]) -> int:
+    """
+    Count from a 1-in-SAMPLE_MODULUS sample and extrapolate, or exactly when the sample holds
+    too few matches for a stable extrapolation.
+
+    `run_count` takes the sample modulus, or None for an exact run.
+    """
+    sampled = run_count(SAMPLE_MODULUS)
+    if sampled >= MIN_SAMPLED_MATCHES:
+        return sampled * SAMPLE_MODULUS
+    return run_count(None)
+
+
+def count_matching_persons(team: Team, filter: Optional[Filter], database: Database, query_type: str) -> int:
+    """Count the persons a filter matches, or every person on the team when filter is None."""
+    return sampled_or_exact_count(
+        lambda sample_modulus: _run_person_count(team, filter, database, query_type, sample_modulus)
+    )
+
+
+def _run_person_count(
+    team: Team, filter: Optional[Filter], database: Database, query_type: str, sample_modulus: Optional[int]
+) -> int:
+    query = build_person_count_query(team, filter, sample_modulus=sample_modulus)
+    response = execute_hogql_query(
+        query=query,
+        team=team,
+        query_type=query_type,
+        # Pin v1 for the unfiltered count, as count_persons_seen_so_far does: the v2 persons
+        # path pushes the executor's default LIMIT into its dedup subquery, which caps a bare
+        # count() at ~101 for teams pinned to v2 (see #87323). The team_id condition already
+        # sends these queries down the persons table's prefilter path, which emits the v1 dedup
+        # whatever this says, so the pin is only load-bearing if that path is ever lost. A
+        # filtered count reads person properties, where v2 is the faster shape, so it keeps the
+        # automatic choice.
+        modifiers=HogQLQueryModifiers(personsArgMaxVersion=PersonsArgMaxVersion.V1) if filter is None else None,
+        context=HogQLContext(team_id=team.pk, database=database),
+        settings=count_settings(sample_modulus),
+    )
+    return response.results[0][0] if response.results else 0
+
+
+def build_person_count_query(team: Team, filter: Optional[Filter], sample_modulus: Optional[int]) -> ast.SelectQuery:
+    where_exprs: list[ast.Expr] = [
+        ast.CompareOperation(
+            op=ast.CompareOperationOp.Eq,
+            left=ast.Field(chain=["persons", "team_id"]),
+            right=ast.Constant(value=team.pk),
+        )
+    ]
+    if sample_modulus is not None:
+        where_exprs.append(person_sample_predicate(sample_modulus))
+    if filter is not None:
+        where_exprs.append(property_to_expr(filter.property_groups, team, scope="person"))
+
+    # A filter can add a one-to-many join: a `distinct_id` person property resolves through
+    # persons.pdi, which gives a person one row per distinct id. So a filtered count dedups on
+    # the person id. The unfiltered total joins nothing, so it keeps the plain count() and
+    # avoids a uniqExact state over every person on the team.
+    if filter is None:
+        count_expr: ast.Expr = ast.Call(name="count", args=[])
+    else:
+        count_expr = ast.Call(name="count", distinct=True, args=[ast.Field(chain=["persons", "id"])])
+
+    return ast.SelectQuery(
+        select=[count_expr],
+        select_from=ast.JoinExpr(table=ast.Field(chain=["persons"])),
+        where=ast.And(exprs=where_exprs),
+    )
+
+
+def person_sample_predicate(modulus: int) -> ast.Expr:
+    # Built as modulo(...) calls, not the % operator: WhereClauseExtractor fails safe to
+    # "no prefilter" on ArithmeticOperation nodes, and the memory bound depends on this
+    # predicate reaching the raw person scan inside the persons lazy table.
+    return sample_predicate(ast.Field(chain=["persons", "id"]), modulus)
+
+
+def sample_predicate(key: ast.Expr, modulus: int) -> ast.Expr:
+    """Keep 1 in `modulus` of the values `key` takes, spread evenly by a hash of the key."""
+    return ast.CompareOperation(
+        op=ast.CompareOperationOp.Eq,
+        left=ast.Call(
+            name="modulo",
+            args=[
+                ast.Call(name="cityHash64", args=[key]),
+                ast.Constant(value=modulus),
+            ],
+        ),
+        right=ast.Constant(value=0),
+    )

products/feature_flags/backend/test/test_blast_radius_v2.py

@@ -0,0 +1,112 @@
+from posthog.test.base import APIBaseTest, ClickhouseTestMixin, _create_person, flush_persons_and_events
+from unittest.mock import patch
+
+from django.test import SimpleTestCase
+
+from parameterized import parameterized
+
+from posthog.models.team.team import Team
+
+from products.feature_flags.backend.user_blast_radius import (
+    BlastRadiusResult,
+    get_person_blast_radius_v2,
+    get_user_blast_radius,
+    use_blast_radius_query_v2,
+)
+
+FILTERS = {"properties": [{"key": "subscribed", "type": "person", "value": ["true"], "operator": "exact"}]}
+
+GATE = "products.feature_flags.backend.api.feature_flag.use_blast_radius_query_v2"
+
+
+class TestBlastRadiusQueryV2(ClickhouseTestMixin, APIBaseTest):
+    def _create_persons(self, subscribed_flags: list[bool]) -> None:
+        for i, subscribed in enumerate(subscribed_flags, start=1):
+            _create_person(
+                team=self.team,
+                distinct_ids=[f"user-{i}"],
+                properties={"subscribed": "true" if subscribed else "false"},
+            )
+        flush_persons_and_events()
+
+    @parameterized.expand(
+        [
+            ("matching a person property", FILTERS, 3, 4),
+            ("with no conditions", {}, 4, 4),
+        ]
+    )
+    def test_sized_condition_matches_the_v1_counts(self, _name, filters, affected, total):
+        # A handful of persons never clears MIN_SAMPLED_MATCHES under the default sample, so
+        # this exercises the exact fallback and pins both counts to the v1 result.
+        self._create_persons([True, True, True, False])
+
+        v1_result = get_user_blast_radius(self.team, filters)
+        result = get_person_blast_radius_v2(self.team, filters)
+
+        assert (result.affected, result.total) == (affected, total)
+        assert (result.affected, result.total) == (v1_result.affected, v1_result.total)
+
+    def test_sampled_path_extrapolates_by_modulus(self):
+        self._create_persons([True, True, False])
+
+        # Modulus 1 samples everyone and MIN_SAMPLED_MATCHES 0 forces the sampled branch,
+        # so the extrapolated result must equal the exact count.
+        with (
+            patch("products.feature_flags.backend.person_sampling.SAMPLE_MODULUS", 1),
+            patch("products.feature_flags.backend.person_sampling.MIN_SAMPLED_MATCHES", 0),
+        ):
+            result = get_person_blast_radius_v2(self.team, FILTERS)
+
+        assert (result.affected, result.total) == (2, 3)
+
+    @parameterized.expand(
+        [
+            ("person condition, gate on", True, None, True),
+            ("person condition, gate off", False, None, False),
+            # The sampled count only covers person audiences, so a group condition stays on v1.
+            ("group condition, gate on", True, 0, False),
+        ]
+    )
+    def test_endpoint_uses_the_sampled_count_only_for_a_gated_person_condition(
+        self, _name, gate_on, group_type_index, expects_v2
+    ):
+        body: dict = {"condition": {"properties": []}}
+        if group_type_index is not None:
+            body["group_type_index"] = group_type_index
+
+        with (
+            patch(GATE, return_value=gate_on),
+            patch("products.feature_flags.backend.api.feature_flag.get_person_blast_radius_v2") as sampled,
+            patch("products.feature_flags.backend.api.feature_flag.get_user_blast_radius") as v1,
+        ):
+            sampled.return_value = BlastRadiusResult(affected=6400, total=64000)
+            v1.return_value = BlastRadiusResult(affected=4, total=10)
+
+            response = self.client.post(
+                f"/api/projects/{self.team.id}/feature_flags/user_blast_radius",
+                body,
+            )
+
+        assert response.status_code == 200, response.json()
+        assert response.json()["affected"] == (6400 if expects_v2 else 4)
+        assert sampled.called is expects_v2
+
+
+class TestBlastRadiusGate(SimpleTestCase):
+    def test_the_gate_evaluates_locally_and_captures_nothing(self):
+        with patch(
+            "products.feature_flags.backend.user_blast_radius.feature_enabled_or_false", return_value=True
+        ) as feature_enabled:
+            assert use_blast_radius_query_v2(Team(pk=42)) is True
+
+        kwargs = feature_enabled.call_args.kwargs
+        # A remote evaluation would put a flags fetch in front of the sizing request the
+        # release-conditions editor sends per condition group, and an event capture would
+        # bill each of those requests as product usage.
+        assert kwargs["only_evaluate_locally"] is True
+        assert kwargs["send_feature_flag_events"] is False
+        # The flag targets the project group, and local evaluation can only resolve that
+        # condition from the ids carried on the call. Without them the gate reads False
+        # for every team, whatever the rollout says.
+        assert kwargs["groups"] == {"project": "42"}
+        assert kwargs["group_properties"] == {"project": {"id": "42"}}

products/feature_flags/backend/test/test_person_sampling.py

@@ -0,0 +1,64 @@
+from collections.abc import Callable
+from typing import Optional
+
+from django.test import SimpleTestCase
+
+from parameterized import parameterized
+
+from products.feature_flags.backend.person_sampling import (
+    MIN_SAMPLED_MATCHES,
+    SAMPLE_MODULUS,
+    bounded_memory_settings,
+    count_settings,
+    sampled_or_exact_count,
+)
+
+
+class TestPersonSamplingSettings(SimpleTestCase):
+    @parameterized.expand(
+        [
+            ("memory bounded", bounded_memory_settings),
+            ("sampled count", lambda: count_settings(64)),
+            ("exact count", lambda: count_settings(None)),
+        ]
+    )
+    def test_settings_reject_a_partial_result_on_timeout(self, _name, build_settings):
+        # Left unset, the mode inherits the cluster profile, which may return whatever the
+        # query aggregated before the timeout. The sampled count would then multiply that
+        # partial tally by the modulus and report it as an audience size.
+        assert build_settings().timeout_overflow_mode == "throw"
+
+
+class TestSampledOrExactCount(SimpleTestCase):
+    # Distinct from every extrapolation below, so the returned value names the branch that ran.
+    EXACT_COUNT = 7
+
+    def _stub_count(self, sampled: int) -> tuple[Callable[[Optional[int]], int], list[Optional[int]]]:
+        moduli: list[Optional[int]] = []
+
+        def run_count(sample_modulus: Optional[int]) -> int:
+            moduli.append(sample_modulus)
+            return sampled if sample_modulus is not None else self.EXACT_COUNT
+
+        return run_count, moduli
+
+    @parameterized.expand(
+        [
+            ("above the threshold", MIN_SAMPLED_MATCHES + 1),
+            # The boundary itself: `>` in place of `>=` reruns exact here and fails.
+            ("at the threshold", MIN_SAMPLED_MATCHES),
+        ]
+    )
+    def test_a_large_enough_sample_is_extrapolated_without_an_exact_rerun(self, _name, sampled):
+        run_count, moduli = self._stub_count(sampled)
+
+        # Returning the sample unmultiplied would under-report the audience by the modulus.
+        assert sampled_or_exact_count(run_count) == sampled * SAMPLE_MODULUS
+        assert moduli == [SAMPLE_MODULUS]
+
+    def test_a_small_sample_is_discarded_for_an_exact_count(self):
+        run_count, moduli = self._stub_count(MIN_SAMPLED_MATCHES - 1)
+
+        assert sampled_or_exact_count(run_count) == self.EXACT_COUNT
+        # The probe still runs and is paid for before the exact count replaces it.
+        assert moduli == [SAMPLE_MODULUS, None]