fix: pull population counts directly from enrollments table

jetstream/analysis.py

@@ -678,6 +678,37 @@ def counts(
             )
             return StatisticResultCollection.model_validate([])
 
+    @dask.delayed
+    def population_subset_table(
+        self,
+        segment: str,
+        analysis_basis: AnalysisBasis,
+    ) -> DataFrame:
+        """Pulls the full enrolled population for this segment/analysis basis directly
+        from the enrollments table.
+
+        Unlike subset_metric_table, this doesn't filter out rows based on any
+        particular metric's null values, so it's safe to use for population counts.
+        """
+        query = self._create_population_subset_query(segment, analysis_basis)
+
+        logger.debug(f"population_subset_table: {segment}, {analysis_basis}\n{query}")
+
+        try:
+            results: DataFrame = self.bigquery.execute(query).to_dataframe()
+        except GoogleAPICallError as e:
+            logger.exception(
+                str(e),
+                extra={
+                    "experiment": self.config.experiment.normandy_slug,
+                    "analysis_basis": analysis_basis,
+                    "segment": segment,
+                },
+            )
+            return None
+
+        return results
+
     @dask.delayed
     def subset_metric_table(
         self,
@@ -805,6 +836,13 @@ def _create_subset_metric_table_query_univariate(
         """
         )
 
+        query += self._basis_and_segment_filter(segment, analysis_basis)
+
+        return query
+
+    def _basis_and_segment_filter(self, segment: str, analysis_basis: AnalysisBasis) -> str:
+        """Builds the analysis-basis and segment WHERE clauses shared by metric-table
+        and population-table subset queries."""
         if analysis_basis == AnalysisBasis.ENROLLMENTS:
             basis_filter = """enrollment_date IS NOT NULL"""
         elif analysis_basis == AnalysisBasis.EXPOSURES:
@@ -815,16 +853,35 @@ def _create_subset_metric_table_query_univariate(
                 + f"Allowed values are: {[AnalysisBasis.ENROLLMENTS, AnalysisBasis.EXPOSURES]}"
             )
 
-        query += basis_filter
+        filter_clause = basis_filter
 
         if segment != "all":
-            segment_filter = dedent(
+            filter_clause += dedent(
                 f"""
             AND m.{segment} = TRUE"""
             )
-            query += segment_filter
 
-        return query
+        return filter_clause
+
+    def _create_population_subset_query(
+        self,
+        segment: str,
+        analysis_basis: AnalysisBasis,
+    ) -> str:
+        """Creates a SQL query string to pull the full enrolled population for a
+        segment/analysis basis directly from the enrollments table, independent of
+        any particular metric's null values."""
+        normalized_slug = bq_normalize_name(self.config.experiment.normandy_slug)
+        enrollments_table_name = f"{self.project}.{self.dataset}.enrollments_{normalized_slug}"
+
+        query = dedent(
+            f"""
+        SELECT branch
+        FROM `{enrollments_table_name}` m
+        WHERE """
+        )
+
+        return query + self._basis_and_segment_filter(segment, analysis_basis)
 
     def _covariate_table_metric_name(
         self,
@@ -1484,7 +1541,9 @@ def run(
                             ).model_dump(warnings=False)
 
                         segment_results.root += self.counts(
-                            segment_data, segment, analysis_basis
+                            self.population_subset_table(segment, analysis_basis),
+                            segment,
+                            analysis_basis,
                         ).model_dump(warnings=False)
 
                     # done with analysis_basis: publish metrics view
@@ -1625,7 +1684,9 @@ def run(
 
                             if segment not in counted_segments:
                                 segment_results.root += self.counts(
-                                    segment_data, segment, analysis_basis
+                                    self.population_subset_table(segment, analysis_basis),
+                                    segment,
+                                    analysis_basis,
                                 ).model_dump(warnings=False)
                                 counted_segments.add(segment)
 

jetstream/tests/test_analysis.py

@@ -778,6 +778,52 @@ def test_create_subset_metric_table_query_univariate_unsupported_analysis_basis(
         )
 
 
+def test_create_population_subset_query_basic(experiments):
+    expected_query = dedent(
+        """
+    SELECT branch
+    FROM `spam.eggs.enrollments_normandy_test_slug` m
+    WHERE enrollment_date IS NOT NULL"""
+    )
+
+    actual_query = _empty_analysis(experiments)._create_population_subset_query(
+        "all", AnalysisBasis.ENROLLMENTS
+    )
+
+    assert expected_query == actual_query
+
+
+def test_create_population_subset_query_segment(experiments):
+    expected_query = dedent(
+        """
+    SELECT branch
+    FROM `spam.eggs.enrollments_normandy_test_slug` m
+    WHERE enrollment_date IS NOT NULL
+    AND m.mysegment = TRUE"""
+    )
+
+    actual_query = _empty_analysis(experiments)._create_population_subset_query(
+        "mysegment", AnalysisBasis.ENROLLMENTS
+    )
+
+    assert expected_query == actual_query
+
+
+def test_create_population_subset_query_exposures(experiments):
+    expected_query = dedent(
+        """
+    SELECT branch
+    FROM `spam.eggs.enrollments_normandy_test_slug` m
+    WHERE enrollment_date IS NOT NULL AND m.exposure_date IS NOT NULL"""
+    )
+
+    actual_query = _empty_analysis(experiments)._create_population_subset_query(
+        "all", AnalysisBasis.EXPOSURES
+    )
+
+    assert expected_query == actual_query
+
+
 def test_create_subset_metric_table_query_covariate_unsupported_analysis_basis(
     experiments, monkeypatch
 ):
@@ -1376,6 +1422,39 @@ def test_subset_metric_table_returns_none_on_google_api_error(experiments, monke
     assert "simulated subset error" in caplog.text
 
 
+def test_population_subset_table_ignores_metric_null_values(experiments, monkeypatch):
+    """population_subset_table's query has no per-metric null filter, unlike
+    subset_metric_table, so its row count reflects the full enrolled population
+    regardless of which metric's data happens to be sparse."""
+    mock_bq = MagicMock()
+    monkeypatch.setattr("jetstream.analysis.BigQueryClient", Mock(return_value=mock_bq))
+
+    _empty_analysis(experiments).population_subset_table("all", AnalysisBasis.ENROLLMENTS).compute(
+        scheduler="synchronous"
+    )
+
+    executed_query = mock_bq.execute.call_args[0][0]
+    assert "enrollments_normandy_test_slug" in executed_query
+    assert "IS NOT NULL AND" not in executed_query
+
+
+def test_population_subset_table_returns_none_on_google_api_error(experiments, monkeypatch, caplog):
+    """population_subset_table returns None (not raises) on GoogleAPICallError."""
+    mock_bq = MagicMock()
+    mock_bq.execute.side_effect = GoogleAPICallError("simulated population error")
+    monkeypatch.setattr("jetstream.analysis.BigQueryClient", Mock(return_value=mock_bq))
+
+    analysis = _empty_analysis(experiments)
+
+    with caplog.at_level(logging.ERROR):
+        result = analysis.population_subset_table("all", AnalysisBasis.ENROLLMENTS).compute(
+            scheduler="synchronous"
+        )
+
+    assert result is None
+    assert "simulated population error" in caplog.text
+
+
 def test_counts_returns_empty_for_none_segment_data(experiments):
     """counts returns an empty StatisticResultCollection when segment_data is None."""
     result = (