From 476b867a456ae1fb37ae5a5eec16a1b21969bc41 Mon Sep 17 00:00:00 2001 From: Shruti Patel Date: Wed, 1 Apr 2026 08:58:12 -0700 Subject: [PATCH] Fix pandas 3.0 compatibility in MarginalEffectsPlot _prepare_data Summary: Fix pandas 3.0 compatibility in MarginalEffectsPlot _prepare_data Similar changes done in D91825190. Under pandas 3.0, `future.infer_string=True` is the default, which causes `StringDtype` inference on arm parameter DataFrames. This breaks the downstream `marginal_effects()` call in `statstools.py`, producing a `KeyError` on the result DataFrame's column selection. Apply the `_use_object_dtype_for_strings` decorator to `_prepare_data`, matching the pattern already used at 11 other callsites in the Ax codebase (e.g. `_extract_arm_data`, `_extract_observation_data` in `data_utils.py`). Fixes `test_marginal_effects::test_compute_adhoc` failure in GitHub Actions CI where the latest pandas (3.0+) is installed from PyPI. Differential Revision: D99107237 --- ax/analysis/plotly/marginal_effects.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/ax/analysis/plotly/marginal_effects.py b/ax/analysis/plotly/marginal_effects.py index 24902e5fee7..0f6b800c1b8 100644 --- a/ax/analysis/plotly/marginal_effects.py +++ b/ax/analysis/plotly/marginal_effects.py @@ -14,6 +14,7 @@ from ax.analysis.plotly.plotly_analysis import create_plotly_analysis_card from ax.analysis.utils import extract_relevant_adapter, validate_experiment from ax.core.analysis_card import AnalysisCardGroup +from ax.core.data_utils import _use_object_dtype_for_strings from ax.core.experiment import Experiment from ax.core.parameter import ChoiceParameter from ax.generation_strategy.generation_strategy import GenerationStrategy @@ -195,6 +196,7 @@ def compute_marginal_effects_adhoc( ) +@_use_object_dtype_for_strings def _prepare_data( metric: str, adapter: Adapter,