REFACTOR-#7427: Track engine and storage format at the query compiler level.

sfc-gh-mvashishtha · sfc-gh-mvashishtha · commit 052e82e73fb6 · 2025-01-27T16:01:04.000-08:00
Signed-off-by: sfc-gh-mvashishtha &lt;mahesh.vashishtha@snowflake.com&gt;
diff --git a/modin/conftest.py b/modin/conftest.py
@@ -165,6 +165,9 @@ class TestQC(BaseQueryCompiler):
     def __init__(self, modin_frame):
         self._modin_frame = modin_frame
 
+    storage_format = property(lambda self: "Base")
+    engine = property(lambda self: "Python")
+
     def finalize(self):
         self._modin_frame.finalize()
 
diff --git a/modin/core/dataframe/algebra/binary.py b/modin/core/dataframe/algebra/binary.py
@@ -393,7 +393,7 @@ def caller(
                             and query_compiler.columns.equals(other.columns)
                         ):
                             shape_hint = "column"
-                    return query_compiler.__constructor__(
+                    return query_compiler._constructor(
                         query_compiler._modin_frame.broadcast_apply(
                             axis,
                             lambda left, right: func(
@@ -417,7 +417,7 @@ def caller(
                             and query_compiler.columns.equals(other.columns)
                         ):
                             shape_hint = "column"
-                    return query_compiler.__constructor__(
+                    return query_compiler._constructor(
                         query_compiler._modin_frame.n_ary_op(
                             lambda x, y: func(x, y, *args, **kwargs),
                             [other._modin_frame],
@@ -453,7 +453,7 @@ def caller(
                         dtypes=dtypes,
                         lazy=True,
                     )
-                return query_compiler.__constructor__(
+                return query_compiler._constructor(
                     new_modin_frame, shape_hint=shape_hint
                 )
 
diff --git a/modin/core/dataframe/algebra/fold.py b/modin/core/dataframe/algebra/fold.py
@@ -82,7 +82,7 @@ def caller(
                 A new query compiler representing the result of executing the
                 function.
             """
-            return query_compiler.__constructor__(
+            return query_compiler._constructor(
                 query_compiler._modin_frame.fold(
                     cls.validate_axis(fold_axis),
                     lambda x: fold_function(x, *args, **kwargs),
diff --git a/modin/core/dataframe/algebra/groupby.py b/modin/core/dataframe/algebra/groupby.py
@@ -446,7 +446,7 @@ def caller(
             new_index=new_index,
         )
 
-        result = query_compiler.__constructor__(new_modin_frame)
+        result = query_compiler._constructor(new_modin_frame)
         return result
 
     @classmethod
diff --git a/modin/core/dataframe/algebra/map.py b/modin/core/dataframe/algebra/map.py
@@ -60,7 +60,7 @@ def caller(
         ) -> PandasQueryCompiler:
             """Execute Map function against passed query compiler."""
             shape_hint = call_kwds.pop("shape_hint", None) or query_compiler._shape_hint
-            return query_compiler.__constructor__(
+            return query_compiler._constructor(
                 query_compiler._modin_frame.map(
                     lambda x: function(x, *args, **kwargs), *call_args, **call_kwds
                 ),
diff --git a/modin/core/dataframe/algebra/reduce.py b/modin/core/dataframe/algebra/reduce.py
@@ -60,7 +60,7 @@ def caller(
         ) -> PandasQueryCompiler:
             """Execute Reduce function against passed query compiler."""
             _axis = kwargs.get("axis") if axis is None else axis
-            return query_compiler.__constructor__(
+            return query_compiler._constructor(
                 query_compiler._modin_frame.reduce(
                     cls.validate_axis(_axis),
                     lambda x: reduce_function(x, *args, **kwargs),
diff --git a/modin/core/dataframe/algebra/tree_reduce.py b/modin/core/dataframe/algebra/tree_reduce.py
@@ -70,7 +70,7 @@ def caller(
             if compute_dtypes and query_compiler.frame_has_materialized_dtypes:
                 new_dtypes = str(compute_dtypes(query_compiler.dtypes, *args, **kwargs))
 
-            return query_compiler.__constructor__(
+            return query_compiler._constructor(
                 query_compiler._modin_frame.tree_reduce(
                     cls.validate_axis(_axis),
                     lambda x: map_function(x, *args, **kwargs),
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -33,7 +33,6 @@
 from pandas.core.indexes.api import Index, RangeIndex
 
 from modin.config import (
-    Engine,
     IsRayCluster,
     MinColumnPartitionSize,
     MinRowPartitionSize,
@@ -1677,7 +1676,9 @@ def copy(self):
         )
 
     @lazy_metadata_decorator(apply_axis="both")
-    def astype(self, col_dtypes, errors: str = "raise"):
+    def astype(
+        self, col_dtypes, engine: str, storage_format: str, errors: str = "raise"
+    ):
         """
         Convert the columns dtypes to given dtypes.
 
@@ -1707,7 +1708,7 @@ def astype(self, col_dtypes, errors: str = "raise"):
                         new_dtypes = self_dtypes.copy()
                     # Update the new dtype series to the proper pandas dtype
                     new_dtype = pandas.api.types.pandas_dtype(dtype)
-                    if Engine.get() == "Dask" and hasattr(dtype, "_is_materialized"):
+                    if engine == "Dask" and hasattr(dtype, "_is_materialized"):
                         # FIXME: https://github.com/dask/distributed/issues/8585
                         _ = dtype._materialize_categories()
 
@@ -1736,7 +1737,7 @@ def astype_builder(df):
             if not (col_dtypes == self_dtypes).all():
                 new_dtypes = self_dtypes.copy()
                 new_dtype = pandas.api.types.pandas_dtype(col_dtypes)
-                if Engine.get() == "Dask" and hasattr(new_dtype, "_is_materialized"):
+                if engine == "Dask" and hasattr(new_dtype, "_is_materialized"):
                     # FIXME: https://github.com/dask/distributed/issues/8585
                     _ = new_dtype._materialize_categories()
                 if isinstance(new_dtype, pandas.CategoricalDtype):
diff --git a/modin/core/execution/dispatching/factories/factories.py b/modin/core/execution/dispatching/factories/factories.py
@@ -629,6 +629,8 @@ def _to_pickle_glob(cls, *args, **kwargs):
         **kwargs : kwargs
             Arguments to the writer method.
         """
+        # TODO(https://github.com/modin-project/modin/issues/7429): Use
+        # frame-level execution instead of the global, default execution.
         current_execution = get_current_execution()
         if current_execution not in supported_executions:
             raise NotImplementedError(
diff --git a/modin/core/execution/modin_aqp.py b/modin/core/execution/modin_aqp.py
@@ -91,6 +91,8 @@ def call_progress_bar(result_parts, line_no):
 
     threading.Thread(target=_show_time_updates, args=(progress_bars[pbar_id],)).start()
 
+    # TODO(https://github.com/modin-project/modin/issues/7429): Use
+    # frame-level engine config.
     modin_engine = Engine.get()
     engine_wrapper = None
     if modin_engine == "Ray":
diff --git a/modin/core/execution/utils.py b/modin/core/execution/utils.py
@@ -42,6 +42,9 @@ def remote_function(func, ignore_defaults=False):
 
 # Check if the function already exists to avoid circular imports
 elif "remote_function" not in dir():
+    # TODO(https://github.com/modin-project/modin/issues/7429): Use
+    # frame-level engine config.
+
     from modin.config import Engine
 
     if Engine.get() == "Ray":
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
@@ -150,6 +150,16 @@ class BaseQueryCompiler(
     _modin_frame: PandasDataframe
     _shape_hint: Optional[str]
 
+    @property
+    @abc.abstractmethod
+    def storage_format(self) -> str:
+        pass
+
+    @property
+    @abc.abstractmethod
+    def engine(self) -> str:
+        pass
+
     def __wrap_in_qc(self, obj):
         """
         Wrap `obj` in query compiler.
@@ -7019,7 +7029,7 @@ def repartition(self, axis=None):
 
         new_query_compiler = self
         for _ax in axes:
-            new_query_compiler = new_query_compiler.__constructor__(
+            new_query_compiler = new_query_compiler._constructor(
                 new_query_compiler._modin_frame.apply_full_axis(
                     _ax,
                     lambda df: df,
diff --git a/modin/core/storage_formats/pandas/aggregations.py b/modin/core/storage_formats/pandas/aggregations.py
@@ -97,7 +97,7 @@ def corr_method(
                 new_columns=new_columns,
                 new_dtypes=new_dtypes,
             )
-            return qc.__constructor__(result)
+            return qc._constructor(result)
 
         return corr_method
 
diff --git a/modin/core/storage_formats/pandas/groupby.py b/modin/core/storage_formats/pandas/groupby.py
@@ -366,7 +366,7 @@ def applyier(df, other):  # pragma: no cover
 
             return result
 
-        result = qc.__constructor__(
+        result = qc._constructor(
             to_group._modin_frame.broadcast_apply_full_axis(
                 axis=0, func=applyier, other=keys_columns._modin_frame
             )
@@ -438,7 +438,7 @@ def range_partition_impl(
         if to_unstack is None:
             result = result.transpose()
 
-        return qc.__constructor__(result)
+        return qc._constructor(result)
 
     @staticmethod
     def _pivot_table_from_groupby(
diff --git a/modin/core/storage_formats/pandas/merge.py b/modin/core/storage_formats/pandas/merge.py
@@ -87,7 +87,7 @@ def func(left, right):
             suffixes=kwargs.get("suffixes", ("_x", "_y")),
         )
 
-        return left.__constructor__(
+        return left._constructor(
             left._modin_frame._apply_func_to_range_partitioning_broadcast(
                 right._modin_frame,
                 func=func,
@@ -195,7 +195,7 @@ def map_func(
             ):
                 left = left.repartition(axis=0)
 
-            new_left = left.__constructor__(
+            new_left = left._constructor(
                 left._modin_frame.broadcast_apply_full_axis(
                     axis=1,
                     func=map_func,
diff --git a/modin/core/storage_formats/pandas/native_query_compiler.py b/modin/core/storage_formats/pandas/native_query_compiler.py
@@ -598,6 +598,9 @@ def __init__(self, pandas_frame, shape_hint: Optional[str] = None):
         self._modin_frame = pandas_frame
         self._shape_hint = shape_hint
 
+    storage_format = property(lambda self: "Pandas")
+    engine = property(lambda self: "Native")
+
     def execute(self):
         pass
 
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
diff --git a/modin/distributed/dataframe/pandas/partitions.py b/modin/distributed/dataframe/pandas/partitions.py
diff --git a/modin/error_message.py b/modin/error_message.py
diff --git a/modin/experimental/batch/pipeline.py b/modin/experimental/batch/pipeline.py
diff --git a/modin/tests/core/storage_formats/pandas/test_internals.py b/modin/tests/core/storage_formats/pandas/test_internals.py
diff --git a/modin/tests/pandas/native_df_mode/test_indexing.py b/modin/tests/pandas/native_df_mode/test_indexing.py
diff --git a/modin/tests/pandas/utils.py b/modin/tests/pandas/utils.py

Original file line number	Diff line number	Diff line change
`@@ -446,7 +446,7 @@ def caller(`
`446`	`446`	`new_index=new_index,`
`447`	`447`	`)`
`448`	`448`
`449`		`- result = query_compiler.__constructor__(new_modin_frame)`
	`449`	`+ result = query_compiler._constructor(new_modin_frame)`
`450`	`450`	`return result`
`451`	`451`
`452`	`452`	`@classmethod`