diff --git a/CHANGELOG.md b/CHANGELOG.md index bd2deaeb..1c0422ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 when the extension is absent or `DIFF_DIFF_BACKEND=python`; kernel-side validation errors fall back to the numpy engine. Affects every MAP consumer (TWFE, SunAbraham, Bacon, Wooldridge, DiD/MPD `absorb=`, and the survey replicate refits). +- **Opt-in chunked dispatch for the Rust `demean_map` kernel** (internal + `DIFF_DIFF_DEMEAN_CHUNK_COLS` env knob; default behavior unchanged - single full-width + dispatch). When set to a positive integer, the wrapper feeds the kernel balanced + near-equal column blocks, capping the kernel's transient input/working copies on wide + absorbed designs for memory-constrained runs (a width near the machine's core count + measured best). Results are identical by construction - each column's MAP loop is fully + independent, so chunking changes neither demeaned values nor iteration counts (locked by + exact-equality tests vs single-call dispatch). Shipped opt-in rather than default-on: + measured end-to-end on the FE-absorption lane, the fit-level peak RSS on the widest + workload is dominated by the downstream solver phase, so default chunking would cost + ~2-7% wall-clock for only ~5-12% peak-RSS reduction (details + corrected memory + attribution in `docs/performance-plan.md`). ### Changed - **FE-absorption demeaning rewritten: factorize-once + `np.bincount` method of alternating diff --git a/TODO.md b/TODO.md index 7d54b009..ca0e0fcf 100644 --- a/TODO.md +++ b/TODO.md @@ -46,7 +46,7 @@ generic sparse-FE, QR+SVD rank-detection redundancy, `check_finite` bypass — m | Issue | Location | Origin | Effort | Priority | |-------|----------|--------|--------|----------| -| SunAbraham-family fits are now SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel. The parked "QR+SVD redundancy is minimal vs the SVD" note predates these measurements - the pivoted QR alone is now ~20% of the whole fit. Candidates: reuse the rank-detection QR factors for the solve, or a Cholesky fast path for well-conditioned tall-skinny designs. Measured attribution in docs/performance-plan.md (FE-absorption baseline). | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Heavy | Medium | +| Rust-backend `solve_ols` peak memory on wide designs: the faer path holds an owned copy of the stacked design plus SVD workspace, dominating fit-level peak RSS on wide absorbed fits (firm_churn-class: ~19-21 GB rust-backend vs ~13.4 GB numpy-backend for the same fit; measured PR-D 2026-07, corrected attribution in docs/performance-plan.md - the demean kernel's transients were NOT the dominator, and opt-in chunked dispatch caps them anyway). Candidates: borrow-view instead of owned copy into faer, workspace reuse, or routing wide designs to the numpy gelsd path. Interim escape hatch: DIFF_DIFF_BACKEND=python. | `rust/src/linalg.rs`, `diff_diff/linalg.py::solve_ols` | PR-D probe | Heavy | Medium | | `ImputationDiD` dense `(A0'A0).toarray()` scales `O((U+T+K)^2)` — OOM risk on large panels (only triggers when the sparse solver fails). Needs an alternative dense fallback or richer sparse strategy. | `imputation.py` | #141 | Heavy | Medium | | CR2 Bell-McCaffrey DOF uses a naive `O(n²k)` per-coefficient loop over cluster pairs; Pustejovsky-Tipton (2018) Appendix B has a scores-based formulation avoiding the full `n×n` `M`. Switch when a user hits a large-`n` cluster-robust design. | `linalg.py::_compute_cr2_bm` | Phase 1a | Heavy | Low | | Rust-backend HC2: the Rust path only supports HC1; HC2 and CR2 Bell-McCaffrey fall through to NumPy. Noticeable for large-`n` fits. | `rust/src/linalg.rs` | Phase 1a | Mid | Low | @@ -124,6 +124,7 @@ Doable in principle, but no current caller and/or explicitly out of paper scope. | Issue | Location | PR | Priority | |-------|----------|----|----------| | `StackedDiD` survey re-resolution intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2). The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation LANDED (#226 shared helpers `ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD is deliberately NOT on that path (control units are duplicated across sub-experiments, so it re-resolves at stacked granularity rather than collapsing to one row per unit). The residual is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization. Post-filter re-resolution / metadata-recompute unification across the three estimators was assessed and is not warranted — they use genuinely different mechanisms and already delegate to shared `_resolve_survey_for_fit` / `compute_survey_metadata`. | `stacked_did.py` | #226 | Low | +| SunAbraham-family fits are SOLVER-bound after the v3.6.x demeaning speedups: on the county-class shape (3.1k units x 60 months, ~100 interaction columns) `solve_ols` = ~60% of fit time (faer SVD 1.24s + pivoted-QR rank detection 0.70s of 3.27s), and the share grows further under the Rust demean kernel; the pivoted QR alone is ~20% of the whole fit (measured attribution in docs/performance-plan.md). Candidate fix is QR-reuse (one factorization for rank detection + solve). **Parked per the 2026-07 correctness-first decision**: the change perturbs every coefficient in the library at the last digits (golden/parity/bit-identity recapture) and removes the SVD's independent singular-value truncation - a second, different near-collinearity detector that is deliberate defense-in-depth for the no-silent-failures contract (the FE-span guard episode showed borderline cases are real). Re-open only on practitioner demand for wide event studies; any implementation gates on fixest/R end-to-end parity, not just internal consistency. | `linalg.py::solve_ols`, `linalg.py::_detect_rank_deficiency` | PR-C attribution | Medium | | dCDH parity-test SE/CI assertions only cover pure-direction scenarios; mixed-direction SE comparison is structurally apples-to-oranges (cell-count vs obs-count weighting). | `test_chaisemartin_dhaultfoeuille_parity.py` | #294 | Low | | `HeterogeneousAdoptionDiD` survey-design API consolidation (**scheduled: next minor bump**): drop the deprecated `survey=` / `weights=` kwargs on all 8 HAD surfaces; only `survey_design=` remains. Also fold the legacy back-end `weights=` routing into the unified `_resolve_survey_for_fit` path. DeprecationWarning has shipped; the removal is ~50 LoC gated on the semver bump. | `had.py`, `had_pretests.py` | next minor | Medium | | `HeterogeneousAdoptionDiD` joint cross-horizon covariance / sup-t bands: per-horizon SEs use independent sandwiches (paper-faithful pointwise CIs per Pierce-Schott Fig 2). Follow-ups (low demand): IF-based stacking for joint cross-horizon inference; analytical H×H covariance on the weighted ES path; a sup-t band on the unweighted ES path. | `had.py::_fit_event_study` | Phase 2b / 4.5 B | Low | diff --git a/diff_diff/utils.py b/diff_diff/utils.py index 1e13a7bf..fc0c0a3b 100644 --- a/diff_diff/utils.py +++ b/diff_diff/utils.py @@ -2,6 +2,7 @@ Utility functions for difference-in-differences estimation. """ +import os import warnings from dataclasses import dataclass, field from typing import Any, Dict, Iterable, List, Optional, Tuple @@ -13,15 +14,14 @@ # Import Rust backend if available (from _backend to avoid circular imports) from diff_diff._backend import ( HAS_RUST_BACKEND, - _rust_demean_map, - _rust_project_simplex, - _rust_sdid_unit_weights, - _rust_compute_time_weights, _rust_compute_noise_level, + _rust_compute_time_weights, + _rust_demean_map, _rust_sc_weight_fw, - _rust_sc_weight_fw_with_convergence, _rust_sc_weight_fw_weighted, _rust_sc_weight_fw_weighted_with_convergence, + _rust_sc_weight_fw_with_convergence, + _rust_sdid_unit_weights, ) from diff_diff.linalg import compute_robust_vcov as _compute_robust_vcov_linalg from diff_diff.linalg import solve_ols as _solve_ols_linalg @@ -2707,6 +2707,41 @@ def _demean_map_numpy( return demeaned, iters +# Opt-in column-block width for Rust demean_map kernel calls (None = single +# dispatch, the default). The kernel holds an owned copy of its input block +# plus the result, and chunking caps those transients at the block width while +# leaving per-column results bit-identical: every column's MAP loop is fully +# independent (no cross-column arithmetic), so partitioning changes neither +# values nor iteration counts. Measured on the 2.4M x 130 firm-panel workload +# (2026-07): dispatch-level transients shrink as designed, but the fit's peak +# RSS is dominated by the downstream solver phase, so end-to-end peak dropped +# only ~5-12% while wall-clock rose ~2-7% - hence OFF by default; the env +# knob remains for memory-constrained runs (a block width near the machine's +# core count measured best: each chunk is one full parallel wave). Peak-RSS +# measurements need repeated runs under matched machine state - macOS memory +# compression deflates single-run ru_maxrss readings under ambient pressure. +_DEMEAN_MAP_CHUNK_COLS: Optional[int] = None + + +def _resolve_demean_chunk_cols() -> Optional[int]: + """Opt-in chunk width for the Rust kernel dispatch (None = unchunked). + + ``DIFF_DIFF_DEMEAN_CHUNK_COLS`` (positive integer) enables chunking; + invalid or non-positive values fall back silently to the module default, + mirroring ``DIFF_DIFF_BACKEND``'s convention. Unlike the backend switch + this is read PER CALL, not at import - deliberate, so benchmarks and tests + can A/B chunked vs unchunked dispatch within one process/build. + """ + raw = os.environ.get("DIFF_DIFF_DEMEAN_CHUNK_COLS") + if raw is None: + return _DEMEAN_MAP_CHUNK_COLS + try: + value = int(raw) + except ValueError: + return _DEMEAN_MAP_CHUNK_COLS + return value if value > 0 else _DEMEAN_MAP_CHUNK_COLS + + def _demean_map_rust( x_cols: List[np.ndarray], codes_list: List[np.ndarray], @@ -2715,36 +2750,61 @@ def _demean_map_rust( tol: float, max_iter: int, ) -> Optional[Tuple[List[np.ndarray], List[int]]]: - """Marshal to the Rust ``demean_map`` kernel. + """Marshal to the Rust ``demean_map`` kernel in column chunks. Returns None to signal "use the canonical numpy engine" (kernel absent, degenerate shapes, or a deliberate kernel-side validation error). Dtypes are coerced explicitly BEFORE the call - never rely on exception handling for dtype mismatches. + + Variables are dispatched in blocks of ``_resolve_demean_chunk_cols()`` + columns to bound peak memory (see ``_DEMEAN_MAP_CHUNK_COLS``); the codes + matrix and weights are built once and shared across chunks. Chunking is + exact partitioning - per-column outputs and iteration counts are identical + to a single-call dispatch. """ if _rust_demean_map is None: return None if not x_cols or x_cols[0].shape[0] == 0: return None - x_mat = np.ascontiguousarray(np.column_stack(x_cols), dtype=np.float64) codes_mat = np.ascontiguousarray(np.column_stack(codes_list), dtype=np.int64) w = None if weights is None else np.ascontiguousarray(weights, dtype=np.float64) - try: - out, iters = _rust_demean_map( # type: ignore[misc] - x_mat, - codes_mat, - [int(g) for g in n_groups_list], - w, - float(tol), - int(max_iter), - ) - except ValueError as e: - if "demean_map" in str(e): - # deliberate kernel-side validation marker -> numpy fallback - return None - raise - # F-order result: per-column views are contiguous, no extra copy - return [out[:, j] for j in range(out.shape[1])], [int(i) for i in iters] + n_groups = [int(g) for g in n_groups_list] + k = len(x_cols) + chunk = _resolve_demean_chunk_cols() + if chunk is None: + chunk = k # default: single dispatch, no partitioning + # Balanced partition into ceil(k / chunk) near-equal blocks. A naive + # fixed-stride split leaves a small remainder block (e.g. 130 columns at + # width 32 -> 4x32 + 2); that tail runs nearly serial across the rayon + # pool, and one slow-converging column in it stalls the whole dispatch. + n_chunks = -(-k // chunk) + bounds = [round(i * k / n_chunks) for i in range(n_chunks + 1)] + demeaned: List[np.ndarray] = [] + iters_all: List[int] = [] + for start, stop in zip(bounds, bounds[1:]): + x_mat = np.ascontiguousarray(np.column_stack(x_cols[start:stop]), dtype=np.float64) + try: + out, iters = _rust_demean_map( # type: ignore[misc] + x_mat, + codes_mat, + n_groups, + w, + float(tol), + int(max_iter), + ) + except ValueError as e: + if "demean_map" in str(e): + # deliberate kernel-side validation marker -> numpy fallback + # (kernel validation is chunk-invariant, so this fires on the + # first chunk; the whole call falls back, same contract as an + # unchunked dispatch) + return None + raise + # F-order result: per-column views are contiguous, no extra copy + demeaned.extend(out[:, j] for j in range(out.shape[1])) + iters_all.extend(int(i) for i in iters) + return demeaned, iters_all def demean_by_groups( diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 176e7431..972cb88f 100644 --- a/docs/doc-deps.yaml +++ b/docs/doc-deps.yaml @@ -904,7 +904,7 @@ sources: - path: docs/methodology/REGISTRY.md section: "Absorbed Fixed Effects with Survey Weights" type: methodology - note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation, rust demean_map dispatch" + note: "demean_by_groups / within_transform: MAP convergence contract (max_iter/tol), bincount accumulation numerics, NaN-group-key guard, snap_absorbed_regressors FE-spanned snap + exact LSMR span confirmation, rust demean_map dispatch (chunked, DIFF_DIFF_DEMEAN_CHUNK_COLS)" - path: docs/practitioner_getting_started.rst section: "Step 4: Check Whether the Result Is Trustworthy" type: user_guide diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 33ab7fd3..20eee050 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -4351,6 +4351,11 @@ unequal selection probabilities). tests assert iteration-count equality plus `assert_allclose` at atol=1e-12 (never a bit-identity claim). `DIFF_DIFF_BACKEND=python` disables the kernel; any kernel-side validation error falls back to the numpy engine. + The dispatch can partition variables into balanced column blocks + (`DIFF_DIFF_DEMEAN_CHUNK_COLS`, internal opt-in knob, OFF by default) to + bound the kernel's transient memory; per-column results and iteration + counts are unchanged by construction (each column's MAP loop is fully + independent - no cross-column arithmetic). - **Edge case:** NaN in an absorbed group column raises a `ValueError` naming the column. `pd.factorize` codes NaN keys as -1, which would otherwise silently index the last group's mean; the prior pandas behavior was itself silently bad diff --git a/docs/performance-plan.md b/docs/performance-plan.md index ec546a7b..ad2b8538 100644 --- a/docs/performance-plan.md +++ b/docs/performance-plan.md @@ -76,9 +76,22 @@ submit gate), tail_stress 31.74s -> 11.37s (2.79x), geo_experiment 0.976s -> 0.484s (2.02x, now ahead of the pyfixest yardstick's 0.623s), scanner 1.59x, survey 1.44x, county 1.36x, guard_small unregressed. Identity vs the committed after-baselines: 1e-13-1e-16 on every scenario incl. tail_stress. -Known trade-off: the kernel holds one owned input copy plus the result -(firm_churn peak RSS 13.4 GB -> 21.0 GB); chunking variables through the -kernel is the noted follow-up if that bites a real workload. +Known trade-off: the rust-backend firm_churn fit peaks at ~21 GB RSS vs +~13.4 GB under the numpy backend. **Corrected attribution (2026-07, PR-D +probes):** the peak is dominated by the downstream SOLVER phase (the Rust +faer path holds an owned copy of the ~2.5 GB stacked design plus SVD +workspace), not by the demean kernel's dispatch marshalling as this +paragraph originally claimed. Measured evidence: an isolated width-16 +chunked dispatch cuts the kernel's transients to near-numpy footprint, yet +the end-to-end fit still peaks ~19-21 GB (only ~5-12% below unchunked) at a ++2-7% wall-clock cost - so chunked dispatch shipped as an opt-in env knob +(`DIFF_DIFF_DEMEAN_CHUNK_COLS`, default off) rather than default-on, and +the solver-phase peak is tracked as its own TODO row. `DIFF_DIFF_BACKEND= +python` remains the OOM escape hatch. Measurement note: single-run +`ru_maxrss` on macOS is unreliable under ambient memory pressure (the +compressor deflates resident peaks - one probe read 15.1 GB for a +configuration that reproducibly peaks ~19-20 GB); gate memory claims on +repeated runs under matched machine state. ### FE-absorption suite results diff --git a/tests/test_rust_backend.py b/tests/test_rust_backend.py index 72511c70..23759ad4 100644 --- a/tests/test_rust_backend.py +++ b/tests/test_rust_backend.py @@ -3157,3 +3157,176 @@ def test_stale_symbol_none_falls_back_to_numpy(self, monkeypatch): ) is None ) + + @staticmethod + def _counting_kernel(monkeypatch): + """Wrap the kernel symbol to count invocations (proves the chunked + path actually ran, per-chunk).""" + import diff_diff.utils as utils_mod + + calls = {"kernel": 0} + orig = utils_mod._rust_demean_map + + def counting(*a, **kw): + calls["kernel"] += 1 + return orig(*a, **kw) + + monkeypatch.setattr(utils_mod, "_rust_demean_map", counting) + return calls + + @pytest.mark.parametrize("weighted", [False, True]) + def test_chunked_dispatch_exactly_equals_single_call(self, monkeypatch, weighted): + """Chunking is exact partitioning: per-column outputs are IDENTICAL + (assert_array_equal - same code path, no cross-backend caveat) and + iteration counts equal, vs both single-chunk rust and the numpy + engine.""" + import diff_diff.utils as utils_mod + from diff_diff.utils import _demean_map_numpy, _demean_map_rust + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + x_cols, codes_list, n_groups, w = self._fixture("unbalanced", seed=13, k=8) + weights = w if weighted else None + + monkeypatch.setattr(utils_mod, "_DEMEAN_MAP_CHUNK_COLS", 1000) + single = _demean_map_rust(x_cols, codes_list, n_groups, weights, 1e-10, 10_000) + assert single is not None + + calls = self._counting_kernel(monkeypatch) + monkeypatch.setattr(utils_mod, "_DEMEAN_MAP_CHUNK_COLS", 3) + chunked = _demean_map_rust(x_cols, codes_list, n_groups, weights, 1e-10, 10_000) + assert chunked is not None + assert calls["kernel"] == 3 # ceil(8 / 3) + + assert chunked[1] == single[1] + for c_col, s_col in zip(chunked[0], single[0]): + np.testing.assert_array_equal(c_col, s_col) + + numpy_res = _demean_map_numpy(x_cols, codes_list, n_groups, weights, 1e-10, 10_000) + assert chunked[1] == numpy_res[1] + for c_col, p_col in zip(chunked[0], numpy_res[0]): + np.testing.assert_allclose(c_col, p_col, rtol=0, atol=1e-12) + + @pytest.mark.parametrize("k", [1, 4, 5]) # below / at / above the chunk boundary + def test_chunk_boundaries(self, monkeypatch, k): + """k <= chunk is a single kernel call; k = chunk+1 exercises the + two-chunk boundary (balanced partition: 5 -> 2+3). All cases match + the numpy engine.""" + import math + + import diff_diff.utils as utils_mod + from diff_diff.utils import _demean_map_numpy, _demean_map_rust + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + monkeypatch.setattr(utils_mod, "_DEMEAN_MAP_CHUNK_COLS", 4) + x_cols, codes_list, n_groups, _ = self._fixture("balanced", seed=14, k=k) + calls = self._counting_kernel(monkeypatch) + rust = _demean_map_rust(x_cols, codes_list, n_groups, None, 1e-10, 10_000) + assert rust is not None + assert calls["kernel"] == math.ceil(k / 4) + numpy_res = _demean_map_numpy(x_cols, codes_list, n_groups, None, 1e-10, 10_000) + assert rust[1] == numpy_res[1] + for r_col, p_col in zip(rust[0], numpy_res[0]): + np.testing.assert_allclose(r_col, p_col, rtol=0, atol=1e-12) + + def test_nonconverged_variable_in_second_chunk_still_named(self, monkeypatch): + """The non-convergence warning names a variable whose column lands in + the SECOND chunk (iters aggregation preserves variable order).""" + import diff_diff.utils as utils_mod + from diff_diff.utils import demean_by_groups + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + monkeypatch.setattr(utils_mod, "_DEMEAN_MAP_CHUNK_COLS", 2) + rng = np.random.default_rng(15) + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(30), 6), + "period": np.tile(np.arange(6), 30), + } + ) + df = df[rng.random(len(df)) > 0.3].reset_index(drop=True) + for name in ["x1", "x2", "y"]: # 'y' is column 3 -> chunk 2 of 2 + df[name] = rng.normal(size=len(df)) + calls = self._counting_kernel(monkeypatch) + with pytest.warns(UserWarning, match=r"\['x1', 'x2', 'y'\].*did not converge"): + demean_by_groups( + df, ["x1", "x2", "y"], ["unit", "period"], suffix="_dm", max_iter=1, tol=1e-15 + ) + assert calls["kernel"] == 2 # ceil(3 / 2): multi-chunk path exercised + + def test_estimator_level_chunked_att_parity(self, monkeypatch): + """DiD(absorb=) ATT/SE with chunk=2 identical to the default chunk.""" + import diff_diff.utils as utils_mod + from diff_diff import DifferenceInDifferences + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + rng = np.random.default_rng(16) + n_units, n_periods = 60, 8 + df = pd.DataFrame( + { + "unit": np.repeat(np.arange(n_units), n_periods), + "period": np.tile(np.arange(n_periods), n_units), + } + ) + df["treated"] = (df["unit"] < 30).astype(int) + df["post"] = (df["period"] >= 4).astype(int) + df["x1"] = rng.normal(size=len(df)) + df["x2"] = rng.normal(size=len(df)) + df["y"] = ( + 1.5 * df["treated"] * df["post"] + + 0.5 * df["x1"] + - 0.3 * df["x2"] + + rng.normal(0, 0.5, len(df)) + ) + + def fit(): + return DifferenceInDifferences().fit( + df, + outcome="y", + treatment="treated", + time="post", + absorb=["unit", "period"], + covariates=["x1", "x2"], + ) + + r_default = fit() + monkeypatch.setattr(utils_mod, "_DEMEAN_MAP_CHUNK_COLS", 2) + r_chunked = fit() + np.testing.assert_allclose(r_chunked.att, r_default.att, rtol=0, atol=1e-12) + np.testing.assert_allclose(r_chunked.se, r_default.se, rtol=0, atol=1e-12) + + def test_default_is_single_dispatch(self, monkeypatch): + """With the env unset and the default constant (None), a k=8 dispatch + makes exactly ONE kernel call - chunking is opt-in.""" + from diff_diff.utils import _demean_map_rust + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + x_cols, codes_list, n_groups, _ = self._fixture("balanced", seed=17, k=8) + calls = self._counting_kernel(monkeypatch) + result = _demean_map_rust(x_cols, codes_list, n_groups, None, 1e-10, 10_000) + assert result is not None + assert calls["kernel"] == 1 + + +class TestDemeanChunkResolver: + """Pure-Python env-override logic - runs regardless of the Rust build.""" + + def test_default_when_unset_is_none(self, monkeypatch): + """Chunking is OPT-IN: env unset -> None -> single dispatch.""" + import diff_diff.utils as utils_mod + + monkeypatch.delenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", raising=False) + assert utils_mod._DEMEAN_MAP_CHUNK_COLS is None + assert utils_mod._resolve_demean_chunk_cols() is None + + def test_valid_override_honored(self, monkeypatch): + import diff_diff.utils as utils_mod + + monkeypatch.setenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", "7") + assert utils_mod._resolve_demean_chunk_cols() == 7 + + @pytest.mark.parametrize("bad", ["abc", "0", "-4", "", "3.5"]) + def test_invalid_values_fall_back_to_default(self, monkeypatch, bad): + import diff_diff.utils as utils_mod + + monkeypatch.setenv("DIFF_DIFF_DEMEAN_CHUNK_COLS", bad) + assert utils_mod._resolve_demean_chunk_cols() is utils_mod._DEMEAN_MAP_CHUNK_COLS