Checkpoint

Author: sfc-gh-jkew

modin/core/dataframe/algebra/default2pandas/default.py

@@ -202,9 +202,9 @@ def args_cast(self, *args, **kwargs):
 
             Cast all Modin objects that function arguments contain to its pandas representation.
             """
-            args = try_cast_to_pandas(args)
-            kwargs = try_cast_to_pandas(kwargs)
-            return wrapper(self, *args, **kwargs)
+            args1 = try_cast_to_pandas(args)
+            kwargs1 = try_cast_to_pandas(kwargs)
+            return wrapper(self, *args1, **kwargs1)
 
         return args_cast
 

modin/core/storage_formats/base/query_compiler.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import abc
+from enum import Enum
 import warnings
 from functools import cached_property
 from typing import TYPE_CHECKING, Hashable, List, Literal, Optional
@@ -107,6 +108,43 @@ def axis_setter(self, labels):
     return axis_setter
 
 
+class QCCoercionCost(Enum):
+    '''
+    Coercion costs between query compilers can be expressed
+    as integers in the range -1 to 1000, where 1000 is
+    considered impossible. Since coercsion costs can be a
+    function of many variables ( dataset size, partitioning, 
+    network throughput, and query time ) we define a set range
+    of cost values to simplify comparisons between two query
+    compilers / engines in a unified way.
+    
+    COST_UNKNOWN means we do not know the cost associated with changing
+    query compilers.
+    
+    COST_ZERO means there is no cost associated, or that the query compilers
+    are the same.
+    
+    COST_IMPOSSIBLE means the coercion is effectively impossible, which can
+    occur if the target system is unable to store the data as a result
+    of the coercion.
+    '''
+    COST_UNKNOWN = -1
+    COST_ZERO = 0
+    COST_LOW = 250
+    COST_MEDIUM = 500
+    COST_HIGH = 750
+    COST_IMPOSSIBLE = 1000
+
+    def validate_coercsion_cost(cost:QCCoercionCost):
+        if int(cost) < int(QCCoercionCost.COST_UNKNOWN) or int(cost) > int(QCCoercionCost.COST_IMPOSSIBLE):
+            raise ValueError("Query compiler coercsion cost out of range") 
+        
+    def __int__(self):
+        return self.value
+    
+    def __add__(self, other) -> int:
+        return int(self) + int(other)
+
 # FIXME: many of the BaseQueryCompiler methods are hiding actual arguments
 # by using *args and **kwargs. They should be spread into actual parameters.
 # Currently actual arguments are placed in the methods docstrings, but since they're
@@ -247,6 +285,18 @@ def default_to_pandas(self, pandas_op, *args, **kwargs) -> Self:
             return [self.__wrap_in_qc(obj) for obj in result]
         return self.__wrap_in_qc(result)
 
+    def qc_engine_switch_cost(self, other_qc) -> dict[type, int]:
+        '''
+        Coercion costs to and from other_qc
+        
+        Returns a map of type to QCCoercionCost, where type is the type we are casting to.
+        This provides a mechanism for the query compilers to provide information to
+        modin on the cost of moving data to another query compiler ( or the other way ).
+        '''
+        if isinstance(type(self), type(other_qc)):
+            return  {type(self): QCCoercionCost.COST_ZERO}
+        return {}
+        
     # Abstract Methods and Fields: Must implement in children classes
     # In some cases, there you may be able to use the same implementation for
     # some of these abstract methods, but for the sake of generality they are

modin/core/storage_formats/pandas/query_compiler_caster.py

@@ -21,16 +21,70 @@
 
 import functools
 import inspect
+from itertools import combinations
 from types import FunctionType, MethodType
 from typing import Any, Dict, Tuple, TypeVar
 
 from pandas.core.indexes.frozen import FrozenList
 
-from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler
+from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler, QCCoercionCost
 
 Fn = TypeVar("Fn", bound=Any)
 
-
+class QueryCompilerCasterCalculator:
+    
+    def __init__(self):
+        self._caster_costing_map = {}
+        self._data_cls_map = {}
+        self._qc_list = []
+        self._qc_cls_list = []
+        self._result_type = None
+    
+    def add_query_compiler(self, query_compiler):
+        if isinstance(query_compiler, type):
+            # class
+            qc_type = query_compiler
+        else:
+            # instance
+            qc_type = type(query_compiler)
+            self._qc_list.append(query_compiler)
+            self._data_cls_map[qc_type] = query_compiler._modin_frame
+        self._qc_cls_list.append(qc_type)
+    
+    def calculate(self):
+        if self._result_type is not None:
+            return self._result_type
+        if len(self._qc_cls_list) == 1:
+            return self._qc_cls_list[0]
+        if len(self._qc_cls_list) == 0:
+            raise ValueError("No query compilers registered")
+        
+        for (qc_1, qc_2) in combinations(self._qc_list, 2):
+            costs_1 = qc_1.qc_engine_switch_cost(qc_2)
+            costs_2 = qc_2.qc_engine_switch_cost(qc_1)
+            self._add_cost_data(costs_1)
+            self._add_cost_data(costs_2)
+        
+        min_value = min(self._caster_costing_map.values())
+        for key, value in self._caster_costing_map.items():
+            if min_value == value:
+                self._result_type = key
+                break
+        return self._result_type
+            
+    def _add_cost_data(self, costs:dict):
+        for k, v in costs.items():
+            # filter out any extranious query compilers not in this operation
+            if k in self._qc_cls_list:
+                QCCoercionCost.validate_coercsion_cost(v)
+                # Adds the costs associated with all coercions to a type, k
+                self._caster_costing_map[k] = v + self._caster_costing_map[k] if k in self._caster_costing_map else v
+    
+    def result_data_frame(self):
+        qc_type = self.calculate()
+        return self._data_cls_map[qc_type]
+        
+    
 class QueryCompilerCaster:
     """Cast all query compiler arguments of the member function to current query compiler."""
 
@@ -55,7 +109,9 @@ def __init_subclass__(
         apply_argument_cast(cls)
 
 
-def cast_nested_args_to_current_qc_type(arguments, current_qc):
+def visit_nested_args(arguments, 
+                      current_qc:BaseQueryCompiler, 
+                      fn:callable):
     """
     Cast all arguments in nested fashion to current query compiler.
 
