Thanks for looking into this, @aman-coder03.

As I commented on #855, converting the dataframe into numpy is not ideal for Polars, as it doesn't benefit from its performant features. Instead, adding support for native Pandas and Polars DataFrames is recommended. It will require careful inspection of the current indexing to make it comparable to the indexing on DataFrames.

thanks for the feedback @jeongyoonlee

i understand the concern....converting to NumPy at the boundary means Polars users lose all the performance benefits(lazy evaluation, zero-copy operations, columnar efficiency) that they came for in the first place

for the native approach, my plan would be...

• add thin abstraction layer in utils.py with helpers like filter_by_mask(), concat_cols(), and get_values() that dispatch to the correct pandas/polars/numpy operation based on the input type
• audit every indexing pattern across all five learner files(tlearner.py, slearner.py, xlearner.py, rlearner.py, drlearner.py) then replace them with these helpers
• key patterns to handle are: boolean mask filtering (X[mask]), column concatenation (np.hstack), and in-place mutation(Polars is immutable)
• keep convert_pd_to_np() only at the final NumPy-only boundaries(sklearn's cross_val_predict in the R-learner), since sklearn doesn't accept Polars natively

Before I go ahead and rewrite, a couple of questions to make sure i am heading in the right direction...

1. should the output of predict() remain a NumPy array, or would you like an option to return a Polars DataFrame when the input was Polars?
2. for sklearn calls like cross_val_predict (R-learner) and propensity model fitting, converting to NumPy at that specific boundary seems unavoidable....does that approach work for you, or do you have a preferred alternative?

Happy to update the PR once we are aligned

Thanks for the detailed plan, @aman-coder03. One important correction before you start: converting to NumPy at the sklearn boundary isn't necessary.

• scikit-learn ≥ 1.4 accepts Polars (and pandas) DataFrames natively via the DataFrame Interchange Protocol (release notes); causalml pins scikit-learn>=1.6.0.
• XGBoost ≥ 3.1 accepts pl.DataFrame and pl.LazyFrame directly (docs).
• LightGBM nominally accepts pl.DataFrame but has a known sklearn-API bug (issue #6849) — document as a caveat.

So the rewrite scope collapses to causalml's own indexing/concat/mutation — not the model calls.

Answers to your questions:

1. predict() output type. Keep returning NumPy. The entire underlying stack does the same: scikit-learn's .predict() returns NumPy regardless of input (only transformers honor set_output("polars"), not predictors), XGBoost's .predict() returns NumPy (or cupy.ndarray on GPU) even when fed a pl.DataFrame via its zero-copy Arrow path, and LightGBM's .predict() returns NumPy as documented. Matching that convention keeps causalml's API consistent with sklearn/xgboost/lightgbm and preserves backward compatibility. A return_type="polars" opt-in can be added later if users ask.

2. NumPy conversion at sklearn boundaries. Skip it — pass Polars/pandas straight through. Zero materialization overhead, and dtype/feature-name plumbing is handled by the ML library.

Suggested phasing:

• Phase 1: Remove convert_pd_to_np() in favor of native pandas DataFrame support across all meta-learners. No Polars yet — this isolates the indexing-polymorphism work and gives a clean baseline.
• Phase 2: Add Polars support to one estimator (e.g. BaseTLearner) as a reference implementation with equivalence tests (NumPy / pandas / Polars produce identical te).
• Phase 3+: Extend Polars to the remaining learners (S/X/R/DR), one PR each.

A few practical notes for the audit:

• Enumerate the actual patterns first. Grep all five learners for indexing/concat/mutation (X[mask], np.hstack, .iloc, in-place assignment) and design the minimal helper set from real usage, not guesses.
• Polars is immutable — anywhere current code does X[mask] = value or accumulates by index assignment, you'll need an explicit rewrite (with_columns, when/then/otherwise), not just a dispatch helper.
• Add polars to test deps in Phase 2. The current PR's tests don't run in CI because Polars isn't installed there — pytest.importorskip("polars") skips the whole file silently.

Happy to review Phase 1 once it's ready.

thanks for the detailed feedback @jeongyoonlee
I understand the phasing suggestion, but I've intentionally tried to complete the full Polars support in this single PR....covering all five learners (T/S/X/R/DR) together. My reasoning was that the indexing patterns are largely the same across learners, so the helper set (filter_mask, filter_index, prepend_column, etc.) naturally fell out of auditing all of them at once rather than discovering gaps PR by PR.

i'm happy to reorganize if you'd prefer the phased approach...I can split this into separate PRs for Phase 1 (pandas cleanup) and Phase 2/3 (Polars per learner). But if the current implementation looks sound to you, I'd love to land it as is to avoid rebasing overhead