Repository: modin-project/modin (python)

Issue:
BUG: Extensions overrides of module-level re-exports does not dispatch to the appropriate backend
### Modin version checks

- [x] I have checked that this issue has not already been reported.

- [x] I have confirmed this bug exists on the latest released version of Modin.

- [x] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-main-branch).)


### Reproducible Example

```python
from modin.pandas.api.extensions import register_pd_accessor
from modin.config import Backend
register_pd_accessor("date_range", backend="Pandas")(lambda i: print("fake override"))
Backend.put("Pandas")
pd.date_range(start=1, periods=2, freq="1M")
```

### Issue Description

When a module-level attribute lookup is performed, the module's `__getattr__` is not invoked if the property is already defined in the module's namespace. As a result, top-level functions and classes that are re-exported from pandas (such as `pd.date_range` and `pd.Index`) are not affected by the extensions dispatch system.

### Expected Behavior

If an extension is specified for one of these methods, it should dispatch the method to the appropriate backend.

### Error Logs

<details>

```python-traceback

Replace this line with the error backtrace (if applicable).

```

</details>


### Installed Versions

<details>

INSTALLED VERSIONS
------------------
commit                : a68dab4ad71e15010f684c237bf1961662eecde5
python                : 3.10.13.final.0
python-bits           : 64
OS                    : Darwin
OS-release            : 24.4.0
Version               : Darwin Kernel Version 24.4.0: Fri Apr 11 18:33:39 PDT 2025; root:xnu-11417.101.15~117/RELEASE_ARM64_T6020
machine               : arm64
processor             : arm
byteorder             : little
LC_ALL                : None
LANG                  : en_US.UTF-8
LOCALE                : en_US.UTF-8

Modin dependencies
------------------
modin                 : 0.32.0+43.ga68dab4a
ray                   : 2.34.0
dask                  : 2024.8.1
distributed           : 2024.8.1

pandas dependencies
-------------------
pandas                : 2.2.2
numpy                 : 1.26.4
pytz                  : 2023.3.post1
dateutil              : 2.8.2
setuptools            : 68.0.0
pip                   : 23.3
Cython                : None
pytest                : 8.3.2
hypothesis            : None
sphinx                : 5.3.0
blosc                 : None
feather               : None
xlsxwriter            : None
lxml.etree            : 5.3.0
html5lib              : None
pymysql               : None
psycopg2              : 2.9.9
jinja2                : 3.1.4
IPython               : 8.17.2
pandas_datareader     : None
adbc-driver-postgresql: None
adbc-driver-sqlite    : None
bs4                   : 4.12.2
bottleneck            : None
dataframe-api-compat  : None
fastparquet           : 2024.5.0
fsspec                : 2024.6.1
gcsfs                 : None
matplotlib            : 3.9.2
numba                 : None
numexpr               : 2.10.1
odfpy                 : None
openpyxl              : 3.1.5
pandas_gbq            : 0.23.1
pyarrow               : 17.0.0
pyreadstat            : None
python-calamine       : None
pyxlsb                : None
s3fs                  : 2024.6.1
scipy                 : 1.14.1
sqlalchemy            : 2.0.32
tables                : 3.10.1
tabulate              : 0.9.0
xarray                : 2024.7.0
xlrd                  : 2.0.1
zstandard             : None
tzdata                : 2023.3
qtpy                  : None
pyqt5                 : None
</details>


Pull request description:
FIX-#7528: Dispatch module-level extensions to the correct backend
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

