Use positional indices for cell membership, add duplicate-index test

igerber · claude · igerber · commit 1bbdfa882b74 · 2026-04-07T11:01:02.000-04:00
Replace label-based index lookup with stable positional row tracking
via _row_pos column, so duplicate DataFrame indices cannot break or
mis-map cell aggregation. Add regression test verifying identical
results with duplicated vs clean indices.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/diff_diff/prep.py b/diff_diff/prep.py
@@ -1532,14 +1532,15 @@ def aggregate_survey(
     y_arrays: Dict[str, np.ndarray] = {var: data[var].values.astype(np.float64) for var in all_vars}
 
     # --- Per-cell computation ---
-    grouped = data.groupby(by_cols, sort=True)
+    # Use stable positional indices (safe with duplicate DataFrame indices)
+    row_positions = np.arange(n_total)
+    grouped = data.assign(_row_pos=row_positions).groupby(by_cols, sort=True)
     rows: List[Dict[str, Any]] = []
     srs_cells: List[str] = []
     zero_var_cells: List[str] = []
 
     for cell_key, cell_df in grouped:
-        cell_idx = np.array(cell_df.index)
-        pos_idx = data.index.get_indexer(cell_idx)
+        pos_idx = cell_df["_row_pos"].values
 
         # Boolean mask for full-design domain estimation
         cell_mask = np.zeros(n_total, dtype=bool)
diff --git a/tests/test_prep.py b/tests/test_prep.py
@@ -2452,6 +2452,38 @@ def test_error_all_missing_grouping_keys(self, design):
                 survey_design=design_simple,
             )
 
+    def test_duplicate_index(self):
+        """Duplicate DataFrame indices do not break aggregation."""
+        rng = np.random.RandomState(77)
+        n = 40
+        data = pd.DataFrame(
+            {
+                "geo": np.repeat(["A", "B"], n // 2),
+                "time": np.tile(np.repeat([0, 1], n // 4), 2),
+                "wt": np.ones(n),
+                "y": rng.normal(10, 2, n),
+            }
+        )
+        # Create duplicate indices (e.g., from concat without reset_index)
+        data.index = list(range(n // 2)) * 2  # 0..19, 0..19
+
+        design = SurveyDesign(weights="wt")
+        panel_dup, _ = aggregate_survey(
+            data, by=["geo", "time"], outcomes="y", survey_design=design
+        )
+
+        # Compare against clean-index version
+        data_clean = data.reset_index(drop=True)
+        panel_clean, _ = aggregate_survey(
+            data_clean, by=["geo", "time"], outcomes="y", survey_design=design
+        )
+
+        # Results should be identical
+        np.testing.assert_allclose(
+            panel_dup["y_mean"].values, panel_clean["y_mean"].values, rtol=1e-12
+        )
+        np.testing.assert_allclose(panel_dup["y_se"].values, panel_clean["y_se"].values, rtol=1e-12)
+
     def test_domain_estimation_preserves_full_design(self):
         """Full-design domain estimation accounts for PSUs outside the cell.
 
