Repository: modin-project/modin

Issue:
Precompute dtypes for `sum` operation



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/algebra/tree_reduce.py b/modin/core/dataframe/algebra/tree_reduce.py
--- a/modin/core/dataframe/algebra/tree_reduce.py
+++ b/modin/core/dataframe/algebra/tree_reduce.py
@@ -20,7 +20,9 @@ class TreeReduce(Operator):
     """Builder class for TreeReduce operator."""
 
     @classmethod
-    def register(cls, map_function, reduce_function=None, axis=None):
+    def register(
+        cls, map_function, reduce_function=None, axis=None, compute_dtypes=None
+    ):
         """
         Build TreeReduce operator.
 
@@ -32,6 +34,8 @@ def register(cls, map_function, reduce_function=None, axis=None):
             Source reduce function.
         axis : int, optional
             Specifies axis to apply function along.
+        compute_dtypes : callable(pandas.Series, *func_args, **func_kwargs) -> np.dtype, optional
+            Callable for computing dtypes.
 
         Returns
         -------
@@ -45,11 +49,17 @@ def register(cls, map_function, reduce_function=None, axis=None):
         def caller(query_compiler, *args, **kwargs):
             """Execute TreeReduce function against passed query compiler."""
             _axis = kwargs.get("axis") if axis is None else axis
+
+            new_dtypes = None
+            if compute_dtypes and query_compiler._modin_frame.has_materialized_dtypes:
+                new_dtypes = str(compute_dtypes(query_compiler.dtypes, *args, **kwargs))
+
             return query_compiler.__constructor__(
                 query_compiler._modin_frame.tree_reduce(
                     cls.validate_axis(_axis),
                     lambda x: map_function(x, *args, **kwargs),
                     lambda y: reduce_function(y, *args, **kwargs),
+                    dtypes=new_dtypes,
                 )
             )
 
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -1936,7 +1936,9 @@ def tree_reduce(
         reduce_parts = self._partition_mgr_cls.map_axis_partitions(
             axis.value, map_parts, reduce_func
         )
-        return self._compute_tree_reduce_metadata(axis.value, reduce_parts)
+        return self._compute_tree_reduce_metadata(
+            axis.value, reduce_parts, dtypes=dtypes
+        )
 
     @lazy_metadata_decorator(apply_axis=None)
     def map(
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -852,7 +852,14 @@ def is_series_like(self):
 
     # TreeReduce operations
     count = TreeReduce.register(pandas.DataFrame.count, pandas.DataFrame.sum)
-    sum = TreeReduce.register(pandas.DataFrame.sum)
+
+    def _dtypes_sum(dtypes: pandas.Series, *func_args, **func_kwargs):  # noqa: GL08
+        # The common type evaluation for `TreeReduce` operator may differ depending
+        # on the pandas function, so it's better to pass a evaluation function that
+        # should be defined for each Modin's function.
+        return find_common_type(dtypes.tolist())
+
+    sum = TreeReduce.register(pandas.DataFrame.sum, compute_dtypes=_dtypes_sum)
     prod = TreeReduce.register(pandas.DataFrame.prod)
     any = TreeReduce.register(pandas.DataFrame.any, pandas.DataFrame.any)
     all = TreeReduce.register(pandas.DataFrame.all, pandas.DataFrame.all)


Test patch: