FEAT-#7472: Add an option register dataframe and series accessors with a particular backend. (#7473)

Add an option to associate each base, dataframe, and series extension with a particular
backend. Also, add `register_base_extension()`.

Wrap each base, dataframe and series method in a dispatcher that checks
the backend of its first argument and then dispatches to the
corresponding backend's extension, if it exists. For non-method
extensions, including properties and scalars (like ints), we have to
extend `__getatribute__`, `__getattr__`, `__delattr__`, and
`__setattr__` to implement the extension.

In the future, the dispatcher could choose an appropriate backend when
it has inputs from multiple backends. It could then convert all
arguments to the correct backend and then choose the extension matching
that 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 #7472
- [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 -->

---------

Signed-off-by: sfc-gh-mvashishtha <[email protected]>
Co-authored-by: Jonathan Shi <[email protected]>

Author: sfc-gh-mvashishtha

.github/workflows/ci.yml

@@ -157,6 +157,7 @@ jobs:
       - run: python -m pytest modin/tests/test_dataframe_api_standard.py
       - run: python -m pytest modin/tests/test_logging.py
       - run: python -m pytest modin/tests/test_metrics.py
+      - run: python -m pytest modin/tests/pandas/extensions
       - uses: ./.github/actions/upload-coverage
 
   test-defaults:

modin/conftest.py

@@ -239,7 +239,9 @@ def get_unique_base_execution():
     execution_name = f"{format_name}On{engine_name}"
 
     # Dynamically building all the required classes to form a new execution
-    base_qc = type(format_name, (TestQC,), {})
+    base_qc = type(
+        format_name, (TestQC,), {"get_backend": (lambda self: execution_name)}
+    )
     base_io = type(
         f"{execution_name}IO", (BaseOnPythonIO,), {"query_compiler_cls": base_qc}
     )
@@ -252,7 +254,7 @@ def get_unique_base_execution():
     # Setting up the new execution
     setattr(factories, f"{execution_name}Factory", base_factory)
     Backend.register_backend(
-        "execution_name", Execution(engine=engine_name, storage_format=format_name)
+        execution_name, Execution(engine=engine_name, storage_format=format_name)
     )
     old_engine, old_format = modin.set_execution(
         engine=engine_name, storage_format=format_name

modin/core/storage_formats/base/query_compiler.py

@@ -31,6 +31,7 @@
 from pandas._typing import DtypeBackend, IndexLabel, Suffixes
 from pandas.core.dtypes.common import is_number, is_scalar
 
+from modin.config import Backend, Execution
 from modin.core.dataframe.algebra.default2pandas import (
     BinaryDefault,
     CatDefault,
@@ -52,6 +53,7 @@
 from modin.error_message import ErrorMessage
 from modin.logging import ClassLogger
 from modin.logging.config import LogLevel
+from modin.logging.logger_decorator import disable_logging
 from modin.utils import MODIN_UNNAMED_SERIES_LABEL, try_cast_to_pandas
 
 from . import doc_utils
@@ -207,6 +209,23 @@ def _maybe_warn_on_default(self, *, message: str = "", reason: str = "") -> None
         if self._should_warn_on_default_to_pandas:
             ErrorMessage.default_to_pandas(message=message, reason=reason)
 
+    @disable_logging
+    def get_backend(self) -> str:
+        """
+        Get the backend for this query compiler.
+
+        Returns
+        -------
+        str
+            The backend for this query compiler.
+        """
+        return Backend.get_backend_for_execution(
+            Execution(
+                engine=self.engine,
+                storage_format=self.storage_format,
+            )
+        )
+
     @property
     @abc.abstractmethod
     def storage_format(self) -> str:

modin/pandas/api/extensions/__init__.py

@@ -12,12 +12,14 @@
 # governing permissions and limitations under the License.
 
 from .extensions import (
+    register_base_accessor,
     register_dataframe_accessor,
     register_pd_accessor,
     register_series_accessor,
 )
 
 __all__ = [
+    "register_base_accessor",
     "register_dataframe_accessor",
     "register_series_accessor",
     "register_pd_accessor",

modin/pandas/api/extensions/extensions.py

@@ -11,14 +11,40 @@
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
 
-from types import ModuleType
-from typing import Any, Union
+import inspect
+from collections import defaultdict
+from functools import wraps
+from typing import Any, Callable, Optional
 
 import modin.pandas as pd
+from modin.config import Backend
+
+# 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)
+
+_NON_EXTENDABLE_ATTRIBUTES = (
+    # we use these attributes to implement the extension system, so we can't
+    # allow extensions to override them.
+    "__getattribute__",
+    "__setattr__",
+    "__delattr__",
+    "get_backend",
+    "set_backend",
+    "__getattr__",
+    "_get_extension",
+    "_getattribute__from_extension_impl",
+    "_getattr__from_extension_impl",
+    "_query_compiler",
+)
 
 
 def _set_attribute_on_obj(
-    name: str, extensions_dict: dict, obj: Union[pd.DataFrame, pd.Series, ModuleType]
+    name: str, extensions: dict, backend: Optional[str], obj: type
 ):
     """
     Create a new or override existing attribute on obj.
@@ -27,8 +53,11 @@ def _set_attribute_on_obj(
     ----------
     name : str
         The name of the attribute to assign to `obj`.
-    extensions_dict : dict
+    extensions : dict
         The dictionary mapping extension name to `new_attr` (assigned below).
+    backend : Optional[str]
+        The backend to which the accessor applies. If `None`, this accessor
+        will become the default for all backends.
     obj : DataFrame, Series, or modin.pandas
         The object we are assigning the new attribute to.
 
@@ -37,10 +66,12 @@ def _set_attribute_on_obj(
     decorator
         Returns the decorator function.
     """
+    if name in _NON_EXTENDABLE_ATTRIBUTES:
+        raise ValueError(f"Cannot register an extension with the reserved name {name}.")
 
     def decorator(new_attr: Any):
         """
-        The decorator for a function or class to be assigned to name
+        Decorate a function or class to be assigned to the given name.
 
         Parameters
         ----------
@@ -52,14 +83,24 @@ def decorator(new_attr: Any):
         new_attr
             Unmodified new_attr is return from the decorator.
         """
-        extensions_dict[name] = new_attr
-        setattr(obj, name, new_attr)
+        extensions[None if backend is None else Backend.normalize(backend)][
+            name
+        ] = new_attr
+        if callable(new_attr) and name not in dir(obj):
+            # For callable extensions, we add a method to the class that
+            # dispatches to the correct implementation.
+            setattr(
+                obj,
+                name,
+                wrap_method_in_backend_dispatcher(name, new_attr, extensions),
+            )
+            _attrs_to_delete_on_test[obj].append(name)
         return new_attr
 
     return decorator
 
 
-def register_dataframe_accessor(name: str):
+def register_dataframe_accessor(name: str, *, backend: Optional[str] = None):
     """
     Registers a dataframe attribute with the name provided.
 
@@ -88,13 +129,16 @@ def my_new_dataframe_method(*args, **kwargs):
     -------
     decorator
         Returns the decorator function.
+    backend : Optional[str]
+        The backend to which the accessor applies. If ``None``, this accessor
+        will become the default for all backends.
     """
     return _set_attribute_on_obj(
-        name, pd.dataframe._DATAFRAME_EXTENSIONS_, pd.DataFrame
+        name, pd.dataframe._DATAFRAME_EXTENSIONS_, backend, pd.dataframe.DataFrame
     )
 
 
-def register_series_accessor(name: str):
+def register_series_accessor(name: str, *, backend: Optional[str] = None):
     """
     Registers a series attribute with the name provided.
 
@@ -118,13 +162,61 @@ def my_new_series_method(*args, **kwargs):
     ----------
     name : str
         The name of the attribute to assign to Series.
+    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, pd.series._SERIES_EXTENSIONS_, backend=backend, obj=pd.series.Series
+    )
+
+
+def register_base_accessor(name: str, *, backend: Optional[str] = None):
+    """
+    Register a base attribute with the name provided.
+
+    This is a decorator that assigns a new attribute to BasePandasDataset. It can be used
+    with the following syntax:
+
+    ```
+    @register_base_accessor("new_method")
+    def register_base_accessor(*args, **kwargs):
+        # logic goes here
+        return
+    ```
+
+    The new attribute can then be accessed with the name provided:
+
+    ```
+    s.new_method(*my_args, **my_kwargs)
+    ```
+
+    Parameters
+    ----------
+    name : str
+        The name of the attribute to assign to BasePandasDataset.
+    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, pd.series._SERIES_EXTENSIONS_, pd.Series)
+    import modin.pandas.base
+
+    return _set_attribute_on_obj(
+        name,
+        modin.pandas.base._BASE_EXTENSIONS,
+        backend=backend,
+        obj=modin.pandas.base.BasePandasDataset,
+    )
 
 
 def register_pd_accessor(name: str):
@@ -160,4 +252,128 @@ def my_new_pd_function(*args, **kwargs):
     decorator
         Returns the decorator function.
     """
-    return _set_attribute_on_obj(name, pd._PD_EXTENSIONS_, 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.
+
+        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
+
+    return decorator
+
+
+def wrap_method_in_backend_dispatcher(
+    name: str, method: Callable, extensions: EXTENSION_DICT_TYPE
+) -> Callable:
+    """
+    Wraps a method to dispatch to the correct backend implementation.
+
+    This function is a wrapper that is used to dispatch to the correct backend
+    implementation of a method.
+
+    Parameters
+    ----------
+    name : str
+        The name of the method being wrapped.
+    method : Callable
+        The method being wrapped.
+    extensions : EXTENSION_DICT_TYPE
+        The extensions dictionary for the class this method is defined on.
+
+    Returns
+    -------
+    Callable
+        Returns the wrapped function.
+    """
+
+    @wraps(method)
+    def method_dispatcher(*args, **kwargs):
+        if len(args) == 0 and len(kwargs) == 0:
+            # Handle some cases like __init__()
+            return method(*args, **kwargs)
+        # TODO(https://github.com/modin-project/modin/issues/7470): this
+        # method may take dataframes and series backed by different backends
+        # as input, e.g. if we are here because of a call like
+        # `pd.DataFrame().set_backend('python_test').merge(pd.DataFrame().set_backend('pandas'))`.
+        # In that case, we should determine which backend to cast to, cast all
+        # arguments to that backend, and then choose the appropriate extension
+        # method, if it exists.
+
+        # Assume that `self` is the first argument.
+        self = args[0]
+        remaining_args = args[1:]
+        if (
+            hasattr(self, "_query_compiler")
+            and self.get_backend() in extensions
+            and name in extensions[self.get_backend()]
+        ):
+            # If `self` is using a query compiler whose backend has an
+            # extension for this method, use that extension.
+            return extensions[self.get_backend()][name](self, *remaining_args, **kwargs)
+        else:
+            # Otherwise, use the default implementation.
+            if name not in extensions[None]:
+                raise AttributeError(
+                    f"{type(self).__name__} object has no attribute {name}"
+                )
+            return extensions[None][name](self, *remaining_args, **kwargs)
+
+    return method_dispatcher
+
+
+def wrap_class_methods_in_backend_dispatcher(
+    extensions: EXTENSION_DICT_TYPE,
+) -> Callable:
+    """
+    Get a function that can wrap a class's instance methods so that they dispatch to the correct backend.
+
+    Parameters
+    ----------
+    extensions : EXTENSION_DICT_TYPE
+        The extension dictionary for the class.
+
+    Returns
+    -------
+    Callable
+        The class wrapper.
+    """
+
+    def wrap_methods(cls: type):
+        # We want to avoid wrapping synonyms like __add__() and add() with
+        # different wrappers, so keep a dict mapping methods we've wrapped
+        # to their wrapped versions.
+        already_seen_to_wrapped: dict[Callable, Callable] = {}
+        for method_name, method_value in inspect.getmembers(
+            cls, predicate=inspect.isfunction
+        ):
+            if method_value in already_seen_to_wrapped:
+                setattr(cls, method_name, already_seen_to_wrapped[method_value])
+                continue
+            elif method_name not in _NON_EXTENDABLE_ATTRIBUTES:
+                extensions[None][method_name] = method_value
+                wrapped = wrap_method_in_backend_dispatcher(
+                    method_name, method_value, extensions
+                )
+                if method_name not in cls.__dict__:
+                    # If this class's method comes from a superclass (i.e.
+                    # it's not in cls.__dict__), mark it so that
+                    # modin.utils._inherit_docstrings knows that the method
+                    # must get its docstrings from its superclass.
+                    wrapped._wrapped_superclass_method = method_value
+                setattr(cls, method_name, wrapped)
+                already_seen_to_wrapped[method_value] = getattr(cls, method_name)
+        return cls
+
+    return wrap_methods