@@ -70,33 +126,25 @@ def cast_nested_args_to_current_qc_type(arguments, current_qc):
         Returns args and kwargs with all query compilers casted to current_qc.
     """
 
-    def cast_arg_to_current_qc(arg):
-        current_qc_type = type(current_qc)
-        if isinstance(arg, BaseQueryCompiler) and not isinstance(arg, current_qc_type):
-            data_cls = current_qc._modin_frame
-            return current_qc_type.from_pandas(arg.to_pandas(), data_cls)
-        else:
-            return arg
-
     imutable_types = (FrozenList, tuple)
     if isinstance(arguments, imutable_types):
         args_type = type(arguments)
         arguments = list(arguments)
-        arguments = cast_nested_args_to_current_qc_type(arguments, current_qc)
+        arguments = visit_nested_args(arguments, current_qc, fn)
 
         return args_type(arguments)
     if isinstance(arguments, list):
         for i in range(len(arguments)):
             if isinstance(arguments[i], (list, dict)):
-                cast_nested_args_to_current_qc_type(arguments[i], current_qc)
+                visit_nested_args(arguments[i], current_qc, fn)
             else:
-                arguments[i] = cast_arg_to_current_qc(arguments[i])
+                arguments[i] = fn(arguments[i])
     elif isinstance(arguments, dict):
         for key in arguments:
             if isinstance(arguments[key], (list, dict)):
-                cast_nested_args_to_current_qc_type(arguments[key], current_qc)
+                visit_nested_args(arguments[key], current_qc, fn)
             else:
-                arguments[key] = cast_arg_to_current_qc(arguments[key])
+                arguments[key] = fn(arguments[key])
     return arguments
 
 
@@ -116,6 +164,9 @@ def apply_argument_cast(obj: Fn) -> Fn:
     if isinstance(obj, type):
         all_attrs = dict(inspect.getmembers(obj))
         all_attrs.pop("__abstractmethods__")
+        all_attrs.pop("__init__")
+        all_attrs.pop("qc_engine_switch_cost")
+        all_attrs.pop("from_pandas")
 
         # This is required because inspect converts class methods to member functions
         current_class_attrs = vars(obj)
@@ -150,10 +201,57 @@ def cast_args(*args: Tuple, **kwargs: Dict) -> Any:
         -------
         Any
         """
+        if len(args) == 0 and len(kwargs) == 0:
+            return
+        print(f"Adding wrapper {obj}\n")
         current_qc = args[0]
+        calculator = QueryCompilerCasterCalculator()
+        calculator.add_query_compiler(current_qc)
+
+        def arg_needs_casting(arg):
+            current_qc_type = type(current_qc)
+            if not isinstance(arg, BaseQueryCompiler):
+                return False
+            if isinstance(arg, current_qc_type):
+                return False
+            return True
+        
+        def register_query_compilers(arg):
+            if not arg_needs_casting(arg):
+                return arg
+            calculator.add_query_compiler(arg)
+            return arg
+        
+        def cast_to_qc(arg):
+            if not arg_needs_casting(arg):
+                return arg
+            qc_type = calculator.calculate()
+            if qc_type == None or qc_type == type(arg):
+                return arg
+            frame_data = calculator.result_data_frame()
+            result = qc_type.from_pandas(arg.to_pandas(), frame_data)
+            return result
+        
+            
         if isinstance(current_qc, BaseQueryCompiler):