In this PR, the extensions system tracks objects that were re-exported from pandas, and either wraps them in `wrap_free_function_in_argument_caster` if it is a function, or removes them from the namespace and dispatches via `__getattr__` if it is a class. Technically the `__getattr__` approach can be used for all objects, but this remains more consistent with how we handle dispatch for top-level functions that modin itself defines.

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #7528 <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Patch:
diff --git a/modin/pandas/__init__.py b/modin/pandas/__init__.py
index 08fd15ae..b629167b 100644
--- a/modin/pandas/__init__.py
+++ b/modin/pandas/__init__.py
@@ -31,69 +31,85 @@ if (
 # to not pollute namespace
 del version
 
+
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
-    from pandas import (
-        eval,
-        factorize,
-        test,
-        date_range,
-        period_range,
-        Index,
-        MultiIndex,
-        CategoricalIndex,
-        bdate_range,
-        DatetimeIndex,
-        Timedelta,
-        Timestamp,
-        set_eng_float_format,
-        options,
-        describe_option,
-        set_option,
-        get_option,
-        reset_option,
-        option_context,
-        NaT,
-        PeriodIndex,
-        Categorical,
-        Interval,
-        UInt8Dtype,
-        UInt16Dtype,
-        UInt32Dtype,
-        UInt64Dtype,
-        SparseDtype,
-        Int8Dtype,
-        Int16Dtype,
-        Int32Dtype,
-        Int64Dtype,
-        StringDtype,
-        BooleanDtype,
-        CategoricalDtype,
-        DatetimeTZDtype,
-        IntervalDtype,
-        PeriodDtype,
-        RangeIndex,
-        TimedeltaIndex,
-        IntervalIndex,
-        IndexSlice,
-        Grouper,
-        array,
-        Period,
-        DateOffset,
-        timedelta_range,
-        infer_freq,
-        interval_range,
-        ExcelWriter,
-        NamedAgg,
-        NA,
-        api,
-        ArrowDtype,
-        Flags,
-        Float32Dtype,
-        Float64Dtype,
-        from_dummies,
-        testing,
+    import inspect
+
+    from modin.core.storage_formats.pandas.query_compiler_caster import (
+        wrap_free_function_in_argument_caster,
+    )
+
+    # To allow the extensions system to override these methods, we must wrap all objects re-exported
+    # from pandas in a backend dispatcher.
+    _reexport_list = (
+        "eval",
+        "factorize",
+        "test",
+        "date_range",
+        "period_range",
+        "Index",
+        "MultiIndex",
+        "CategoricalIndex",
+        "bdate_range",
+        "DatetimeIndex",
+        "Timedelta",
+        "Timestamp",
+        "set_eng_float_format",
+        "options",
+        "describe_option",
+        "set_option",
+        "get_option",
+        "reset_option",
+        "option_context",
+        "NaT",
+        "PeriodIndex",
+        "Categorical",
+        "Interval",
+        "UInt8Dtype",
+        "UInt16Dtype",
+        "UInt32Dtype",
+        "UInt64Dtype",
+        "SparseDtype",
+        "Int8Dtype",
+        "Int16Dtype",
+        "Int32Dtype",
+        "Int64Dtype",
+        "StringDtype",
+        "BooleanDtype",
+        "CategoricalDtype",
+        "DatetimeTZDtype",
+        "IntervalDtype",
+        "PeriodDtype",
+        "RangeIndex",
+        "TimedeltaIndex",
+        "IntervalIndex",
+        "IndexSlice",
+        "Grouper",
+        "array",
+        "Period",
+        "DateOffset",
+        "timedelta_range",
+        "infer_freq",
+        "interval_range",
+        "ExcelWriter",
+        "NamedAgg",
+        "NA",
+        "api",
+        "ArrowDtype",
+        "Flags",
+        "Float32Dtype",
+        "Float64Dtype",
+        "from_dummies",
+        "testing",
     )
+    for name in _reexport_list:
+        item = getattr(pandas, name)
+        if inspect.isfunction(item):
+            # Note that this is applied to only functions, not classes.
+            item = wrap_free_function_in_argument_caster(name)(item)
+        globals()[name] = item
+    del inspect, item, _reexport_list, name, wrap_free_function_in_argument_caster
 
 import os
 
diff --git a/modin/pandas/api/extensions/extensions.py b/modin/pandas/api/extensions/extensions.py
index 4f851bb4..cbcca65f 100644
--- a/modin/pandas/api/extensions/extensions.py
+++ b/modin/pandas/api/extensions/extensions.py
@@ -11,9 +11,10 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
+import inspect
 from collections import defaultdict
 from types import MethodType, ModuleType
-from typing import Any, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import modin.pandas as pd
 from modin.config import Backend
@@ -24,7 +25,12 @@ from modin.core.storage_formats.pandas.query_compiler_caster import (
     wrap_function_in_argument_caster,
 )
 
-_attrs_to_delete_on_test = defaultdict(list)
+_attrs_to_delete_on_test = defaultdict(set)
+
+# Track a dict of module-level classes that are re-exported from pandas that may need to dynamically
+# change when overridden by the extensions system, such as pd.Index.
+# See register_pd_accessor for details.
+_reexport_classes: Dict[str, Any] = {}
 
 
 def _set_attribute_on_obj(
@@ -32,6 +38,7 @@ def _set_attribute_on_obj(
     extensions: EXTENSION_DICT_TYPE,
     backend: Optional[str],
     obj: Union[type, ModuleType],
+    set_reexport: bool = False,
 ):
     """
     Create a new or override existing attribute on obj.
@@ -47,6 +54,8 @@ def _set_attribute_on_obj(
         will become the default for all backends.
     obj : DataFrame, Series, or modin.pandas
         The object we are assigning the new attribute to.
+    set_reexport : bool, default False
+        If True, register the original property in `_reexport_classes`.
 
     Returns
     -------
@@ -70,12 +79,33 @@ def _set_attribute_on_obj(
         new_attr
             Unmodified new_attr is return from the decorator.
         """
+        # Module-level functions are resolved by `wrap_free_function_in_argument_caster`, which dynamically
+        # identifies the appropriate backend to use. We cannot apply this wrapper to classes in order
+        # to preserve the vailidity of `isinstance` checks, and instead must force __getattr__ to directly
+        # return the correct class.
+        # Because the module-level __getattr__ function is not called if the object is found in the namespace,
+        # any overrides from the extensions system must `delattr` the attribute to force any future lookups
+        # to hit this code path.
+        # We cannot do this by omitting those exports at module initialization time because the
+        # __getattr__ codepath performs a call to Backend.get() that assumes the presence of an engine;
+        # in an extensions system that may reference types like pd.Timestamp/pd.Index before registering
+        # itself as an engine, this will cause errors.
+        if set_reexport:
+            original_attr = getattr(pd, name)
+            _reexport_classes[name] = original_attr
+            delattr(pd, name)
         extensions[None if backend is None else Backend.normalize(backend)][
             name
         ] = new_attr
-        if callable(new_attr) and name not in dir(obj):
+        if (
+            callable(new_attr)
+            and name not in dir(obj)
+            and not inspect.isclass(new_attr)
+        ):
             # For callable extensions, we add a method to `obj`'s namespace that
             # dispatches to the correct implementation.
+            # If the extension is a class like pd.Index, do not add a wrapper and let
+            # the getattr dispatcher choose the correct item.
             setattr(
                 obj,
                 name,
@@ -89,7 +119,9 @@ def _set_attribute_on_obj(
                     name=name,
                 ),
             )
-            _attrs_to_delete_on_test[obj].append(name)
+            # "Free" functions are permanently kept in the wrapper, so no need to clear them in tests.
+            if obj is not pd:
+                _attrs_to_delete_on_test[obj].add(name)
         return new_attr
 
     return decorator
@@ -97,7 +129,7 @@ def _set_attribute_on_obj(
 
 def register_dataframe_accessor(name: str, *, backend: Optional[str] = None):
     """
-    Registers a dataframe attribute with the name provided.
+    Register a dataframe attribute with the name provided.
 
     This is a decorator that assigns a new attribute to DataFrame. It can be used
     with the following syntax:
@@ -135,7 +167,7 @@ def register_dataframe_accessor(name: str, *, backend: Optional[str] = None):
 
 def register_series_accessor(name: str, *, backend: Optional[str] = None):
     """
-    Registers a series attribute with the name provided.
+    Register a series attribute with the name provided.
 
     This is a decorator that assigns a new attribute to Series. It can be used
     with the following syntax:
@@ -216,7 +248,7 @@ def register_base_accessor(name: str, *, backend: Optional[str] = None):
 
 def register_pd_accessor(name: str, *, backend: Optional[str] = None):
     """
-    Registers a pd namespace attribute with the name provided.
+    Register a pd namespace attribute with the name provided.
 
     This is a decorator that assigns a new attribute to modin.pandas. It can be used
     with the following syntax:
@@ -250,14 +282,19 @@ def register_pd_accessor(name: str, *, backend: Optional[str] = None):
     decorator
         Returns the decorator function.
     """
+    set_reexport = name not in _GENERAL_EXTENSIONS[backend] and name in dir(pd)
     return _set_attribute_on_obj(
-        name=name, extensions=_GENERAL_EXTENSIONS, backend=backend, obj=pd
+        name=name,
+        extensions=_GENERAL_EXTENSIONS,
+        backend=backend,
+        obj=pd,
+        set_reexport=set_reexport,
     )
 
 
 def __getattr___impl(name: str):
     """
-    Override __getatttr__ on the modin.pandas module to enable extensions.
+    Override __getattr__ on the modin.pandas module to enable extensions.
 
     Note that python only falls back to this function if the attribute is not
     found in this module's namespace.
@@ -273,7 +310,6 @@ def __getattr___impl(name: str):
         Returns the extension attribute, if it exists, otherwise returns the attribute
         imported in this file.
     """
-
     from modin.config import Backend
 
     backend = Backend.get()
@@ -281,5 +317,7 @@ def __getattr___impl(name: str):
         return _GENERAL_EXTENSIONS[backend][name]
     elif name in _GENERAL_EXTENSIONS[None]:
         return _GENERAL_EXTENSIONS[None][name]
+    elif name in _reexport_classes:
+        return _reexport_classes[name]
     else:
         raise AttributeError(f"module 'modin.pandas' has no attribute '{name}'")
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
index 6f180535..8130a907 100644
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -36,6 +36,7 @@ from typing import (
 
 import numpy as np
 import pandas
+from pandas import Categorical
 from pandas._libs import lib
 from pandas._typing import (
     CompressionOptions,
@@ -62,7 +63,6 @@ from modin.config import PersistentPickle
 from modin.core.storage_formats.pandas.query_compiler_caster import EXTENSION_DICT_TYPE
 from modin.error_message import ErrorMessage
 from modin.logging import disable_logging
-from modin.pandas import Categorical
 from modin.pandas.io import from_non_pandas, from_pandas, to_pandas
 from modin.utils import (
     MODIN_UNNAMED_SERIES_LABEL,


Test patch:
diff --git a/modin/tests/pandas/extensions/test_pd_extensions.py b/modin/tests/pandas/extensions/test_pd_extensions.py
index 2d8f0a7c..be28d25a 100644
--- a/modin/tests/pandas/extensions/test_pd_extensions.py
+++ b/modin/tests/pandas/extensions/test_pd_extensions.py
@@ -167,3 +167,21 @@ class TestRegisterForOneBackend:
             pd.concat([modin_on_python_df, modin_on_pandas_df])
             == "python_concat_result"
         )
+
+    def test_index_class_override(self):
+        class FakeIndex:
+            def __init__(self, _values):
+                pass
+
+            def fake_method(self) -> str:
+                return "python_fake_index"
+
+        register_pd_accessor("Index", backend="Python_Test")(FakeIndex)
+
+        with config_context(Backend="Pandas"):
+            # Should return an actual native pandas index object
+            df_equals(pd.Index([1]).to_series(), pd.Series([1], index=[1]))
+
+        with config_context(Backend="Python_Test"):
+            # Should just return a string
+            assert pd.Index([1]).fake_method() == "python_fake_index"
diff --git a/modin/tests/pandas/extensions/test_series_extensions.py b/modin/tests/pandas/extensions/test_series_extensions.py
index 33f717fd..55912a92 100644
--- a/modin/tests/pandas/extensions/test_series_extensions.py
+++ b/modin/tests/pandas/extensions/test_series_extensions.py
@@ -255,3 +255,22 @@ def test_disallowed_extensions(Backend1, non_extendable_attribute_name):
         register_series_accessor(name=non_extendable_attribute_name, backend=Backend1)(
             "unused_value"
         )
+
+
+def test_wrapped_extension(Backend1):
+    """
+    Tests using the extensions system to overwrite a method with a wrapped version of the original method
+    obtained via getattr.
+    Because the QueryCompilerCaster ABC automatically wraps all methods with a dispatch to the appropriate
+    backend, we must use the __wrapped__ property of the originally-defined attribute to avoid
+    infinite recursion.
+    """
+    original_item = pd.Series.item.__wrapped__
+
+    @register_series_accessor(name="item", backend=Backend1)
+    def item_implementation(self):
+        return (original_item(self) + 2) * 5
+
+    series = pd.Series([3])
+    assert series.item() == 3
+    assert series.set_backend(Backend1).item() == 25
diff --git a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
index 7257bd2e..c7ab640c 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -825,6 +825,7 @@ class TestSwitchBackendPreOp:
                 ).move_to("Small_Data_Local")
             )
         )
+        pandas_read_json.__name__ = "read_json"
         cloud_read_json = mock.Mock(
             wraps=(
                 lambda *args, **kwargs: _GENERAL_EXTENSIONS[None]["read_json"](
@@ -832,6 +833,7 @@ class TestSwitchBackendPreOp:
                 ).move_to("Big_Data_Cloud")
             )
         )
+        cloud_read_json.__name__ = "read_json"
 
         register_pd_accessor("read_json", backend="Small_Data_Local")(pandas_read_json)
         register_pd_accessor("read_json", backend="Big_Data_Cloud")(cloud_read_json)