tests

Author: madsbk

python/cudf_polars/cudf_polars/experimental/dask_registers.py

@@ -132,6 +132,7 @@ def _(header: ColumnHeader, frames: tuple[memoryview, memoryview]) -> Column:
 
     @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
@@ -143,4 +144,5 @@ def _(x: Column) -> int:
 
     @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)

python/cudf_polars/tests/experimental/test_dask_sizeof.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import pyarrow as pa
+import pytest
+from dask.sizeof import sizeof
+
+import pylibcudf as plc
+
+from cudf_polars.containers import DataFrame
+from cudf_polars.experimental.dask_registers import register
+
+# Must register sizeof dispatch before running tests
+register()
+
+
+@pytest.mark.parametrize(
+    "arrow_tbl, size",
+    [
+        (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),
+    ],
+)
+def test_dask_sizeof(arrow_tbl, size):
+    plc_tbl = plc.interop.from_arrow(arrow_tbl)
+    df = DataFrame.from_table(plc_tbl, names=arrow_tbl.column_names)
+    assert sizeof(df) == size