analytics-toolkit 1.3.7.3__tar.gz → 1.3.8.2__tar.gz

{analytics_toolkit-1.3.7.3 → analytics_toolkit-1.3.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: analytics-toolkit
-Version: 1.3.7.3
+Version: 1.3.8.2
 Summary: Shared Python utilities for SQL, Excel, and date helpers.
 Author: analytics_toolkit contributors
 License: MIT License
@@ -75,7 +75,7 @@ Python toolkit for AB-test analysis, SQL workflows, Excel reports, and date help
 
 [Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
 
-**Version:** `1.3.7.3`<br>
+**Version:** `1.3.8.2`<br>
 **Depends:** Python (`>=3.8,<3.15`)<br>
 **Imports:** [clickhouse-connect](https://pypi.org/project/clickhouse-connect/) (`>=0.5.14,<1`), [lz4](https://pypi.org/project/lz4/) (`>=4.3.2,<5`), [numpy](https://pypi.org/project/numpy/) (`>=1.24.2,<2`), [openpyxl](https://pypi.org/project/openpyxl/) (`>=3.1.1,<4`), [orjson](https://pypi.org/project/orjson/) (`>=3.8.7,<4`), [pandas](https://pypi.org/project/pandas/) (`>=1.4.4,<3`), [psycopg2-binary](https://pypi.org/project/psycopg2-binary/) (`>=2.9.5,<3`), [python-dateutil](https://pypi.org/project/python-dateutil/) (`>=2.8.2,<3`), [pytz](https://pypi.org/project/pytz/) (`>=2022.7`), [requests](https://pypi.org/project/requests/) (`>=2.28.2,<3`), [scipy](https://pypi.org/project/scipy/) (`>=1.10.1,<2`), [sqlglot](https://pypi.org/project/sqlglot/) (`>=20,<31`), [sqlparse](https://pypi.org/project/sqlparse/) (`>=0.4.3,<1`), [tqdm](https://pypi.org/project/tqdm/) (`>=4.65.0,<5`), [trino](https://pypi.org/project/trino/) (`>=0.320,<1`), [zstandard](https://pypi.org/project/zstandard/) (`>=0.20.0,<1`)<br>
 **Suggests:** [apache-airflow](https://pypi.org/project/apache-airflow/) (`>=2.4,<3`; optional extra `airflow`)<br>
@@ -107,13 +107,111 @@ pip install git+https://github.com/Karapsin/analytics_toolkit.git
 - `analytics_toolkit.dates`: date and period helpers.
 - `analytics_toolkit.general`: shared logging and file path helpers.
 
-## SQL Import Policy
+## SQL Workflows
 
-Supported SQL imports are `from analytics_toolkit import sql` or
-`import analytics_toolkit.sql as sql`. Deep imports under
-`analytics_toolkit.sql.*` are internal implementation details and may change;
-call SQL helpers through the `sql` facade, for example `sql.create_sql_table(...)`
-or `sql.transfer(...)`.
+[All SQL functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/sql/functions/index.md)
+
+[SQL module guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/sql/index.md)
+
+`sql.transfer` streams query results between configured SQL backends, with
+batching, retries, and table creation or replacement handled by one call.
+
+```python
+from analytics_toolkit import sql
+
+rows = sql.transfer(
+    from_db="trino",
+    to_db="gp",
+    from_sql="select user_id, order_id, amount from iceberg.analytics.orders",
+    to_table="sandbox.orders_copy",
+    write_mode="replace",
+    batch_size=50_000,
+    progress=True,
+)
+```
+
+Aliases can point to the same backend type, so Greenplum-to-Greenplum transfers
+work the same way.
+
+```python
+rows = sql.transfer(
+    from_db="gp_sales",
+    to_db="gp_finance",
+    from_sql="select user_id, order_id, amount from mart.sales_orders",
+    to_table="finance.sales_orders_copy",
+    write_mode="replace",
+    batch_size=50_000,
+)
+```
+
+- `sql.read`: run a query and return a dataframe.
+- `sql.execute`: run DDL or DML without returning a dataframe.
+- `sql.execute_read`: run setup SQL and return the final result as a dataframe.
+- `sql.load_df`: load a pandas dataframe into a configured backend table.
+- `sql.transfer`: move rows from a source query to a target table across backends.
+
+## AB Metrics
+
+[All AB functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/ab_utils/functions/index.md)
+
+[AB utilities guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/ab_utils/index.md)
+
+`compute_test_metrics` compares experiment groups across mean and ratio metrics,
+with optional CUPED statistics and bootstrap multiple-comparison adjustment.
+
+```python
+from analytics_toolkit.ab_utils import compute_test_metrics
+
+result = compute_test_metrics(
+    experiment_df,
+    group="group_name",
+    control="control",
+    user_id="user_id",
+    ratio_metrics=[
+        {"name": "ctr", "numerator": "clicks", "denominator": "views"},
+    ],
+    pre_exp_metrics_df=pre_experiment_df,
+    multiple_comparisons_adjustment=True,
+    multiple_comparisons_adjustment_resamples=1000,
+)
+```
+
+Example output with CUPED and bootstrap columns enabled:
+
+| metric_type | group_1 | group_2 | metric_name | n0 | n1 | outliers_cutoff | outliers_n_control | outliers_n_test | metric_control | metric_test | variance_control | variance_test | delta_abs | delta_relative | mde_abs | mde_relative | s.e. | p-value | s.e. CUPED | p-value CUPED | mde_abs CUPED | mde_relative CUPED | s.e. bootstrap | bootstrap_adj_p |
+| --- | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| mean | test | control | revenue | 10000 | 10050 | 250.0 | 3 | 4 | 12.40 | 13.10 | 45.20 | 47.80 | 0.70 | 0.056 | 0.42 | 0.034 | 0.15 | 0.003 | 0.11 | 0.001 | 0.31 | 0.025 | 0.14 | 0.012 |
+| ratio | test | control | ctr | 10000 | 10050 | 1.0 | 0 | 0 | 0.082 | 0.087 | 0.0009 | 0.0010 | 0.005 | 0.061 | 0.003 | 0.037 | 0.001 | 0.008 | 0.001 | 0.006 | 0.002 | 0.024 | 0.001 | 0.019 |
+
+## Date Helpers
+
+[All date functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/dates/functions/index.md)
+
+[Date helpers guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/dates/index.md)
+
+Date helpers cover reporting ranges, period boundaries, offsets, and stable
+string formatting for SQL and filenames.
+
+```python
+from analytics_toolkit.dates import add_days, first_day, gen_dates_list, last_day
+
+report_days = gen_dates_list("2026-06-01", "2026-06-07")
+# ["2026-06-01", "2026-06-02", "2026-06-03", "2026-06-04", "2026-06-05", "2026-06-06", "2026-06-07"]
+
+month_start = first_day("2026-06-08", "month")
+# "2026-06-01"
+
+month_end = last_day("2026-06-08", "month")
+# "2026-06-30"
+
+next_run = add_days("2026-06-08", 1)
+# "2026-06-09"
+```
+
+- `gen_dates_list`: build daily, weekly, monthly, or quarterly sequences.
+- `first_day` / `last_day`: get week, month, or quarter boundaries.
+- `add_days`, `add_weeks`, `add_months`, `add_quarters`: shift dates.
+- `sanitize_date`: convert a date to compact `YYYYMMDD` text.
 
 ## Documentation
 

analytics_toolkit-1.3.8.2/README.md

@@ -0,0 +1,153 @@
+# analytics_toolkit
+
+Python toolkit for AB-test analysis, SQL workflows, Excel reports, and date helpers.
+
+[Quick Start Guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/QUICK_START.md)
+
+[Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
+
+**Version:** `1.3.8.2`<br>
+**Depends:** Python (`>=3.8,<3.15`)<br>
+**Imports:** [clickhouse-connect](https://pypi.org/project/clickhouse-connect/) (`>=0.5.14,<1`), [lz4](https://pypi.org/project/lz4/) (`>=4.3.2,<5`), [numpy](https://pypi.org/project/numpy/) (`>=1.24.2,<2`), [openpyxl](https://pypi.org/project/openpyxl/) (`>=3.1.1,<4`), [orjson](https://pypi.org/project/orjson/) (`>=3.8.7,<4`), [pandas](https://pypi.org/project/pandas/) (`>=1.4.4,<3`), [psycopg2-binary](https://pypi.org/project/psycopg2-binary/) (`>=2.9.5,<3`), [python-dateutil](https://pypi.org/project/python-dateutil/) (`>=2.8.2,<3`), [pytz](https://pypi.org/project/pytz/) (`>=2022.7`), [requests](https://pypi.org/project/requests/) (`>=2.28.2,<3`), [scipy](https://pypi.org/project/scipy/) (`>=1.10.1,<2`), [sqlglot](https://pypi.org/project/sqlglot/) (`>=20,<31`), [sqlparse](https://pypi.org/project/sqlparse/) (`>=0.4.3,<1`), [tqdm](https://pypi.org/project/tqdm/) (`>=4.65.0,<5`), [trino](https://pypi.org/project/trino/) (`>=0.320,<1`), [zstandard](https://pypi.org/project/zstandard/) (`>=0.20.0,<1`)<br>
+**Suggests:** [apache-airflow](https://pypi.org/project/apache-airflow/) (`>=2.4,<3`; optional extra `airflow`)<br>
+**Install:** `pip install analytics-toolkit`<br>
+**PyPI:** [pypi.org/project/analytics-toolkit](https://pypi.org/project/analytics-toolkit/)<br>
+**License:** MIT<br>
+**Source:** [github.com/Karapsin/analytics_toolkit](https://github.com/Karapsin/analytics_toolkit)<br>
+**Issues:** [GitHub Issues](https://github.com/Karapsin/analytics_toolkit/issues)
+
+## Installation
+
+From PyPI:
+
+```bash
+pip install analytics-toolkit
+```
+
+From GitHub:
+
+```bash
+pip install git+https://github.com/Karapsin/analytics_toolkit.git
+```
+
+## Areas
+
+- `analytics_toolkit.ab_utils`: AB-test metric comparison helpers.
+- `analytics_toolkit.sql`: SQL read, execute, load, and transfer helpers.
+- `analytics_toolkit.excel`: Excel report helpers for long-format dataframes.
+- `analytics_toolkit.dates`: date and period helpers.
+- `analytics_toolkit.general`: shared logging and file path helpers.
+
+## SQL Workflows
+
+[All SQL functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/sql/functions/index.md)
+
+[SQL module guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/sql/index.md)
+
+`sql.transfer` streams query results between configured SQL backends, with
+batching, retries, and table creation or replacement handled by one call.
+
+```python
+from analytics_toolkit import sql
+
+rows = sql.transfer(
+    from_db="trino",
+    to_db="gp",
+    from_sql="select user_id, order_id, amount from iceberg.analytics.orders",
+    to_table="sandbox.orders_copy",
+    write_mode="replace",
+    batch_size=50_000,
+    progress=True,
+)
+```
+
+Aliases can point to the same backend type, so Greenplum-to-Greenplum transfers
+work the same way.
+
+```python
+rows = sql.transfer(
+    from_db="gp_sales",
+    to_db="gp_finance",
+    from_sql="select user_id, order_id, amount from mart.sales_orders",
+    to_table="finance.sales_orders_copy",
+    write_mode="replace",
+    batch_size=50_000,
+)
+```
+
+- `sql.read`: run a query and return a dataframe.
+- `sql.execute`: run DDL or DML without returning a dataframe.
+- `sql.execute_read`: run setup SQL and return the final result as a dataframe.
+- `sql.load_df`: load a pandas dataframe into a configured backend table.
+- `sql.transfer`: move rows from a source query to a target table across backends.
+
+## AB Metrics
+
+[All AB functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/ab_utils/functions/index.md)
+
+[AB utilities guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/ab_utils/index.md)
+
+`compute_test_metrics` compares experiment groups across mean and ratio metrics,
+with optional CUPED statistics and bootstrap multiple-comparison adjustment.
+
+```python
+from analytics_toolkit.ab_utils import compute_test_metrics
+
+result = compute_test_metrics(
+    experiment_df,
+    group="group_name",
+    control="control",
+    user_id="user_id",
+    ratio_metrics=[
+        {"name": "ctr", "numerator": "clicks", "denominator": "views"},
+    ],
+    pre_exp_metrics_df=pre_experiment_df,
+    multiple_comparisons_adjustment=True,
+    multiple_comparisons_adjustment_resamples=1000,
+)
+```
+
+Example output with CUPED and bootstrap columns enabled:
+
+| metric_type | group_1 | group_2 | metric_name | n0 | n1 | outliers_cutoff | outliers_n_control | outliers_n_test | metric_control | metric_test | variance_control | variance_test | delta_abs | delta_relative | mde_abs | mde_relative | s.e. | p-value | s.e. CUPED | p-value CUPED | mde_abs CUPED | mde_relative CUPED | s.e. bootstrap | bootstrap_adj_p |
+| --- | --- | --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |
+| mean | test | control | revenue | 10000 | 10050 | 250.0 | 3 | 4 | 12.40 | 13.10 | 45.20 | 47.80 | 0.70 | 0.056 | 0.42 | 0.034 | 0.15 | 0.003 | 0.11 | 0.001 | 0.31 | 0.025 | 0.14 | 0.012 |
+| ratio | test | control | ctr | 10000 | 10050 | 1.0 | 0 | 0 | 0.082 | 0.087 | 0.0009 | 0.0010 | 0.005 | 0.061 | 0.003 | 0.037 | 0.001 | 0.008 | 0.001 | 0.006 | 0.002 | 0.024 | 0.001 | 0.019 |
+
+## Date Helpers
+
+[All date functions](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/dates/functions/index.md)
+
+[Date helpers guide](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/dates/index.md)
+
+Date helpers cover reporting ranges, period boundaries, offsets, and stable
+string formatting for SQL and filenames.
+
+```python
+from analytics_toolkit.dates import add_days, first_day, gen_dates_list, last_day
+
+report_days = gen_dates_list("2026-06-01", "2026-06-07")
+# ["2026-06-01", "2026-06-02", "2026-06-03", "2026-06-04", "2026-06-05", "2026-06-06", "2026-06-07"]
+
+month_start = first_day("2026-06-08", "month")
+# "2026-06-01"
+
+month_end = last_day("2026-06-08", "month")
+# "2026-06-30"
+
+next_run = add_days("2026-06-08", 1)
+# "2026-06-09"
+```
+
+- `gen_dates_list`: build daily, weekly, monthly, or quarterly sequences.
+- `first_day` / `last_day`: get week, month, or quarter boundaries.
+- `add_days`, `add_weeks`, `add_months`, `add_quarters`: shift dates.
+- `sanitize_date`: convert a date to compact `YYYYMMDD` text.
+
+## Documentation
+
+- [Documentation Overview](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/README.md)
+- [Quick Start](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/QUICK_START.md)
+- [Airflow SQL Manual](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/AIRFLOW_SQL_MANUAL.md)
+- [Module documentation](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/modules/README.md)
+- [Changelog](https://github.com/Karapsin/analytics_toolkit/blob/main/docs/CHANGELOG.md)

{analytics_toolkit-1.3.7.3 → analytics_toolkit-1.3.8.2}/analytics_toolkit/ab_utils/api.py

@@ -6,10 +6,15 @@ from analytics_toolkit.general import time_print
 
 from .bootstrap import _apply_multiple_comparisons_adjustment
 from .constants import DEFAULT_ALPHA, DEFAULT_POWER
-from .cuped import _compute_cuped_statistics
+from .cuped import _compute_cuped_statistics, _prepare_cuped_context
 from .outliers import _build_outlier_context
 from .ratio import _normalize_ratio_metrics
-from .rows import _build_comparisons, _build_metric_definitions, _build_metric_row
+from .rows import (
+    _build_comparisons,
+    _build_metric_definitions,
+    _build_metric_row,
+    _prepare_metric_context,
+)
 from .stats import _compute_mde_from_standard_error, _safe_relative
 from .validation import (
     _validate_input_columns,
@@ -105,6 +110,15 @@ def compute_test_metrics(
     ratio_specs = _normalize_ratio_metrics(df, ratio_metrics, reserved_columns={group, user_id})
     comparisons = _build_comparisons(group_names, control, test_vs_test=test_vs_test)
     metric_definitions = _build_metric_definitions(metric_columns, ratio_specs)
+    group_values = df[group].to_numpy()
+    group_masks = {group_name: group_values == group_name for group_name in group_names}
+    comparison_frames = {
+        (test_group, baseline_group): df.loc[
+            group_masks[test_group] | group_masks[baseline_group],
+            [user_id, group],
+        ].copy()
+        for test_group, baseline_group in comparisons
+    }
     time_print(
         "compute_test_metrics: setup complete "
         f"groups={len(group_names)} comparisons={len(comparisons)} "
@@ -129,6 +143,26 @@ def compute_test_metrics(
             )
     time_print("compute_test_metrics: outlier contexts complete")
 
+    prepared_metric_contexts: dict[str, dict[str, object]] = {}
+    prepared_cuped_contexts: dict[str, dict[str, object]] = {}
+    for metric_definition in metric_definitions:
+        metric_key = str(metric_definition["metric_key"])
+        outlier_context = metric_definition.get("_outlier_context")
+        prepared_metric_contexts[metric_key] = _prepare_metric_context(
+            df=df,
+            metric_definition=metric_definition,
+            outlier_context=outlier_context,
+        )
+        if pre_exp_metrics_df is not None:
+            prepared_cuped_contexts[metric_key] = _prepare_cuped_context(
+                df=df,
+                pre_exp_metrics_df=pre_exp_metrics_df,
+                user_id_column=user_id,
+                metric_definition=metric_definition,
+                outlier_context=outlier_context,
+                pre_outlier_context=metric_definition.get("_pre_outlier_context"),
+            )
+
     rows: list[dict[str, object]] = []
     for test_group, baseline_group in comparisons:
         time_print(f"compute_test_metrics: comparison {test_group} vs {baseline_group}")
@@ -146,6 +180,8 @@ def compute_test_metrics(
                 mde_alpha=mde_alpha,
                 mde_power=mde_power,
                 outlier_context=outlier_context,
+                prepared_metric_context=prepared_metric_contexts.get(metric_name),
+                group_masks=group_masks,
             )
             row["_comparison_key"] = (test_group, baseline_group)
             if pre_exp_metrics_df is not None:
@@ -160,6 +196,10 @@ def compute_test_metrics(
                     metric_definition=metric_definition,
                     outlier_context=outlier_context,
                     pre_outlier_context=metric_definition.get("_pre_outlier_context"),
+                    prepared_cuped_context={
+                        **prepared_cuped_contexts[metric_name],
+                        "comparison_frame": comparison_frames[(test_group, baseline_group)],
+                    },
                 )
                 row["mde_abs CUPED"] = _compute_mde_from_standard_error(
                     standard_error=row["s.e. CUPED"],

{analytics_toolkit-1.3.7.3 → analytics_toolkit-1.3.8.2}/analytics_toolkit/ab_utils/cuped.py

@@ -41,29 +41,68 @@ def _compute_cuped_p_value(
     return p_value
 
 
-def _compute_cuped_statistics(
+def _prepare_cuped_context(
     df: pd.DataFrame,
     pre_exp_metrics_df: pd.DataFrame,
-    group_column: str,
     user_id_column: str,
-    baseline_group: str,
-    test_group: str,
     metric_definition: dict[str, object],
     outlier_context: dict[str, object] | None = None,
     pre_outlier_context: dict[str, object] | None = None,
-) -> tuple[float, float]:
-    cuped_frame, reason = _build_cuped_frame(
+) -> dict[str, object]:
+    exp_values, exp_error = _build_metric_values_by_user(
         df=df,
-        pre_exp_metrics_df=pre_exp_metrics_df,
         user_id_column=user_id_column,
-        group_column=group_column,
-        baseline_group=baseline_group,
-        test_group=test_group,
         metric_definition=metric_definition,
+        value_column="metric_exp",
         outlier_context=outlier_context,
-        pre_outlier_context=pre_outlier_context,
     )
+    pre_values, pre_error = _build_metric_values_by_user(
+        df=pre_exp_metrics_df,
+        user_id_column=user_id_column,
+        metric_definition=metric_definition,
+        value_column="metric_pre",
+        outlier_context=pre_outlier_context,
+    )
+    return {
+        "exp_values": exp_values.set_index(user_id_column) if exp_error is None else None,
+        "pre_values": pre_values.set_index(user_id_column) if pre_error is None else None,
+        "exp_error": exp_error,
+        "pre_error": pre_error,
+    }
+
+
+def _compute_cuped_statistics(
+    df: pd.DataFrame,
+    pre_exp_metrics_df: pd.DataFrame,
+    group_column: str,
+    user_id_column: str,
+    baseline_group: str,
+    test_group: str,
+    metric_definition: dict[str, object],
+    outlier_context: dict[str, object] | None = None,
+    pre_outlier_context: dict[str, object] | None = None,
+    prepared_cuped_context: dict[str, object] | None = None,
+) -> tuple[float, float]:
     metric_name = str(metric_definition["metric_key"])
+    if prepared_cuped_context is None:
+        cuped_frame, reason = _build_cuped_frame(
+            df=df,
+            pre_exp_metrics_df=pre_exp_metrics_df,
+            user_id_column=user_id_column,
+            group_column=group_column,
+            baseline_group=baseline_group,
+            test_group=test_group,
+            metric_definition=metric_definition,
+            outlier_context=outlier_context,
+            pre_outlier_context=pre_outlier_context,
+        )
+    else:
+        cuped_frame, reason = _build_cuped_frame_from_prepared_context(
+            comparison_frame=prepared_cuped_context["comparison_frame"],
+            user_id_column=user_id_column,
+            prepared_cuped_context=prepared_cuped_context,
+        )
+
     if reason is not None:
         warnings.warn(
             (
@@ -93,6 +132,26 @@ def _compute_cuped_statistics(
     return p_value, standard_error
 
 
+def _build_cuped_frame_from_prepared_context(
+    comparison_frame: pd.DataFrame,
+    user_id_column: str,
+    prepared_cuped_context: dict[str, object],
+) -> tuple[pd.DataFrame | None, str | None]:
+    exp_values = prepared_cuped_context.get("exp_values")
+    if exp_values is None:
+        return None, f"experiment metric values are unavailable: {prepared_cuped_context.get('exp_error')}"
+    pre_values = prepared_cuped_context.get("pre_values")
+    if pre_values is None:
+        return None, f"pre-experiment metric values are unavailable: {prepared_cuped_context.get('pre_error')}"
+
+    cuped_frame = comparison_frame.copy()
+    cuped_frame = cuped_frame.join(exp_values, on=user_id_column).join(pre_values, on=user_id_column)
+    cuped_frame = cuped_frame.dropna(subset=["metric_exp", "metric_pre"]).reset_index(drop=True)
+    if cuped_frame.empty:
+        return None, "no overlapping non-missing experiment/pre-experiment observations"
+    return cuped_frame, None
+
+
 def _build_cuped_frame(
     df: pd.DataFrame,
     pre_exp_metrics_df: pd.DataFrame,

{analytics_toolkit-1.3.7.3 → analytics_toolkit-1.3.8.2}/analytics_toolkit/ab_utils/metrics.py

@@ -29,6 +29,7 @@ from .cuped import (
     _compute_cuped_p_value_from_frame,
     _compute_cuped_statistics,
     _compute_cuped_statistics_from_frame,
+    _prepare_cuped_context,
 )
 from .outliers import (
     _apply_outliers_to_agg_ratio_components,
@@ -59,6 +60,7 @@ from .rows import (
     _build_metric_definitions,
     _build_metric_row,
     _build_ratio_metric_row,
+    _prepare_metric_context,
 )
 from .stats import (
     _both_present,
@@ -110,6 +112,8 @@ __all__ = [
     "_build_ratio_valid_mask",
     "_build_ratio_valid_mask_from_arrays",
     "_build_value_outlier_mask",
+    "_prepare_cuped_context",
+    "_prepare_metric_context",
     "_compute_agg_ratio_diff_standard_error",
     "_compute_agg_ratio_group_stats",
     "_compute_agg_ratio_group_stats_arrays",

{analytics_toolkit-1.3.7.3 → analytics_toolkit-1.3.8.2}/analytics_toolkit/ab_utils/parallel.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextvars
+import inspect
 from collections.abc import Mapping
 from contextlib import ExitStack
 from concurrent.futures import Future, ThreadPoolExecutor, as_completed
@@ -20,6 +21,9 @@ from .api import compute_test_metrics
 async_sql = sql_facade.async_sql
 
 _SQL_DATAFRAME_FIELDS = frozenset({"df", "pre_exp_df", "pre_exp_metrics_df"})
+_COMPUTE_TEST_METRICS_FIELDS = frozenset(
+    inspect.signature(compute_test_metrics).parameters
+)
 _DEFAULT_HARD_CONCURRENCY_CAP = 10
 _CONCURRENCY_STATE: contextvars.ContextVar["_ConcurrencyState | None"] = (
     contextvars.ContextVar("analytics_toolkit_ab_parallel_concurrency", default=None)
@@ -127,9 +131,11 @@ def parallel_compute_metrics_from_sql(
     soft_concurrency_cap: int | None = None,
     hard_concurrency_cap: int = _DEFAULT_HARD_CONCURRENCY_CAP,
     progress: bool = False,
+    **metric_defaults: Any,
 ) -> dict[str, pd.DataFrame | str]:
     """Load SQL-backed task dataframes, then run ``parallel_compute_metrics``."""
     _validate_start_comment("start_comment", start_comment)
+    metric_defaults = _validate_metric_defaults(metric_defaults)
     task_defs = _validate_sql_tasks(tasks)
     _validate_concurrency(concurrency)
     _validate_optional_soft_concurrency_cap(soft_concurrency_cap)
@@ -191,7 +197,8 @@ def parallel_compute_metrics_from_sql(
             sql_failures[name] = df_result
             continue
 
-        metric_kwargs = dict(kwargs)
+        metric_kwargs = dict(metric_defaults)
+        metric_kwargs.update(kwargs)
         metric_kwargs["df"] = df_result
         if pre_exp_sql is not None:
             pre_exp_df_result = sql_results[_sql_read_task_name(name, "pre_exp_sql")]
@@ -424,6 +431,38 @@ def _validate_sql_task_spec(
     return name, kwargs, sql, pre_exp_sql, start_comment
 
 
+def _validate_metric_defaults(metric_defaults: Mapping[str, Any]) -> dict[str, Any]:
+    if not metric_defaults:
+        return {}
+
+    dataframe_fields = sorted(_SQL_DATAFRAME_FIELDS.intersection(metric_defaults))
+    if dataframe_fields:
+        fields = ", ".join(dataframe_fields)
+        raise ValueError(
+            "parallel_compute_metrics_from_sql top-level metric defaults cannot "
+            f"define SQL-backed dataframe field(s): {fields}."
+        )
+
+    supported_fields = _COMPUTE_TEST_METRICS_FIELDS - {
+        "df",
+        "pre_exp_metrics_df",
+    }
+    unexpected_fields = sorted(set(metric_defaults) - supported_fields)
+    if unexpected_fields:
+        quoted_fields = ", ".join(repr(field) for field in unexpected_fields)
+        if len(unexpected_fields) == 1:
+            raise TypeError(
+                "parallel_compute_metrics_from_sql() got an unexpected keyword "
+                f"argument {quoted_fields}"
+            )
+        raise TypeError(
+            "parallel_compute_metrics_from_sql() got unexpected keyword "
+            f"arguments: {quoted_fields}"
+        )
+
+    return dict(metric_defaults)
+
+
 def _make_sql_read_task(
     *,
     name: str,
@@ -435,7 +474,7 @@ def _make_sql_read_task(
     task = {
         "name": _sql_read_task_name(name, field),
         "type": "read",
-        "connection_type": db,
+        "db_key": db,
         "query": query,
     }
     if start_comment is not _MISSING: