pylibcudf.Column: add device_buffer_size and register a dask.sizeof function for cudf-polars Column and DataFrame

python/cudf_polars/cudf_polars/containers/column.py

@@ -58,6 +58,28 @@ def __init__(
         self.name = name
         self.set_sorted(is_sorted=is_sorted, order=order, null_order=null_order)
 
+    @functools.cached_property
+    def device_buffer_size(self) -> int:
+        """
+        The total size of the device buffers used by the Column.
+
+        Returns
+        -------
+        Number of bytes.
+        """
+
+        def _pylibcudf_device_size(col: plc.Column) -> int:
+            ret = 0
+            if col.data() is not None:
+                ret += col.data().nbytes
+            if col.null_mask() is not None:
+                ret += col.null_mask().nbytes
+            if col.children() is not None:
+                ret += sum(_pylibcudf_device_size(c) for c in col.children())
+            return ret
+
+        return _pylibcudf_device_size(self.obj)
+
     @classmethod
     def deserialize(
         cls, header: ColumnHeader, frames: tuple[memoryview, plc.gpumemoryview]

python/cudf_polars/cudf_polars/containers/dataframe.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import functools
 from functools import cached_property
 from typing import TYPE_CHECKING, cast
 
@@ -49,6 +50,17 @@ def __init__(self, columns: Iterable[Column]) -> None:
         self.column_map = {c.name: c for c in self.columns}
         self.table = plc.Table([c.obj for c in self.columns])
 
+    @functools.cached_property
+    def device_buffer_size(self) -> int:
+        """
+        The total size of the device buffers used by the Table.
+
+        Returns
+        -------
+        Number of bytes.
+        """
+        return sum(c.device_buffer_size for c in self.columns)
+
     def copy(self) -> Self:
         """Return a shallow copy of self."""
         return type(self)(c.copy() for c in self.columns)

python/cudf_polars/cudf_polars/experimental/dask_registers.py

@@ -130,19 +130,7 @@ def _(header: ColumnHeader, frames: tuple[memoryview, memoryview]) -> Column:
             frames = frames[0], plc.gpumemoryview(rmm.DeviceBuffer.to_device(frames[1]))
             return Column.deserialize(header, frames)
 
-    @sizeof_dispatch.register(Column)
-    def _(x: Column) -> int:
-        """The total size of the device buffers used by the Column."""
-        ret = 0
-        if x.obj.data() is not None:
-            ret += x.obj.data().nbytes
-        if x.obj.null_mask() is not None:
-            ret += x.obj.null_mask().nbytes
-        if x.obj.children() is not None:
-            ret += sum(sizeof_dispatch(c) for c in x.obj.children())
-        return ret
-
-    @sizeof_dispatch.register(DataFrame)
-    def _(x: DataFrame) -> int:
-        """The total size of the device buffers used by the DataFrame."""
-        return sum(sizeof_dispatch(c) for c in x.columns)
+    @sizeof_dispatch.register((Column, DataFrame))
+    def _(x: Column | DataFrame) -> int:
+        """The total size of the device buffers used by the DataFrame or Column."""
+        return x.device_buffer_size

python/cudf_polars/tests/experimental/test_dask_sizeof.py

@@ -22,7 +22,7 @@
         (pa.table([]), 0),
         (pa.table({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}), 9 * 8),
         (pa.table({"a": [1, 2, 3]}), 3 * 8),
-        (pa.table({"a": [1], "b": [2], "c": [3]}), 3 * 8),
+        (pa.table({"a": ["a"], "b": ["bc"]}), 2 * 8 + 3),
         (pa.table({"a": [1, 2, None]}), 88),
     ],
 )