Enable pd.Timestamp objects to be picklable when cudf.pandas is active

python/cudf/cudf/pandas/_wrappers/pandas.py

@@ -1,7 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-
+import copyreg
+import pickle
 import sys
 
 import pandas as pd
@@ -1304,3 +1305,50 @@ def _df_query_method(self, *args, local_dict=None, global_dict=None, **kwargs):
         _Unusable,
         typ,
     )
+
+# timestamps and timedeltas are not proxied and "real" pandas types
+# are currently not picklable when the module accelerator is enabled.
+def _unpickle_timestamp(pickled_args):
+    from cudf.pandas.module_accelerator import disable_module_accelerator
+
+    with disable_module_accelerator():
+        unpickler, args = pickle.loads(pickled_args)
+    ts = unpickler(*args)
+    return ts
+
+
+def _reduce_timestamp(ts):
+    from cudf.pandas.module_accelerator import disable_module_accelerator
+
+    with disable_module_accelerator():
+        # args can contain objects that are unpicklable
+        # when the module accelerator is disabled
+        # (freq is of a proxy type):
+        pickled_args = pickle.dumps(ts.__reduce__())
+
+    return _unpickle_timestamp, (pickled_args,)
+
+
+def _unpickle_timedelta(pickled_args):
+    from cudf.pandas.module_accelerator import disable_module_accelerator
+
+    with disable_module_accelerator():
+        unpickler, args = pickle.loads(pickled_args)
+    ts = unpickler(*args)
+    return ts
+
+
+def _reduce_timedelta(ts):
+    from cudf.pandas.module_accelerator import disable_module_accelerator
+
+    with disable_module_accelerator():
+        # args can contain objects that are unpicklable
+        # when the module accelerator is disabled
+        # (freq is of a proxy type):
+        pickled_args = pickle.dumps(ts.__reduce__())
+
+    return _unpickle_timedelta, (pickled_args,)
+
+
+copyreg.dispatch_table[pd.Timestamp] = _reduce_timestamp
+copyreg.dispatch_table[pd.Timedelta] = _reduce_timedelta

python/cudf/cudf_pandas_tests/test_cudf_pandas.py

@@ -1076,6 +1076,10 @@ def test_np_array_of_timestamps():
         xpd.Index(["a", 2, 3]),
         # Other types
         xpd.tseries.offsets.BDay(5),
+        xpd.Timestamp("2001-01-01"),
+        xpd.Timestamp("2001-01-01", freq="D"),
+        xpd.Timedelta("1 days"),
+        xpd.Timedelta(1, "D"),
     ],
 )
 def test_pickle(obj):