diff --git a/python/pylibcudf/pylibcudf/column.pxd b/python/pylibcudf/pylibcudf/column.pxd index 8cb3b7d4093f..4fdc9ad08f49 100644 --- a/python/pylibcudf/pylibcudf/column.pxd +++ b/python/pylibcudf/pylibcudf/column.pxd @@ -53,6 +53,7 @@ cdef class Column: # _children: List[Column] list _children size_type _num_children + object __weakref__ cdef column_view view(self) cdef mutable_column_view mutable_view(self) diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pxd b/python/pylibcudf/pylibcudf/gpumemoryview.pxd index 87f52bb8bc58..08be0a8b4771 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pxd +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport uint64_t, uintptr_t @@ -9,3 +9,4 @@ cdef class gpumemoryview: cdef readonly object obj cdef readonly dict cai cdef readonly uint64_t nbytes + cdef object __weakref__ diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pyi b/python/pylibcudf/pylibcudf/gpumemoryview.pyi index 10d8d3aec052..81c9a4779839 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pyi +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pyi @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from collections.abc import Mapping @@ -9,6 +9,7 @@ class gpumemoryview: @property def __cuda_array_interface__(self) -> Mapping[str, Any]: ... def __len__(self) -> int: ... + def byte_slice(self, s: slice) -> gpumemoryview: ... @property def ptr(self) -> int: ... @property diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pyx b/python/pylibcudf/pylibcudf/gpumemoryview.pyx index 03af09faca47..0312198e2a84 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pyx +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pyx @@ -1,13 +1,25 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stddef cimport size_t +from libc.stdint cimport uintptr_t, uint64_t import functools import operator from .types cimport DataType, size_of, type_id +cdef gpumemoryview _slice(gpumemoryview parent, uintptr_t ptr, uint64_t nbytes): + cdef gpumemoryview v = gpumemoryview.__new__(gpumemoryview) + v.ptr = ptr + v.nbytes = nbytes + v.obj = parent + # always returns a raw byte view regardless of the source dtype. + # TODO: Need to propagate stream from parent.cai if present + v.cai = {"data": (ptr, parent.cai["data"][1]), "shape": (nbytes,), "typestr": "|u1", "version": 3} + return v + + __all__ = ["gpumemoryview"] @@ -59,6 +71,7 @@ cdef class gpumemoryview: self.obj = obj self.cai = cai # TODO: Need to respect readonly + # TODO: Need to synchronize on stream if present in cai self.ptr = cai["data"][0] # Compute the buffer size. @@ -83,6 +96,40 @@ cdef class gpumemoryview: return self.nbytes def __len__(self): - return self.obj.__cuda_array_interface__["shape"][0] + return self.cai["shape"][0] + + def byte_slice(self, s): + """Return a byte-range sub-view of this buffer. + + Parameters + ---------- + s : slice + Byte-based slice. + + Returns + ------- + gpumemoryview + A ``|u1`` view of the requested byte range. The returned view + holds a reference to the parent buffer, keeping it alive. + + Raises + ------ + TypeError + If ``s`` is not a slice. + ValueError + If the slice step is not 1. Out-of-range or reversed ranges + return a zero-length view rather than raising. + """ + if not isinstance(s, slice): + raise TypeError( + f"byte_slice requires a slice, not {type(s).__name__}" + ) + start, stop, step = s.indices(self.nbytes) + if step != 1: + raise ValueError("byte_slice only supports step=1 slices") + length = stop - start + if length <= 0: + return _slice(self, self.ptr + start, 0) + return _slice(self, self.ptr + start, length) __hash__ = None diff --git a/python/pylibcudf/tests/test_gpumemoryview.py b/python/pylibcudf/tests/test_gpumemoryview.py index 98a04422563b..9bd0f6c840ab 100644 --- a/python/pylibcudf/tests/test_gpumemoryview.py +++ b/python/pylibcudf/tests/test_gpumemoryview.py @@ -1,7 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import itertools +import weakref import numpy as np import pytest @@ -63,3 +64,40 @@ def test_len(np_array, stream): assert len(gpumemview) == len(np_array_view) assert gpumemview.nbytes == np_array.nbytes + + +@pytest.mark.parametrize( + "s", + [ + slice(1, 3), + slice(None, 2), + slice(3, None), + slice(2, 2), + slice(0, 10000), + ], +) +def test_slice(np_array, s): + gv = plc.Column.from_array(np_array.view("u1")).data() + result = plc.Column.from_array(gv.byte_slice(s)).to_pylist() + assert result == np_array.view("u1")[s].tolist() + + +def test_slice_fails(np_array): + gv = plc.Column.from_array(np_array.view("u1")).data() + with pytest.raises(TypeError, match="requires a slice"): + gv.byte_slice(0) + with pytest.raises(ValueError, match="step=1"): + gv.byte_slice(slice(None, None, 2)) + + +def test_slice_keeps_parent_alive(): + col = plc.Column.from_array(np.arange(10, dtype="u1")) + gv = col.data() + col_ref = weakref.ref(col) + gv_ref = weakref.ref(gv) + s = gv.byte_slice(slice(2, 5)) + del col, gv + assert col_ref() is None + assert gv_ref() is not None + del s + assert gv_ref() is None