Validate grouping columns for NaN, fix examples, add missing-key tests

Address remaining P1/P3 from AI review:
- Reject NaN in by columns before groupby (P1)
- Make docstring/RST examples illustrative pseudocode (P3)
- Add tests for partial and all-NaN grouping keys (P3)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: igerber

diff_diff/prep.py

@@ -1466,10 +1466,13 @@ def aggregate_survey(
     ...     microdata, by=["state", "year"],
     ...     outcomes="smoking_rate", survey_design=design,
     ... )
-    >>> result = DifferenceInDifferences().fit(
-    ...     panel, outcome="smoking_rate_mean",
-    ...     treatment="treated", time="year", survey_design=stage2,
-    ... )
+    >>> # Add treatment/time indicators at the panel level, then fit:
+    >>> # panel["treated"] = ...  # e.g., from policy adoption data
+    >>> # panel["post"] = (panel["year"] >= treatment_year).astype(int)
+    >>> # result = DifferenceInDifferences().fit(
+    >>> #     panel, outcome="smoking_rate_mean",
+    >>> #     treatment="treated", time="post", survey_design=stage2,
+    >>> # )
     """
     import warnings
     from dataclasses import replace
@@ -1508,6 +1511,15 @@ def aggregate_survey(
     if data.empty:
         raise ValueError("data must be non-empty")
 
+    # --- Validate grouping columns have no missing values ---
+    by_missing = data[by_cols].isna().any()
+    cols_with_na = list(by_missing[by_missing].index)
+    if cols_with_na:
+        raise ValueError(
+            f"Missing values in grouping column(s): {cols_with_na}. "
+            f"Drop or fill NaN values before calling aggregate_survey()."
+        )
+
     # --- Resolve design once on full data ---
     effective_design = (
         replace(survey_design, lonely_psu=lonely_psu) if lonely_psu else survey_design

docs/api/prep.rst

@@ -284,13 +284,13 @@ Example
    #   cell_n, cell_n_eff, srs_fallback
 
    # stage2 is pre-configured: aweights + state-level clustering
-   result = DifferenceInDifferences().fit(
-       panel,
-       outcome="smoking_rate_mean",
-       treatment="treated",
-       time="year",
-       survey_design=stage2,
-   )
+   # Add treatment/time indicators at the panel level, then fit:
+   # panel["treated"] = ...  # from policy adoption data
+   # panel["post"] = (panel["year"] >= treatment_year).astype(int)
+   # result = DifferenceInDifferences().fit(
+   #     panel, outcome="smoking_rate_mean",
+   #     treatment="treated", time="post", survey_design=stage2,
+   # )
 
 Data Validation
 ---------------

tests/test_prep.py

@@ -2421,6 +2421,37 @@ def test_error_empty_data(self, design):
                 survey_design=design,
             )
 
+    def test_error_missing_grouping_keys(self, micro_data, design):
+        """NaN in grouping columns raises ValueError."""
+        data = micro_data.copy()
+        data.loc[0, "state"] = np.nan
+        with pytest.raises(ValueError, match="Missing values in grouping column"):
+            aggregate_survey(
+                data,
+                by=["state", "year"],
+                outcomes="y",
+                survey_design=design,
+            )
+
+    def test_error_all_missing_grouping_keys(self, design):
+        """All-NaN grouping column raises ValueError."""
+        data = pd.DataFrame(
+            {
+                "state": [np.nan] * 10,
+                "year": np.ones(10, dtype=int),
+                "y": np.random.RandomState(1).normal(0, 1, 10),
+                "wt": np.ones(10),
+            }
+        )
+        design_simple = SurveyDesign(weights="wt")
+        with pytest.raises(ValueError, match="Missing values in grouping column"):
+            aggregate_survey(
+                data,
+                by=["state", "year"],
+                outcomes="y",
+                survey_design=design_simple,
+            )
+
     def test_domain_estimation_preserves_full_design(self):
         """Full-design domain estimation accounts for PSUs outside the cell.