Pandas: Upgrade to 3.0 and fix compatibility issues (#4838)

Summary:
Pandas 3.0 introduces several breaking changes (https://pandas.pydata.org/docs/whatsnew/v3.0.0.html#other-api-changes) that required fixes across the codebase:

1. StringDtype inference: Pandas 3.0 infers string columns to use StringDtype instead of object dtype. This breaks DataFrame comparisons since our Data class expects object dtype (defined in COLUMN_DATA_TYPES). Fixed by setting `pd.options.future.infer_string = False` in modules that construct DataFrames.

2. Deprecated inplace on replace(): The `inplace=True` parameter on Series.replace() is deprecated. Changed to assignment pattern: `df["col"] = df["col"].replace(...)`.

3. read_json() no longer accepts strings: pd.read_json() now requires file paths or file-like objects, not raw JSON strings. Wrapped JSON strings with StringIO().

4. Read-only arrays from DataFrame.to_numpy(): Arrays returned by to_numpy() are now read-only views. Added .copy() before passing to torch.from_numpy() which requires writable arrays.

5. Test dtype mismatches: Tests comparing DataFrames failed because manually constructed expected DataFrames had different dtypes than production code. Fixed by wrapping expected DataFrames with Data(df=...).df to ensure consistent dtype casting.

Pull Request resolved: #4838

Reviewed By: mgrange1998

Differential Revision: D91825190

Pulled By: CristianLara

fbshipit-source-id: 242e097f1f3c90c4c9bb18cb7dc68ea4609adecd

Author: CristianLara

ax/adapter/data_utils.py

@@ -15,13 +15,15 @@
 
 from __future__ import annotations
 
+import functools
 import warnings
 from collections.abc import Iterable
 from copy import deepcopy
 from dataclasses import dataclass, InitVar
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
+import pandas as pd
 from ax.core.data import Data, MAP_KEY
 from ax.core.experiment import Experiment
 from ax.core.map_metric import MapMetric
@@ -351,6 +353,32 @@ def extract_experiment_data(
     return ExperimentData(arm_data=arm_data, observation_data=observation_data)
 
 
+def _use_object_dtype_for_strings(
+    func: Callable[..., Any],
+) -> Callable[..., Any]:
+    """Decorator to disable pandas 3.0 StringDtype inference.
+
+    This ensures string columns like arm_name keep object dtype to match
+    Data.COLUMN_DATA_TYPES. See: https://pandas.pydata.org/docs/whatsnew/v3.0.0.html
+
+    On older pandas versions that don't have the future.infer_string option,
+    this decorator is a no-op since the StringDtype inference doesn't exist.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Check if the future.infer_string option exists (pandas 3.0+)
+        if hasattr(pd.options, "future") and hasattr(pd.options.future, "infer_string"):
+            with pd.option_context("future.infer_string", False):
+                return func(*args, **kwargs)
+        else:
+            # Older pandas version - no StringDtype inference to disable
+            return func(*args, **kwargs)
+
+    return wrapper
+
+
+@_use_object_dtype_for_strings
 def _extract_arm_data(experiment: Experiment) -> DataFrame:
     """Extract a dataframe containing the trial index, arm name,
     parameterizations, and metadata from the given experiment.
@@ -383,6 +411,7 @@ def _extract_arm_data(experiment: Experiment) -> DataFrame:
     return df
 
 
+@_use_object_dtype_for_strings
 def _extract_observation_data(
     experiment: Experiment,
     data_loader_config: DataLoaderConfig,

ax/adapter/tests/test_data_utils.py

@@ -11,7 +11,11 @@
 from unittest import mock
 
 import numpy as np
-from ax.adapter.data_utils import DataLoaderConfig, extract_experiment_data
+from ax.adapter.data_utils import (
+    _use_object_dtype_for_strings,
+    DataLoaderConfig,
+    extract_experiment_data,
+)
 from ax.adapter.registry import Generators
 from ax.core.data import Data, MAP_KEY
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
@@ -97,6 +101,7 @@ def test_extract_experiment_data_empty(self) -> None:
             )
             self.assertEqual(experiment_data, experiment_data)
 
+    @_use_object_dtype_for_strings
     def test_extract_experiment_data_non_map(self) -> None:
         # This is a 2 objective experiment with 2 trials, 1 arm each.
         observations = [[0.1, 1.0], [0.2, 2.0]]
@@ -248,6 +253,7 @@ def test_extract_experiment_data_non_map(self) -> None:
             )
         )
 
+    @_use_object_dtype_for_strings
     def test_extract_experiment_data_map(self) -> None:
         exp = get_branin_experiment_with_timestamp_map_metric(with_trials_and_data=True)
         t_0_metric = 55.602112642270264
@@ -261,7 +267,8 @@ def test_extract_experiment_data_map(self) -> None:
         expected_arm_df = DataFrame(
             [{"x1": 0.0, "x2": 0.0}, {"x1": 1.0, "x2": 1.0}],
             index=MultiIndex.from_tuples(
-                [(0, "0_0"), (1, "1_0")], names=["trial_index", "arm_name"]
+                [(0, "0_0"), (1, "1_0")],
+                names=["trial_index", "arm_name"],
             ),
         )
         assert_frame_equal(
@@ -359,6 +366,7 @@ def test_extract_experiment_data_map(self) -> None:
         # Check equality with self.
         self.assertEqual(experiment_data, experiment_data)
 
+    @_use_object_dtype_for_strings
     def test_extract_experiment_data_multiple_map(self) -> None:
         # Checks that multiple map metrics are correctly normalized.
         # Using a custom Data input to simplify testing.
@@ -467,6 +475,7 @@ def test_extract_experiment_data_batch_trials(self) -> None:
         for df in [experiment_data.arm_data, experiment_data.observation_data]:
             self.assertEqual(set(df.index.get_level_values("arm_name")), expected_arms)
 
+    @_use_object_dtype_for_strings
     def test_extract_experiment_data_with_metadata_columns(self) -> None:
         # Tests the case where the Data.df includes additional columns,
         # such as start_time and end_time, besides the usual required columns.
@@ -522,7 +531,7 @@ def test_extract_experiment_data_with_metadata_columns(self) -> None:
                 names=["trial_index", "arm_name"],
             ),
             columns=MultiIndex.from_tuples(
-                tuples=[
+                [
                     ("mean", "branin_a"),
                     ("mean", "branin_b"),
                     ("sem", "branin_a"),

ax/adapter/transforms/tests/test_cast_transform.py

@@ -11,7 +11,11 @@
 
 import numpy as np
 from ax.adapter.base import DataLoaderConfig
-from ax.adapter.data_utils import ExperimentData, extract_experiment_data
+from ax.adapter.data_utils import (
+    _use_object_dtype_for_strings,
+    ExperimentData,
+    extract_experiment_data,
+)
 from ax.adapter.transforms.cast import Cast
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.parameter import (
@@ -449,6 +453,7 @@ def test_transform_experiment_data_flatten_with_missing_columns(self) -> None:
         )
         self.assertEqual(set(transformed.arm_data.columns), expected_columns)
 
+    @_use_object_dtype_for_strings
     def test_transform_experiment_data_cast(self) -> None:
         # Test for casting to the correct data type and dropping of Nones.
         experiment = get_experiment_with_observations(
@@ -495,6 +500,7 @@ def test_transform_experiment_data_cast(self) -> None:
         ]
         assert_frame_equal(transformed.observation_data, expected_obs_data)
 
+    @_use_object_dtype_for_strings
     def test_transform_experiment_data_cast_map_data(self) -> None:
         # Check that indexing for removal of NaNs works correctly with data that
         # has a "step" column.

ax/adapter/transforms/tests/test_choice_encode_transform.py

@@ -9,7 +9,7 @@
 from copy import deepcopy
 
 from ax.adapter.base import DataLoaderConfig
-from ax.adapter.data_utils import extract_experiment_data
+from ax.adapter.data_utils import _use_object_dtype_for_strings, extract_experiment_data
 from ax.adapter.transforms.choice_encode import (
     ChoiceToNumericChoice,
     OrderedChoiceToIntegerRange,
@@ -258,6 +258,7 @@ def test_hss_dependents_are_preserved(self) -> None:
         self.assertEqual(hss.parameters["x2"].parameter_type, ParameterType.INT)
         self.assertEqual(hss.parameters["x2"].dependents, {0: [], 1: ["x3"]})
 
+    @_use_object_dtype_for_strings
     def test_transform_experiment_data(self) -> None:
         parameterizations = [
             {"x": 2.2, "a": 2, "b": 10.0, "c": 10.0, "d": "r", "e": "q"},

ax/adapter/transforms/tests/test_one_hot_transform.py

@@ -9,7 +9,7 @@
 from copy import deepcopy
 
 from ax.adapter.base import DataLoaderConfig
-from ax.adapter.data_utils import extract_experiment_data
+from ax.adapter.data_utils import _use_object_dtype_for_strings, extract_experiment_data
 from ax.adapter.transforms.one_hot import OH_PARAM_INFIX, OneHot
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import (
@@ -226,6 +226,7 @@ def test_heterogeneous_search_space(self) -> None:
         untf_obs = self.t.untransform_observation_features(obs_ft)
         self.assertFalse(any(obs.parameters.get("b") == "b" for obs in untf_obs))
 
+    @_use_object_dtype_for_strings
     def test_transform_experiment_data(self) -> None:
         parameterizations = [
             {"x": 2.2, "a": 2, "b": "b", "c": False, "d": 10.0},

ax/adapter/transforms/tests/test_unit_x_transform.py

@@ -9,7 +9,7 @@
 from copy import deepcopy
 
 from ax.adapter.base import DataLoaderConfig
-from ax.adapter.data_utils import extract_experiment_data
+from ax.adapter.data_utils import _use_object_dtype_for_strings, extract_experiment_data
 from ax.adapter.transforms.unit_x import UnitX
 from ax.core.observation import ObservationFeatures
 from ax.core.parameter import ChoiceParameter, ParameterType, RangeParameter
@@ -205,6 +205,7 @@ def test_TransformNewSearchSpace(self) -> None:
         t.transform_search_space(new_search_space_with_target)
         self.assertEqual(new_search_space_with_target.parameters["x"].target_value, 0.5)
 
+    @_use_object_dtype_for_strings
     def test_transform_experiment_data(self) -> None:
         parameterizations = [
             {"x": 1.0, "y": 1.5, "z": 1.0, "a": 1, "b": "b"},