Skip to content

Polars support (#901) follow-ups: non-blocking polish + bootstrap CI test coverage #920

Description

@jeongyoonlee

Follow-ups from the review of #901 (native Polars support for meta-learners), which is being merged now that the one merge blocker (_fit_bootstrap_clone crashing on a native X) is resolved. Everything below is non-blocking polish plus a test-coverage gap — safe to batch into a single small PR.

1. S-learner predict rebuilds the augmented matrices inside the per-group loop

causalml/inference/meta/slearner.py (BaseSRegressor.predict ~L138-139; BaseSClassifier.predict likewise) builds X_new_c = prepend_column(0.0, X) / X_new_t = prepend_column(1.0, X) inside for group in self.t_groups:. These depend only on X and are identical every iteration → 2·N full-matrix copies instead of 2. Hoist above the loop (restores the prior "build once" optimization).

2. X-learner fit target re-filters instead of reusing the numpy slice

causalml/inference/meta/xlearner.py:162 (and :592 in BaseXClassifier) fits the outcome model with filter_mask(y_filt, w == 1), a redundant second filter that yields a native pandas/polars Series, while y_filt_np = to_numpy(y_filt) is already computed one line below. Use y_filt_np[w == 1] — consistent with the rest of the method and avoids relying on polars-Series→numpy coercion at the sklearn boundary.

3. Test: use shared constants

tests/test_polars_support.py:33 hardcodes RANDOM_STATE = 42; per repo convention use RANDOM_SEED from tests/const.py.

4. Test gap: store_bootstrapsreturn_ci with a DataFrame X

The _fit_bootstrap_clone fix (the merge blocker) is CI-invisible — no test drives fit(store_bootstraps=True)predict(return_ci=True) with a pandas or polars DataFrame X, which is the exact path that was crashing. Add one, plus a DR-classifier polars case (still untested), to lock the fixed path.

Context: full review thread on #901 (2026-07-01 round).

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions