Add support for pylibcudf.DataType serialization

python/pylibcudf/pylibcudf/tests/test_serialize.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+import pickle
+
+import pytest
+
+from pylibcudf import DataType
+from pylibcudf.types import TypeId
+
+
+@pytest.fixture(params=list(TypeId))
+def dtype(request):
+    tid = request.param
+    if tid in {TypeId.DECIMAL32, TypeId.DECIMAL64, TypeId.DECIMAL128}:
+        scale = 5
+    else:
+        scale = 0
+    return DataType(tid, scale)
+
+
+def test_reduce(dtype):
+    (typ, (tid, scale)) = dtype.__reduce__()
+    assert typ is DataType
+    assert tid == dtype.id()
+    assert scale == dtype.scale()
+
+
+def test_pickle(dtype):
+    assert dtype == pickle.loads(pickle.dumps(dtype))

python/pylibcudf/pylibcudf/types.pyx

@@ -76,6 +76,9 @@ cdef class DataType:
     def __hash__(self):
         return hash((self.c_obj.id(), self.c_obj.scale()))
 
+    def __reduce__(self):
+        return (type(self), (self.c_obj.id(), self.c_obj.scale()))
+
     @staticmethod
     cdef DataType from_libcudf(data_type dt):
         """Create a DataType from a libcudf data_type.