diff --git a/jetstream/analysis.py b/jetstream/analysis.py index c8b3af39..8205b68d 100644 --- a/jetstream/analysis.py +++ b/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) diff --git a/jetstream/tests/test_analysis.py b/jetstream/tests/test_analysis.py index 7ab7a409..a13e42a9 100644 --- a/jetstream/tests/test_analysis.py +++ b/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 = (