Fix dtype infer in DataFrame arithmetic on datetime consts (mars-project#2879)

Author: wjsi

mars/core/graph/tests/test_graph.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import pytest
 
 from .... import tensor as mt
@@ -103,4 +104,17 @@ def test_to_dot():
     graph = arr2.build_graph(fuse_enabled=False, tile=True)
 
     dot = str(graph.to_dot(trunc_key=5))
-    assert all(str(n.op.key)[5] in dot for n in graph) is True
+    try:
+        assert all(str(n.op.key)[5] in dot for n in graph) is True
+    except AssertionError:
+        graph_reprs = []
+        for n in graph:
+            graph_reprs.append(
+                f"{n.op.key} -> {[succ.op.key for succ in graph.successors(n)]}"
+            )
+        logging.error(
+            "Unexpected error in test_to_dot.\ndot = %r\ngraph_repr: %r",
+            dot,
+            "\n".join(graph_reprs),
+        )
+        raise

mars/dataframe/arithmetic/tests/test_arithmetic.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import itertools
 import operator
 from dataclasses import dataclass
@@ -1550,3 +1551,19 @@ def test_arithmetic_lazy_chunk_meta():
     pd.testing.assert_index_equal(chunk.index_value.to_pandas(), pd.RangeIndex(3))
     assert chunk._FIELD_VALUES.get("_columns_value") is None
     pd.testing.assert_index_equal(chunk.columns_value.to_pandas(), pd.RangeIndex(3))
+
+
+def test_datetime_arithmetic():
+    data1 = (
+        pd.Series([pd.Timedelta(days=d) for d in range(10)]) + datetime.datetime.now()
+    )
+    s1 = from_pandas_series(data1)
+
+    assert (s1 + pd.Timedelta(days=10)).dtype == (data1 + pd.Timedelta(days=10)).dtype
+    assert (s1 + datetime.timedelta(days=10)).dtype == (
+        data1 + datetime.timedelta(days=10)
+    ).dtype
+    assert (s1 - pd.Timestamp.now()).dtype == (data1 - pd.Timestamp.now()).dtype
+    assert (s1 - datetime.datetime.now()).dtype == (
+        data1 - datetime.datetime.now()
+    ).dtype

mars/services/scheduling/worker/quota.py

@@ -334,7 +334,9 @@ async def update_mem_stats(self):
             await self._process_requests()
         self._last_memory_available = cur_mem_available
         self._report_quota_info()
-        self.ref().update_mem_stats.tell_delay(delay=self._refresh_time)
+        self._stat_refresh_task = self.ref().update_mem_stats.tell_delay(
+            delay=self._refresh_time
+        )
 
     async def _has_space(self, delta: int):
         if self._hard_limit is None:

mars/utils.py

@@ -16,6 +16,7 @@
 
 import asyncio
 import dataclasses
+import datetime
 import enum
 import functools
 import importlib
@@ -1033,6 +1034,10 @@ def is_object_dtype(dtype: np.dtype) -> bool:
 def get_dtype(dtype: Union[np.dtype, pd.api.extensions.ExtensionDtype]):
     if pd.api.types.is_extension_array_dtype(dtype):
         return dtype
+    elif dtype is pd.Timestamp or dtype is datetime.datetime:
+        return np.dtype("datetime64[ns]")
+    elif dtype is pd.Timedelta or dtype is datetime.timedelta:
+        return np.dtype("timedelta64[ns]")
     else:
         return np.dtype(dtype)