Repository: modin-project/modin
Pull request #7531 (2025-04-25 22:15:59)

FEAT-#7511: Support pre-operation switch for init by passing arguments to cost functions.


Patch:
diff --git a/modin/core/io/io.py b/modin/core/io/io.py
index a492988c..442d42a2 100644
--- a/modin/core/io/io.py
+++ b/modin/core/io/io.py
@@ -82,7 +82,7 @@ class BaseIO:
         return None
 
     @classmethod
-    def from_pandas(cls, df):
+    def from_pandas(cls, df) -> BaseQueryCompiler:
         """
         Create a Modin `query_compiler` from a `pandas.DataFrame`.
 
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
index 048ef049..8431a114 100644
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -23,7 +23,8 @@ import abc
 import warnings
 from enum import IntEnum
 from functools import cached_property
-from typing import TYPE_CHECKING, Hashable, List, Literal, Optional
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Hashable, List, Literal, Optional
 
 import numpy as np
 import pandas
@@ -318,7 +319,11 @@ class BaseQueryCompiler(
 
     @disable_logging
     def move_to_cost(
-        self, other_qc_type: type, api_cls_name: Optional[str], operation: str
+        self,
+        other_qc_type: type,
+        api_cls_name: Optional[str],
+        operation: str,
+        arguments: MappingProxyType[str, Any],
     ) -> int:
         """
         Return the coercion costs of this qc to other_qc type.
@@ -340,6 +345,8 @@ class BaseQueryCompiler(
         operation : str
             The operation being performed which can be used as a consideration
             for the costing analysis.
+        arguments : MappingProxyType[str, Any]
+            The arguments to the operation.
 
         Returns
         -------
@@ -397,7 +404,12 @@ class BaseQueryCompiler(
         return int(total_cost)
 
     @disable_logging
-    def stay_cost(self, api_cls_name: Optional[str], operation: str) -> Optional[int]:
+    def stay_cost(
+        self,
+        api_cls_name: Optional[str],
+        operation: str,
+        arguments: MappingProxyType[str, Any],
+    ) -> Optional[int]:
         """
         Return the "opportunity cost" of not moving the data.
 
@@ -424,6 +436,8 @@ class BaseQueryCompiler(
         operation : str, default: None
             The operation being performed which can be used as a consideration
             for the costing analysis.
+        arguments : MappingProxyType[str, Any]
+            The arguments to the operation.
 
         Returns
         -------
@@ -440,7 +454,11 @@ class BaseQueryCompiler(
     @disable_logging
     @classmethod
     def move_to_me_cost(
-        cls, other_qc: BaseQueryCompiler, api_cls_name: Optional[str], operation: str
+        cls,
+        other_qc: BaseQueryCompiler,
+        api_cls_name: Optional[str],
+        operation: str,
+        arguments: MappingProxyType[str, Any],
     ) -> Optional[int]:
         """
         Return the execution and hidden coercion costs from other_qc.
@@ -466,6 +484,8 @@ class BaseQueryCompiler(
         operation : str
             The operation being performed which can be used as a consideration
             for the costing analysis.
+        arguments : MappingProxyType[str, Any]
+            The arguments to the operation.
 
         Returns
         -------
diff --git a/modin/core/storage_formats/base/query_compiler_calculator.py b/modin/core/storage_formats/base/query_compiler_calculator.py
index 00d7d1b5..5b15274f 100644
--- a/modin/core/storage_formats/base/query_compiler_calculator.py
+++ b/modin/core/storage_formats/base/query_compiler_calculator.py
@@ -20,7 +20,8 @@ all query compilers to determine the best query compiler to use.
 """
 
 import logging
-from typing import Optional
+from types import MappingProxyType
+from typing import Any, Optional
 
 from modin.core.storage_formats.base.query_compiler import (
     BaseQueryCompiler,
@@ -57,19 +58,25 @@ class BackendCostCalculator:
 
     Parameters
     ----------
+    operation_arguments : MappingProxyType[str, Any]
+        Mapping from operation argument names to their values.
     api_cls_name : str or None
         Representing the class name of the function being called.
-    operation : str or None representing the operation being performed
+    operation : str representing the operation being performed
     """
 
     def __init__(
-        self, api_cls_name: Optional[str] = None, operation: Optional[str] = None
+        self,
+        operation_arguments: MappingProxyType[str, Any],
+        api_cls_name: Optional[str],
+        operation: str,
     ):
         self._backend_data: dict[str, AggregatedBackendData] = {}
         self._qc_list: list[BaseQueryCompiler] = []
         self._result_backend = None
         self._api_cls_name = api_cls_name
         self._op = operation
+        self._operation_arguments = operation_arguments
 
     def add_query_compiler(self, query_compiler: BaseQueryCompiler):
         """
@@ -106,7 +113,12 @@ class BackendCostCalculator:
                 if qc_cls_to not in qc_to_cls_costed:
                     qc_to_cls_costed.add(qc_cls_to)
                     backend_to = qc_to.get_backend()
-                    cost = qc_from.move_to_cost(qc_cls_to, self._api_cls_name, self._op)
+                    cost = qc_from.move_to_cost(
+                        qc_cls_to,
+                        self._api_cls_name,
+                        self._op,
+                        self._operation_arguments,
+                    )
                     if cost is not None:
                         self._add_cost_data(backend_to, cost)
                     else:
@@ -114,7 +126,10 @@ class BackendCostCalculator:
                         # qc_from does not know about qc_to types so we instead
                         # ask the same question but of qc_to.
                         cost = qc_cls_to.move_to_me_cost(
-                            qc_from, self._api_cls_name, self._op
+                            qc_from,
+                            self._api_cls_name,
+                            self._op,
+                            self._operation_arguments,
                         )
                         if cost is not None:
                             self._add_cost_data(backend_to, cost)
diff --git a/modin/core/storage_formats/pandas/native_query_compiler.py b/modin/core/storage_formats/pandas/native_query_compiler.py
index d58b4ee4..207ecf9b 100644
--- a/modin/core/storage_formats/pandas/native_query_compiler.py
+++ b/modin/core/storage_formats/pandas/native_query_compiler.py
@@ -26,7 +26,10 @@ from pandas.core.dtypes.common import is_scalar
 from modin.core.dataframe.base.interchange.dataframe_protocol.dataframe import (
     ProtocolDataframe,
 )
-from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
+from modin.core.storage_formats.base.query_compiler import (
+    BaseQueryCompiler,
+    QCCoercionCost,
+)
 from modin.utils import _inherit_docstrings
 
 _NO_REPARTITION_ON_NATIVE_EXECUTION_EXCEPTION_MESSAGE = (
@@ -90,6 +93,11 @@ class NativeQueryCompiler(BaseQueryCompiler):
         The pandas frame to query with the compiled queries.
     """
 
+    _MAX_SIZE_THIS_ENGINE_CAN_HANDLE = 10_000_000
+    _OPERATION_INITIALIZATION_OVERHEAD = 0
+    _OPERATION_PER_ROW_OVERHEAD = 0
+    _TRANSFER_THRESHOLD = 10_000_000
+
     _modin_frame: pandas.DataFrame
     _should_warn_on_default_to_pandas: bool = False
 
@@ -251,3 +259,22 @@ class NativeQueryCompiler(BaseQueryCompiler):
     @_inherit_docstrings(BaseQueryCompiler.repartition)
     def repartition(self, axis=None):
         raise Exception(_NO_REPARTITION_ON_NATIVE_EXECUTION_EXCEPTION_MESSAGE)
+
+    @classmethod
+    def move_to_me_cost(cls, other_qc, api_cls_name, operation, arguments):
+        if api_cls_name in ("DataFrame", "Series") and operation == "__init__":
+            if (query_compiler := arguments.get("query_compiler")) is not None:
+                # When we create a dataframe or series with a query compiler
+                # input, we should not switch the resulting dataframe or series
+                # to a different backend.
+                return (
+                    QCCoercionCost.COST_ZERO
+                    if isinstance(query_compiler, cls)
+                    else QCCoercionCost.COST_IMPOSSIBLE
+                )
+            else:
+                # Otherwise we're constructing a dataframe or series from
+                # local data, which we assume is easy to store using the native
+                # backend.
+                return QCCoercionCost.COST_LOW
+        return super().move_to_me_cost(other_qc, api_cls_name, operation, arguments)
diff --git a/modin/core/storage_formats/pandas/query_compiler_caster.py b/modin/core/storage_formats/pandas/query_compiler_caster.py
index dd9c9695..288b2586 100644
--- a/modin/core/storage_formats/pandas/query_compiler_caster.py
+++ b/modin/core/storage_formats/pandas/query_compiler_caster.py
@@ -24,7 +24,7 @@ import inspect
 import logging
 from abc import ABC, abstractmethod
 from collections import defaultdict, namedtuple
-from types import FunctionType, MethodType
+from types import FunctionType, MappingProxyType, MethodType
 from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, Union, ValuesView
 
 import pandas
@@ -317,8 +317,9 @@ def apply_argument_cast_to_class(klass: type) -> type:
 
 def _maybe_switch_backend_pre_op(
     function_name: str,
-    qc_list: list[BaseQueryCompiler],
+    input_qc: BaseQueryCompiler,
     class_of_wrapped_fn: Optional[str],
+    arguments: MappingProxyType[str, Any],
 ) -> tuple[str, Callable[[Any], Any]]:
     """
     Possibly switch backend before a function.
@@ -327,11 +328,13 @@ def _maybe_switch_backend_pre_op(
     ----------
     function_name : str
         The name of the function.
-    qc_list : list[BaseQueryCompiler]
-        The list of query compilers that are arguments to the function.
+    input_qc : BaseQueryCompiler
+        The input query compiler.
     class_of_wrapped_fn : Optional[str]
         The name of the class that the function belongs to. `None` for functions
         in the modin.pandas module.
+    arguments : MappingProxyType[str, Any]
+        Mapping from operation argument names to their values.
 
     Returns
     -------
@@ -339,30 +342,20 @@ def _maybe_switch_backend_pre_op(
         A tuple of the new backend and a function that casts all castable arguments
         to the new query compiler type.
     """
-    assert len(qc_list) in (0, 1)
-    if len(qc_list) == 0:
-        # For nullary functions, we need to create a dummy query compiler
-        # to calculate the cost of switching backends.
-        from modin.core.execution.dispatching.factories.dispatcher import (
-            FactoryDispatcher,
-        )
-
-        input_qc = FactoryDispatcher.get_factory().io_cls.from_pandas(
-            pandas.DataFrame()
-        )
-    else:
-        input_qc = qc_list[0]
-    input_backend = Backend.get() if len(qc_list) == 0 else input_qc.get_backend()
+    input_backend = input_qc.get_backend()
     if (
         function_name
         in _CLASS_AND_BACKEND_TO_PRE_OP_SWITCH_METHODS[
-            BackendAndClassName(backend=input_backend, class_name=class_of_wrapped_fn)
+            BackendAndClassName(
+                backend=input_qc.get_backend(), class_name=class_of_wrapped_fn
+            )
         ]
     ):
         result_backend = _get_backend_for_auto_switch(
             input_qc=input_qc,
             class_of_wrapped_fn=class_of_wrapped_fn,
             function_name=function_name,
+            arguments=arguments,
         )
     else:
         result_backend = input_backend
@@ -386,6 +379,7 @@ def _maybe_switch_backend_post_op(
     qc_list: list[BaseQueryCompiler],
     starting_backend: str,
     class_of_wrapped_fn: Optional[str],
+    arguments: MappingProxyType[str, Any],
 ) -> Any:
     """
     Possibly switch the backend of the result of a function.
@@ -409,6 +403,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.
+    arguments : MappingProxyType[str, Any]
+        Mapping from operation argument names to their values.
 
     Returns
     -------
@@ -437,6 +433,7 @@ def _maybe_switch_backend_post_op(
                 input_qc=input_qc,
                 class_of_wrapped_fn=class_of_wrapped_fn,
                 function_name=function_name,
+                arguments=arguments,
             )
         )
     return result
@@ -446,6 +443,7 @@ def _get_backend_for_auto_switch(
     input_qc: BaseQueryCompiler,
     class_of_wrapped_fn: str,
     function_name: str,
+    arguments: MappingProxyType[str, Any],
 ) -> str:
     """
     Get the best backend to switch to.
@@ -464,6 +462,8 @@ def _get_backend_for_auto_switch(
         in the modin.pandas module.
     function_name : str
         The name of the function.
+    arguments : MappingProxyType[str, Any]
+        Mapping from operation argument names to their values.
 
     Returns
     -------
@@ -484,6 +484,7 @@ def _get_backend_for_auto_switch(
     stay_cost = input_qc.stay_cost(
         api_cls_name=class_of_wrapped_fn,
         operation=function_name,
+        arguments=arguments,
     )
     for backend in Backend.get_active_backends():
         if backend in ("Ray", "Unidist", "Dask"):
@@ -502,11 +503,13 @@ def _get_backend_for_auto_switch(
             move_to_class,
             api_cls_name=class_of_wrapped_fn,
             operation=function_name,
+            arguments=arguments,
         )
         other_execute_cost = move_to_class.move_to_me_cost(
             input_qc,
             api_cls_name=class_of_wrapped_fn,
             operation=function_name,
+            arguments=arguments,
         )
         if (
             move_to_cost is not None
@@ -541,6 +544,61 @@ def _get_backend_for_auto_switch(
     return best_backend
 
 
+def _get_extension_for_method(
+    name: str,
+    extensions: EXTENSION_DICT_TYPE,
+    backend: str,
+    args: tuple,
+    wrapping_function_type: Optional[
+        Union[type[classmethod], type[staticmethod], type[MethodType]]
+    ],
+) -> callable:
+    """
+    Get the extension implementation for a method.
+
+    Parameters
+    ----------
+    name : str
+        The name of the method.
+    extensions : EXTENSION_DICT_TYPE
+        The extension dictionary for the modin-API-level object (e.g. class
+        DataFrame or module modin.pandas) that the method belongs to.
+    backend : str
+        The backend to use for this method call.
+    args : tuple
+        The arguments to the method.
+    wrapping_function_type : Union[type[classmethod], type[staticmethod], type[MethodType]]
+        The type of the original function that `f` implements.
+        - `None` means we are wrapping a free function, e.g. pd.concat()
+        - `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.
+
+    Returns
+    -------
+    callable
+        The implementation of the method for the given backend.
+    """
+    if name in extensions[backend]:
+        f_to_apply = extensions[backend][name]
+    else:
+        if name not in extensions[None]:
+            raise AttributeError(
+                (
+                    # 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 = extensions[None][name]
+    return f_to_apply
+
+
 def wrap_function_in_argument_caster(
     klass: Optional[type],
     f: callable,
@@ -613,29 +671,70 @@ def wrap_function_in_argument_caster(
         inplace_update_trackers: list[InplaceUpdateTracker] = []
         # The function name and class name of the function are passed to the calculator as strings
         class_of_wrapped_fn = klass.__name__ if klass is not None else None
-        calculator: BackendCostCalculator = BackendCostCalculator(
-            class_of_wrapped_fn, f.__name__
-        )
+
+        input_query_compilers: list[BaseQueryCompiler] = []
 
         def register_query_compilers(arg):
             if (
                 isinstance(arg, QueryCompilerCaster)
                 and (qc := arg._get_query_compiler()) is not None
             ):
-                calculator.add_query_compiler(qc)
+                input_query_compilers.append(qc)
             elif isinstance(arg, BaseQueryCompiler):
                 # We might get query compiler arguments in __init__()
-                calculator.add_query_compiler(arg)
+                input_query_compilers.append(arg)
             return arg
 
         visit_nested_args(args, register_query_compilers)
         visit_nested_args(kwargs, register_query_compilers)
 
-        if len(calculator._qc_list) < 2:
+        if len(input_query_compilers) == 0:
+            # For nullary functions, we need to create a dummy query compiler
+            # to calculate the cost of switching backends.
+            from modin.core.execution.dispatching.factories.dispatcher import (
+                FactoryDispatcher,
+            )
+
+            input_qc_for_pre_op_switch = (
+                FactoryDispatcher.get_factory().io_cls.from_pandas(pandas.DataFrame())
+            )
+            input_backend = Backend.get()
+        else:
+            input_qc_for_pre_op_switch = input_query_compilers[0]
+            input_backend = input_qc_for_pre_op_switch.get_backend()
+
+        # Bind the arguments using the function implementation for the input
+        # backend. TODO(https://github.com/modin-project/modin/issues/7525):
+        # Ideally every implementation would have the same signature.
+        bound_arguments = inspect.signature(
+            _get_extension_for_method(
+                name=name,
+                extensions=extensions,
+                backend=input_backend,
+                args=args,
+                wrapping_function_type=wrapping_function_type,
+            ),
+        ).bind(*args, **kwargs)
+        bound_arguments.apply_defaults()
+        args_dict = MappingProxyType(bound_arguments.arguments)
+
+        if len(input_query_compilers) < 2:
             result_backend, cast_to_qc = _maybe_switch_backend_pre_op(
-                name, calculator._qc_list, class_of_wrapped_fn=class_of_wrapped_fn
+                name,
+                input_qc=input_qc_for_pre_op_switch,
+                class_of_wrapped_fn=class_of_wrapped_fn,
+                arguments=args_dict,
             )
         else:
+            calculator: BackendCostCalculator = BackendCostCalculator(
+                operation_arguments=args_dict,
+                api_cls_name=class_of_wrapped_fn,
+                operation=name,
+            )
+
+            for qc in input_query_compilers:
+                calculator.add_query_compiler(qc)
+
             result_backend = calculator.calculate()
 
             def cast_to_qc(arg):
@@ -657,23 +756,17 @@ def wrap_function_in_argument_caster(
 
         args = visit_nested_args(args, cast_to_qc)
         kwargs = visit_nested_args(kwargs, cast_to_qc)
-        if name in extensions[result_backend]:
-            f_to_apply = extensions[result_backend][name]
-        else:
-            if name not in extensions[None]:
-                raise AttributeError(
-                    (
-                        # 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 = extensions[None][name]
+
+        # `result_backend` may be different from `input_backend`, so we have to
+        # look up the correct implementation based on `result_backend`.
+        f_to_apply = _get_extension_for_method(
+            name=name,
+            extensions=extensions,
+            backend=result_backend,
+            args=args,
+            wrapping_function_type=wrapping_function_type,
+        )
+
         # We have to set the global Backend correctly for I/O methods like
         # read_json() to use the correct backend.
         with config_context(Backend=result_backend):
@@ -690,9 +783,10 @@ def wrap_function_in_argument_caster(
         return _maybe_switch_backend_post_op(
             result,
             function_name=name,
-            qc_list=calculator._qc_list,
+            qc_list=input_query_compilers,
             starting_backend=result_backend,
             class_of_wrapped_fn=class_of_wrapped_fn,
+            arguments=args_dict,
         )
 
     f_with_argument_casting._wrapped_method_for_casting = f


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 bbb7f892..e925d10a 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -14,6 +14,7 @@
 import contextlib
 import json
 from io import StringIO
+from types import MappingProxyType
 from typing import Iterator
 from unittest import mock
 
@@ -51,15 +52,15 @@ class CalculatorTestQc(NativeQueryCompiler):
     """
 
     @classmethod
-    def move_to_me_cost(cls, other_qc, api_cls_name, operation):
+    def move_to_me_cost(cls, other_qc, api_cls_name, operation, arguments):
         if isinstance(other_qc, cls):
             return QCCoercionCost.COST_ZERO
         return None
 
-    def stay_cost(self, api_cls_name, operation):
+    def stay_cost(self, api_cls_name, operation, arguments):
         return QCCoercionCost.COST_ZERO
 
-    def move_to_cost(self, other_qc_type, api_cls_name, operation):
+    def move_to_cost(self, other_qc_type, api_cls_name, operation, arguments):
         if isinstance(self, other_qc_type):
             return QCCoercionCost.COST_ZERO
         return None
@@ -74,7 +75,7 @@ class CloudQC(CalculatorTestQc):
     def max_cost(self):
         return QCCoercionCost.COST_IMPOSSIBLE
 
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         assert op is not None
         assert api_cls_name in [
             None,
@@ -94,7 +95,7 @@ class CloudQC(CalculatorTestQc):
             OmniscientLazyQC: None,
         }[other_qc_cls]
 
-    def stay_cost(self, api_cls_name, op):
+    def stay_cost(self, api_cls_name, op, arguments):
         return QCCoercionCost.COST_HIGH
 
 
@@ -107,7 +108,7 @@ class ClusterQC(CalculatorTestQc):
     def max_cost(self):
         return QCCoercionCost.COST_HIGH
 
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         return {
             CloudQC: QCCoercionCost.COST_MEDIUM,
             ClusterQC: QCCoercionCost.COST_ZERO,
@@ -126,7 +127,7 @@ class LocalMachineQC(CalculatorTestQc):
     def max_cost(self):
         return QCCoercionCost.COST_MEDIUM
 
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         return {
             CloudQC: QCCoercionCost.COST_MEDIUM,
             ClusterQC: QCCoercionCost.COST_LOW,
@@ -144,7 +145,7 @@ class PicoQC(CalculatorTestQc):
     def max_cost(self):
         return QCCoercionCost.COST_LOW
 
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         return {
             CloudQC: QCCoercionCost.COST_LOW,
             ClusterQC: QCCoercionCost.COST_LOW,
@@ -159,7 +160,7 @@ class AdversarialQC(CalculatorTestQc):
     def get_backend(self):
         return "Adversarial"
 
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         return {
             CloudQC: -1000,
             ClusterQC: 10000,
@@ -174,14 +175,14 @@ class OmniscientEagerQC(CalculatorTestQc):
         return "Eager"
 
     # keep other workloads from getting my workload
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         if OmniscientEagerQC is other_qc_cls:
             return QCCoercionCost.COST_ZERO
         return QCCoercionCost.COST_IMPOSSIBLE
 
     # try to force other workloads to my engine
     @classmethod
-    def move_to_me_cost(cls, other_qc, api_cls_name, operation):
+    def move_to_me_cost(cls, other_qc, api_cls_name, operation, arguments):
         return QCCoercionCost.COST_ZERO
 
 
@@ -192,12 +193,12 @@ class OmniscientLazyQC(CalculatorTestQc):
         return "Lazy"
 
     # encorage other engines to take my workload
-    def move_to_cost(self, other_qc_cls, api_cls_name, op):
+    def move_to_cost(self, other_qc_cls, api_cls_name, op, arguments):
         return QCCoercionCost.COST_ZERO
 
     # try to keep other workloads from getting my workload
     @classmethod
-    def move_to_me_cost(cls, other_qc, api_cls_name, operation):
+    def move_to_me_cost(cls, other_qc, api_cls_name, operation, arguments):
         if isinstance(other_qc, cls):
             return QCCoercionCost.COST_ZERO
         return QCCoercionCost.COST_IMPOSSIBLE
@@ -236,14 +237,33 @@ class CloudForBigDataQC(BaseTestAutoMover):
     def __init__(self, pandas_frame):
         super().__init__(pandas_frame)
 
-    def stay_cost(self, api_cls_name, operation):
+    def stay_cost(self, api_cls_name, operation, arguments):
         if operation == "read_json":
             return QCCoercionCost.COST_IMPOSSIBLE
-        return super().stay_cost(api_cls_name, operation)
+        return super().stay_cost(api_cls_name, operation, arguments)
 
     def get_backend(self) -> str:
         return "Big_Data_Cloud"
 
+    @classmethod
+    def move_to_me_cost(cls, other_qc, api_cls_name, operation, arguments):
+        if api_cls_name in ("DataFrame", "Series") and operation == "__init__":
+            if (query_compiler := arguments.get("query_compiler")) is not None:
+                # When we create a dataframe or series with a query compiler
+                # input, we should not switch the resulting dataframe or series
+                # to a different backend.
+                return (
+                    QCCoercionCost.COST_ZERO
+                    if isinstance(query_compiler, cls)
+                    else QCCoercionCost.COST_IMPOSSIBLE
+                )
+            else:
+                # Moving the in-memory __init__ inputs to the cloud is expensive.
+                return QCCoercionCost.COST_HIGH
+        return super().move_to_me_cost(
+            cls, other_qc, api_cls_name, operation, arguments
+        )
+
 
 class LocalForSmallDataQC(BaseTestAutoMover):
     """Represents a local query compiler that prefers small data."""
@@ -499,7 +519,11 @@ def test_default_to_caller(default_df, default2_df):
 
 
 def test_no_qc_to_calculate():
-    calculator = BackendCostCalculator()
+    calculator = BackendCostCalculator(
+        operation_arguments=MappingProxyType({}),
+        api_cls_name=None,
+        operation="operation0",
+    )
     with pytest.raises(ValueError):
         calculator.calculate()
 
@@ -510,6 +534,7 @@ def test_qc_default_self_cost(default_df, default2_df):
             other_qc_type=type(default2_df._query_compiler),
             api_cls_name=None,
             operation="operation0",
+            arguments=MappingProxyType({}),
         )
         is None
     )
@@ -518,6 +543,7 @@ def test_qc_default_self_cost(default_df, default2_df):
             other_qc_type=type(default_df._query_compiler),
             api_cls_name=None,
             operation="operation0",
+            arguments=MappingProxyType({}),
         )
         is QCCoercionCost.COST_ZERO
     )
@@ -593,10 +619,13 @@ def test_switch_local_to_cloud_with_iloc___setitem__(local_df, cloud_df):
 def test_stay_or_move_evaluation(cloud_df, default_df):
     default_cls = type(default_df._get_query_compiler())
     cloud_cls = type(cloud_df._get_query_compiler())
+    empty_arguments = MappingProxyType({})
 
-    stay_cost = cloud_df._get_query_compiler().stay_cost("Series", "myop")
+    stay_cost = cloud_df._get_query_compiler().stay_cost(
+        "Series", "myop", arguments=empty_arguments
+    )
     move_cost = cloud_df._get_query_compiler().move_to_cost(
-        default_cls, "Series", "myop"
+        default_cls, "Series", "myop", arguments=empty_arguments
     )
     df = cloud_df
     if stay_cost > move_cost:
@@ -604,8 +633,12 @@ def test_stay_or_move_evaluation(cloud_df, default_df):
     else:
         assert False
 
-    stay_cost = df._get_query_compiler().stay_cost("Series", "myop")
-    move_cost = df._get_query_compiler().move_to_cost(cloud_cls, "Series", "myop")
+    stay_cost = df._get_query_compiler().stay_cost(
+        "Series", "myop", arguments=empty_arguments
+    )
+    move_cost = df._get_query_compiler().move_to_cost(
+        cloud_cls, "Series", "myop", arguments=empty_arguments
+    )
     assert stay_cost is not None
     assert move_cost is None
 
@@ -835,3 +868,50 @@ class TestSwitchBackendPreOp:
                 __inplace__=True,
             )
             assert md_df.get_backend() == expected_backend
+
+    @pytest.mark.parametrize(
+        "args, kwargs, expected_backend",
+        (
+            param((), {}, "Small_Data_Local", id="no_args_or_kwargs"),
+            param(([1],), {}, "Small_Data_Local", id="small_list_data_in_arg"),
+            param(
+                (list(range(BIG_DATA_CLOUD_MIN_NUM_ROWS)),),
+                {},
+                "Small_Data_Local",
+                id="big_list_data_in_arg",
+            ),
+            param((), {"data": [1]}, "Small_Data_Local", id="list_data_in_kwarg"),
+            param(
+                (),
+                {"data": pandas.Series([1])},
+                "Small_Data_Local",
+                id="series_data_in_kwarg",
+            ),
+            param(
+                (),
+                {"query_compiler": CloudForBigDataQC(pandas.DataFrame([0, 1, 2]))},
+                "Big_Data_Cloud",
+                id="cloud_query_compiler_in_kwarg",
+            ),
+            param(
+                (),
+                {"query_compiler": LocalForSmallDataQC(pandas.DataFrame([0, 1, 2]))},
+                "Small_Data_Local",
+                id="small_query_compiler_in_kwarg",
+            ),
+        ),
+    )
+    @pytest.mark.parametrize("data_class", [pd.DataFrame, pd.Series])
+    def test___init___with_in_memory_data_uses_native_query_compiler(
+        self, args, kwargs, expected_backend, data_class
+    ):
+        register_function_for_pre_op_switch(
+            class_name=data_class.__name__,
+            method="__init__",
+            backend="Big_Data_Cloud",
+        )
+        with backend_test_context(
+            test_backend="Big_Data_Cloud",
+            choices=("Big_Data_Cloud", "Small_Data_Local"),
+        ):
+            assert data_class(*args, **kwargs).get_backend() == expected_backend