Add struct accessor to dask-cudf

python/cudf/cudf/core/column/column.py

@@ -1345,7 +1345,13 @@ def column_empty(
     dtype = pandas_dtype(dtype)
     children = ()  # type: Tuple[ColumnBase, ...]
 
-    if is_categorical_dtype(dtype):
+    if is_struct_dtype(dtype):
+        data = None
+        children = tuple(
+            column_empty(row_count, dtype.fields[f])
+            for f in dtype.fields.keys()
+        )
+    elif is_categorical_dtype(dtype):
         data = None
         children = (
             build_column(

python/dask_cudf/dask_cudf/accessors.py

@@ -1,5 +1,53 @@
 # Copyright (c) 2021, NVIDIA CORPORATION.
 
+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]
+            return self.d_series.map_partitions(
+                lambda s: s.struct.field(key),
+                meta=self.d_series._meta._constructor([], dtype=typ),
+            )
+        except KeyError as e:
+            if isinstance(key, int):
+                key_list = [dict_key for dict_key in self.d_series._meta.dtype.fields.keys()]
+                typ_key = key_list[key]
+                typ = self.d_series._meta.dtype.fields[typ_key]
+                return self.d_series.map_partitions(
+                    lambda s: s.struct.field(key), 
+                    meta=self.d_series._meta._constructor([], dtype=typ),
+                )
+            else:
+                print('Field "' + str(key) + '" is not found in the set of existing keys.')
+                raise e
+
+
 
 class ListMethods:
     def __init__(self, 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,77 @@ 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}],
+
+        [{'b':3, 'c':4}],
+
+    ],
+)
+def test_struct_field_a(data):
+    expect = Series(data).struct.field('a')
+    ds_got = dgd.from_cudf(Series(data), 2).struct.field('a')
+    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_zero(data):
+    expect = Series(data).struct.field(0)
+    ds_got = dgd.from_cudf(Series(data), 2).struct.field(0)
+    assert_eq(expect, ds_got.compute())

python/dask_cudf/dask_cudf/tests/test_struct.py

@@ -0,0 +1,36 @@
+import pytest
+
+import cudf
+
+import dask_cudf
+
+
+@pytest.mark.parametrize(
+    "data, column",
+    [
+        (
+            {
+                "a": [{"a": [1, 2, 3, 4], "b": "Hello world"}, {}, {"a": []}],
+                "b": [1, 2, 3],
+                "c": ["rapids", "cudf", "hi"],
+            },
+            "a",
+        ),
+        (
+            {"a": [{}, {}, {}], "b": [1, 2, 3], "c": ["rapids", "cudf", "hi"]},
+            "a",
+        ),
+        (
+            {
+                "a": [{}, {}, {}],
+                "b": [{"a": 1}, {"b": 5}, {"c": "Hello"}],
+                "c": ["rapids", "cudf", "hi"],
+            },
+            "b",
+        ),
+    ],
+)
+def test_select_struct(data, column):
+    df = cudf.DataFrame(data)
+    ddf = dask_cudf.from_cudf(df, 2)
+    assert df[column].to_arrow() == ddf[column].compute().to_arrow()