Repository: modin-project/modin (python)

Issue:
FEAT: allow module-level methods like pd.to_datetime() to dispatch to the extension for the correct backend

Pull request description:
FEAT-#7474: Register general functions with a particular backend.
Prior to this commit, general extensions would apply to all backends. Now, we can register an extension that applies to just one backend.

- [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 #7474
- [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/core/storage_formats/pandas/query_compiler_caster.py b/modin/core/storage_formats/pandas/query_compiler_caster.py
index 372c6c3e..6b05aa9a 100644
--- a/modin/core/storage_formats/pandas/query_compiler_caster.py
+++ b/modin/core/storage_formats/pandas/query_compiler_caster.py
@@ -22,7 +22,7 @@ This ensures compatibility between different query compiler classes.
 import functools
 import inspect
 from abc import ABC, abstractmethod
-from collections import namedtuple
+from collections import defaultdict, namedtuple
 from types import FunctionType, MethodType
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, ValuesView
 
@@ -34,11 +34,17 @@ from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
 from modin.core.storage_formats.base.query_compiler_calculator import (
     BackendCostCalculator,
 )
-from modin.error_message import ErrorMessage
 
 Fn = TypeVar("Fn", bound=Any)
 
 
+# This type describes a defaultdict that maps backend name (or `None` for
+# method implementation and not bound to any one extension) to the dictionary of
+# extensions for that backend. The keys of the inner dictionary are the names of
+# the extensions, and the values are the extensions themselves.
+EXTENSION_DICT_TYPE = defaultdict[Optional[str], dict[str, Any]]
+
+
 _NON_EXTENDABLE_ATTRIBUTES = {
     # we use these attributes to implement casting and backend dispatching, so
     # we can't allow extensions to override them.
@@ -234,7 +240,7 @@ def apply_argument_cast_to_class(klass: type) -> type:
                             else MethodType
                         )
                     ),
-                    cls=klass,
+                    extensions=klass._extensions,
                     name=attr_name,
                 )
                 wrapped = (
@@ -264,7 +270,7 @@ def wrap_function_in_argument_caster(
     wrapping_function_type: Optional[
         Union[type[classmethod], type[staticmethod], type[MethodType]]
     ],
-    cls: Optional[type],
+    extensions: EXTENSION_DICT_TYPE,
 ) -> callable:
     """
     Wrap a function so that it casts all castable arguments to a consistent query compiler, and uses the correct extension implementation for methods.
@@ -281,7 +287,7 @@ def wrap_function_in_argument_caster(
         - `classmethod` means we are wrapping a classmethod.
         - `staticmethod` means we are wrapping a staticmethod.
         - `MethodType` means we are wrapping a regular method of a class.
-    cls : Optional[type]
+    extensions : EXTENSION_DICT_TYPE
         The class of the function we are wrapping. This should be None if
         and only if `wrapping_function_type` is None.
 
@@ -290,11 +296,6 @@ def wrap_function_in_argument_caster(
     callable
         The wrapped function.
     """
-    ErrorMessage.catch_bugs_and_request_email(
-        (cls is None and wrapping_function_type is not None)
-        or (cls is not None and wrapping_function_type is None),
-        extra_log="`cls` should be None if and only if `wrapping_function_type` is None",
-    )
 
     @functools.wraps(f)
     def f_with_argument_casting(*args: Tuple, **kwargs: Dict) -> Any:
@@ -312,9 +313,6 @@ def wrap_function_in_argument_caster(
         -------
         Any
         """
-        if len(args) == 0 and len(kwargs) == 0:
-            return
-
         if wrapping_function_type in (classmethod, staticmethod):
             # TODO: currently we don't support any kind of casting or extension
             # for classmethod or staticmethod.
@@ -357,8 +355,7 @@ def wrap_function_in_argument_caster(
             def cast_to_qc(arg):
                 if not (
                     isinstance(arg, QueryCompilerCaster)
-                    and (original_query_compiler := arg._get_query_compiler())
-                    is not None
+                    and arg._get_query_compiler() is not None
                     and arg.get_backend() != result_backend
                 ):
                     return arg
@@ -366,7 +363,7 @@ def wrap_function_in_argument_caster(
                 inplace_update_trackers.append(
                     InplaceUpdateTracker(
                         input_castable=arg,
-                        original_query_compiler=original_query_compiler,
+                        original_query_compiler=cast._get_query_compiler(),
                         new_castable=cast,
                     )
                 )
@@ -376,29 +373,24 @@ def wrap_function_in_argument_caster(
             kwargs = visit_nested_args(kwargs, cast_to_qc)
         else:
             result_backend = Backend.get()
-
-        if wrapping_function_type is None:
-            # TODO(https://github.com/modin-project/modin/issues/7474): dispatch
-            # free functions like pd.concat() to the correct extension.
-            return f(*args, **kwargs)
-
-        # Now we should be dealing with a regular method of a class.
-        ErrorMessage.catch_bugs_and_request_email(
-            wrapping_function_type is not MethodType
-        )
-
-        # When python invokes a method on an object, it passes the object as
-        # the first positional argument.
-        self = args[0]
-        if name in cls._extensions[result_backend]:
-            f_to_apply = cls._extensions[result_backend][name]
+        if name in extensions[result_backend]:
+            f_to_apply = extensions[result_backend][name]
         else:
-            if name not in cls._extensions[None]:
+            if name not in extensions[None]:
                 raise AttributeError(
-                    f"{type(self).__name__} object has no attribute {name}"
+                    (
+                        # When python invokes a method on an object, it passes the object as
+                        # the first positional argument.
+                        (
+                            f"{(type(args[0]).__name__)} object"
+                            if wrapping_function_type is MethodType
+                            else "module 'modin.pandas'"
+                        )
+                        + f" has no attribute {name}"
+                    )
                 )
-            f_to_apply = cls._extensions[None][name]
-        method_result = f_to_apply(*args, **kwargs)
+            f_to_apply = extensions[None][name]
+        result = f_to_apply(*args, **kwargs)
         for (
             original_castable,
             original_qc,
@@ -407,11 +399,14 @@ def wrap_function_in_argument_caster(
             new_qc = new_castable._get_query_compiler()
             if original_qc is not new_qc:
                 new_castable._copy_into(original_castable)
-        return method_result
+        return result
 
     return f_with_argument_casting
 
 
+_GENERAL_EXTENSIONS: EXTENSION_DICT_TYPE = defaultdict(dict)
+
+
 def wrap_free_function_in_argument_caster(name: str) -> callable:
     """
     Get a wrapper for a free function that casts all castable arguments to a consistent query compiler.
@@ -426,9 +421,16 @@ def wrap_free_function_in_argument_caster(name: str) -> callable:
     callable
         A wrapper for a free function that casts all castable arguments to a consistent query compiler.
     """
-    return functools.partial(
-        wrap_function_in_argument_caster,
-        wrapping_function_type=None,
-        cls=None,
-        name=name,
-    )
+
+    def wrapper(f):
+        if name not in _GENERAL_EXTENSIONS[None]:
+            _GENERAL_EXTENSIONS[None][name] = f
+
+        return wrap_function_in_argument_caster(
+            f=f,
+            wrapping_function_type=None,
+            extensions=_GENERAL_EXTENSIONS,
+            name=name,
+        )
+
+    return wrapper
diff --git a/modin/pandas/__init__.py b/modin/pandas/__init__.py
index b7eee337..08fd15ae 100644
--- a/modin/pandas/__init__.py
+++ b/modin/pandas/__init__.py
@@ -27,8 +27,6 @@ if (
         + f" Modin ({__pandas_version__}.X). This may cause undesired side effects!"
     )
 
-# The extensions assigned to this module
-_PD_EXTENSIONS_ = {}
 
 # to not pollute namespace
 del version
@@ -138,6 +136,7 @@ def _initialize_engine(engine_string: str):
 
 
 from modin.pandas import arrays, errors
+from modin.pandas.api.extensions.extensions import __getattr___impl
 from modin.utils import show_versions
 
 from .. import __version__
@@ -195,30 +194,10 @@ from .io import (
 from .plotting import Plotting as plotting
 from .series import Series
 
-
-def __getattr__(name: str):
-    """
-    Overrides getattr on the module to enable extensions.
-
-    Parameters
-    ----------
-    name : str
-        The name of the attribute being retrieved.
-
-    Returns
-    -------
-    Attribute
-        Returns the extension attribute, if it exists, otherwise returns the attribute
-        imported in this file.
-    """
-    try:
-        return _PD_EXTENSIONS_.get(name, globals()[name])
-    except KeyError:
-        raise AttributeError(f"module 'modin.pandas' has no attribute '{name}'")
+__getattr__ = __getattr___impl
 
 
 __all__ = [  # noqa: F405
-    "_PD_EXTENSIONS_",
     "DataFrame",
     "Series",
     "read_csv",
@@ -332,4 +311,5 @@ __all__ = [  # noqa: F405
     "errors",
 ]
 
-del pandas, Parameter
+# Remove these attributes from this module's namespace.
+del pandas, Parameter, __getattr___impl
diff --git a/modin/pandas/api/extensions/extensions.py b/modin/pandas/api/extensions/extensions.py
index 7a0b43f3..23a9b383 100644
--- a/modin/pandas/api/extensions/extensions.py
+++ b/modin/pandas/api/extensions/extensions.py
@@ -12,27 +12,26 @@
 # governing permissions and limitations under the License.
 
 from collections import defaultdict
-from types import MethodType
-from typing import Any, Optional
+from types import MethodType, ModuleType
+from typing import Any, Optional, Union
 
 import modin.pandas as pd
 from modin.config import Backend
 from modin.core.storage_formats.pandas.query_compiler_caster import (
+    _GENERAL_EXTENSIONS,
     _NON_EXTENDABLE_ATTRIBUTES,
+    EXTENSION_DICT_TYPE,
     wrap_function_in_argument_caster,
 )
 
-# This type describes a defaultdict that maps backend name (or `None` for
-# method implementation and not bound to any one extension) to the dictionary of
-# extensions for that backend. The keys of the inner dictionary are the names of
-# the extensions, and the values are the extensions themselves.
-EXTENSION_DICT_TYPE = defaultdict[Optional[str], dict[str, Any]]
-
 _attrs_to_delete_on_test = defaultdict(list)
 
 
 def _set_attribute_on_obj(
-    name: str, extensions: dict, backend: Optional[str], obj: type
+    name: str,
+    extensions: EXTENSION_DICT_TYPE,
+    backend: Optional[str],
+    obj: Union[type, ModuleType],
 ):
     """
     Create a new or override existing attribute on obj.
@@ -41,7 +40,7 @@ def _set_attribute_on_obj(
     ----------
     name : str
         The name of the attribute to assign to `obj`.
-    extensions : dict
+    extensions : EXTENSION_DICT_TYPE
         The dictionary mapping extension name to `new_attr` (assigned below).
     backend : Optional[str]
         The backend to which the accessor applies. If `None`, this accessor
@@ -75,15 +74,17 @@ def _set_attribute_on_obj(
             name
         ] = new_attr
         if callable(new_attr) and name not in dir(obj):
-            # For callable extensions, we add a method to the class that
+            # For callable extensions, we add a method to `obj`'s namespace that
             # dispatches to the correct implementation.
             setattr(
                 obj,
                 name,
                 wrap_function_in_argument_caster(
                     f=new_attr,
-                    wrapping_function_type=MethodType,
-                    cls=obj,
+                    wrapping_function_type=(
+                        MethodType if isinstance(obj, type) else None
+                    ),
+                    extensions=extensions,
                     name=name,
                 ),
             )
@@ -212,7 +213,7 @@ def register_base_accessor(name: str, *, backend: Optional[str] = None):
     )
 
 
-def register_pd_accessor(name: str):
+def register_pd_accessor(name: str, *, backend: Optional[str] = None):
     """
     Registers a pd namespace attribute with the name provided.
 
@@ -239,29 +240,45 @@ def register_pd_accessor(name: str):
     ----------
     name : str
         The name of the attribute to assign to modin.pandas.
+    backend : Optional[str]
+        The backend to which the accessor applies. If ``None``, this accessor
+        will become the default for all backends.
 
     Returns
     -------
     decorator
         Returns the decorator function.
     """
+    return _set_attribute_on_obj(
+        name=name, extensions=_GENERAL_EXTENSIONS, backend=backend, obj=pd
+    )
 
-    def decorator(new_attr: Any):
-        """
-        The decorator for a function or class to be assigned to name
 
-        Parameters
-        ----------
-        new_attr : Any
-            The new attribute to assign to name.
+def __getattr___impl(name: str):
+    """
+    Override __getatttr__ on the modin.pandas module to enable extensions.
 
-        Returns
-        -------
-        new_attr
-            Unmodified new_attr is return from the decorator.
-        """
-        pd._PD_EXTENSIONS_[name] = new_attr
-        setattr(pd, name, new_attr)
-        return new_attr
+    Note that python only falls back to this function if the attribute is not
+    found in this module's namespace.
 
-    return decorator
+    Parameters
+    ----------
+    name : str
+        The name of the attribute being retrieved.
+
+    Returns
+    -------
+    Attribute
+        Returns the extension attribute, if it exists, otherwise returns the attribute
+        imported in this file.
+    """
+
+    from modin.config import Backend
+
+    backend = Backend.get()
+    if name in _GENERAL_EXTENSIONS[backend]:
+        return _GENERAL_EXTENSIONS[backend][name]
+    elif name in _GENERAL_EXTENSIONS[None]:
+        return _GENERAL_EXTENSIONS[None][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 a7fad2e7..6f180535 100644
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -59,10 +59,10 @@ from pandas.util._decorators import doc
 from pandas.util._validators import validate_bool_kwarg
 
 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.api.extensions.extensions import EXTENSION_DICT_TYPE
 from modin.pandas.io import from_non_pandas, from_pandas, to_pandas
 from modin.utils import (
     MODIN_UNNAMED_SERIES_LABEL,
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index 67df6e55..503d7116 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -40,10 +40,12 @@ from pandas.api.types import is_bool, is_list_like
 from pandas.core.dtypes.common import is_bool_dtype, is_integer, is_integer_dtype
 from pandas.core.indexing import IndexingError
 
-from modin.core.storage_formats.pandas.query_compiler_caster import QueryCompilerCaster
+from modin.core.storage_formats.pandas.query_compiler_caster import (
+    EXTENSION_DICT_TYPE,
+    QueryCompilerCaster,
+)
 from modin.error_message import ErrorMessage
 from modin.logging import ClassLogger, disable_logging
-from modin.pandas.api.extensions.extensions import EXTENSION_DICT_TYPE
 from modin.utils import _inherit_docstrings
 
 from .dataframe import DataFrame
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
index 71dea966..892ebaaf 100644
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -33,8 +33,8 @@ from pandas.util._decorators import doc
 from pandas.util._validators import validate_bool_kwarg
 
 from modin.config import PersistentPickle
+from modin.core.storage_formats.pandas.query_compiler_caster import EXTENSION_DICT_TYPE
 from modin.logging import disable_logging
-from modin.pandas.api.extensions.extensions import EXTENSION_DICT_TYPE
 from modin.pandas.io import from_pandas, to_pandas
 from modin.utils import (
     MODIN_UNNAMED_SERIES_LABEL,


Test patch:
diff --git a/modin/conftest.py b/modin/conftest.py
index 536cc1e7..e36ec77e 100644
--- a/modin/conftest.py
+++ b/modin/conftest.py
@@ -14,6 +14,7 @@
 # We turn off mypy type checks in this file because it's not imported anywhere
 # type: ignore
 
+import copy
 import os
 import platform
 import shutil
@@ -56,6 +57,7 @@ import uuid  # noqa: E402
 
 import modin  # noqa: E402
 import modin.config  # noqa: E402
+import modin.pandas as pd  # noqa: E402
 import modin.tests.config  # noqa: E402
 from modin.config import (  # noqa: E402
     AsyncReadMode,
@@ -73,6 +75,9 @@ from modin.core.storage_formats import (  # noqa: E402
     BaseQueryCompiler,
     PandasQueryCompiler,
 )
+from modin.core.storage_formats.pandas.query_compiler_caster import (  # noqa: E402
+    _GENERAL_EXTENSIONS,
+)
 from modin.tests.pandas.utils import (  # noqa: E402
     NROWS,
     _make_csv_file,
@@ -719,3 +724,28 @@ def modify_config(request):
                 key.put(False)
             else:
                 raise e
+
+
+@pytest.fixture(autouse=True)
+def clean_up_extensions():
+
+    original_dataframe_extensions = copy.deepcopy(pd.dataframe.DataFrame._extensions)
+    original_series_extensions = copy.deepcopy(pd.Series._extensions)
+    original_base_extensions = copy.deepcopy(pd.base.BasePandasDataset._extensions)
+    original_general_extensions = copy.deepcopy(_GENERAL_EXTENSIONS)
+    yield
+    pd.dataframe.DataFrame._extensions.clear()
+    pd.dataframe.DataFrame._extensions.update(original_dataframe_extensions)
+    pd.Series._extensions.clear()
+    pd.Series._extensions.update(original_series_extensions)
+    pd.base.BasePandasDataset._extensions.clear()
+    pd.base.BasePandasDataset._extensions.update(original_base_extensions)
+    _GENERAL_EXTENSIONS.clear()
+    _GENERAL_EXTENSIONS.update(original_general_extensions)
+
+    from modin.pandas.api.extensions.extensions import _attrs_to_delete_on_test
+
+    for k, v in _attrs_to_delete_on_test.items():
+        for obj in v:
+            delattr(k, obj)
+    _attrs_to_delete_on_test.clear()
diff --git a/modin/tests/pandas/extensions/conftest.py b/modin/tests/pandas/extensions/conftest.py
index 0573caaa..9b9a0df4 100644
--- a/modin/tests/pandas/extensions/conftest.py
+++ b/modin/tests/pandas/extensions/conftest.py
@@ -11,11 +11,8 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
-import copy
-
 import pytest
 
-import modin.pandas as pd
 from modin.config import Backend, Engine, Execution, StorageFormat
 from modin.core.execution.dispatching.factories import factories
 from modin.core.execution.dispatching.factories.factories import BaseFactory, NativeIO
@@ -39,28 +36,6 @@ class Test1Factory(BaseFactory):
         cls.io_cls = Test1IO
 
 
-@pytest.fixture(autouse=True)
-def clean_up_extensions():
-
-    original_dataframe_extensions = copy.deepcopy(pd.dataframe.DataFrame._extensions)
-    original_series_extensions = copy.deepcopy(pd.Series._extensions)
-    original_base_extensions = copy.deepcopy(pd.base.BasePandasDataset._extensions)
-    yield
-    pd.dataframe.DataFrame._extensions.clear()
-    pd.dataframe.DataFrame._extensions.update(original_dataframe_extensions)
-    pd.Series._extensions.clear()
-    pd.Series._extensions.update(original_series_extensions)
-    pd.base.BasePandasDataset._extensions.clear()
-    pd.base.BasePandasDataset._extensions.update(original_base_extensions)
-
-    from modin.pandas.api.extensions.extensions import _attrs_to_delete_on_test
-
-    for k, v in _attrs_to_delete_on_test.items():
-        for obj in v:
-            delattr(k, obj)
-    _attrs_to_delete_on_test.clear()
-
-
 @pytest.fixture
 def Backend1():
     factories.Test1_Storage_FormatOnTest1_EngineFactory = Test1Factory
diff --git a/modin/tests/pandas/extensions/test_pd_extensions.py b/modin/tests/pandas/extensions/test_pd_extensions.py
index b027ac40..10d1b8f5 100644
--- a/modin/tests/pandas/extensions/test_pd_extensions.py
+++ b/modin/tests/pandas/extensions/test_pd_extensions.py
@@ -11,27 +11,126 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
+import re
+
+import pandas
+import pytest
+
 import modin.pandas as pd
+from modin.config import context as config_context
 from modin.pandas.api.extensions import register_pd_accessor
+from modin.tests.pandas.utils import df_equals, eval_general
+
+
+class TestRegisterForAllBackends:
+    def test_add_new_function(self):
+        expected_string_val = "Some string value"
+        method_name = "new_method"
+
+        @register_pd_accessor(method_name)
+        def my_method_implementation():
+            return expected_string_val
+
+        assert pd.new_method() == expected_string_val
+
+    def test_add_new_non_method(self):
+        expected_val = 4
+        attribute_name = "four"
+        register_pd_accessor(attribute_name)(expected_val)
+        assert pd.four == expected_val
+
+
+class TestRegisterForOneBackend:
+    def test_add_new_function(self):
+        backend = "Pandas"
+        expected_string_val = "Some string value"
+        method_name = "new_method"
+
+        @register_pd_accessor(method_name, backend=backend)
+        def my_method_implementation():
+            return expected_string_val
+
+        with config_context(Backend=backend):
+            assert getattr(pd, method_name)() == expected_string_val
+        with config_context(Backend="Python_Test"):
+            with pytest.raises(
+                AttributeError,
+                match=re.escape(
+                    f"module 'modin.pandas' has no attribute {method_name}"
+                ),
+            ):
+                getattr(pd, method_name)()
+
+    def test_override_function(self):
+        backend = "Pandas"
+        expected_string_val = "Some string value"
+
+        @register_pd_accessor("to_datetime", backend=backend)
+        def my_method_implementation(*args, **kwargs):
+            return expected_string_val
+
+        with config_context(Backend=backend):
+            # Since there are no query compiler inputs to to_datetime(), use
+            # the to_datetime() implementation for Backend.get()
+            assert pd.to_datetime(1) == expected_string_val
+
+        with config_context(Backend="Python_Test"):
+            # There are no query compiler inputs to to_datetime(), and
+            # the current Backend.get() does not have a to_datetime() extension,
+            # so fall back to the default to_datetime() implementation, which
+            # should return the same result as pandas.to_datetime().
+            eval_general(pd, pandas, lambda lib: lib.to_datetime(1))
+
+    def test_add_new_non_method(self):
+        backend = "Pandas"
+        expected_val = 4
+        attribute_name = "four"
+        register_pd_accessor(attribute_name, backend=backend)(expected_val)
+        with config_context(Backend=backend):
+            assert pd.four == expected_val
+        with config_context(Backend="Python_Test"):
+            assert not hasattr(pd, attribute_name)
+
+    def test_to_datetime_dispatches_to_implementation_for_input(self):
+
+        @register_pd_accessor("to_datetime", backend="Pandas")
+        def pandas_to_datetime(*args, **kwargs):
+            return "pandas_to_datetime_result"
+
+        with config_context(Backend="Pandas"):
+            pandas_backend_series = pd.Series(1)
+
+        with config_context(Backend="Python_Test"):
+            python_backend_df = pd.Series(1)
 
+        assert pd.to_datetime(pandas_backend_series) == "pandas_to_datetime_result"
+        df_equals(
+            pd.to_datetime(python_backend_df),
+            pandas.to_datetime(python_backend_df._to_pandas()),
+        )
 
-def test_dataframe_extension_simple_method():
-    expected_string_val = "Some string value"
-    method_name = "new_method"
+    def test_concat_with_two_different_backends(self):
+        with config_context(Backend="Pandas"):
+            modin_on_pandas_df = pd.DataFrame({"a": [1, 2, 3]})
+        with config_context(Backend="Python_Test"):
+            modin_on_python_df = pd.DataFrame({"a": [4, 5, 6]})
 
-    @register_pd_accessor(method_name)
-    def my_method_implementation():
-        return expected_string_val
+        @register_pd_accessor("concat", backend="Pandas")
+        def pandas_concat(*args, **kwargs):
+            return "pandas_concat_result"
 
-    assert method_name in pd._PD_EXTENSIONS_.keys()
-    assert pd._PD_EXTENSIONS_[method_name] is my_method_implementation
-    assert pd.new_method() == expected_string_val
+        @register_pd_accessor("concat", backend="Python_Test")
+        def python_concat(*args, **kwargs):
+            return "python_concat_result"
 
+        # If the backends are different, we dispatch to the concat() override
+        # for the backend of the first argument.
+        assert (
+            pd.concat([modin_on_pandas_df, modin_on_python_df])
+            == "pandas_concat_result"
+        )
 
-def test_dataframe_extension_non_method():
-    expected_val = 4
-    attribute_name = "four"
-    register_pd_accessor(attribute_name)(expected_val)
-    assert attribute_name in pd._PD_EXTENSIONS_.keys()
-    assert pd._PD_EXTENSIONS_[attribute_name] == 4
-    assert pd.four == expected_val
+        assert (
+            pd.concat([modin_on_python_df, modin_on_pandas_df])
+            == "python_concat_result"
+        )
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 3af66cf8..fd81b9ae 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -26,6 +26,7 @@ from modin.core.storage_formats.base.query_compiler_calculator import (
     BackendCostCalculator,
 )
 from modin.core.storage_formats.pandas.native_query_compiler import NativeQueryCompiler
+from modin.pandas.api.extensions import register_pd_accessor
 from modin.tests.pandas.utils import df_equals
 
 
@@ -254,6 +255,17 @@ def test_cast_to_second_backend_with_concat(pico_df, cluster_df):
     assert df3.get_backend() == "Cluster"  # result should be on cluster
 
 
+def test_cast_to_second_backend_with_concat_uses_second_backend_api_override(
+    pico_df, cluster_df
+):
+    register_pd_accessor(name="concat", backend="Cluster")(
+        lambda *args, **kwargs: "custom_concat_result"
+    )
+    assert pd.concat([pico_df, cluster_df], axis=1) == "custom_concat_result"
+    assert pico_df.get_backend() == "Pico"
+    assert cluster_df.get_backend() == "Cluster"
+
+
 def test_moving_pico_to_cluster_in_place_calls_set_backend_only_once_github_issue_7490(
     pico_df, cluster_df
 ):
@@ -279,6 +291,17 @@ def test_cast_to_first_backend(pico_df, cluster_df):
     assert df3.get_backend() == cluster_df.get_backend()  # result should be on cluster
 
 
+def test_cast_to_first_backend_with_concat_uses_first_backend_api_override(
+    pico_df, cluster_df
+):
+    register_pd_accessor(name="concat", backend="Cluster")(
+        lambda *args, **kwargs: "custom_concat_result"
+    )
+    assert pd.concat([cluster_df, pico_df], axis=1) == "custom_concat_result"
+    assert pico_df.get_backend() == "Pico"
+    assert cluster_df.get_backend() == "Cluster"
+
+
 def test_cast_to_first_backend_with___init__(pico_df, cluster_df):
     df3 = pd.DataFrame(
         {