diff --git a/python/cudf/cudf/core/column/column.py b/python/cudf/cudf/core/column/column.py index d52f63a79f5..fc8b9028fa1 100644 --- a/python/cudf/cudf/core/column/column.py +++ b/python/cudf/cudf/core/column/column.py @@ -509,7 +509,9 @@ def slice(self, start: int, stop: int, stride: int = None) -> ColumnBase: return column_empty(0, self.dtype, masked=True) # compute mask slice if stride == 1: - return libcudf.copying.column_slice(self, [start, stop])[0] + return libcudf.copying.column_slice(self, [start, stop])[ + 0 + ]._with_type_metadata(self.dtype) else: # Need to create a gather map for given slice with stride gather_map = arange( @@ -721,6 +723,7 @@ def take( self.as_frame() ._gather(indices, keep_index=keep_index, nullify=nullify) ._as_column() + ._with_type_metadata(self.dtype) ) except RuntimeError as e: if "out of bounds" in str(e): diff --git a/python/cudf/cudf/core/column/struct.py b/python/cudf/cudf/core/column/struct.py index fd63b4de144..c035a5bfb6b 100644 --- a/python/cudf/cudf/core/column/struct.py +++ b/python/cudf/cudf/core/column/struct.py @@ -87,7 +87,7 @@ def __getitem__(self, args): field: value for field, value in zip(self.dtype.fields, result.values()) } - return result._rename_fields(self.dtype.fields.keys()) + return result def __setitem__(self, key, value): if isinstance(value, dict): diff --git a/python/cudf/cudf/tests/test_struct.py b/python/cudf/cudf/tests/test_struct.py index 702a3fe75d6..4e5e9c96146 100644 --- a/python/cudf/cudf/tests/test_struct.py +++ b/python/cudf/cudf/tests/test_struct.py @@ -260,6 +260,18 @@ def test_struct_slice(series, start, end): assert sr[start:end].to_arrow() == expected.to_arrow() +def test_struct_slice_nested_struct(): + data = [ + {"a": {"b": 42, "c": "abc"}}, + {"a": {"b": 42, "c": "hello world"}}, + ] + + got = cudf.Series(data)[0:1] + expect = cudf.Series(data[0:1]) + assert got.__repr__() == expect.__repr__() + assert got.dtype.to_arrow() == expect.dtype.to_arrow() + + @pytest.mark.parametrize( "data", [