-            kwargs = cast_nested_args_to_current_qc_type(kwargs, current_qc)
-            args = cast_nested_args_to_current_qc_type(args, current_qc)
+            visit_nested_args(kwargs, current_qc, register_query_compilers)
+            visit_nested_args(args, current_qc, register_query_compilers)
+            
+            args = visit_nested_args(args, current_qc, cast_to_qc)
+            kwargs = visit_nested_args(kwargs, current_qc, cast_to_qc)
+
+        
+        qc = calculator.calculate()
+
+        if qc == None or qc == type(current_qc):
+            return obj(*args, **kwargs)
+
+        #breakpoint()
+        # we need to cast current_qc to a new query compiler
+        if qc != current_qc:
+            data_cls = current_qc._modin_frame
+            return qc.from_pandas(current_qc.to_pandas(), data_cls)
+        # need to find the new function for obj
         return obj(*args, **kwargs)
 
     return cast_args

modin/tests/pandas/native_df_interoperability/test_compiler_caster.py

@@ -0,0 +1,92 @@
+
+
+import pandas
+from modin.core.storage_formats.base.query_compiler import BaseQueryCompiler, QCCoercionCost
+from modin.core.storage_formats.pandas.native_query_compiler import NativeQueryCompiler
+from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
+from modin.utils import _inherit_docstrings
+
+
+class CloudQC(NativeQueryCompiler):
+    'Represents a cloud-hosted query compiler'
+    def __init__(self, pandas_frame):
+        self._modin_frame = pandas_frame
+        super().__init__(pandas_frame)
+      
+    def qc_engine_switch_cost(self, other_qc):
+        return {CloudQC: QCCoercionCost.COST_ZERO,
+                ClusterQC: QCCoercionCost.COST_MEDIUM,
+                LocalMachineQC: QCCoercionCost.COST_HIGH,
+                PicoQC: QCCoercionCost.COST_IMPOSSIBLE}
+
+class ClusterQC(NativeQueryCompiler):
+    'Represents a local network cluster query compiler'
+    def __init__(self, pandas_frame):
+        self._modin_frame = pandas_frame
+        super().__init__(pandas_frame)
+        
+    def qc_engine_switch_cost(self, other_qc):
+        return {CloudQC: QCCoercionCost.COST_MEDIUM,
+                ClusterQC: QCCoercionCost.COST_ZERO,
+                LocalMachineQC: QCCoercionCost.COST_MEDIUM,
+                PicoQC: QCCoercionCost.COST_HIGH}
+    
+class LocalMachineQC(NativeQueryCompiler):
+    'Represents a local machine query compiler'
+    def __init__(self, pandas_frame):
+        self._modin_frame = pandas_frame
+        super().__init__(pandas_frame)
+        
+    def qc_engine_switch_cost(self, other_qc):
+        return {CloudQC: QCCoercionCost.COST_MEDIUM,
+                ClusterQC: QCCoercionCost.COST_LOW,
+                LocalMachineQC: QCCoercionCost.COST_ZERO,
+                PicoQC: QCCoercionCost.COST_MEDIUM}
+
+class PicoQC(NativeQueryCompiler):
+    'Represents a query compiler with very few resources'
+    def __init__(self, pandas_frame):
+        self._modin_frame = pandas_frame
+        super().__init__(pandas_frame)
+        
+    def qc_engine_switch_cost(self, other_qc):
+        return {CloudQC: QCCoercionCost.COST_LOW,
+                ClusterQC: QCCoercionCost.COST_LOW,
+                LocalMachineQC: QCCoercionCost.COST_LOW,
+                PicoQC: QCCoercionCost.COST_ZERO}
+
+def test_two_same_qc_types_noop():
+    df = PicoQC(pandas.DataFrame([0, 1, 2]))
+    df2 = PicoQC(pandas.DataFrame([0, 1, 2]))
+    df3 = df.concat(axis=1, other=df2)
+    assert(type(df3) == type(df2))
+
+def test_two_two_qc_types_rhs():
+    df = PicoQC(pandas.DataFrame([0, 1, 2]))
+    df2 = ClusterQC(pandas.DataFrame([0, 1, 2]))
+    df3 = df.concat(axis=1, other=df2)
+    assert(type(df3) == type(df2))
+
+def test_two_two_qc_types_lhs():
+    df = PicoQC(pandas.DataFrame([0, 1, 2]))
+    df2 = ClusterQC(pandas.DataFrame([0, 1, 2]))
+    df3 = df2.concat(axis=1, other=df)
+    assert(type(df3) == type(df2)) # should move to cluster
+
+def test_three_two_qc_types_rhs():
+    pass
+
+def test_three_two_qc_types_lhs():
+    pass
+
+def test_three_two_qc_types_middle():
+    pass
+
+def test_three_three_qc_types_rhs():
+    pass
+
+def test_three_three_qc_types_lhs():
+    pass
+
+def test_three_three_qc_types_middle():
+    pass