Repository: modin-project/modin (python)

Issue:
FEAT: Support pinning objects to a backend
When automatic backend switching is enabled, users should be able to opt out via a `pin_backend` function on inidvidual objects.

Pull request:
FEAT-#7521: Support pinning objects to a 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?

Adds `is_backend_pinned`, `pin_backend`, and `unpin_backend` methods to the BasePandasDataset, which prevent data from automatically moving to other backends outside explicit `set_backend` or other materialization operations. Pinning an object will prevent all subsequent operations from doing heuristic-based backend switching; a pinned object will remain pinned until an in-place `unpin_backend` or `set_backend` is performed. Pinning an object does not prevent materialization for operations like `__repr__` or `to_pandas()`, which are invoked explicitly.

n-ary operations that have the potential to mix multiple pinned/unpinned backends of different engines have not been tested. In the future, we should update the calculator's behavior to pin to a choice backend if any objects are pinned to the same backend, and error out if multiple objects are pinned to different backends.

- [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 #7521 <!-- 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/core/storage_formats/pandas/query_compiler_caster.py b/modin/core/storage_formats/pandas/query_compiler_caster.py
index 288b2586..88477d37 100644
--- a/modin/core/storage_formats/pandas/query_compiler_caster.py
+++ b/modin/core/storage_formats/pandas/query_compiler_caster.py
@@ -31,7 +31,7 @@ import pandas
 from pandas.core.indexes.frozen import FrozenList
 from typing_extensions import Self
 
-from modin.config import Backend
+from modin.config import AutoSwitchBackend, Backend
 from modin.config import context as config_context
 from modin.core.storage_formats.base.query_compiler import (
     BaseQueryCompiler,
@@ -68,6 +68,11 @@ _NON_EXTENDABLE_ATTRIBUTES = {
     "_query_compiler",
     "_get_query_compiler",
     "_copy_into",
+    "is_backend_pinned",
+    "_set_backend_pinned",
+    "pin_backend",
+    "unpin_backend",
+    "__dict__",
 }
 
 BackendAndClassName = namedtuple("BackendAndClassName", ["backend", "class_name"])
@@ -119,6 +124,70 @@ class QueryCompilerCaster(ABC):
         """
         pass
 
+    @abstractmethod
+    def is_backend_pinned(self) -> bool:
+        """
+        Get whether this object's data is pinned to a particular backend.
+
+        Returns
+        -------
+        bool
+            True if the data is pinned.
+        """
+        pass
+
+    @abstractmethod
+    def _set_backend_pinned(self, pinned: bool, inplace: bool) -> Optional[Self]:
+        """
+        Update whether this object's data is pinned to a particular backend.
+
+        Parameters
+        ----------
+        pinned : bool
+            Whether the data is pinned.
+
+        inplace : bool, default: False
+            Whether to update the object in place.
+
+        Returns
+        -------
+        Optional[Self]
+            The object with the new pin state, if `inplace` is False. Otherwise, None.
+        """
+        pass
+
+    def pin_backend(self, inplace: bool = False) -> Optional[Self]:
+        """
+        Pin the object's underlying data, preventing Modin from automatically moving it to another backend.
+
+        Parameters
+        ----------
+        inplace : bool, default: False
+            Whether to update the object in place.
+
+        Returns
+        -------
+        Optional[Self]
+            The newly-pinned object, if `inplace` is False. Otherwise, None.
+        """
+        return self._set_backend_pinned(True, inplace)
+
+    def unpin_backend(self, inplace: bool = False) -> Optional[Self]:
+        """
+        Unpin the object's underlying data, allowing Modin to automatically move it to another backend.
+
+        Parameters
+        ----------
+        inplace : bool, default: False
+            Whether to update the object in place.
+
+        Returns
+        -------
+        Optional[Self]
+            The newly-unpinned object, if `inplace` is False. Otherwise, None.
+        """
+        return self._set_backend_pinned(False, inplace)
+
     @abstractmethod
     def get_backend(self) -> str:
         """
@@ -379,6 +448,7 @@ def _maybe_switch_backend_post_op(
     qc_list: list[BaseQueryCompiler],
     starting_backend: str,
     class_of_wrapped_fn: Optional[str],
+    pin_backend: bool,
     arguments: MappingProxyType[str, Any],
 ) -> Any:
     """
@@ -403,6 +473,8 @@ def _maybe_switch_backend_post_op(
     class_of_wrapped_fn : Optional[str]
         The name of the class that the function belongs to. `None` for functions
         in the modin.pandas module.
+    pin_backend : bool
+        Whether the result should have its backend pinned, and therefore not moved.
     arguments : MappingProxyType[str, Any]
         Mapping from operation argument names to their values.
 
@@ -411,6 +483,11 @@ def _maybe_switch_backend_post_op(
     Any
         The result of the function, possibly with its backend switched.
     """
+    # If any input QC was pinned, then the output should be as well.
+    if pin_backend:
+        if isinstance(result, QueryCompilerCaster):
+            result.pin_backend(inplace=True)
+        return result
     if (
         # only apply post-operation switch to nullary and unary methods
         len(qc_list) in (0, 1)
@@ -611,6 +688,8 @@ def wrap_function_in_argument_caster(
     """
     Wrap a function so that it casts all castable arguments to a consistent query compiler, and uses the correct extension implementation for methods.
 
+    Also propagates pin behavior across operations.
+
     Parameters
     ----------
     klass : Optional[type]
@@ -674,11 +753,22 @@ def wrap_function_in_argument_caster(
 
         input_query_compilers: list[BaseQueryCompiler] = []
 
+        pin_target_backend = None
+
         def register_query_compilers(arg):
+            nonlocal pin_target_backend
             if (
                 isinstance(arg, QueryCompilerCaster)
                 and (qc := arg._get_query_compiler()) is not None
             ):
+                arg_backend = arg.get_backend()
+                if pin_target_backend is not None:
+                    if arg.is_backend_pinned() and arg_backend != pin_target_backend:
+                        raise ValueError(
+                            f"Cannot combine arguments that are pinned to conflicting backends ({pin_target_backend}, {arg_backend})"
+                        )
+                elif arg.is_backend_pinned():
+                    pin_target_backend = arg_backend
                 input_query_compilers.append(qc)
             elif isinstance(arg, BaseQueryCompiler):
                 # We might get query compiler arguments in __init__()
@@ -688,6 +778,22 @@ def wrap_function_in_argument_caster(
         visit_nested_args(args, register_query_compilers)
         visit_nested_args(kwargs, register_query_compilers)
 
+        # Before determining any automatic switches, we perform the following checks:
+        # 1. If the global AutoSwitchBackend configuration variable is set to False, do not switch.
+        # 2. If there's only one query compiler and it's pinned, do not switch.
+        # 3. If there are multiple query compilers, and at least one is pinned to a particular
+        #    backend, then switch to that backend.
+        # 4. If there are multiple query compilers, at least two of which are pinned to distinct
+        #    backends, raise a ValueError.
+
+        if not AutoSwitchBackend.get() or (
+            len(input_query_compilers) < 2 and pin_target_backend is not None
+        ):
+            result = f(*args, **kwargs)
+            if isinstance(result, QueryCompilerCaster):
+                result._set_backend_pinned(True, inplace=True)
+            return result
+
         if len(input_query_compilers) == 0:
             # For nullary functions, we need to create a dummy query compiler
             # to calculate the cost of switching backends.
@@ -719,6 +825,7 @@ def wrap_function_in_argument_caster(
         args_dict = MappingProxyType(bound_arguments.arguments)
 
         if len(input_query_compilers) < 2:
+            # No need to check should_pin_result() again, since we have already done so above.
             result_backend, cast_to_qc = _maybe_switch_backend_pre_op(
                 name,
                 input_qc=input_qc_for_pre_op_switch,
@@ -735,7 +842,10 @@ def wrap_function_in_argument_caster(
             for qc in input_query_compilers:
                 calculator.add_query_compiler(qc)
 
-            result_backend = calculator.calculate()
+            if pin_target_backend is None:
+                result_backend = calculator.calculate()
+            else:
+                result_backend = pin_target_backend
 
             def cast_to_qc(arg):
                 if not (
@@ -786,6 +896,7 @@ def wrap_function_in_argument_caster(
             qc_list=input_query_compilers,
             starting_backend=result_backend,
             class_of_wrapped_fn=class_of_wrapped_fn,
+            pin_backend=pin_target_backend is not None,
             arguments=args_dict,
         )
 
diff --git a/modin/pandas/base.py b/modin/pandas/base.py
index 888db2a4..cea33bfb 100644
--- a/modin/pandas/base.py
+++ b/modin/pandas/base.py
@@ -113,6 +113,11 @@ _EXTENSION_NO_LOOKUP = {
     "_getattr__from_extension_impl",
     "_get_query_compiler",
     "set_backend",
+    "_pinned",
+    "is_backend_pinned",
+    "_set_backend_pinned",
+    "pin_backend",
+    "unpin_backend",
 }
 
 # Do not lookup certain attributes in columns or index, as they're used for some
@@ -231,6 +236,7 @@ class BasePandasDataset(QueryCompilerCaster, ClassLogger):
     _siblings: list[BasePandasDataset]
 
     _extensions: EXTENSION_DICT_TYPE = EXTENSION_DICT_TYPE(dict)
+    _pinned: bool = False
 
     @cached_property
     def _is_dataframe(self) -> bool:
@@ -4445,6 +4451,50 @@ class BasePandasDataset(QueryCompilerCaster, ClassLogger):
     # namespace for additional Modin functions that are not available in Pandas
     modin: ModinAPI = CachedAccessor("modin", ModinAPI)
 
+    @disable_logging
+    def is_backend_pinned(self) -> bool:
+        """
+        Get whether this object's data is pinned to a particular backend.
+
+        Returns
+        -------
+        bool
+            True if the data is pinned.
+        """
+        return self._pinned
+
+    def _set_backend_pinned(
+        self, pinned: bool, inplace: bool = False
+    ) -> Optional[Self]:
+        """
+        Update whether this object's data is pinned to a particular backend.
+
+        Parameters
+        ----------
+        pinned : bool
+            Whether the data is pinned.
+
+        inplace : bool, default: False
+            Whether to update the object in place.
+
+        Returns
+        -------
+        Optional[Self]
+            The object with the new pin state, if `inplace` is False. Otherwise, None.
+        """
+        change = (self.is_backend_pinned() and not pinned) or (
+            not self.is_backend_pinned() and pinned
+        )
+        if inplace:
+            self._pinned = pinned
+            return None
+        else:
+            if change:
+                new_obj = self.__constructor__(query_compiler=self._query_compiler)
+                new_obj._pinned = pinned
+                return new_obj
+            return self
+
     @doc(SET_BACKEND_DOC, class_name=__qualname__)
     def set_backend(self, backend: str, inplace: bool = False) -> Optional[Self]:
         # TODO(https://github.com/modin-project/modin/issues/7467): refactor
@@ -4481,6 +4531,8 @@ class BasePandasDataset(QueryCompilerCaster, ClassLogger):
             pass
         if inplace:
             self._update_inplace(query_compiler)
+            # Always unpin after an explicit set_backend operation
+            self._pinned = False
             return None
         else:
             return self.__constructor__(query_compiler=query_compiler)
diff --git a/modin/pandas/indexing.py b/modin/pandas/indexing.py
index 72a2c5f0..9902befb 100644
--- a/modin/pandas/indexing.py
+++ b/modin/pandas/indexing.py
@@ -292,6 +292,45 @@ class _LocationIndexerBase(QueryCompilerCaster, ClassLogger):
     qc: BaseQueryCompiler
     _extensions: EXTENSION_DICT_TYPE = EXTENSION_DICT_TYPE(dict)
 
+    def is_backend_pinned(self) -> bool:
+        """
+        Get whether this object's data is pinned to a particular backend.
+
+        Returns
+        -------
+        bool
+            True if the data is pinned.
+        """
+        return self.df.is_backend_pinned()
+
+    def _set_backend_pinned(self, pinned: bool, inplace: bool = False):
+        """
+        Update whether this object's data is pinned to a particular backend.
+
+        Parameters
+        ----------
+        pinned : bool
+            Whether the data is pinned.
+
+        inplace : bool, default: False
+            Whether to update the object in place.
+
+        Returns
+        -------
+        Optional[Self]
+            The object with the new pin state, if `inplace` is False. Otherwise, None.
+        """
+        change = (self.is_backend_pinned() and not pinned) or (
+            not self.is_backend_pinned() and pinned
+        )
+        if not change:
+            return None if inplace else self
+        result = type(self)(self.df._set_backend_pinned(pinned))
+        if inplace:
+            result._copy_into(self)
+            return None
+        return result
+
     @disable_logging
     @_inherit_docstrings(QueryCompilerCaster.set_backend)
     def set_backend(self, backend, inplace: bool = False):
@@ -316,6 +355,7 @@ class _LocationIndexerBase(QueryCompilerCaster, ClassLogger):
     def _copy_into(self, other: Series):
         other.qc = self.df._query_compiler
         other.df._update_inplace(new_query_compiler=self.df._query_compiler)
+        other.df._set_backend_pinned(self.is_backend_pinned())
         return None
 
     def __init__(self, modin_df: Union[DataFrame, Series]):


Test patch:
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 e925d10a..92d37f81 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -23,6 +23,7 @@ import pytest
 from pytest import param
 
 import modin.pandas as pd
+from modin.config import context as config_context
 from modin.config.envvars import Backend, Engine, Execution
 from modin.core.execution.dispatching.factories import factories
 from modin.core.execution.dispatching.factories.factories import BaseFactory
@@ -605,12 +606,15 @@ def test_setitem_in_place_with_self_switching_backend(cloud_df, local_df):
     assert cloud_df.get_backend() == "Cloud"
 
 
-def test_switch_local_to_cloud_with_iloc___setitem__(local_df, cloud_df):
+@pytest.mark.parametrize("pin_local", [True, False], ids=["pinned", "unpinned"])
+def test_switch_local_to_cloud_with_iloc___setitem__(local_df, cloud_df, pin_local):
+    if pin_local:
+        local_df = local_df.pin_backend()
     local_df.iloc[:, 0] = cloud_df.iloc[:, 0] + 1
     expected_pandas = local_df._to_pandas()
     expected_pandas.iloc[:, 0] = cloud_df._to_pandas().iloc[:, 0] + 1
     df_equals(local_df, expected_pandas)
-    assert local_df.get_backend() == "Cloud"
+    assert local_df.get_backend() == "Local_machine" if pin_local else "Cloud"
 
 
 # Outlines a future generic function for determining when to stay
@@ -691,6 +695,49 @@ class TestSwitchBackendPostOpDependingOnDataSize:
             assert df.get_backend() == "Big_Data_Cloud"
             assert df.sum().get_backend() == "Small_Data_Local"
 
+    def test_agg_pinned(self):
+        # The operation in test_agg would naturally cause an automatic switch, but the
+        # absence of AutoSwitchBackend or the presence of a pin on the frame prevent this
+        # switch from happening.
+        with backend_test_context(
+            test_backend="Big_Data_Cloud",
+            choices=("Big_Data_Cloud", "Small_Data_Local"),
+        ):
+            register_function_for_post_op_switch(
+                class_name="DataFrame", backend="Big_Data_Cloud", method="sum"
+            )
+            # No pin or config, should switch
+            df = pd.DataFrame([[1, 2], [3, 4]])
+            assert df.get_backend() == "Big_Data_Cloud"
+            assert df.sum().get_backend() == "Small_Data_Local"
+            # config set to false, should not switch
+            with config_context(AutoSwitchBackend=False):
+                df = pd.DataFrame([[1, 2], [3, 4]])
+                assert df.get_backend() == "Big_Data_Cloud"
+                assert df.sum().get_backend() == "Big_Data_Cloud"
+            # no config, but data is pinned
+            df = pd.DataFrame([[1, 2], [3, 4]]).pin_backend()
+            assert df.get_backend() == "Big_Data_Cloud"
+            assert df.sum().get_backend() == "Big_Data_Cloud"
+            # a frame-level pin remains valid across a transformation
+            df_copy = df + 1
+            assert df_copy.get_backend() == "Big_Data_Cloud"
+            assert df_copy.sum().get_backend() == "Big_Data_Cloud"
+            # unpinning df allows a switch again
+            df = df.unpin_backend()
+            assert df.get_backend() == "Big_Data_Cloud"
+            assert df.sum().get_backend() == "Small_Data_Local"
+            df_copy = df + 1
+            assert df_copy.get_backend() == "Big_Data_Cloud"
+            assert df_copy.sum().get_backend() == "Small_Data_Local"
+            # check in-place pin/unpin operations
+            df.pin_backend(inplace=True)
+            assert df.get_backend() == "Big_Data_Cloud"
+            assert df.sum().get_backend() == "Big_Data_Cloud"
+            df.unpin_backend(inplace=True)
+            assert df.get_backend() == "Big_Data_Cloud"
+            assert df.sum().get_backend() == "Small_Data_Local"
+
 
 class TestSwitchBackendPreOp:
     @pytest.mark.parametrize(
@@ -869,6 +916,58 @@ class TestSwitchBackendPreOp:
             )
             assert md_df.get_backend() == expected_backend
 
+    def test_iloc_pinned(self):
+        # The operation in test_iloc would naturally cause an automatic switch, but the
+        # absence of AutoSwitchBackend or the presence of a pin on the frame prevent this
+        # switch from happening.
+        data_size = BIG_DATA_CLOUD_MIN_NUM_ROWS - 1
+        with backend_test_context(
+            test_backend="Big_Data_Cloud",
+            choices=("Big_Data_Cloud", "Small_Data_Local"),
+        ):
+            register_function_for_pre_op_switch(
+                class_name="_iLocIndexer",
+                backend="Big_Data_Cloud",
+                method="__setitem__",
+            )
+            # No pin or config, should switch
+            df = pd.DataFrame(list(range(data_size)))
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.iloc[(0, 0)] = -1
+            assert df.get_backend() == "Small_Data_Local"
+            # config set to false, should not switch
+            with config_context(AutoSwitchBackend=False):
+                df = pd.DataFrame(list(range(data_size)))
+                assert df.get_backend() == "Big_Data_Cloud"
+                df.iloc[(0, 0)] = -2
+                assert df.get_backend() == "Big_Data_Cloud"
+            # no config, but data is pinned
+            df = pd.DataFrame(list(range(data_size))).pin_backend()
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.iloc[(0, 0)] = -3
+            assert df.get_backend() == "Big_Data_Cloud"
+            # a frame-level pin remains valid across a transformation
+            df_copy = df + 1
+            assert df_copy.get_backend() == "Big_Data_Cloud"
+            df_copy.iloc[(0, 0)] = -4
+            assert df_copy.get_backend() == "Big_Data_Cloud"
+            # unpinning df allows a switch again
+            df.unpin_backend(inplace=True)
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.iloc[(0, 0)] = -5
+            assert df.get_backend() == "Small_Data_Local"
+            # An in-place set_backend operation clears the pin
+            df.move_to("Big_Data_Cloud", inplace=True)
+            # check in-place pin/unpin operations
+            df.pin_backend(inplace=True)
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.iloc[(0, 0)] = -6
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.unpin_backend(inplace=True)
+            assert df.get_backend() == "Big_Data_Cloud"
+            df.iloc[(0, 0)] = -7
+            assert df.get_backend() == "Small_Data_Local"
+
     @pytest.mark.parametrize(
         "args, kwargs, expected_backend",
         (
@@ -915,3 +1014,76 @@ class TestSwitchBackendPreOp:
             choices=("Big_Data_Cloud", "Small_Data_Local"),
         ):
             assert data_class(*args, **kwargs).get_backend() == expected_backend
+
+
+def test_move_to_clears_pin():
+    # Pin status is reset to false after a set_backend call
+    with backend_test_context(
+        test_backend="Big_Data_Cloud",
+        choices=("Big_Data_Cloud", "Small_Data_Local"),
+    ):
+        df = pd.DataFrame(list(range(10)))
+        # in-place
+        df.pin_backend(inplace=True)
+        assert df.is_backend_pinned()
+        df.move_to("Small_Data_Local", inplace=True)
+        assert not df.is_backend_pinned()
+        # not in-place
+        intermediate = df.pin_backend().move_to("Big_Data_Cloud")
+        assert not intermediate.is_backend_pinned()
+        assert intermediate.pin_backend().is_backend_pinned()
+
+
+@pytest.mark.parametrize(
+    "pin_backends, expected_backend",
+    [
+        param(
+            [("Small_Data_Local", False), ("Big_Data_Cloud", False)],
+            "Small_Data_Local",
+            id="no_pin",
+        ),  # no backend pinned
+        param(
+            [("Small_Data_Local", True), ("Big_Data_Cloud", False)],
+            "Small_Data_Local",
+            id="one_pin",
+        ),  # one backend is pinned, so move there
+        param(
+            [
+                ("Big_Data_Cloud", False),
+                ("Small_Data_Local", True),
+                ("Small_Data_Local", True),
+            ],
+            "Small_Data_Local",
+            id="two_pin",
+        ),  # two identical pinned backends
+        param(
+            [("Small_Data_Local", True), ("Big_Data_Cloud", True)],
+            None,
+            id="conflict_pin",
+        ),  # conflicting pins raises ValueError
+    ],
+)
+def test_concat_with_pin(pin_backends, expected_backend):
+    with backend_test_context(
+        test_backend="Big_Data_Cloud",
+        choices=("Big_Data_Cloud", "Small_Data_Local"),
+    ):
+        dfs = [
+            pd.DataFrame([1] * 10).move_to(backend)._set_backend_pinned(should_pin)
+            for backend, should_pin in pin_backends
+        ]
+        if expected_backend is None:
+            with pytest.raises(
+                ValueError,
+                match="Cannot combine arguments that are pinned to conflicting backends",
+            ):
+                pd.concat(dfs)
+        else:
+            result = pd.concat(dfs)
+            assert result.is_backend_pinned() == any(
+                df.is_backend_pinned() for df in dfs
+            )
+            assert result.get_backend() == expected_backend
+            df_equals(
+                result, pandas.concat([pandas.DataFrame([1] * 10)] * len(pin_backends))
+            )
diff --git a/modin/tests/pandas/test_api.py b/modin/tests/pandas/test_api.py
index cbc92c9c..8f90dd65 100644
--- a/modin/tests/pandas/test_api.py
+++ b/modin/tests/pandas/test_api.py
@@ -28,6 +28,10 @@ _MODIN_EXTRA_ATTRIBUTES = (
     "set_backend",
     # move_to - set storage and engine backend for the current DataFrame
     "move_to",
+    # is_backend_pinned, pin_backend, unpin_backend - change automatic switching behavior
+    "is_backend_pinned",
+    "pin_backend",
+    "unpin_backend",
 )