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,31 @@ def _df_query_method(self, *args, local_dict=None, global_dict=None, **kwargs):
         _Unusable,
         typ,
     )
+
+# timestamps and timedeltas are not proxied, but non-proxied
+# pandas types are currently not picklable. Thus, we define
+# custom reducer/unpicker functions for these types:
+def _reduce_obj(obj):
+    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(obj.__reduce__())
+
+    return _unpickle_obj, (pickled_args,)
+
+
+def _unpickle_obj(pickled_args):
+    from cudf.pandas.module_accelerator import disable_module_accelerator
+
+    with disable_module_accelerator():
+        unpickler, args = pickle.loads(pickled_args)
+    obj = unpickler(*args)
+    return obj
+
+
+copyreg.dispatch_table[pd.Timestamp] = _reduce_obj
+# same reducer/unpickler can be used for Timedelta:
+copyreg.dispatch_table[pd.Timedelta] = _reduce_obj

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):