Add struct accessor to dask-cudf

python/dask_cudf/dask_cudf/accessors.py

@@ -1,6 +1,62 @@
 # Copyright (c) 2021, NVIDIA CORPORATION.
 
 
+# from cudf.python.dask_cudf.dask_cudf.core import Index
+
+
+class StructMethods:
+    def __init__(self, d_series):
+        self.d_series = d_series
+
+    def field(self, key):
+        """
+        Extract children of the specified struct column
+        in the Series
+        Parameters
+        ----------
+        key: int or str
+            index/position or field name of the respective
+            struct column
+        Returns
+        -------
+        Series
+        Examples
+        --------
+        >>> s = cudf.Series([{'a': 1, 'b': 2}, {'a': 3, 'b': 4}])
+        >>> ds = dask_cudf.from_cudf(s, 2)
+        >>> ds.struct.field(0).compute()
+        0    1
+        1    3
+        dtype: int64
+        >>> ds.struct.field('a').compute()
+        0    1
+        1    3
+        dtype: int64
+        """
+        try:
+            typ = self.d_series._meta.dtype.fields[key]
+
+        except KeyError:
+            try:
+                if isinstance(key, int):
+                    typ = self.d_series._meta.dtype.fields[
+                        list(self.d_series._meta.dtype.fields)[key]
+                    ]
+                else:
+                    raise KeyError(
+                        f"Field '{key}' is not in the set of existing keys"
+                    )
+            except TypeError:
+                raise IndexError(
+                    f"Index '{key}' is greater than the number of fields"
+                )
+
+        return self.d_series.map_partitions(
+            lambda s: s.struct.field(key),
+            meta=self.d_series._meta._constructor([], dtype=typ),
+        )
+
+
 class ListMethods:
     def __init__(self, d_series):
         self.d_series = d_series

python/dask_cudf/dask_cudf/core.py

@@ -27,7 +27,7 @@
 from cudf import _lib as libcudf
 
 from dask_cudf import sorting
-from dask_cudf.accessors import ListMethods
+from dask_cudf.accessors import ListMethods, StructMethods
 
 DASK_VERSION = LooseVersion(dask.__version__)
 
@@ -414,6 +414,10 @@ def groupby(self, *args, **kwargs):
     def list(self):
         return ListMethods(self)
 
+    @property
+    def struct(self):
+        return StructMethods(self)
+
 
 class Index(Series, dd.core.Index):
     _partition_type = cudf.Index

python/dask_cudf/dask_cudf/tests/test_accessor.py

@@ -438,3 +438,88 @@ def test_sorting(data, ascending, na_position, ignore_index):
         .reset_index(drop=True)
     )
     assert_eq(expect, got)
+
+
+#############################################################################
+#                            Struct Accessor                                #
+#############################################################################
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [{"a": 5, "b": 10}, {"a": 3, "b": 7}, {"a": -3, "b": 11}],
+        [{"a": None, "b": 1}, {"a": None, "b": 0}, {"a": -3, "b": None}],
+        [{"a": 1, "b": 2}],
+        [{"b": 3, "c": 4}],
+    ],
+)
+def test_create_struct_series(data):
+    expect = pd.Series(data)
+    ds_got = dgd.from_cudf(Series(data), 2)
+    assert_eq(expect, ds_got.compute())
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [{"a": 5, "b": 10}, {"a": 3, "b": 7}, {"a": -3, "b": 11}],
+        [{"a": None, "b": 1}, {"a": None, "b": 0}, {"a": -3, "b": None}],
+        [{"a": 1, "b": 2}],
+    ],
+)
+def test_struct_field_str(data):
+    for test_key in ["a", "b"]:
+        expect = Series(data).struct.field(test_key)
+        ds_got = dgd.from_cudf(Series(data), 2).struct.field(test_key)
+        assert_eq(expect, ds_got.compute())
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [{"a": 5, "b": 10}, {"a": 3, "b": 7}, {"a": -3, "b": 11}],
+        [{"a": None, "b": 1}, {"a": None, "b": 0}, {"a": -3, "b": None}],
+        [{"a": 1, "b": 2}],
+        [{"b": 3, "c": 4}],
+    ],
+)
+def test_struct_field_integer(data):
+    for test_key in [0, 1]:
+        expect = Series(data).struct.field(test_key)
+        ds_got = dgd.from_cudf(Series(data), 2).struct.field(test_key)
+        assert_eq(expect, ds_got.compute())
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [{"a": 5, "b": 10}, {"a": 3, "b": 7}, {"a": -3, "b": 11}],
+        [{"a": None, "b": 1}, {"a": None, "b": 0}, {"a": -3, "b": None}],
+        [{"a": 1, "b": 2}],
+        [{"b": 3, "c": 4}],
+    ],
+)
+def test_dask_struct_field_Key_Error(data):
+    got = dgd.from_cudf(Series(data), 2)
+
+    # import pdb; pdb.set_trace()
+    with pytest.raises(KeyError):
+        got.struct.field("notakey").compute()
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        [{"a": 5, "b": 10}, {"a": 3, "b": 7}, {"a": -3, "b": 11}],
+        [{"a": None, "b": 1}, {"a": None, "b": 0}, {"a": -3, "b": None}],
+        [{"a": 1, "b": 2}],
+        [{"b": 3, "c": 4}],
+    ],
+)
+def test_dask_struct_field_Int_Error(data):
+    # breakpoint()
+    got = dgd.from_cudf(Series(data), 2)
+
+    with pytest.raises(IndexError):
+        got.struct.field(1000).compute()