diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index 4b6cc014a71e..64468754dcef 100644 --- a/docs/cudf/source/conf.py +++ b/docs/cudf/source/conf.py @@ -652,6 +652,13 @@ def on_missing_reference(app, env, node, contnode): ("py:class", "Value"), ("py:class", "polars.lazyframe.frame.LazyFrame"), ("py:class", "cudf_polars.engine.persisted_result.PersistedBackend"), + # pylibcudf typing aliases rendered as bare names in autodoc signatures. + ("py:class", "ColumnNameSpec"), + ("py:class", "CudaStreamLike"), + ("py:class", "Datasource"), + ("py:class", "Span"), + ("py:class", "SupportsArrayInterface"), + ("py:class", "SupportsCudaArrayInterface"), ] # Temporarily disable nitpick warnings for pandas: https://github.com/pandas-dev/pandas/issues/64584 nitpick_ignore_regex = [ diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index a903039d8c95..8bf0bbcc1cbf 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1814,8 +1814,6 @@ def pytest_unconfigure(config): "tests/groupby/test_missing.py::test_groupby_column_index_name_lost_fill_funcs[bfill]": "AssertionError: Index are different", "tests/groupby/test_missing.py::test_groupby_column_index_name_lost_fill_funcs[ffill]": "AssertionError: Index are different", "tests/groupby/test_missing.py::test_indices_with_missing": "TODO: Add a reason for failure", - "tests/groupby/test_numeric_only.py::TestNumericOnly::test_extrema[max]": "Failed: DID NOT RAISE ", - "tests/groupby/test_numeric_only.py::TestNumericOnly::test_extrema[min]": "Failed: DID NOT RAISE ", "tests/groupby/test_reductions.py::test_basic_aggregations[float32]": "AssertionError: Attributes of Series are different", "tests/groupby/test_reductions.py::test_basic_aggregations[int32]": "AssertionError: Attributes of Series are different", "tests/groupby/test_reductions.py::test_groupby_mean_no_overflow": "TODO: Add a reason for failure", @@ -5243,6 +5241,13 @@ def pytest_unconfigure(config): "tests/window/moments/test_moments_consistency_rolling.py::test_rolling_apply_consistency_sum[all_data7-rolling_consistency_cases0-False-sum]": "pandas xfails, but xpasses with cudf.pandas", } +# Keep keys in alphabetical order +NODEIDS_THAT_MAY_FAIL = { + "tests/groupby/test_numeric_only.py::TestNumericOnly::test_extrema[max]": "Environment-sensitive TypeError expectation", + "tests/groupby/test_numeric_only.py::TestNumericOnly::test_extrema[min]": "Environment-sensitive TypeError expectation", + "tests/io/test_spss.py::test_spss_metadata": "pandas 3.0.3 metadata expectation is incompatible with pyreadstat 1.3.6", +} + def pytest_configure(config): config.addinivalue_line( @@ -5272,5 +5277,9 @@ def pytest_collection_modifyitems(session, config, items): ) ) is not None: item.add_marker(pytest.mark.skip(reason=reason)) + elif ( + reason := NODEIDS_THAT_MAY_FAIL.get(item.nodeid, None) + ) is not None: + item.add_marker(pytest.mark.xfail(reason=reason, strict=False)) elif (reason := NODEIDS_THAT_FAIL.get(item.nodeid, None)) is not None: item.add_marker(pytest.mark.xfail(reason=reason)) diff --git a/python/pylibcudf/pylibcudf/__init__.py b/python/pylibcudf/pylibcudf/__init__.py index 3df8e83addab..afce8bdea681 100644 --- a/python/pylibcudf/pylibcudf/__init__.py +++ b/python/pylibcudf/pylibcudf/__init__.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # If libcudf was installed as a wheel, we must request it to load the library symbols. @@ -52,6 +52,7 @@ transform, transpose, types, + typing, unary, utilities, utils, @@ -110,6 +111,7 @@ "transform", "transpose", "types", + "typing", "unary", "utilities", "utils", diff --git a/python/pylibcudf/pylibcudf/aggregation.pyx b/python/pylibcudf/pylibcudf/aggregation.pyx index d7e7571bca2b..b7e6b5559354 100644 --- a/python/pylibcudf/pylibcudf/aggregation.pyx +++ b/python/pylibcudf/pylibcudf/aggregation.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cython.operator cimport dereference @@ -428,7 +428,7 @@ cpdef Aggregation median(): return Aggregation.from_libcudf(move(make_median_aggregation[aggregation]())) -cpdef Aggregation quantile(list quantiles, interpolation interp = interpolation.LINEAR): +cpdef Aggregation quantile(list quantiles: list[float], interpolation interp = interpolation.LINEAR): """Create a quantile aggregation. For details, see :cpp:func:`make_quantile_aggregation`. diff --git a/python/pylibcudf/pylibcudf/binaryop.pyx b/python/pylibcudf/pylibcudf/binaryop.pyx index ea72b2d3d4e8..8ec2777a848e 100644 --- a/python/pylibcudf/pylibcudf/binaryop.pyx +++ b/python/pylibcudf/pylibcudf/binaryop.pyx @@ -22,6 +22,10 @@ from .scalar cimport Scalar from .types cimport DataType from .utils cimport _get_stream, _get_memory_resource from cuda.bindings.cyruntime cimport cudaStream_t +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike __all__ = ["BinaryOperator", "binary_operation", "is_supported_operation"] @@ -30,7 +34,7 @@ cpdef Column binary_operation( RightBinaryOperand rhs, binary_operator op, DataType output_type, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a binary operation between a column and another column or scalar. diff --git a/python/pylibcudf/pylibcudf/column.pyx b/python/pylibcudf/pylibcudf/column.pyx index f7f2a59925e0..3ffe35ee212b 100644 --- a/python/pylibcudf/pylibcudf/column.pyx +++ b/python/pylibcudf/pylibcudf/column.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Iterable, Sequence from cython.operator cimport dereference from cpython.exc cimport PyErr_Occurred @@ -58,15 +59,19 @@ from .traits cimport ( from .types cimport DataType, size_of, type_id from pylibcudf.types import TypeId from .utils cimport _get_stream, _get_memory_resource +from pylibcudf.typing import SupportsArrayInterface, SupportsCudaArrayInterface +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from .gpumemoryview import _datatype_from_dtype_desc from ._interop_helpers import ArrowLike, ColumnMetadata, _ObjectWithArrowMetadata +from itertools import accumulate import array import functools import operator -from itertools import accumulate -from typing import Iterable from cuda.bindings.cyruntime cimport cudaStream_t try: @@ -77,7 +82,12 @@ except ImportError as e: pa_err = e -__all__ = ["Column", "ListsColumnView", "StructsColumnView", "is_c_contiguous"] +__all__ = [ + "Column", + "ListsColumnView", + "StructsColumnView", + "is_c_contiguous", +] cdef is_iterable(obj): @@ -157,7 +167,7 @@ class ArrayInterfaceWrapper: self.__array_interface__ = iface -cdef gpumemoryview _copy_array_to_device(object buf, object stream=None): +cdef gpumemoryview _copy_array_to_device(object buf, object stream: CudaStreamLike | None = None): """ Copy a host-side array.array buffer to device memory. @@ -346,7 +356,7 @@ cdef class Column: def __init__( self, DataType data_type not None, size_type size, object data, object mask, size_type null_count, size_type offset, - children, bint validate=True + children: Iterable[Column], bint validate=True ): children = list(children) if not all(isinstance(c, Column) for c in children): @@ -406,7 +416,7 @@ cdef class Column: def from_arrow( obj: ArrowLike, dtype: DataType | None = None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ) -> ArrowLike: """ @@ -604,7 +614,7 @@ cdef class Column: DeviceBuffer buff, DataType dtype, size_type size, - children, + children: Iterable[Column], ): """ Create a Column from an RMM DeviceBuffer. @@ -813,7 +823,7 @@ cdef class Column: def from_scalar( Scalar slr, size_type size, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a Column from a Scalar. @@ -846,7 +856,7 @@ cdef class Column: ) return Column.from_libcudf(move(c_result), _stream, mr) - cpdef Scalar to_scalar(self, object stream=None, DeviceMemoryResource mr=None): + cpdef Scalar to_scalar(self, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Return the first value of 1-element column as a Scalar. @@ -882,7 +892,7 @@ cdef class Column: def all_null_like( Column like, size_type size, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create an all null column from a template. @@ -921,7 +931,7 @@ cdef class Column: tuple shape, DataType dtype, Column base=None, - object stream=None, + object stream: CudaStreamLike | None = None, ): """ Construct a list Column from a gpumemoryview and array @@ -974,7 +984,11 @@ cdef class Column: return nested @classmethod - def from_array_interface(cls, obj, object stream=None): + def from_array_interface( + cls, + obj: SupportsArrayInterface, + object stream: CudaStreamLike | None = None, + ): """ Create a Column from an object implementing the NumPy Array Interface. @@ -1028,7 +1042,11 @@ cdef class Column: ) @classmethod - def from_cuda_array_interface(cls, obj, object stream=None): + def from_cuda_array_interface( + cls, + obj: SupportsCudaArrayInterface, + object stream: CudaStreamLike | None = None, + ): """ Create a Column from an object implementing the CUDA Array Interface. @@ -1067,7 +1085,11 @@ cdef class Column: ) @classmethod - def from_array(cls, obj, object stream=None): + def from_array( + cls, + obj: SupportsCudaArrayInterface | SupportsArrayInterface, + object stream: CudaStreamLike | None = None, + ): """ Create a Column from any object which supports the NumPy or CUDA array interface. @@ -1113,7 +1135,7 @@ cdef class Column: def from_iterable_of_py( obj: Iterable, dtype: DataType | None = None, - object stream=None + object stream: CudaStreamLike | None = None ) -> Column: """ Create a Column from a Python iterable of scalar values or nested iterables. @@ -1399,7 +1421,7 @@ cdef class Column: """The children of the column.""" return self._children - cpdef Column copy(self, object stream=None, DeviceMemoryResource mr=None): + cpdef Column copy(self, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Create a copy of the column.""" cdef unique_ptr[column] c_result cdef Stream _stream = _get_stream(stream) @@ -1461,7 +1483,7 @@ cdef class Column: return PyCapsule_New(raw_schema_ptr, 'arrow_schema', _release_schema) - def _to_host_array(self, object stream): + def _to_host_array(self, object stream: CudaStreamLike): cdef ArrowArray* raw_host_array_ptr cdef Stream _stream = _get_stream(stream) cdef cudaStream_t _cs = _stream.view().value() @@ -1532,7 +1554,7 @@ cdef class ListsColumnView: """ return lists_column_view(self._column.view()) - cpdef Column get_sliced_child(self, object stream=None): + cpdef Column get_sliced_child(self, object stream: CudaStreamLike | None = None): """ Get the list elements child properly sliced to match parent's view. @@ -1570,7 +1592,7 @@ cdef class StructsColumnView: """ return structs_column_view(self._column.view()) - cpdef Column get_sliced_child(self, int index, object stream=None): + cpdef Column get_sliced_child(self, int index, object stream: CudaStreamLike | None = None): """ Get the struct elements child properly sliced to match parent's view. diff --git a/python/pylibcudf/pylibcudf/column_factories.pyx b/python/pylibcudf/pylibcudf/column_factories.pyx index 45d590f41060..7b16d0d57126 100644 --- a/python/pylibcudf/pylibcudf/column_factories.pyx +++ b/python/pylibcudf/pylibcudf/column_factories.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.memory cimport unique_ptr from libcpp.utility cimport move @@ -20,6 +20,10 @@ from .types cimport DataType, type_id from .types import MaskState, TypeId from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -34,8 +38,8 @@ __all__ = [ ] cpdef Column make_empty_column( - MakeEmptyColumnOperand type_or_id, - object stream=None, + MakeEmptyColumnOperand type_or_id: DataType | TypeId, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Creates an empty column of the specified type. @@ -83,7 +87,7 @@ cpdef Column make_numeric_column( DataType type_, size_type size, MaskArg mstate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Creates an empty numeric column. @@ -122,7 +126,7 @@ cpdef Column make_fixed_point_column( DataType type_, size_type size, MaskArg mstate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): @@ -158,7 +162,7 @@ cpdef Column make_timestamp_column( DataType type_, size_type size, MaskArg mstate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): @@ -194,7 +198,7 @@ cpdef Column make_duration_column( DataType type_, size_type size, MaskArg mstate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): @@ -230,7 +234,7 @@ cpdef Column make_fixed_width_column( DataType type_, size_type size, MaskArg mstate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): @@ -264,7 +268,7 @@ cpdef Column make_fixed_width_column( cpdef Column make_empty_lists_column( DataType child_type, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Creates an empty column of the specified type. diff --git a/python/pylibcudf/pylibcudf/concatenate.pyx b/python/pylibcudf/pylibcudf/concatenate.pyx index d9bf0af16655..d81654c1a14c 100644 --- a/python/pylibcudf/pylibcudf/concatenate.pyx +++ b/python/pylibcudf/pylibcudf/concatenate.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from libcpp.memory cimport unique_ptr from libcpp.utility cimport move from libcpp.vector cimport vector @@ -16,11 +17,19 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["concatenate"] -cpdef concatenate(objects, object stream=None, DeviceMemoryResource mr=None): +cpdef concatenate( + objects: Sequence[Column] | Sequence[Table], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, +): """Concatenate columns or tables. Parameters diff --git a/python/pylibcudf/pylibcudf/contiguous_split.pyx b/python/pylibcudf/pylibcudf/contiguous_split.pyx index 1f5cb5c52f89..de665b691f9c 100644 --- a/python/pylibcudf/pylibcudf/contiguous_split.pyx +++ b/python/pylibcudf/pylibcudf/contiguous_split.pyx @@ -37,6 +37,10 @@ from .gpumemoryview cimport gpumemoryview from .table cimport Table from .span import is_span from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -165,7 +169,7 @@ cdef class ChunkedPack: def create( Table input, size_t user_buffer_size, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource temp_mr=None, ): """ @@ -319,7 +323,7 @@ cdef class ChunkedPack: ) -cpdef PackedColumns pack(Table input, object stream=None, DeviceMemoryResource mr=None): +cpdef PackedColumns pack(Table input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Deep-copy a table into a serialized contiguous memory format. Later use `unpack` or `unpack_from_memoryviews` to unpack the serialized @@ -361,7 +365,7 @@ cpdef PackedColumns pack(Table input, object stream=None, DeviceMemoryResource m return PackedColumns.from_libcudf(move(pack), _stream, mr) -cpdef Table unpack(PackedColumns input, object stream=None): +cpdef Table unpack(PackedColumns input, object stream: CudaStreamLike | None = None): """Deserialize the result of `pack`. Copies the result of a serialized table into a table. @@ -390,7 +394,7 @@ cpdef Table unpack(PackedColumns input, object stream=None): cpdef Table unpack_from_memoryviews( memoryview metadata, object gpu_data, - object stream=None, + object stream: CudaStreamLike | None = None, ): """Deserialize the result of `pack`. diff --git a/python/pylibcudf/pylibcudf/copying.pyx b/python/pylibcudf/pylibcudf/copying.pyx index 22289c325835..9b09536d314d 100644 --- a/python/pylibcudf/pylibcudf/copying.pyx +++ b/python/pylibcudf/pylibcudf/copying.pyx @@ -40,6 +40,10 @@ from .column cimport Column from .scalar cimport Scalar from .table cimport Table from .utils cimport _as_vector, _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -65,7 +69,7 @@ cpdef Table gather( Table source_table, Column gather_map, out_of_bounds_policy bounds_policy, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Select rows from source_table according to the provided gather_map. @@ -117,7 +121,7 @@ cpdef Table scatter( TableOrListOfScalars source, Column scatter_map, Table target_table, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Scatter from source into target_table according to scatter_map. @@ -195,7 +199,7 @@ cpdef Table scatter( cpdef ColumnOrTable empty_like( - ColumnOrTable input, object stream=None, DeviceMemoryResource mr=None + ColumnOrTable input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Create an empty column or table with the same type as ``input``. @@ -236,7 +240,7 @@ cpdef Column allocate_like( Column input_column, mask_allocation_policy policy, size=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Allocate a column with the same type as input_column. @@ -286,7 +290,7 @@ cpdef Column copy_range_in_place( size_type input_begin, size_type input_end, size_type target_begin, - object stream=None + object stream: CudaStreamLike | None = None ): """Copy a range of elements from input_column to target_column. @@ -344,7 +348,7 @@ cpdef Column copy_range( size_type input_begin, size_type input_end, size_type target_begin, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Copy a range of elements from input_column to target_column. @@ -404,7 +408,7 @@ cpdef Column shift( Column input, size_type offset, Scalar fill_value, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Shift the elements of input by offset. @@ -451,7 +455,11 @@ cpdef Column shift( return Column.from_libcudf(move(c_result), _stream, mr) -cpdef list slice(ColumnOrTable input, list indices, object stream=None): +cpdef list slice( + ColumnOrTable input, + list indices: list[int], + object stream: CudaStreamLike | None = None, +): """Slice input according to indices. For details on the implementation, see :cpp:func:`slice`. @@ -508,7 +516,11 @@ cpdef list slice(ColumnOrTable input, list indices, object stream=None): ] -cpdef list split(ColumnOrTable input, list splits, object stream=None): +cpdef list split( + ColumnOrTable input, + list splits: list[int], + object stream: CudaStreamLike | None = None, +): """Split input into multiple. For details on the implementation, see :cpp:func:`split`. @@ -561,7 +573,7 @@ cpdef Column copy_if_else( LeftCopyIfElseOperand lhs, RightCopyIfElseOperand rhs, Column boolean_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Copy elements from lhs or rhs into a new column according to boolean_mask. @@ -656,7 +668,7 @@ cpdef Table boolean_mask_scatter( TableOrListOfScalars input, Table target, Column boolean_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Scatter rows from input into target according to boolean_mask. @@ -733,7 +745,7 @@ cpdef Table boolean_mask_scatter( cpdef Scalar get_element( Column input_column, size_type index, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Get the element at index from input_column. diff --git a/python/pylibcudf/pylibcudf/datetime.pyx b/python/pylibcudf/pylibcudf/datetime.pyx index 91b6ef83037e..38737cc45393 100644 --- a/python/pylibcudf/pylibcudf/datetime.pyx +++ b/python/pylibcudf/pylibcudf/datetime.pyx @@ -31,6 +31,10 @@ from rmm.pylibrmm.stream cimport Stream from .column cimport Column from .scalar cimport Scalar from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -51,7 +55,7 @@ __all__ = [ cpdef Column extract_datetime_component( Column input, datetime_component component, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -89,7 +93,7 @@ cpdef Column extract_datetime_component( cpdef Column ceil_datetimes( Column input, rounding_frequency freq, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -125,7 +129,7 @@ cpdef Column ceil_datetimes( cpdef Column floor_datetimes( Column input, rounding_frequency freq, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -161,7 +165,7 @@ cpdef Column floor_datetimes( cpdef Column round_datetimes( Column input, rounding_frequency freq, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -197,7 +201,7 @@ cpdef Column round_datetimes( cpdef Column add_calendrical_months( Column input, ColumnOrScalar months, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -247,7 +251,7 @@ cpdef Column add_calendrical_months( return Column.from_libcudf(move(result), _stream, mr) cpdef Column day_of_year( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Computes the day number since the start of @@ -280,7 +284,7 @@ cpdef Column day_of_year( return Column.from_libcudf(move(result), _stream, mr) cpdef Column is_leap_year( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Check if the year of the given date is a leap year. @@ -312,7 +316,7 @@ cpdef Column is_leap_year( return Column.from_libcudf(move(result), _stream, mr) cpdef Column last_day_of_month( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Computes the last day of the month. @@ -344,7 +348,7 @@ cpdef Column last_day_of_month( return Column.from_libcudf(move(result), _stream, mr) cpdef Column extract_quarter( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns the quarter (ie. a value from {1, 2, 3, 4}) @@ -376,7 +380,7 @@ cpdef Column extract_quarter( return Column.from_libcudf(move(result), _stream, mr) cpdef Column days_in_month( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Extract the number of days in the month. diff --git a/python/pylibcudf/pylibcudf/experimental/_join_streams.pyx b/python/pylibcudf/pylibcudf/experimental/_join_streams.pyx index d9efcb19ed9c..783aa2bce584 100644 --- a/python/pylibcudf/pylibcudf/experimental/_join_streams.pyx +++ b/python/pylibcudf/pylibcudf/experimental/_join_streams.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cuda.bindings.cyruntime cimport cudaStream_t @@ -11,12 +11,18 @@ from pylibcudf.libcudf.utilities.span cimport host_span from rmm.pylibrmm.stream cimport Stream from ..utils cimport _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike __all__ = ["join_streams"] -cpdef void join_streams(list streams, object stream): +cpdef void join_streams( + list streams: list[CudaStreamLike], object stream: CudaStreamLike +): """Synchronize a stream to an event on a set of streams. This function synchronizes the joined stream with the waited-on streams diff --git a/python/pylibcudf/pylibcudf/expressions.pyx b/python/pylibcudf/pylibcudf/expressions.pyx index 26fad2bc7b9f..1c7fbaf93a75 100644 --- a/python/pylibcudf/pylibcudf/expressions.pyx +++ b/python/pylibcudf/pylibcudf/expressions.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import ast import functools @@ -476,7 +476,7 @@ class ExpressionTransformer(ast.NodeVisitor): @functools.lru_cache(256) -def to_expression(str expr, tuple column_names): +def to_expression(str expr, tuple column_names: tuple[str, ...]): """ Create an expression for `pylibcudf.transform.compute_column`. diff --git a/python/pylibcudf/pylibcudf/filling.pyx b/python/pylibcudf/pylibcudf/filling.pyx index 4a11a33ea8a4..300f90e5f71b 100644 --- a/python/pylibcudf/pylibcudf/filling.pyx +++ b/python/pylibcudf/pylibcudf/filling.pyx @@ -26,6 +26,10 @@ from .column cimport Column from .scalar cimport Scalar from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -42,7 +46,7 @@ cpdef Column fill( size_type begin, size_type end, Scalar value, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): @@ -94,7 +98,7 @@ cpdef void fill_in_place( size_type begin, size_type end, Scalar value, - object stream=None, + object stream: CudaStreamLike | None = None, ): """Fill destination column in place from begin to end with value. @@ -137,7 +141,7 @@ cpdef Column sequence( size_type size, Scalar init, Scalar step, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a sequence column of size ``size`` with initial value ``init`` and step @@ -183,7 +187,7 @@ cpdef Column sequence( cpdef Table repeat( Table input_table, ColumnOrSize count, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Repeat rows of a Table. @@ -245,7 +249,7 @@ cpdef Column calendrical_month_sequence( size_type n, Scalar init, size_type months, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): diff --git a/python/pylibcudf/pylibcudf/gpumemoryview.pyx b/python/pylibcudf/pylibcudf/gpumemoryview.pyx index 0312198e2a84..34cedb5331ea 100644 --- a/python/pylibcudf/pylibcudf/gpumemoryview.pyx +++ b/python/pylibcudf/pylibcudf/gpumemoryview.pyx @@ -3,8 +3,11 @@ from libc.stddef cimport size_t from libc.stdint cimport uintptr_t, uint64_t +from collections.abc import Mapping +import cython import functools import operator +from typing import Any from .types cimport DataType, size_of, type_id @@ -83,7 +86,7 @@ cdef class gpumemoryview: self.nbytes = functools.reduce(operator.mul, cai["shape"]) * itemsize @property - def __cuda_array_interface__(self): + def __cuda_array_interface__(self) -> Mapping[str, Any]: return self.cai @property @@ -95,10 +98,11 @@ cdef class gpumemoryview: """ return self.nbytes - def __len__(self): + def __len__(self) -> int: return self.cai["shape"][0] - def byte_slice(self, s): + @cython.annotation_typing(False) + def byte_slice(self, s: slice) -> gpumemoryview: """Return a byte-range sub-view of this buffer. Parameters diff --git a/python/pylibcudf/pylibcudf/groupby.pyx b/python/pylibcudf/pylibcudf/groupby.pyx index 1646639bbb83..acaaac857677 100644 --- a/python/pylibcudf/pylibcudf/groupby.pyx +++ b/python/pylibcudf/pylibcudf/groupby.pyx @@ -30,6 +30,13 @@ from .types cimport null_order, null_policy, order, sorted from .utils cimport _as_vector, _get_stream, _get_memory_resource from cuda.bindings.cyruntime cimport cudaStream_t +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.replace import ReplacePolicy + from pylibcudf.scalar import Scalar + from pylibcudf.typing import CudaStreamLike + __all__ = ["GroupBy", "GroupByRequest"] @@ -49,7 +56,7 @@ cdef class GroupByRequest: aggregations : List[Aggregation] The list of aggregations to perform. """ - def __init__(self, Column values, list aggregations): + def __init__(self, Column values, list aggregations: list[Aggregation]): self._values = values self._aggregations = aggregations @@ -163,7 +170,10 @@ cdef class GroupBy: return group_keys, results cpdef tuple aggregate( - self, list requests, object stream=None, DeviceMemoryResource mr=None + self, + list requests: list[GroupByRequest], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, ): """Compute aggregations on columns. @@ -204,7 +214,10 @@ cdef class GroupBy: return GroupBy._parse_outputs(move(c_res), _stream, mr) cpdef tuple scan( - self, list requests, object stream=None, DeviceMemoryResource mr=None + self, + list requests: list[GroupByRequest], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, ): """Compute scans on columns. @@ -246,9 +259,9 @@ cdef class GroupBy: cpdef tuple shift( self, Table values, - list offset, - list fill_values, - object stream=None, + list offset: list[int], + list fill_values: list[Scalar], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Compute shifts on columns. @@ -299,8 +312,8 @@ cdef class GroupBy: cpdef tuple replace_nulls( self, Table value, - list replace_policies, - object stream=None, + list replace_policies: list[ReplacePolicy], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replace nulls in columns. @@ -341,7 +354,7 @@ cdef class GroupBy: ) cpdef tuple get_groups( - self, Table values=None, object stream=None, DeviceMemoryResource mr=None + self, Table values=None, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Get the grouped keys and values labels for each row. diff --git a/python/pylibcudf/pylibcudf/hashing.pyx b/python/pylibcudf/pylibcudf/hashing.pyx index fb2fcad3c82a..ac866828c36e 100644 --- a/python/pylibcudf/pylibcudf/hashing.pyx +++ b/python/pylibcudf/pylibcudf/hashing.pyx @@ -25,6 +25,10 @@ from rmm.pylibrmm.stream cimport Stream from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -46,7 +50,7 @@ LIBCUDF_DEFAULT_HASH_SEED = DEFAULT_HASH_SEED cpdef Column murmurhash3_x86_32( Table input, uint32_t seed=DEFAULT_HASH_SEED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the MurmurHash3 32-bit hash value of each row in the given table. @@ -86,7 +90,7 @@ cpdef Column murmurhash3_x86_32( cpdef Table murmurhash3_x64_128( Table input, uint64_t seed=DEFAULT_HASH_SEED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the MurmurHash3 64-bit hash value of each row in the given table. @@ -126,7 +130,7 @@ cpdef Table murmurhash3_x64_128( cpdef Column xxhash_32( Table input, uint32_t seed=DEFAULT_HASH_SEED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the xxHash 32-bit hash value of each row in the given table. @@ -167,7 +171,7 @@ cpdef Column xxhash_32( cpdef Column xxhash_64( Table input, uint64_t seed=DEFAULT_HASH_SEED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the xxHash 64-bit hash value of each row in the given table. @@ -207,7 +211,7 @@ cpdef Column xxhash_64( cpdef Column md5( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the MD5 hash value of each row in the given table. @@ -241,7 +245,7 @@ cpdef Column md5( cpdef Column sha1( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the SHA-1 hash value of each row in the given table. @@ -274,7 +278,7 @@ cpdef Column sha1( cpdef Column sha224( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the SHA-224 hash value of each row in the given table. @@ -307,7 +311,7 @@ cpdef Column sha224( cpdef Column sha256( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the SHA-256 hash value of each row in the given table. @@ -340,7 +344,7 @@ cpdef Column sha256( cpdef Column sha384( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the SHA-384 hash value of each row in the given table. @@ -373,7 +377,7 @@ cpdef Column sha384( cpdef Column sha512( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the SHA-512 hash value of each row in the given table. diff --git a/python/pylibcudf/pylibcudf/interop.pyx b/python/pylibcudf/pylibcudf/interop.pyx index b31792f13b2f..afd1c2ea6b64 100644 --- a/python/pylibcudf/pylibcudf/interop.pyx +++ b/python/pylibcudf/pylibcudf/interop.pyx @@ -23,6 +23,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from ._interop_helpers import ColumnMetadata from cuda.bindings.cyruntime cimport cudaStream_t @@ -35,7 +39,7 @@ __all__ = [ cpdef Table from_dlpack( - object managed_tensor, object stream=None, DeviceMemoryResource mr=None + object managed_tensor, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Convert a DLPack DLTensor into a cudf table. @@ -82,7 +86,7 @@ cpdef Table from_dlpack( return result -cpdef object to_dlpack(Table input, object stream=None, DeviceMemoryResource mr=None): +cpdef object to_dlpack(Table input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Convert a cudf table into a DLPack DLTensor. diff --git a/python/pylibcudf/pylibcudf/io/avro.pyx b/python/pylibcudf/pylibcudf/io/avro.pyx index f2bd021cdde2..7343fe46559e 100644 --- a/python/pylibcudf/pylibcudf/io/avro.pyx +++ b/python/pylibcudf/pylibcudf/io/avro.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.string cimport string @@ -19,6 +19,10 @@ from pylibcudf.libcudf.io.avro cimport ( from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike __all__ = ["read_avro", "AvroReaderOptions", "AvroReaderOptionsBuilder"] @@ -53,7 +57,7 @@ cdef class AvroReaderOptions: avro_builder.source = source return avro_builder - cpdef void set_columns(self, list col_names): + cpdef void set_columns(self, list col_names: list[str]): """ Set names of the column to be read. @@ -89,7 +93,7 @@ cdef class AvroReaderOptions: cdef class AvroReaderOptionsBuilder: - cpdef AvroReaderOptionsBuilder columns(self, list col_names): + cpdef AvroReaderOptionsBuilder columns(self, list col_names: list[str]): """ Set names of the column to be read. @@ -153,7 +157,7 @@ cdef class AvroReaderOptionsBuilder: cpdef TableWithMetadata read_avro( AvroReaderOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/io/csv.pyx b/python/pylibcudf/pylibcudf/io/csv.pyx index 16824f2b9f24..71de5671ecf9 100644 --- a/python/pylibcudf/pylibcudf/io/csv.pyx +++ b/python/pylibcudf/pylibcudf/io/csv.pyx @@ -34,6 +34,10 @@ from pylibcudf.table cimport Table from pylibcudf.types cimport DataType from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike __all__ = [ "read_csv", @@ -87,7 +91,7 @@ cdef class CsvReaderOptions: """ self.c_obj.set_header(header) - cpdef void set_names(self, list col_names): + cpdef void set_names(self, list col_names: list[str]): """ Sets names of the column. @@ -120,7 +124,7 @@ cdef class CsvReaderOptions: """ self.c_obj.set_prefix(prefix.encode()) - cpdef void set_use_cols_indexes(self, list col_indices): + cpdef void set_use_cols_indexes(self, list col_indices: list[int]): """ Sets indexes of columns to read. @@ -138,7 +142,7 @@ cdef class CsvReaderOptions: vec.push_back(i) self.c_obj.set_use_cols_indexes(vec) - cpdef void set_use_cols_names(self, list col_names): + cpdef void set_use_cols_names(self, list col_names: list[str]): """ Sets names of the columns to be read. @@ -201,7 +205,7 @@ cdef class CsvReaderOptions: """ self.c_obj.set_comment(ord(comment)) - cpdef void set_parse_dates(self, list val): + cpdef void set_parse_dates(self, list val: list[int | str]): """ Sets indexes or names of columns to read as datetime. @@ -227,7 +231,7 @@ cdef class CsvReaderOptions: self.c_obj.set_parse_dates(vec_str) self.c_obj.set_parse_dates(vec_int) - cpdef void set_parse_hex(self, list val): + cpdef void set_parse_hex(self, list val: list[int | str]): """ Sets indexes or names of columns to parse as hexadecimal. @@ -254,7 +258,7 @@ cdef class CsvReaderOptions: self.c_obj.set_parse_hex(vec_str) self.c_obj.set_parse_hex(vec_int) - cpdef void set_dtypes(self, object types): + cpdef void set_dtypes(self, object types: dict[str, DataType] | list[DataType]): """ Sets per-column types. @@ -281,7 +285,7 @@ cdef class CsvReaderOptions: else: raise TypeError("Must pass an dict or list") - cpdef void set_true_values(self, list true_values): + cpdef void set_true_values(self, list true_values: list[str]): """ Sets additional values to recognize as boolean true values. @@ -299,7 +303,7 @@ cdef class CsvReaderOptions: vec.push_back(val.encode()) self.c_obj.set_true_values(vec) - cpdef void set_false_values(self, list false_values): + cpdef void set_false_values(self, list false_values: list[str]): """ Sets additional values to recognize as boolean false values. @@ -317,7 +321,7 @@ cdef class CsvReaderOptions: vec.push_back(val.encode()) self.c_obj.set_false_values(vec) - cpdef void set_na_values(self, list na_values): + cpdef void set_na_values(self, list na_values: list[str]): """ Sets additional values to recognize as null values. @@ -673,7 +677,7 @@ cdef class CsvReaderOptionsBuilder: cpdef TableWithMetadata read_csv( CsvReaderOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -902,7 +906,7 @@ cdef class CsvWriterOptionsBuilder: cpdef void write_csv( CsvWriterOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, ): """ Write to CSV format. diff --git a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx index 664eb489428c..dfdaabfe5b89 100644 --- a/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx +++ b/python/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyx @@ -31,6 +31,11 @@ from pylibcudf.libcudf.io.types cimport table_with_metadata from pylibcudf.libcudf.types cimport size_type from pylibcudf.libcudf.utilities.span cimport device_span, host_span from pylibcudf.utils cimport _get_memory_resource, _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing_extensions import Buffer + from pylibcudf.typing import CudaStreamLike from pylibcudf.span import is_span from pylibcudf.io.parquet_metadata import FileMetaData @@ -81,14 +86,20 @@ cdef class HybridScanReader: >>> row_groups = reader.all_row_groups(options) """ - def __init__(self, const uint8_t[::1] footer_bytes, ParquetReaderOptions options): + def __init__( + self, + const uint8_t[::1] footer_bytes: Buffer, + ParquetReaderOptions options, + ): self.c_obj = make_unique[cpp_hybrid_scan_reader]( host_span[const_uint8_t](&footer_bytes[0], len(footer_bytes)), options.c_obj ) @staticmethod - def from_parquet_metadata(c_FileMetaData metadata, ParquetReaderOptions options): + def from_parquet_metadata( + c_FileMetaData metadata, ParquetReaderOptions options + ) -> HybridScanReader: """Create a HybridScanReader from pre-populated metadata. Parameters @@ -109,7 +120,7 @@ cdef class HybridScanReader: ) return reader - def parquet_metadata(self): + def parquet_metadata(self) -> FileMetaData: """Get the Parquet file footer metadata. Returns @@ -119,7 +130,7 @@ cdef class HybridScanReader: """ return c_FileMetaData.from_cpp(self.c_obj.get()[0].parquet_metadata()) - def page_index_byte_range(self): + def page_index_byte_range(self) -> ByteRangeInfo: """Get the byte range of the page index. Returns @@ -130,7 +141,9 @@ cdef class HybridScanReader: cdef byte_range_info info = self.c_obj.get()[0].page_index_byte_range() return ByteRangeInfo(info.offset(), info.size()) - def setup_page_index(self, const uint8_t[::1] page_index_bytes): + def setup_page_index( + self, const uint8_t[::1] page_index_bytes: Buffer + ) -> None: """Setup the page index within the Parquet file metadata. Parameters @@ -142,7 +155,7 @@ cdef class HybridScanReader: host_span[const_uint8_t](&page_index_bytes[0], len(page_index_bytes)) ) - def all_row_groups(self, ParquetReaderOptions options): + def all_row_groups(self, ParquetReaderOptions options) -> list[int]: """Get all available row groups from the parquet file. Parameters @@ -160,7 +173,9 @@ cdef class HybridScanReader: ) return list(row_groups) - def total_rows_in_row_groups(self, list row_group_indices): + def total_rows_in_row_groups( + self, list row_group_indices: list[int] + ) -> int: """Get the total number of top-level rows in the row groups. Parameters @@ -178,7 +193,7 @@ cdef class HybridScanReader: std_span[const_size_type](indices_vec.data(), indices_vec.size()) ) - def reset_column_selection(self): + def reset_column_selection(self) -> None: """Reset the column selection state. Resets the internal column selection state forcing re-selection of columns in @@ -188,10 +203,10 @@ cdef class HybridScanReader: def filter_row_groups_with_stats( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options, - object stream=None - ): + object stream: CudaStreamLike | None = None + ) -> list[int]: """Filter row groups using column chunk statistics. Parameters @@ -223,9 +238,9 @@ cdef class HybridScanReader: def secondary_filters_byte_ranges( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options - ): + ) -> tuple[list[ByteRangeInfo], list[ByteRangeInfo]]: """Get byte ranges of bloom filters and dictionary pages. Parameters @@ -258,10 +273,10 @@ cdef class HybridScanReader: def filter_row_groups_with_dictionary_pages( self, list dictionary_page_data, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options, - object stream=None - ): + object stream: CudaStreamLike | None = None + ) -> list[int]: """Filter row groups using column chunk dictionary pages. Parameters @@ -301,10 +316,10 @@ cdef class HybridScanReader: def filter_row_groups_with_bloom_filters( self, list bloom_filter_data, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options, - object stream=None - ): + object stream: CudaStreamLike | None = None + ) -> list[int]: """Filter row groups using column chunk bloom filters. Parameters @@ -343,11 +358,11 @@ cdef class HybridScanReader: def build_row_mask_with_page_index_stats( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> Column: """Build a boolean column indicating surviving rows from page stats. Parameters @@ -380,9 +395,9 @@ cdef class HybridScanReader: def filter_column_chunks_byte_ranges( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options - ): + ) -> list[ByteRangeInfo]: """Get byte ranges of column chunks of filter columns. Parameters @@ -407,14 +422,14 @@ cdef class HybridScanReader: def materialize_filter_columns( self, - list row_group_indices, + list row_group_indices: list[int], list column_chunk_data, Column row_mask, cpp_use_data_page_mask mask_data_pages, ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> TableWithMetadata: """Materialize filter columns and update the row mask. Parameters @@ -464,9 +479,9 @@ cdef class HybridScanReader: def payload_column_chunks_byte_ranges( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options - ): + ) -> list[ByteRangeInfo]: """Get byte ranges of column chunks of payload columns. Parameters @@ -491,14 +506,14 @@ cdef class HybridScanReader: def materialize_payload_columns( self, - list row_group_indices, + list row_group_indices: list[int], list column_chunk_data, Column row_mask, cpp_use_data_page_mask mask_data_pages, ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> TableWithMetadata: """Materialize payload columns and apply the row mask. Parameters @@ -548,9 +563,9 @@ cdef class HybridScanReader: def all_column_chunks_byte_ranges( self, - list row_group_indices, + list row_group_indices: list[int], ParquetReaderOptions options - ): + ) -> list[ByteRangeInfo]: """Get byte ranges of column chunks of all columns. Parameters @@ -575,12 +590,12 @@ cdef class HybridScanReader: def materialize_all_columns( self, - list row_group_indices, + list row_group_indices: list[int], list column_chunk_data, ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> TableWithMetadata: """Materialize all columns. Parameters @@ -624,14 +639,14 @@ cdef class HybridScanReader: self, size_t chunk_read_limit, size_t pass_read_limit, - list row_group_indices, + list row_group_indices: list[int], Column row_mask, cpp_use_data_page_mask mask_data_pages, list column_chunk_data, ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> None: """Setup chunking information for filter columns. Parameters @@ -682,7 +697,7 @@ cdef class HybridScanReader: def materialize_filter_columns_chunk( self, Column row_mask - ): + ) -> TableWithMetadata: """Materialize a chunk of filter columns. Parameters @@ -707,14 +722,14 @@ cdef class HybridScanReader: self, size_t chunk_read_limit, size_t pass_read_limit, - list row_group_indices, + list row_group_indices: list[int], Column row_mask, cpp_use_data_page_mask mask_data_pages, list column_chunk_data, ParquetReaderOptions options, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None - ): + ) -> None: """Setup chunking information for payload columns. Parameters @@ -765,7 +780,7 @@ cdef class HybridScanReader: def materialize_payload_columns_chunk( self, Column row_mask, - ): + ) -> TableWithMetadata: """Materialize a chunk of payload columns. Parameters @@ -788,9 +803,9 @@ cdef class HybridScanReader: def construct_row_group_passes( self, - list row_group_indices, + list row_group_indices: list[int], size_t pass_read_limit, - ): + ) -> list[list[int]]: """Partition row groups into passes such that the GPU memory required to materialize a pass is bounded by the specified limit. @@ -824,7 +839,7 @@ cdef class HybridScanReader: pass_read_limit ) - def has_next_table_chunk(self): + def has_next_table_chunk(self) -> bool: """Check if there is any parquet data left to read. Returns diff --git a/python/pylibcudf/pylibcudf/io/json.pyx b/python/pylibcudf/pylibcudf/io/json.pyx index a1c716f18a66..b38ee649339b 100644 --- a/python/pylibcudf/pylibcudf/io/json.pyx +++ b/python/pylibcudf/pylibcudf/io/json.pyx @@ -7,6 +7,8 @@ from libcpp.string cimport string from libcpp.utility cimport move from libcpp.vector cimport vector +from typing import TypeAlias + from rmm.pylibrmm.stream cimport Stream from pylibcudf.concatenate cimport concatenate @@ -46,6 +48,10 @@ from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.types cimport DataType from pylibcudf.utils cimport _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cython.operator import dereference @@ -63,6 +69,8 @@ __all__ = [ "JsonWriterOptionsBuilder" ] +NameAndType: TypeAlias = tuple[str, DataType, list["NameAndType"]] + cdef map[string, schema_element] _generate_schema_map(list dtypes): cdef map[string, schema_element] schema_map cdef schema_element s_elem @@ -177,7 +185,7 @@ cdef class JsonReaderOptions: json_builder.source = source return json_builder - cpdef void set_dtypes(self, list types): + cpdef void set_dtypes(self, list types: list[DataType] | list[NameAndType]): """ Set data types for columns to be read. @@ -330,7 +338,7 @@ cdef class JsonReaderOptions: cpdef void allow_nonnumeric_numbers(self, bool val): self.c_obj.allow_nonnumeric_numbers(val) - cpdef void set_na_values(self, list vals): + cpdef void set_na_values(self, list vals: list[str]): cdef vector[string] vec for val in vals: if isinstance(val, str): @@ -706,7 +714,7 @@ cdef class JsonReaderOptionsBuilder: cpdef tuple chunked_read_json( JsonReaderOptions options, int chunk_size=100_000_000, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr = None, ): """ @@ -775,7 +783,7 @@ cpdef tuple chunked_read_json( cpdef TableWithMetadata read_json( JsonReaderOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr = None ): """ @@ -814,7 +822,7 @@ cpdef TableWithMetadata read_json_from_string_column( list dtypes = None, compression_type compression = compression_type.NONE, json_recovery_mode_t recovery_mode = json_recovery_mode_t.RECOVER_WITH_NULL, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr = None ): """ @@ -1096,7 +1104,7 @@ cdef class JsonWriterOptionsBuilder: return json_options -cpdef void write_json(JsonWriterOptions options, object stream = None): +cpdef void write_json(JsonWriterOptions options, object stream: CudaStreamLike | None = None): """ Writes a set of columns to JSON format. diff --git a/python/pylibcudf/pylibcudf/io/orc.pyx b/python/pylibcudf/pylibcudf/io/orc.pyx index a25151bd0365..b91473d79541 100644 --- a/python/pylibcudf/pylibcudf/io/orc.pyx +++ b/python/pylibcudf/pylibcudf/io/orc.pyx @@ -56,6 +56,10 @@ from pylibcudf.types cimport DataType from pylibcudf.variant cimport get_if, holds_alternative from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike __all__ = [ @@ -324,7 +328,7 @@ cdef class OrcReaderOptions: """ self.c_obj.set_skip_rows(skip_rows) - cpdef void set_stripes(self, list stripes): + cpdef void set_stripes(self, list stripes: list[list[int]]): """ Sets list of stripes to read for each input source. @@ -346,7 +350,7 @@ cdef class OrcReaderOptions: vec.clear() self.c_obj.set_stripes(c_stripes) - cpdef void set_decimal128_columns(self, list val): + cpdef void set_decimal128_columns(self, list val: list[str]): """ Set columns that should be read as 128-bit Decimal. @@ -382,7 +386,7 @@ cdef class OrcReaderOptions: """ self.c_obj.set_timestamp_type(type_.c_obj) - cpdef void set_columns(self, list col_names): + cpdef void set_columns(self, list col_names: list[str]): """ Sets names of the column to read. @@ -446,7 +450,7 @@ cdef class OrcReaderOptionsBuilder: cpdef TableWithMetadata read_orc( - OrcReaderOptions options, object stream = None, DeviceMemoryResource mr=None + OrcReaderOptions options, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Read from ORC format. @@ -477,7 +481,7 @@ cpdef TableWithMetadata read_orc( cpdef ParsedOrcStatistics read_parsed_orc_statistics( SourceInfo source_info, - object stream=None + object stream: CudaStreamLike | None = None ): """ Read ORC statistics from a source. @@ -623,7 +627,7 @@ cdef class OrcWriterOptionsBuilder: self.c_obj.enable_statistics(val) return self - cpdef OrcWriterOptionsBuilder key_value_metadata(self, dict kvm): + cpdef OrcWriterOptionsBuilder key_value_metadata(self, dict kvm: dict[str, str]): """ Sets Key-Value footer metadata. @@ -670,7 +674,7 @@ cdef class OrcWriterOptionsBuilder: return orc_options -cpdef void write_orc(OrcWriterOptions options, object stream = None): +cpdef void write_orc(OrcWriterOptions options, object stream: CudaStreamLike | None = None): """ Write to ORC format. @@ -726,7 +730,10 @@ cdef class OrcChunkedWriter: self.c_obj.get()[0].write(c_table) @staticmethod - def from_options(ChunkedOrcWriterOptions options, object stream = None): + def from_options( + ChunkedOrcWriterOptions options, + object stream: CudaStreamLike | None = None, + ) -> OrcChunkedWriter: """ Creates a chunked ORC writer from options @@ -861,7 +868,7 @@ cdef class ChunkedOrcWriterOptionsBuilder: cpdef ChunkedOrcWriterOptionsBuilder key_value_metadata( self, - dict kvm + dict kvm: dict[str, str] ): """ Sets Key-Value footer metadata. diff --git a/python/pylibcudf/pylibcudf/io/parquet.pxd b/python/pylibcudf/pylibcudf/io/parquet.pxd index 58b186647286..3178c67dec11 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pxd +++ b/python/pylibcudf/pylibcudf/io/parquet.pxd @@ -85,7 +85,7 @@ cdef class ChunkedParquetReader: cpdef TableWithMetadata read_chunk(self, DeviceMemoryResource mr=*) -cpdef read_parquet( +cpdef TableWithMetadata read_parquet( ParquetReaderOptions options, object stream = *, DeviceMemoryResource mr=*, diff --git a/python/pylibcudf/pylibcudf/io/parquet.pyx b/python/pylibcudf/pylibcudf/io/parquet.pyx index d287d10900ec..109454b95120 100644 --- a/python/pylibcudf/pylibcudf/io/parquet.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet.pyx @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Mapping, Sequence from cython.operator cimport dereference import warnings @@ -50,6 +51,10 @@ from pylibcudf.libcudf.table.table_view cimport table_view from pylibcudf.libcudf.types cimport size_type, type_id from pylibcudf.table cimport Table from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -140,7 +145,7 @@ cdef class ParquetReaderOptions: parquet_builder.source = source return parquet_builder - cpdef void set_row_groups(self, list row_groups): + cpdef void set_row_groups(self, list row_groups: list[list[int]]): """ Sets list of individual row groups to read. @@ -210,7 +215,7 @@ cdef class ParquetReaderOptions: """ self.c_obj.set_skip_rows(skip_rows) - cpdef void set_columns(self, list col_names): + cpdef void set_columns(self, list col_names: list[str]): """ Sets names of the columns to be read. Deprecated and will be removed in a future version. Use set_column_names instead. @@ -227,7 +232,7 @@ cdef class ParquetReaderOptions: _warn_deprecated("set_columns", "set_column_names") self.set_column_names(col_names) - cpdef void set_column_names(self, list col_names): + cpdef void set_column_names(self, list col_names: list[str]): """ Sets names of the columns to be read. @@ -245,7 +250,7 @@ cdef class ParquetReaderOptions: vec.push_back(str(name).encode()) self.c_obj.set_column_names(vec) - cpdef void set_column_indices(self, list col_indices): + cpdef void set_column_indices(self, list col_indices: list[int]): """ Sets indices of the top-level columns to be read. @@ -263,7 +268,7 @@ cdef class ParquetReaderOptions: vec.push_back(idx) self.c_obj.set_column_indices(vec) - cpdef void set_column_field_ids(self, list column_field_ids): + cpdef void set_column_field_ids(self, list column_field_ids: list[int]): """ Sets Parquet field IDs of the columns/fields to be read. @@ -441,7 +446,7 @@ cdef class ParquetReaderOptionsBuilder: self.c_obj.filter(dereference(filter.c_obj)) return self - cpdef ParquetReaderOptionsBuilder columns(self, list col_names): + cpdef ParquetReaderOptionsBuilder columns(self, list col_names: list[str]): """ Sets names of the columns to be read. Deprecated and will be removed in a future version. Use column_names instead. @@ -458,7 +463,7 @@ cdef class ParquetReaderOptionsBuilder: _warn_deprecated("columns", "column_names") return self.column_names(col_names) - cpdef ParquetReaderOptionsBuilder column_names(self, list col_names): + cpdef ParquetReaderOptionsBuilder column_names(self, list col_names: list[str]): """ Sets names of the columns to be read. @@ -477,7 +482,7 @@ cdef class ParquetReaderOptionsBuilder: self.c_obj.column_names(vec) return self - cpdef ParquetReaderOptionsBuilder column_indices(self, list col_indices): + cpdef ParquetReaderOptionsBuilder column_indices(self, list col_indices: list[int]): """ Sets indices of the top-level columns to be read. @@ -496,7 +501,7 @@ cdef class ParquetReaderOptionsBuilder: self.c_obj.column_indices(vec) return self - cpdef ParquetReaderOptionsBuilder column_field_ids(self, list column_field_ids): + cpdef ParquetReaderOptionsBuilder column_field_ids(self, list column_field_ids: list[int]): """ Sets Parquet field IDs of the columns/fields to be read. @@ -602,11 +607,11 @@ cdef class ChunkedParquetReader: def __init__( self, ParquetReaderOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr = None, size_t chunk_read_limit=0, size_t pass_read_limit=1024000000, - object parquet_metadatas=None, + object parquet_metadatas: Sequence[FileMetaData] | None = None, ): self._stream = _get_stream(stream) self.mr = _get_memory_resource(mr) @@ -681,11 +686,11 @@ cdef class ChunkedParquetReader: return TableWithMetadata.from_libcudf(c_result, self._stream, mr) -cpdef read_parquet( +cpdef TableWithMetadata read_parquet( ParquetReaderOptions options, - object stream = None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, - object parquet_metadatas=None, + object parquet_metadatas: Sequence[FileMetaData] | None = None, ): """ Read from Parquet format. @@ -784,7 +789,10 @@ cdef class ChunkedParquetWriter: self.c_obj.get()[0].write(c_table, partitions) @staticmethod - def from_options(ChunkedParquetWriterOptions options, object stream = None): + def from_options( + ChunkedParquetWriterOptions options, + object stream: CudaStreamLike | None = None, + ) -> ChunkedParquetWriter: """ Creates a chunked Parquet writer from options @@ -858,7 +866,9 @@ cdef class ChunkedParquetWriterOptionsBuilder: self.c_obj.metadata(metadata.c_obj) return self - cpdef ChunkedParquetWriterOptionsBuilder key_value_metadata(self, metadata): + cpdef ChunkedParquetWriterOptionsBuilder key_value_metadata( + self, metadata: Sequence[Mapping[str, str]] + ): """ Sets Key-Value footer metadata. @@ -1047,7 +1057,7 @@ cdef class ParquetWriterOptions: bldr.sink_ref = sink return bldr - cpdef void set_partitions(self, list partitions): + cpdef void set_partitions(self, list partitions: list[PartitionInfo]): """ Sets partitions. @@ -1069,7 +1079,7 @@ cdef class ParquetWriterOptions: self.c_obj.set_partitions(c_partions) - cpdef void set_column_chunks_file_paths(self, file_paths): + cpdef void set_column_chunks_file_paths(self, file_paths: Sequence[str]): """ Sets column chunks file path to be set in the raw output metadata. @@ -1178,7 +1188,9 @@ cdef class ParquetWriterOptionsBuilder: self.c_obj.metadata(metadata.c_obj) return self - cpdef ParquetWriterOptionsBuilder key_value_metadata(self, metadata): + cpdef ParquetWriterOptionsBuilder key_value_metadata( + self, metadata: Sequence[Mapping[str, str]] + ): """ Sets Key-Value footer metadata. @@ -1380,7 +1392,7 @@ cdef class ParquetWriterOptionsBuilder: return parquet_options -cpdef memoryview write_parquet(ParquetWriterOptions options, object stream = None): +cpdef memoryview write_parquet(ParquetWriterOptions options, object stream: CudaStreamLike | None = None): """ Writes a set of columns to parquet format. diff --git a/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx b/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx index f8aca270642e..8a4362153eb6 100644 --- a/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx +++ b/python/pylibcudf/pylibcudf/io/parquet_metadata.pyx @@ -24,6 +24,11 @@ from pylibcudf.libcudf.io.parquet_schema cimport ( from pylibcudf.libcudf.utilities.span cimport host_span from pylibcudf.types cimport DataType +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing_extensions import Buffer + ctypedef const unique_ptr[datasource] const_unique_ptr_datasource @@ -284,17 +289,17 @@ cdef class SortingColumn: return result @property - def column_idx(self): + def column_idx(self) -> int: """Column index (within the row group).""" return self.c_obj.column_idx @property - def descending(self): + def descending(self) -> bool: """Whether this column is sorted in descending order.""" return self.c_obj.descending @property - def nulls_first(self): + def nulls_first(self) -> bool: """Whether null values are ordered before non-null values.""" return self.c_obj.nulls_first @@ -312,42 +317,42 @@ cdef class ColumnChunk: return result @property - def file_path(self): + def file_path(self) -> str: """Relative file path for this column chunk.""" return self.c_obj.file_path.decode("utf-8") @property - def file_offset(self): + def file_offset(self) -> int: """Deprecated byte offset to column metadata.""" return self.c_obj.file_offset @property - def offset_index_offset(self): + def offset_index_offset(self) -> int: """File offset of the chunk's OffsetIndex.""" return self.c_obj.offset_index_offset @property - def offset_index_length(self): + def offset_index_length(self) -> int: """Size of the chunk's OffsetIndex, in bytes.""" return self.c_obj.offset_index_length @property - def column_index_offset(self): + def column_index_offset(self) -> int: """File offset of the chunk's ColumnIndex.""" return self.c_obj.column_index_offset @property - def column_index_length(self): + def column_index_length(self) -> int: """Size of the chunk's ColumnIndex, in bytes.""" return self.c_obj.column_index_length @property - def schema_idx(self): + def schema_idx(self) -> int: """Derived index in the flattened schema.""" return self.c_obj.schema_idx @property - def meta_data(self): + def meta_data(self) -> ColumnChunkMetaData: """Column metadata for this chunk.""" return ColumnChunkMetaData.from_cpp(self.c_obj.meta_data) @@ -367,23 +372,23 @@ cdef class ColumnChunkMetaData: return result @property - def path_in_schema(self): + def path_in_schema(self) -> list[str]: """Column path components in the flattened schema.""" cdef string path return [path.decode("utf-8") for path in self.c_obj.path_in_schema] @property - def num_values(self): + def num_values(self) -> int: """Number of values in this chunk.""" return self.c_obj.num_values @property - def total_uncompressed_size(self): + def total_uncompressed_size(self) -> int: """Total uncompressed page bytes for this chunk.""" return self.c_obj.total_uncompressed_size @property - def total_compressed_size(self): + def total_compressed_size(self) -> int: """Total compressed page bytes for this chunk.""" return self.c_obj.total_compressed_size @@ -401,7 +406,7 @@ cdef class RowGroup: return result @property - def columns(self): + def columns(self) -> list[ColumnChunk]: """Column chunk metadata for each column in this row group.""" cdef cpp_ColumnChunk column_chunk return [ @@ -409,17 +414,17 @@ cdef class RowGroup: ] @property - def total_byte_size(self): + def total_byte_size(self) -> int: """Total uncompressed byte size in this row group.""" return self.c_obj.total_byte_size @property - def num_rows(self): + def num_rows(self) -> int: """Number of rows in this row group.""" return self.c_obj.num_rows @property - def sorting_columns(self): + def sorting_columns(self) -> list[SortingColumn] | None: """Optional row sort order metadata.""" cdef cpp_SortingColumn sorting_column if not self.c_obj.sorting_columns.has_value(): @@ -430,21 +435,21 @@ cdef class RowGroup: ] @property - def file_offset(self): + def file_offset(self) -> int | None: """Optional byte offset to first page in this row group.""" if not self.c_obj.file_offset.has_value(): return None return self.c_obj.file_offset.value() @property - def total_compressed_size(self): + def total_compressed_size(self) -> int | None: """Optional total compressed bytes for this row group.""" if not self.c_obj.total_compressed_size.has_value(): return None return self.c_obj.total_compressed_size.value() @property - def ordinal(self): + def ordinal(self) -> int | None: """Optional row group ordinal within the file.""" if not self.c_obj.ordinal.has_value(): return None @@ -473,28 +478,28 @@ cdef class FileMetaData: return result @property - def version(self): + def version(self) -> int: """Get the file format version.""" return self.c_obj.version @property - def num_rows(self): + def num_rows(self) -> int: """Get the total number of rows.""" return self.c_obj.num_rows @property - def created_by(self): + def created_by(self) -> str: """Get the application that created the file.""" return self.c_obj.created_by.decode("utf-8") @property - def row_groups(self): + def row_groups(self) -> list[RowGroup]: """Get row group metadata in this file.""" cdef cpp_RowGroup row_group return [RowGroup.from_cpp(row_group) for row_group in self.c_obj.row_groups] @property - def row_group_num_rows(self): + def row_group_num_rows(self) -> list[int]: """ Get row counts for each row group in this file. @@ -516,7 +521,7 @@ cdef class FileMetaData: return [self.c_obj.row_groups[i].num_rows for i in range(n)] @property - def columnchunk_metadata(self): + def columnchunk_metadata(self) -> dict[str, list[int]]: """ Get a map of dotted column paths to lists of `total_uncompressed_size` values from every column chunk in @@ -563,7 +568,9 @@ cdef class FileMetaData: return result @classmethod - def from_bytes(cls, const uint8_t[::1] footer_bytes): + def from_bytes( + cls, const uint8_t[::1] footer_bytes: Buffer + ) -> FileMetaData: """Build ``FileMetaData`` from parquet footer bytes. Parameters diff --git a/python/pylibcudf/pylibcudf/io/text.pyx b/python/pylibcudf/pylibcudf/io/text.pyx index be15701a4d88..b2103d42d2bb 100644 --- a/python/pylibcudf/pylibcudf/io/text.pyx +++ b/python/pylibcudf/pylibcudf/io/text.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cython.operator cimport dereference @@ -9,6 +9,10 @@ from libcpp.utility cimport move from pylibcudf.column cimport Column from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.stream cimport Stream from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from pylibcudf.libcudf.column.column cimport column @@ -70,7 +74,7 @@ cdef class ParseOptions: def __init__( self, *, - byte_range=None, + byte_range: tuple[int, int] | list[int] | None = None, strip_delimiters=False, ): self.c_options = cpp_text.parse_options() @@ -194,7 +198,7 @@ cpdef Column multibyte_split( DataChunkSource source, str delimiter, ParseOptions options=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/io/timezone.pyx b/python/pylibcudf/pylibcudf/io/timezone.pyx index 033ed15a1ba6..696e9bdb4164 100644 --- a/python/pylibcudf/pylibcudf/io/timezone.pyx +++ b/python/pylibcudf/pylibcudf/io/timezone.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.memory cimport unique_ptr @@ -11,6 +11,10 @@ from pylibcudf.libcudf.io.timezone cimport ( from pylibcudf.libcudf.table.table cimport table from ..utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from ..table cimport Table from rmm.pylibrmm.stream cimport Stream from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -19,7 +23,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["make_timezone_transition_table"] cpdef Table make_timezone_transition_table( - str tzif_dir, str timezone_name, object stream=None, DeviceMemoryResource mr=None, + str tzif_dir, str timezone_name, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ Creates a transition table to convert ORC timestamps to UTC. diff --git a/python/pylibcudf/pylibcudf/io/types.pyx b/python/pylibcudf/pylibcudf/io/types.pyx index 5a56c02706c1..02dbdd80abf4 100644 --- a/python/pylibcudf/pylibcudf/io/types.pyx +++ b/python/pylibcudf/pylibcudf/io/types.pyx @@ -35,11 +35,13 @@ from pylibcudf.span import is_span from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource +from collections.abc import Mapping, Sequence import codecs import errno import io import os import re +from typing import TYPE_CHECKING, Any, TypeAlias from pylibcudf.libcudf.io.json import \ json_recovery_mode_t as JSONRecoveryMode # no-cython-lint @@ -51,6 +53,10 @@ from pylibcudf.libcudf.io.types import ( statistics_freq as StatisticsFreq, # no-cython-lint ) +if TYPE_CHECKING: + from pylibcudf.column import Column + from pylibcudf.span import Span + __all__ = [ "ColumnEncoding", "ColumnInMetadata", @@ -67,6 +73,9 @@ __all__ = [ "TableWithMetadata", ] +ColumnNameSpec: TypeAlias = tuple[str, list["ColumnNameSpec"]] +ChildNameSpec: TypeAlias = Mapping[str, "ChildNameSpec"] + cdef class PartitionInfo: """ Information used while writing partitioned datasets. @@ -298,7 +307,7 @@ cdef class TableInputMetadata: self.c_obj = table_input_metadata(table.view()) @property - def column_metadata(self): + def column_metadata(self) -> list[ColumnInMetadata]: return [ ColumnInMetadata.from_libcudf(&self.c_obj.column_metadata[i], self) for i in range(self.c_obj.column_metadata.size()) @@ -322,7 +331,7 @@ cdef class TableWithMetadata: [("id", []), ("name", [("first", []), ("last", [])])] """ - def __init__(self, Table tbl, list column_names): + def __init__(self, Table tbl, list column_names: list[ColumnNameSpec]): self.tbl = tbl self.metadata.schema_info = self._make_column_info(column_names) @@ -347,7 +356,7 @@ cdef class TableWithMetadata: return col_name_infos @property - def columns(self): + def columns(self) -> tuple[Column, ...]: """ Return a tuple containing the columns of the table """ @@ -360,7 +369,9 @@ cdef class TableWithMetadata: names.append((child, grandchildren)) return names - def column_names(self, include_children=False): + def column_names( + self, include_children=False + ) -> list[str] | list[ColumnNameSpec]: """ Return a list containing the column names of the table """ @@ -377,7 +388,7 @@ cdef class TableWithMetadata: return names @property - def child_names(self): + def child_names(self) -> ChildNameSpec: """ Return a dictionary mapping the names of columns with children to the names of their child columns. Columns without children @@ -409,7 +420,7 @@ cdef class TableWithMetadata: return out @property - def per_file_user_data(self): + def per_file_user_data(self) -> list[Mapping[bytes, bytes]]: """ Returns a list containing a dict containing file-format specific metadata, @@ -418,7 +429,7 @@ cdef class TableWithMetadata: return self.metadata.per_file_user_data @property - def num_rows_per_source(self): + def num_rows_per_source(self) -> list[int]: """ Returns a list containing the number of rows for each file being read in. @@ -427,7 +438,7 @@ cdef class TableWithMetadata: # The following functions are currently only for Parquet reader @property - def num_input_row_groups(self): + def num_input_row_groups(self) -> int: """ Returns the total number of input Parquet row groups across all data sources. @@ -435,7 +446,7 @@ cdef class TableWithMetadata: return self.metadata.num_input_row_groups @property - def num_row_groups_after_stats_filter(self): + def num_row_groups_after_stats_filter(self) -> int | None: """ Returns the number of remaining Parquet row groups after stats filter. None if no filtering done. @@ -445,7 +456,7 @@ cdef class TableWithMetadata: return None @property - def num_row_groups_after_bloom_filter(self): + def num_row_groups_after_bloom_filter(self) -> int | None: """ Returns the number of remaining Parquet row groups after bloom filter. None if no filtering done. @@ -471,7 +482,7 @@ cdef class FilepathSource: Known file size in bytes. Omit to query size via KvikIO (HEAD for remote URLs). """ - def __init__(self, path, size=None): + def __init__(self, path: str | os.PathLike[Any], size: int | None = None): self.path = os.fspath(path) self.size = size @@ -498,7 +509,19 @@ cdef class SourceInfo: If an empty list, constructs an empty SourceInfo. """ - def __init__(self, sources): + def __init__( + self, + sources: ( + Sequence[str] + | Sequence[os.PathLike[Any]] + | Sequence[FilepathSource] + | Sequence[Datasource] + | Sequence[io.StringIO] + | Sequence[bytes] + | Sequence[io.BytesIO] + | Sequence[Span] + ), + ): if not sources: self.c_obj = move(source_info()) return @@ -662,7 +685,14 @@ cdef class SinkInfo: (that are not all io.IOBase instances) will raise a ValueError. """ - def __init__(self, list sinks): + def __init__( + self, + list sinks: ( + list[str] + | list[os.PathLike[Any]] + | list[io.IOBase] + ), + ): cdef vector[data_sink *] data_sinks cdef vector[string] paths diff --git a/python/pylibcudf/pylibcudf/join.pyx b/python/pylibcudf/pylibcudf/join.pyx index d0f90f777ea2..b3d72c496102 100644 --- a/python/pylibcudf/pylibcudf/join.pyx +++ b/python/pylibcudf/pylibcudf/join.pyx @@ -21,6 +21,10 @@ from .column cimport Column from .expressions cimport Expression from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -64,7 +68,7 @@ cpdef tuple inner_join( Table left_keys, Table right_keys, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform an inner join between two tables. @@ -112,7 +116,7 @@ cpdef tuple left_join( Table left_keys, Table right_keys, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a left join between two tables. @@ -160,7 +164,7 @@ cpdef tuple full_join( Table left_keys, Table right_keys, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a full join between two tables. @@ -208,7 +212,7 @@ cpdef Column left_semi_join( Table left_keys, Table right_keys, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a left semi join between two tables. @@ -259,7 +263,7 @@ cpdef Column left_anti_join( Table left_keys, Table right_keys, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a left anti join between two tables. @@ -307,7 +311,7 @@ cpdef Column left_anti_join( cpdef Table cross_join( - Table left, Table right, object stream=None, DeviceMemoryResource mr=None + Table left, Table right, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Perform a cross join on two tables. @@ -348,7 +352,7 @@ cpdef tuple conditional_inner_join( Table left, Table right, Expression binary_predicate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a conditional inner join between two tables. @@ -398,7 +402,7 @@ cpdef tuple conditional_left_join( Table left, Table right, Expression binary_predicate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a conditional left join between two tables. @@ -448,7 +452,7 @@ cpdef tuple conditional_full_join( Table left, Table right, Expression binary_predicate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a conditional full join between two tables. @@ -496,7 +500,7 @@ cpdef Column conditional_left_semi_join( Table left, Table right, Expression binary_predicate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a conditional left semi join between two tables. @@ -542,7 +546,7 @@ cpdef Column conditional_left_anti_join( Table left, Table right, Expression binary_predicate, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a conditional left anti join between two tables. @@ -591,7 +595,7 @@ cpdef tuple mixed_inner_join( Table right_conditional, Expression binary_predicate, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a mixed inner join between two tables. @@ -655,7 +659,7 @@ cpdef tuple mixed_left_join( Table right_conditional, Expression binary_predicate, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a mixed left join between two tables. @@ -719,7 +723,7 @@ cpdef tuple mixed_full_join( Table right_conditional, Expression binary_predicate, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a mixed full join between two tables. @@ -783,7 +787,7 @@ cpdef Column mixed_left_semi_join( Table right_conditional, Expression binary_predicate, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a mixed left semi join between two tables. @@ -841,7 +845,7 @@ cpdef Column mixed_left_anti_join( Table right_conditional, Expression binary_predicate, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a mixed left anti join between two tables. @@ -910,7 +914,7 @@ cdef class FilteredJoin: Table right, null_equality compare_nulls=null_equality.EQUAL, double load_factor=0.5, - object stream=None, + object stream: CudaStreamLike | None = None, ): """ Construct a filtered hash join object for subsequent probe calls. @@ -944,7 +948,7 @@ cdef class FilteredJoin: def semi_join( self, Table left, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -985,7 +989,7 @@ cdef class FilteredJoin: def anti_join( self, Table left, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/json.pyx b/python/pylibcudf/pylibcudf/json.pyx index 7d71f2a9dca4..ca1222f5385f 100644 --- a/python/pylibcudf/pylibcudf/json.pyx +++ b/python/pylibcudf/pylibcudf/json.pyx @@ -16,6 +16,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["GetJsonObjectOptions", "get_json_object"] @@ -122,7 +126,7 @@ cpdef Column get_json_object( Column col, Scalar json_path, GetJsonObjectOptions options=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/labeling.pyx b/python/pylibcudf/pylibcudf/labeling.pyx index 9d52e308f7c4..1dc6ba43c339 100644 --- a/python/pylibcudf/pylibcudf/labeling.pyx +++ b/python/pylibcudf/pylibcudf/labeling.pyx @@ -15,6 +15,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["Inclusive", "label_bins"] @@ -25,7 +29,7 @@ cpdef Column label_bins( inclusive left_inclusive, Column right_edges, inclusive right_inclusive, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Labels elements based on membership in the specified bins. diff --git a/python/pylibcudf/pylibcudf/lists.pyx b/python/pylibcudf/pylibcudf/lists.pyx index 71996340e5db..930cacec73dc 100644 --- a/python/pylibcudf/pylibcudf/lists.pyx +++ b/python/pylibcudf/pylibcudf/lists.pyx @@ -59,6 +59,10 @@ from .column cimport Column, ListsColumnView from .scalar cimport Scalar from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -88,7 +92,7 @@ __all__ = [ cpdef Table explode_outer( Table input, size_type explode_column_idx, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Explode a column of lists into rows. @@ -126,7 +130,7 @@ cpdef Table explode_outer( cpdef Column concatenate_rows( Table input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Concatenate multiple lists columns into a single lists column row-wise. @@ -163,7 +167,7 @@ cpdef Column concatenate_rows( cpdef Column concatenate_list_elements( Column input, concatenate_null_policy null_policy, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Concatenate multiple lists on the same row into a single list. @@ -200,7 +204,7 @@ cpdef Column concatenate_list_elements( cpdef Column contains( Column input, ColumnOrScalar search_key, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column of bool values indicating whether @@ -257,7 +261,7 @@ cpdef Column contains( cpdef Column contains_nulls( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column of bool values indicating whether @@ -297,7 +301,7 @@ cpdef Column index_of( Column input, ColumnOrScalar search_key, duplicate_find_option find_option, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column of index values indicating the position of a search @@ -352,7 +356,7 @@ cpdef Column index_of( cpdef Column reverse( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Reverse the element order within each list of the input column. @@ -388,7 +392,7 @@ cpdef Column segmented_gather( Column input, Column gather_map_list, out_of_bounds_policy bounds_policy=out_of_bounds_policy.DONT_CHECK, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column with elements gathered based on the indices in gather_map_list @@ -443,7 +447,7 @@ cpdef Column segmented_gather( cpdef Column extract_list_element( Column input, ColumnOrSizeType index, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column of extracted list elements. @@ -486,7 +490,7 @@ cpdef Column extract_list_element( cpdef Column count_elements( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Count the number of rows in each @@ -525,7 +529,7 @@ cpdef Column sequences( Column starts, Column sizes, Column steps = None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a lists column in which each row contains a sequence of @@ -583,7 +587,7 @@ cpdef Column sort_lists( order sort_order, null_order na_position, bool stable = False, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sort the elements within a list in each row of a list column. @@ -641,7 +645,7 @@ cpdef Column difference_distinct( Column rhs, null_equality nulls_equal=null_equality.EQUAL, nan_equality nans_equal=nan_equality.ALL_EQUAL, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column of index values indicating the position of a search @@ -692,7 +696,7 @@ cpdef Column have_overlap( Column rhs, null_equality nulls_equal=null_equality.EQUAL, nan_equality nans_equal=nan_equality.ALL_EQUAL, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Check if lists at each row of the given lists columns overlap. @@ -742,7 +746,7 @@ cpdef Column intersect_distinct( Column rhs, null_equality nulls_equal=null_equality.EQUAL, nan_equality nans_equal=nan_equality.ALL_EQUAL, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a lists column of distinct elements common to two input lists columns. @@ -792,7 +796,7 @@ cpdef Column union_distinct( Column rhs, null_equality nulls_equal=null_equality.EQUAL, nan_equality nans_equal=nan_equality.ALL_EQUAL, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a lists column of distinct elements found in @@ -841,7 +845,7 @@ cpdef Column union_distinct( cpdef Column apply_boolean_mask( Column input, Column boolean_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters elements in each row of the input lists column using a boolean mask @@ -885,7 +889,7 @@ cpdef Column apply_boolean_mask( cpdef Column apply_deletion_mask( Column input, Column deletion_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters elements in each row of the input lists column using a deletion mask. @@ -928,7 +932,7 @@ cpdef Column distinct( Column input, null_equality nulls_equal, nan_equality nans_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a new list column without duplicate elements in each list. diff --git a/python/pylibcudf/pylibcudf/merge.pyx b/python/pylibcudf/pylibcudf/merge.pyx index 3c0cd93a342e..82c075b57d6a 100644 --- a/python/pylibcudf/pylibcudf/merge.pyx +++ b/python/pylibcudf/pylibcudf/merge.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.memory cimport unique_ptr @@ -16,14 +16,20 @@ from .table cimport Table from .utils cimport _get_stream, _get_memory_resource from cuda.bindings.cyruntime cimport cudaStream_t +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.types import NullOrder, Order + from pylibcudf.typing import CudaStreamLike + __all__ = ["merge"] cpdef Table merge ( - list tables_to_merge, - list key_cols, - list column_order, - list null_precedence, - object stream=None, + list tables_to_merge: list[Table], + list key_cols: list[int], + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Merge a set of sorted tables. diff --git a/python/pylibcudf/pylibcudf/null_mask.pyx b/python/pylibcudf/pylibcudf/null_mask.pyx index bdb92b66bb02..ef121704ea1a 100644 --- a/python/pylibcudf/pylibcudf/null_mask.pyx +++ b/python/pylibcudf/pylibcudf/null_mask.pyx @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from libc.stdint cimport uintptr_t from libcpp.memory cimport make_unique from libcpp.pair cimport pair @@ -21,6 +22,10 @@ from .span import is_span as py_is_span from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -43,7 +48,7 @@ cdef DeviceBuffer buffer_to_python( cpdef DeviceBuffer copy_bitmask( Column col, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Copies ``col``'s bitmask into a ``DeviceBuffer``. @@ -81,7 +86,7 @@ cpdef DeviceBuffer copy_bitmask_from_bitmask( object bitmask, size_type begin_bit, size_type end_bit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Copies a portion of a bitmask into a ``DeviceBuffer``. @@ -154,7 +159,7 @@ cpdef size_t bitmask_allocation_size_bytes(size_type number_of_bits): cpdef DeviceBuffer create_null_mask( size_type size, mask_state state = mask_state.UNINITIALIZED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Creates a ``DeviceBuffer`` for use as a null value indicator bitmask of a @@ -192,7 +197,11 @@ cpdef DeviceBuffer create_null_mask( return buffer_to_python(move(db), _stream, mr) -cpdef tuple bitmask_and(columns, object stream=None, DeviceMemoryResource mr=None): +cpdef tuple bitmask_and( + columns: Sequence[Column], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, +): """Performs bitwise AND of the bitmasks of a list of columns. For details, see :cpp:func:`bitmask_and`. @@ -226,7 +235,11 @@ cpdef tuple bitmask_and(columns, object stream=None, DeviceMemoryResource mr=Non return buffer_to_python(move(c_result.first), _stream, mr), c_result.second -cpdef tuple bitmask_or(columns, object stream=None, DeviceMemoryResource mr=None): +cpdef tuple bitmask_or( + columns: Sequence[Column], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, +): """Performs bitwise OR of the bitmasks of a list of columns. For details, see :cpp:func:`bitmask_or`. @@ -262,7 +275,7 @@ cpdef size_type null_count( object bitmask, size_type start, size_type stop, - object stream=None + object stream: CudaStreamLike | None = None ): """Given a validity bitmask, counts the number of null elements. @@ -304,7 +317,7 @@ cpdef size_type index_of_first_set_bit( object bitmask, size_type start, size_type stop, - object stream=None + object stream: CudaStreamLike | None = None ): """Given a validity bitmask, returns the index of the first valid element relative to ``start``. diff --git a/python/pylibcudf/pylibcudf/nvtext/byte_pair_encode.pyx b/python/pylibcudf/pylibcudf/nvtext/byte_pair_encode.pyx index 0db5a345f43c..8a20d0ebf41c 100644 --- a/python/pylibcudf/pylibcudf/nvtext/byte_pair_encode.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/byte_pair_encode.pyx @@ -17,6 +17,10 @@ from pylibcudf.libcudf.scalar.scalar_factories cimport ( ) from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -31,7 +35,7 @@ cdef class BPEMergePairs: def __cinit__( self, Column merge_pairs, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): cdef column_view c_pairs = merge_pairs.view() @@ -49,7 +53,7 @@ cpdef Column byte_pair_encoding( Column input, BPEMergePairs merge_pairs, Scalar separator=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/deduplicate.pyx b/python/pylibcudf/pylibcudf/nvtext/deduplicate.pyx index 1dbbe3e066e6..55e86c128eac 100644 --- a/python/pylibcudf/pylibcudf/nvtext/deduplicate.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/deduplicate.pyx @@ -16,6 +16,10 @@ from pylibcudf.libcudf.nvtext.deduplicate cimport ( ) from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.librmm.device_buffer cimport device_buffer from rmm.pylibrmm.stream cimport Stream @@ -43,7 +47,7 @@ cdef Column _column_from_suffix_array( cpdef Column build_suffix_array( - Column input, size_type min_width, object stream=None, DeviceMemoryResource mr=None + Column input, size_type min_width, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Builds a suffix array for the input strings column. @@ -85,7 +89,7 @@ cpdef Column resolve_duplicates( Column input, Column indices, size_type min_width, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -133,7 +137,7 @@ cpdef Column resolve_duplicates_pair( Column input2, Column indices2, size_type min_width, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/edit_distance.pyx b/python/pylibcudf/pylibcudf/nvtext/edit_distance.pyx index 06ae6a8366dc..03d45f82697a 100644 --- a/python/pylibcudf/pylibcudf/nvtext/edit_distance.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/edit_distance.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.memory cimport unique_ptr @@ -14,6 +14,10 @@ from rmm.pylibrmm.stream cimport Stream from ..column cimport Column from ..utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["edit_distance"] @@ -21,7 +25,7 @@ __all__ = ["edit_distance"] cpdef Column edit_distance( Column input, Column targets, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyx b/python/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyx index 6d70751a5a06..b4778f9077e4 100644 --- a/python/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stdint cimport uint32_t @@ -16,6 +16,10 @@ from pylibcudf.libcudf.scalar.scalar cimport string_scalar from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -30,7 +34,7 @@ cpdef Column generate_ngrams( Column input, size_type ngrams, Scalar separator, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -75,7 +79,7 @@ cpdef Column generate_ngrams( cpdef Column generate_character_ngrams( Column input, size_type ngrams = 2, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -117,7 +121,7 @@ cpdef Column hash_character_ngrams( Column input, size_type ngrams, uint32_t seed, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/jaccard.pyx b/python/pylibcudf/pylibcudf/nvtext/jaccard.pyx index 24a343e4508a..57129757f16d 100644 --- a/python/pylibcudf/pylibcudf/nvtext/jaccard.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/jaccard.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libcpp.memory cimport unique_ptr @@ -11,6 +11,10 @@ from pylibcudf.libcudf.nvtext.jaccard cimport ( ) from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -21,7 +25,7 @@ cpdef Column jaccard_index( Column input1, Column input2, size_type width, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/minhash.pyx b/python/pylibcudf/pylibcudf/nvtext/minhash.pyx index 035c4370acd2..92c9cbfd081f 100644 --- a/python/pylibcudf/pylibcudf/nvtext/minhash.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/minhash.pyx @@ -15,6 +15,10 @@ from pylibcudf.libcudf.nvtext.minhash cimport ( ) from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -32,7 +36,7 @@ cpdef Column minhash( Column a, Column b, size_type width, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ @@ -86,7 +90,7 @@ cpdef Column minhash64( Column a, Column b, size_type width, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ @@ -142,7 +146,7 @@ cpdef Column minhash_ngrams( uint32_t seed, Column a, Column b, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ @@ -199,7 +203,7 @@ cpdef Column minhash64_ngrams( uint64_t seed, Column a, Column b, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/ngrams_tokenize.pyx b/python/pylibcudf/pylibcudf/nvtext/ngrams_tokenize.pyx index 7dfba5c6b989..fccae239a514 100644 --- a/python/pylibcudf/pylibcudf/nvtext/ngrams_tokenize.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/ngrams_tokenize.pyx @@ -14,6 +14,10 @@ from pylibcudf.libcudf.scalar.scalar cimport string_scalar from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -25,7 +29,7 @@ cpdef Column ngrams_tokenize( size_type ngrams, Scalar delimiter, Scalar separator, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/normalize.pyx b/python/pylibcudf/pylibcudf/nvtext/normalize.pyx index 889792d67383..ec083549ccee 100644 --- a/python/pylibcudf/pylibcudf/nvtext/normalize.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/normalize.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.nvtext cimport normalize as cpp_normalize from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -29,7 +33,7 @@ cdef class CharacterNormalizer: self, bool do_lower_case, Column tokens, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): cdef column_view c_tokens = tokens.view() @@ -49,7 +53,7 @@ cdef class CharacterNormalizer: __hash__ = None cpdef Column normalize_spaces( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new strings column by normalizing the whitespace in @@ -86,7 +90,7 @@ cpdef Column normalize_spaces( cpdef Column normalize_characters( Column input, CharacterNormalizer normalizer, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/replace.pyx b/python/pylibcudf/pylibcudf/nvtext/replace.pyx index 70fec3d96640..3c5131cda0ce 100644 --- a/python/pylibcudf/pylibcudf/nvtext/replace.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/replace.pyx @@ -18,6 +18,10 @@ from pylibcudf.libcudf.scalar.scalar_factories cimport ( from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -29,7 +33,7 @@ cpdef Column replace_tokens( Column targets, Column replacements, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -84,7 +88,7 @@ cpdef Column filter_tokens( size_type min_token_length, Scalar replacement=None, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/stemmer.pyx b/python/pylibcudf/pylibcudf/nvtext/stemmer.pyx index cb5c5f69cc23..44d20377dea1 100644 --- a/python/pylibcudf/pylibcudf/nvtext/stemmer.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/stemmer.pyx @@ -15,6 +15,10 @@ from pylibcudf.libcudf.nvtext.stemmer cimport ( from pylibcudf.libcudf.types cimport size_type from pylibcudf.nvtext.stemmer cimport ColumnOrSize from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from pylibcudf.libcudf.nvtext.stemmer import letter_type as LetterType # no-cython-lint from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -27,7 +31,7 @@ cpdef Column is_letter( Column input, bool check_vowels, ColumnOrSize indices, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -77,7 +81,7 @@ cpdef Column is_letter( cpdef Column porter_stemmer_measure( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns the Porter Stemmer measurements of a strings column. diff --git a/python/pylibcudf/pylibcudf/nvtext/tokenize.pyx b/python/pylibcudf/pylibcudf/nvtext/tokenize.pyx index 2459760a9d4c..fa660fb80660 100644 --- a/python/pylibcudf/pylibcudf/nvtext/tokenize.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/tokenize.pyx @@ -22,6 +22,10 @@ from pylibcudf.libcudf.scalar.scalar_factories cimport ( from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -42,7 +46,7 @@ cdef class TokenizeVocabulary: For details, see :cpp:class:`cudf::nvtext::tokenize_vocabulary`. """ - def __cinit__(self, Column vocab, object stream=None, DeviceMemoryResource mr=None): + def __cinit__(self, Column vocab, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): cdef column_view c_vocab = vocab.view() cdef Stream _stream = _get_stream(stream) cdef cudaStream_t _cs = _stream.view().value() @@ -55,7 +59,7 @@ cdef class TokenizeVocabulary: cpdef Column tokenize_scalar( Column input, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -100,7 +104,7 @@ cpdef Column tokenize_scalar( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column tokenize_column( - Column input, Column delimiters, object stream=None, DeviceMemoryResource mr=None + Column input, Column delimiters, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a single column of strings by tokenizing the input @@ -142,7 +146,7 @@ cpdef Column tokenize_column( cpdef Column count_tokens_scalar( Column input, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -187,7 +191,7 @@ cpdef Column count_tokens_scalar( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column count_tokens_column( - Column input, Column delimiters, object stream=None, DeviceMemoryResource mr=None + Column input, Column delimiters, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns the number of tokens in each string of a strings column @@ -227,7 +231,7 @@ cpdef Column count_tokens_column( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column character_tokenize( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a single column of strings by converting @@ -261,7 +265,7 @@ cpdef Column detokenize( Column input, Column row_indices, Scalar separator=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -314,7 +318,7 @@ cpdef Column tokenize_with_vocabulary( TokenizeVocabulary vocabulary, Scalar delimiter, size_type default_id=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/nvtext/wordpiece_tokenize.pyx b/python/pylibcudf/pylibcudf/nvtext/wordpiece_tokenize.pyx index 815139152b59..13b148e40c3b 100644 --- a/python/pylibcudf/pylibcudf/nvtext/wordpiece_tokenize.pyx +++ b/python/pylibcudf/pylibcudf/nvtext/wordpiece_tokenize.pyx @@ -13,6 +13,10 @@ from pylibcudf.libcudf.nvtext.wordpiece_tokenize cimport ( ) from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -30,7 +34,7 @@ cdef class WordPieceVocabulary: def __cinit__( self, Column vocab, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): cdef column_view c_vocab = vocab.view() @@ -48,7 +52,7 @@ cpdef Column wordpiece_tokenize( Column input, WordPieceVocabulary vocabulary, size_type max_words_per_row, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/partitioning.pyx b/python/pylibcudf/pylibcudf/partitioning.pyx index 88f7f9de086e..bf2f1c001312 100644 --- a/python/pylibcudf/pylibcudf/partitioning.pyx +++ b/python/pylibcudf/pylibcudf/partitioning.pyx @@ -17,6 +17,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -32,7 +36,7 @@ cpdef tuple[Table, list] hash_partition( int num_partitions, cpp_partitioning.hash_id hash_function = cpp_partitioning.hash_id.HASH_MURMUR3, uint32_t seed = cpp_partitioning.DEFAULT_HASH_SEED, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -102,7 +106,7 @@ cpdef tuple[Table, list] partition( Table t, Column partition_map, int num_partitions, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -155,7 +159,7 @@ cpdef tuple[Table, list] round_robin_partition( Table input, int num_partitions, int start_partition=0, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/quantiles.pyx b/python/pylibcudf/pylibcudf/quantiles.pyx index 6f28720836ef..0651f3fea22b 100644 --- a/python/pylibcudf/pylibcudf/quantiles.pyx +++ b/python/pylibcudf/pylibcudf/quantiles.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Iterable from libcpp cimport bool from libcpp.memory cimport unique_ptr from libcpp.utility cimport move @@ -21,17 +22,21 @@ from .column cimport Column from .table cimport Table from .types cimport interpolation from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["quantile", "quantiles"] cpdef Column quantile( Column input, - vector[double] q, + vector[double] q: Iterable[float], interpolation interp = interpolation.LINEAR, Column ordered_indices = None, bool exact=True, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes quantiles with interpolation. @@ -97,12 +102,12 @@ cpdef Column quantile( cpdef Table quantiles( Table input, - vector[double] q, + vector[double] q: Iterable[float], interpolation interp = interpolation.NEAREST, sorted is_input_sorted = sorted.NO, list column_order = None, list null_precedence = None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes row quantiles with interpolation. diff --git a/python/pylibcudf/pylibcudf/reduce.pyx b/python/pylibcudf/pylibcudf/reduce.pyx index b2bdfa720872..4e6d0aec5607 100644 --- a/python/pylibcudf/pylibcudf/reduce.pyx +++ b/python/pylibcudf/pylibcudf/reduce.pyx @@ -31,6 +31,10 @@ from .column cimport Column from .scalar cimport Scalar from .types cimport DataType from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from pylibcudf.libcudf.reduce import scan_type as ScanType # no-cython-lint from cuda.bindings.cyruntime cimport cudaStream_t @@ -50,7 +54,7 @@ cpdef Scalar reduce( Aggregation agg, DataType data_type, Scalar init=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a reduction on a column @@ -111,7 +115,7 @@ cpdef Column scan( Column col, Aggregation agg, scan_type inclusive, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a scan on a column @@ -156,7 +160,7 @@ cpdef Column scan( return Column.from_libcudf(move(result), _stream, mr) -cpdef tuple minmax(Column col, object stream=None, DeviceMemoryResource mr=None): +cpdef tuple minmax(Column col, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Compute the minimum and maximum of a column For details, see ``cudf::minmax`` documentation. @@ -215,7 +219,7 @@ cpdef size_type unique_count( Column source, null_policy null_handling, nan_policy nan_handling, - object stream=None + object stream: CudaStreamLike | None = None ): """Returns the number of unique consecutive elements in the input column. @@ -255,7 +259,7 @@ cpdef size_type distinct_count( Column source, null_policy null_handling, nan_policy nan_handling, - object stream=None + object stream: CudaStreamLike | None = None ): """Returns the number of distinct elements in the input column. @@ -289,7 +293,7 @@ cpdef size_type distinct_count( cpdef size_type unique_count_table( Table source, null_equality nulls_equal, - object stream=None + object stream: CudaStreamLike | None = None ): """Returns the number of unique consecutive rows in the input table. @@ -323,7 +327,7 @@ cpdef size_type unique_count_table( cpdef size_type distinct_count_table( Table source, null_equality nulls_equal, - object stream=None + object stream: CudaStreamLike | None = None ): """Returns the number of distinct rows in the input table. diff --git a/python/pylibcudf/pylibcudf/replace.pyx b/python/pylibcudf/pylibcudf/replace.pyx index e411ae645062..604a3cefa40c 100644 --- a/python/pylibcudf/pylibcudf/replace.pyx +++ b/python/pylibcudf/pylibcudf/replace.pyx @@ -19,6 +19,10 @@ from pylibcudf.libcudf.replace import \ from .column cimport Column from .scalar cimport Scalar from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -33,7 +37,7 @@ __all__ = [ cpdef Column replace_nulls( Column source_column, ReplacementType replacement, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replace nulls in source_column. @@ -130,7 +134,7 @@ cpdef Column find_and_replace_all( Column source_column, Column values_to_replace, Column replacement_values, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replace all occurrences of values_to_replace with replacement_values. @@ -182,7 +186,7 @@ cpdef Column clamp( Scalar hi, Scalar lo_replace=None, Scalar hi_replace=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Clamp the values in source_column to the range [lo, hi]. @@ -248,7 +252,7 @@ cpdef Column clamp( cpdef Column normalize_nans_and_zeros( Column source_column, bool inplace=False, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Normalize NaNs and zeros in source_column. diff --git a/python/pylibcudf/pylibcudf/reshape.pyx b/python/pylibcudf/pylibcudf/reshape.pyx index b1b60f4444dd..211bda522ab4 100644 --- a/python/pylibcudf/pylibcudf/reshape.pyx +++ b/python/pylibcudf/pylibcudf/reshape.pyx @@ -25,12 +25,16 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["interleave_columns", "tile", "table_to_array"] cpdef Column interleave_columns( - Table source_table, object stream=None, DeviceMemoryResource mr=None + Table source_table, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Interleave columns of a table into a single column. @@ -73,7 +77,7 @@ cpdef Column interleave_columns( cpdef Table tile( Table source_table, size_type count, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Repeats the rows from input table count times to form a new table. @@ -114,7 +118,7 @@ cpdef void table_to_array( Table input_table, uintptr_t ptr, size_t size, - object stream=None + object stream: CudaStreamLike | None = None ): """ Copy a table into a preallocated column-major device array. diff --git a/python/pylibcudf/pylibcudf/rolling.pyx b/python/pylibcudf/pylibcudf/rolling.pyx index d12a40592983..b6e5e12f2635 100644 --- a/python/pylibcudf/pylibcudf/rolling.pyx +++ b/python/pylibcudf/pylibcudf/rolling.pyx @@ -22,6 +22,10 @@ from .column cimport Column from .scalar cimport Scalar from .types cimport DataType from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t @@ -127,8 +131,8 @@ cpdef Table grouped_range_rolling_window( null_order null_order, PrecedingRangeWindowType preceding, FollowingRangeWindowType following, - list requests, - object stream=None, + list requests: list[RollingRequest], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -194,7 +198,7 @@ cpdef Column rolling_window( WindowType following_window, size_type min_periods, Aggregation agg, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a rolling window operation on a column @@ -290,7 +294,7 @@ cpdef tuple make_range_windows( null_order null_order, PrecedingRangeWindowType preceding, FollowingRangeWindowType following, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/round.pyx b/python/pylibcudf/pylibcudf/round.pyx index a3de7add7533..7a6585c36a08 100644 --- a/python/pylibcudf/pylibcudf/round.pyx +++ b/python/pylibcudf/pylibcudf/round.pyx @@ -20,6 +20,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["RoundingMethod", "round"] @@ -28,7 +32,7 @@ cpdef Column round( Column source, int32_t decimal_places = 0, rounding_method round_method = rounding_method.HALF_UP, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Rounds all the values in a column to the specified number of decimal places. @@ -81,7 +85,7 @@ cpdef Column round_decimal( Column source, int32_t decimal_places = 0, rounding_method round_method = rounding_method.HALF_UP, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Rounds all the values in a column to the specified number of decimal places. diff --git a/python/pylibcudf/pylibcudf/scalar.pyx b/python/pylibcudf/pylibcudf/scalar.pyx index 1b80a62bed96..2a663104a8d6 100644 --- a/python/pylibcudf/pylibcudf/scalar.pyx +++ b/python/pylibcudf/pylibcudf/scalar.pyx @@ -64,6 +64,10 @@ from .column cimport Column from .traits cimport is_floating_point from .types cimport DataType from .utils cimport _get_memory_resource, _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from functools import singledispatch from ._interop_helpers import ArrowLike, ColumnMetadata @@ -153,7 +157,7 @@ cdef class Scalar: """The type of data in the column.""" return self._data_type - cpdef bool is_valid(self, object stream = None): + cpdef bool is_valid(self, object stream: CudaStreamLike | None = None): """True if the scalar is valid, false if not""" cdef Stream _stream = _get_stream(stream) cdef cudaStream_t _cs = _stream.view().value() diff --git a/python/pylibcudf/pylibcudf/search.pyx b/python/pylibcudf/pylibcudf/search.pyx index d67229729d7b..d76d427c0647 100644 --- a/python/pylibcudf/pylibcudf/search.pyx +++ b/python/pylibcudf/pylibcudf/search.pyx @@ -17,14 +17,20 @@ from .table cimport Table from .utils cimport _get_stream, _get_memory_resource from cuda.bindings.cyruntime cimport cudaStream_t +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.types import NullOrder, Order + from pylibcudf.typing import CudaStreamLike + __all__ = ["contains", "lower_bound", "upper_bound"] cpdef Column lower_bound( Table haystack, Table needles, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Find smallest indices in haystack where needles may be inserted to retain order. @@ -76,9 +82,9 @@ cpdef Column lower_bound( cpdef Column upper_bound( Table haystack, Table needles, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Find largest indices in haystack where needles may be inserted to retain order. @@ -128,7 +134,7 @@ cpdef Column upper_bound( cpdef Column contains( - Column haystack, Column needles, object stream=None, DeviceMemoryResource mr=None + Column haystack, Column needles, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Check whether needles are present in haystack. diff --git a/python/pylibcudf/pylibcudf/sorting.pyx b/python/pylibcudf/pylibcudf/sorting.pyx index eac72a4cdda6..fa6bc5bda3b5 100644 --- a/python/pylibcudf/pylibcudf/sorting.pyx +++ b/python/pylibcudf/pylibcudf/sorting.pyx @@ -19,6 +19,12 @@ from .table cimport Table from .utils cimport _get_stream, _get_memory_resource from cuda.bindings.cyruntime cimport cudaStream_t +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.types import NullOrder, Order + from pylibcudf.typing import CudaStreamLike + __all__ = [ "is_sorted", "rank", @@ -36,9 +42,9 @@ __all__ = [ cpdef Column sorted_order( Table source_table, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the row indices required to sort the table. @@ -81,9 +87,9 @@ cpdef Column sorted_order( cpdef Column stable_sorted_order( Table source_table, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the row indices required to sort the table, @@ -132,7 +138,7 @@ cpdef Column rank( null_policy null_handling, null_order null_precedence, bool percentage, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Computes the rank of each element in the column. @@ -181,7 +187,10 @@ cpdef Column rank( cpdef bool is_sorted( - Table tbl, list column_order, list null_precedence, object stream=None + Table tbl, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, ): """Checks if the table is sorted. @@ -223,9 +232,9 @@ cpdef Table segmented_sort_by_key( Table values, Table keys, Column segment_offsets, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table by key, within segments. @@ -278,9 +287,9 @@ cpdef Table stable_segmented_sort_by_key( Table values, Table keys, Column segment_offsets, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table by key preserving order of equal elements, @@ -333,9 +342,9 @@ cpdef Table stable_segmented_sort_by_key( cpdef Table sort_by_key( Table values, Table keys, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table by key. @@ -383,9 +392,9 @@ cpdef Table sort_by_key( cpdef Table stable_sort_by_key( Table values, Table keys, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table by key preserving order of equal elements. @@ -432,9 +441,9 @@ cpdef Table stable_sort_by_key( cpdef Table sort( Table source_table, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table. @@ -477,9 +486,9 @@ cpdef Table sort( cpdef Table stable_sort( Table source_table, - list column_order, - list null_precedence, - object stream=None, + list column_order: list[Order], + list null_precedence: list[NullOrder], + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Sorts the table preserving order of equal elements. @@ -524,7 +533,7 @@ cpdef Column top_k( Column col, size_type k, order sort_order = order.DESCENDING, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -569,7 +578,7 @@ cpdef Column top_k_order( Column col, size_type k, order sort_order = order.DESCENDING, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/stream_compaction.pyx b/python/pylibcudf/pylibcudf/stream_compaction.pyx index 457c82f63a7e..9d64484c9605 100644 --- a/python/pylibcudf/pylibcudf/stream_compaction.pyx +++ b/python/pylibcudf/pylibcudf/stream_compaction.pyx @@ -26,6 +26,10 @@ from .column cimport Column from .expressions cimport Expression from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -43,9 +47,9 @@ __all__ = [ cpdef Table drop_nulls( Table source_table, - list keys, + list keys: list[int], size_type keep_threshold, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters out rows from the input table based on the presence of nulls. @@ -83,9 +87,9 @@ cpdef Table drop_nulls( cpdef Table drop_nans( Table source_table, - list keys, + list keys: list[int], size_type keep_threshold, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters out rows from the input table based on the presence of NaNs. @@ -124,7 +128,7 @@ cpdef Table drop_nans( cpdef Table apply_boolean_mask( Table source_table, Column boolean_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters out rows from the input table based on a boolean mask. @@ -161,7 +165,7 @@ cpdef Table apply_boolean_mask( cpdef Table apply_deletion_mask( Table source_table, Column deletion_mask, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters out rows from the input table using a deletion mask. @@ -197,10 +201,10 @@ cpdef Table apply_deletion_mask( cpdef Table unique( Table input, - list keys, + list keys: list[int], duplicate_keep_option keep, null_equality nulls_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filter duplicate consecutive rows from the input table. @@ -246,11 +250,11 @@ cpdef Table unique( cpdef Table distinct( Table input, - list keys, + list keys: list[int], duplicate_keep_option keep, null_equality nulls_equal, nan_equality nans_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Get the distinct rows from the input table. @@ -297,7 +301,7 @@ cpdef Column distinct_indices( duplicate_keep_option keep, null_equality nulls_equal, nan_equality nans_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Get the indices of the distinct rows from the input table. @@ -336,11 +340,11 @@ cpdef Column distinct_indices( cpdef Table stable_distinct( Table input, - list keys, + list keys: list[int], duplicate_keep_option keep, null_equality nulls_equal, nan_equality nans_equal, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Get the distinct rows from the input table, preserving input order. @@ -386,7 +390,7 @@ cpdef Table filter( Table predicate_table, Expression predicate_expr, Table filter_table, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Filters a table using a predicate expression. diff --git a/python/pylibcudf/pylibcudf/strings/attributes.pyx b/python/pylibcudf/pylibcudf/strings/attributes.pyx index ead494d047f3..5d6597a1a3f7 100644 --- a/python/pylibcudf/pylibcudf/strings/attributes.pyx +++ b/python/pylibcudf/pylibcudf/strings/attributes.pyx @@ -8,6 +8,10 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport attributes as cpp_attributes from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -15,7 +19,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["code_points", "count_bytes", "count_characters"] cpdef Column count_characters( - Column source_strings, object stream=None, DeviceMemoryResource mr=None + Column source_strings, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a column containing character lengths of each string @@ -47,7 +51,7 @@ cpdef Column count_characters( cpdef Column count_bytes( - Column source_strings, object stream=None, DeviceMemoryResource mr=None + Column source_strings, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a column containing byte lengths of each string @@ -79,7 +83,7 @@ cpdef Column count_bytes( cpdef Column code_points( - Column source_strings, object stream=None, DeviceMemoryResource mr=None + Column source_strings, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Creates a numeric column with code point values (integers) diff --git a/python/pylibcudf/pylibcudf/strings/capitalize.pyx b/python/pylibcudf/pylibcudf/strings/capitalize.pyx index 2c5683980191..0c07c6d6baa1 100644 --- a/python/pylibcudf/pylibcudf/strings/capitalize.pyx +++ b/python/pylibcudf/pylibcudf/strings/capitalize.pyx @@ -14,6 +14,10 @@ from pylibcudf.libcudf.strings cimport capitalize as cpp_capitalize from pylibcudf.scalar cimport Scalar from pylibcudf.strings.char_types cimport string_character_types from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -25,7 +29,7 @@ __all__ = ["capitalize", "is_title", "title"] cpdef Column capitalize( Column input, Scalar delimiters=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, # TODO: default scalar values # https://github.com/rapidsai/cudf/issues/15505 @@ -76,7 +80,7 @@ cpdef Column capitalize( cpdef Column title( Column input, string_character_types sequence_type=string_character_types.ALPHA, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Modifies first character of each word to upper-case and lower-cases @@ -109,7 +113,7 @@ cpdef Column title( return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column is_title(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_title(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Checks if the strings in the input column are title formatted. For details, see :cpp:func:`is_title`. diff --git a/python/pylibcudf/pylibcudf/strings/case.pyx b/python/pylibcudf/pylibcudf/strings/case.pyx index 5a122faacb16..111af660b230 100644 --- a/python/pylibcudf/pylibcudf/strings/case.pyx +++ b/python/pylibcudf/pylibcudf/strings/case.pyx @@ -8,13 +8,17 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport case as cpp_case from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["swapcase", "to_lower", "to_upper"] -cpdef Column to_lower(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column to_lower(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Returns a column of lowercased strings. For details, see :cpp:func:`to_lower`. @@ -43,7 +47,7 @@ cpdef Column to_lower(Column input, object stream=None, DeviceMemoryResource mr= return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column to_upper(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column to_upper(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Returns a column of uppercased strings. For details, see :cpp:func:`to_upper`. @@ -72,7 +76,7 @@ cpdef Column to_upper(Column input, object stream=None, DeviceMemoryResource mr= return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column swapcase(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column swapcase(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Returns a column of strings where the lowercase characters are converted to uppercase and the uppercase characters are converted to lowercase. diff --git a/python/pylibcudf/pylibcudf/strings/char_types.pyx b/python/pylibcudf/pylibcudf/strings/char_types.pyx index d7e155f548a5..10e03fee8552 100644 --- a/python/pylibcudf/pylibcudf/strings/char_types.pyx +++ b/python/pylibcudf/pylibcudf/strings/char_types.pyx @@ -11,6 +11,10 @@ from pylibcudf.libcudf.strings cimport char_types as cpp_char_types from pylibcudf.libcudf.strings.char_types cimport string_character_types from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -29,7 +33,7 @@ cpdef Column all_characters_of_type( Column source_strings, string_character_types types, string_character_types verify_types, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -72,7 +76,7 @@ cpdef Column filter_characters_of_type( string_character_types types_to_remove, Scalar replacement, string_character_types types_to_keep, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/combine.pyx b/python/pylibcudf/pylibcudf/strings/combine.pyx index f147374df11d..a2fabe7352c0 100644 --- a/python/pylibcudf/pylibcudf/strings/combine.pyx +++ b/python/pylibcudf/pylibcudf/strings/combine.pyx @@ -14,6 +14,10 @@ from pylibcudf.scalar cimport Scalar from pylibcudf.table cimport Table from pylibcudf.libcudf.table.table_view cimport table_view from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -38,7 +42,7 @@ cpdef Column concatenate( Scalar narep=None, Scalar col_narep=None, separator_on_nulls separate_nulls=separator_on_nulls.YES, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -134,7 +138,7 @@ cpdef Column join_strings( Column input, Scalar separator, Scalar narep, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -190,7 +194,7 @@ cpdef Column join_list_elements( Scalar string_narep, separator_on_nulls separate_nulls, output_if_empty_list empty_list_policy, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/contains.pyx b/python/pylibcudf/pylibcudf/strings/contains.pyx index 2e2e89de2795..a38a46a2ad92 100644 --- a/python/pylibcudf/pylibcudf/strings/contains.pyx +++ b/python/pylibcudf/pylibcudf/strings/contains.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport contains as cpp_contains from pylibcudf.strings.regex_program cimport RegexProgram from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -19,7 +23,7 @@ __all__ = ["contains_re", "count_re", "like", "matches_re"] cpdef Column contains_re( Column input, RegexProgram prog, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Returns a boolean column identifying rows which match the given @@ -61,7 +65,7 @@ cpdef Column contains_re( cpdef Column count_re( Column input, RegexProgram prog, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Returns the number of times the given regex_program's pattern @@ -101,7 +105,7 @@ cpdef Column count_re( cpdef Column matches_re( Column input, RegexProgram prog, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Returns a boolean column identifying rows which @@ -143,7 +147,7 @@ cpdef Column like( Column input, str pattern, str escape_character=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_booleans.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_booleans.pyx index 2a4f77f4c0e8..020fdb066ae3 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_booleans.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_booleans.pyx @@ -11,6 +11,10 @@ from pylibcudf.libcudf.strings.convert cimport ( ) from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cython.operator import dereference from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -20,7 +24,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["from_booleans", "to_booleans"] cpdef Column to_booleans( - Column input, Scalar true_string, object stream=None, DeviceMemoryResource mr=None + Column input, Scalar true_string, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new bool column by parsing boolean values from the strings @@ -67,7 +71,7 @@ cpdef Column from_booleans( Column booleans, Scalar true_string, Scalar false_string, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_datetime.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_datetime.pyx index 9aab9ae06dce..49569055b503 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_datetime.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_datetime.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings.convert cimport ( convert_datetime as cpp_convert_datetime, ) from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -22,7 +26,7 @@ cpdef Column to_timestamps( Column input, DataType timestamp_type, str format, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -71,7 +75,7 @@ cpdef Column from_timestamps( Column timestamps, str format, Column input_strings_names, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -120,7 +124,7 @@ cpdef Column from_timestamps( cpdef Column is_timestamp( Column input, str format, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_durations.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_durations.pyx index 8068d80a8e92..afe319a3ffdc 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_durations.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_durations.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings.convert cimport ( convert_durations as cpp_convert_durations, ) from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -22,7 +26,7 @@ cpdef Column to_durations( Column input, DataType duration_type, str format, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ @@ -71,7 +75,7 @@ cpdef Column to_durations( cpdef Column from_durations( Column durations, str format=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_fixed_point.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_fixed_point.pyx index abc7d7a7848a..7ee321ffd646 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_fixed_point.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_fixed_point.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings.convert cimport ( ) from pylibcudf.types cimport DataType, type_id from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -18,7 +22,7 @@ __all__ = ["from_fixed_point", "is_fixed_point", "to_fixed_point"] cpdef Column to_fixed_point( - Column input, DataType output_type, object stream=None, DeviceMemoryResource mr=None + Column input, DataType output_type, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new fixed-point column parsing decimal values from the @@ -59,7 +63,7 @@ cpdef Column to_fixed_point( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column from_fixed_point( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new strings column converting the fixed-point values @@ -96,7 +100,7 @@ cpdef Column from_fixed_point( cpdef Column is_fixed_point( Column input, DataType decimal_type=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_floats.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_floats.pyx index cd48c98e9d92..c73db3eda134 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_floats.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_floats.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings.convert cimport ( ) from pylibcudf.types cimport DataType from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -19,7 +23,7 @@ __all__ = ["from_floats", "is_float", "to_floats"] cpdef Column to_floats( Column strings, DataType output_type, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -62,7 +66,7 @@ cpdef Column to_floats( cpdef Column from_floats( - Column floats, object stream=None, DeviceMemoryResource mr=None + Column floats, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new strings column converting the float values from the @@ -97,7 +101,7 @@ cpdef Column from_floats( return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column is_float(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_float(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Returns a boolean column identifying strings in which all characters are valid for conversion to floats. diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_integers.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_integers.pyx index 95db3b8829b5..f0c78e852e8e 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_integers.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_integers.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings.convert cimport ( ) from pylibcudf.types cimport DataType from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -24,7 +28,7 @@ __all__ = [ ] cpdef Column to_integers( - Column input, DataType output_type, object stream=None, DeviceMemoryResource mr=None + Column input, DataType output_type, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new integer numeric column parsing integer values from the @@ -68,7 +72,7 @@ cpdef Column to_integers( cpdef Column from_integers( - Column integers, object stream=None, DeviceMemoryResource mr=None + Column integers, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new strings column converting the integer values from the @@ -110,7 +114,7 @@ cpdef Column from_integers( cpdef Column is_integer( Column input, DataType int_type=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -167,7 +171,7 @@ cpdef Column is_integer( cpdef Column hex_to_integers( - Column input, DataType output_type, object stream=None, DeviceMemoryResource mr=None + Column input, DataType output_type, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new integer numeric column parsing hexadecimal values @@ -210,7 +214,7 @@ cpdef Column hex_to_integers( return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column is_hex(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_hex(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Returns a boolean column identifying strings in which all characters are valid for conversion to integers from hex. @@ -249,7 +253,7 @@ cpdef Column is_hex(Column input, object stream=None, DeviceMemoryResource mr=No cpdef Column integers_to_hex( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a new strings column converting integer columns to hexadecimal diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_ipv4.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_ipv4.pyx index f0a262192b8b..3f3ada06ccb2 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_ipv4.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_ipv4.pyx @@ -7,6 +7,10 @@ from pylibcudf.column cimport Column from pylibcudf.libcudf.column.column cimport column, column_view from pylibcudf.libcudf.strings.convert cimport convert_ipv4 as cpp_convert_ipv4 from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -14,7 +18,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["integers_to_ipv4", "ipv4_to_integers", "is_ipv4"] cpdef Column ipv4_to_integers( - Column input, object stream=None, DeviceMemoryResource mr=None + Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Converts IPv4 addresses into integers. @@ -49,7 +53,7 @@ cpdef Column ipv4_to_integers( cpdef Column integers_to_ipv4( - Column integers, object stream=None, DeviceMemoryResource mr=None + Column integers, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Converts integers into IPv4 addresses as strings. @@ -83,7 +87,7 @@ cpdef Column integers_to_ipv4( return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column is_ipv4(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_ipv4(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Returns a boolean column identifying strings in which all characters are valid for conversion to integers from IPv4 format. diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_lists.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_lists.pyx index 903b83a9ea92..af86e771232b 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_lists.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_lists.pyx @@ -16,6 +16,10 @@ from pylibcudf.libcudf.strings.convert cimport ( from pylibcudf.scalar cimport Scalar from pylibcudf.types cimport type_id from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -28,7 +32,7 @@ cpdef Column format_list_column( Column input, Scalar na_rep=None, Column separators=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/convert/convert_urls.pyx b/python/pylibcudf/pylibcudf/strings/convert/convert_urls.pyx index 5c9d4bf17a79..0cd2ebb96d5d 100644 --- a/python/pylibcudf/pylibcudf/strings/convert/convert_urls.pyx +++ b/python/pylibcudf/pylibcudf/strings/convert/convert_urls.pyx @@ -7,6 +7,10 @@ from pylibcudf.column cimport Column from pylibcudf.libcudf.column.column cimport column, column_view from pylibcudf.libcudf.strings.convert cimport convert_urls as cpp_convert_urls from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -14,7 +18,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["url_decode", "url_encode"] -cpdef Column url_encode(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column url_encode(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Encodes each string using URL encoding. @@ -47,7 +51,7 @@ cpdef Column url_encode(Column input, object stream=None, DeviceMemoryResource m return Column.from_libcudf(move(c_result), _stream, mr) -cpdef Column url_decode(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column url_decode(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """ Decodes each string using URL encoding. diff --git a/python/pylibcudf/pylibcudf/strings/extract.pyx b/python/pylibcudf/pylibcudf/strings/extract.pyx index 3dee27693a37..3ca69488d476 100644 --- a/python/pylibcudf/pylibcudf/strings/extract.pyx +++ b/python/pylibcudf/pylibcudf/strings/extract.pyx @@ -12,6 +12,10 @@ from pylibcudf.strings.regex_program cimport RegexProgram from pylibcudf.table cimport Table from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -19,7 +23,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["extract", "extract_all_record", "extract_single"] cpdef Table extract( - Column input, RegexProgram prog, object stream=None, DeviceMemoryResource mr=None + Column input, RegexProgram prog, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a table of strings columns where each column @@ -59,7 +63,7 @@ cpdef Table extract( cpdef Column extract_all_record( - Column input, RegexProgram prog, object stream=None, DeviceMemoryResource mr=None + Column input, RegexProgram prog, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a lists column of strings where each string column @@ -102,7 +106,7 @@ cpdef Column extract_single( Column input, RegexProgram prog, size_type group, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/find.pyx b/python/pylibcudf/pylibcudf/strings/find.pyx index d9c9e81dea7d..bcd9f86c73d7 100644 --- a/python/pylibcudf/pylibcudf/strings/find.pyx +++ b/python/pylibcudf/pylibcudf/strings/find.pyx @@ -9,6 +9,10 @@ from pylibcudf.libcudf.strings cimport find as cpp_find from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -24,7 +28,7 @@ cpdef Column find( ColumnOrScalar target, size_type start=0, size_type stop=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Returns a column of character position values where the target string is @@ -98,7 +102,7 @@ cpdef Column rfind( Scalar target, size_type start=0, size_type stop=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -146,7 +150,7 @@ cpdef Column rfind( cpdef Column contains( Column input, ColumnOrScalar target, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -212,7 +216,7 @@ cpdef Column contains( cpdef Column starts_with( Column input, ColumnOrScalar target, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -278,7 +282,7 @@ cpdef Column starts_with( cpdef Column ends_with( Column input, ColumnOrScalar target, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/find_multiple.pyx b/python/pylibcudf/pylibcudf/strings/find_multiple.pyx index 3ac87788cd6f..611f8347613c 100644 --- a/python/pylibcudf/pylibcudf/strings/find_multiple.pyx +++ b/python/pylibcudf/pylibcudf/strings/find_multiple.pyx @@ -10,6 +10,10 @@ from pylibcudf.libcudf.strings cimport find_multiple as cpp_find_multiple from pylibcudf.libcudf.table.table cimport table from pylibcudf.table cimport Table from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -19,7 +23,7 @@ __all__ = ["find_multiple", "contains_multiple"] cpdef Column find_multiple( Column input, Column targets, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -62,7 +66,7 @@ cpdef Column find_multiple( cpdef Table contains_multiple( Column input, Column targets, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/findall.pyx b/python/pylibcudf/pylibcudf/strings/findall.pyx index 17099870f3aa..214f2c61b4f5 100644 --- a/python/pylibcudf/pylibcudf/strings/findall.pyx +++ b/python/pylibcudf/pylibcudf/strings/findall.pyx @@ -9,6 +9,10 @@ from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport findall as cpp_findall from pylibcudf.strings.regex_program cimport RegexProgram from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -16,7 +20,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["findall", "find_re"] cpdef Column findall( - Column input, RegexProgram pattern, object stream=None, DeviceMemoryResource mr=None + Column input, RegexProgram pattern, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns a lists column of strings for each matching occurrence using @@ -55,7 +59,7 @@ cpdef Column findall( cpdef Column find_re( - Column input, RegexProgram pattern, object stream=None, DeviceMemoryResource mr=None + Column input, RegexProgram pattern, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Returns character positions where the pattern first matches diff --git a/python/pylibcudf/pylibcudf/strings/padding.pyx b/python/pylibcudf/pylibcudf/strings/padding.pyx index 56949d84eb5f..8aafea761109 100644 --- a/python/pylibcudf/pylibcudf/strings/padding.pyx +++ b/python/pylibcudf/pylibcudf/strings/padding.pyx @@ -9,6 +9,10 @@ from pylibcudf.libcudf.strings cimport padding as cpp_padding from pylibcudf.libcudf.strings.side_type cimport side_type from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -20,7 +24,7 @@ cpdef Column pad( size_type width, side_type side, str fill_char, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -65,7 +69,7 @@ cpdef Column pad( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column zfill( - Column input, size_type width, object stream=None, DeviceMemoryResource mr=None + Column input, size_type width, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Add '0' as padding to the left of each string. @@ -102,7 +106,7 @@ cpdef Column zfill( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column zfill_by_widths( - Column input, Column widths, object stream=None, DeviceMemoryResource mr=None + Column input, Column widths, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Add '0' as padding to the left of each string. diff --git a/python/pylibcudf/pylibcudf/strings/repeat.pyx b/python/pylibcudf/pylibcudf/strings/repeat.pyx index 3f3fd1f0310a..fb5f12415d92 100644 --- a/python/pylibcudf/pylibcudf/strings/repeat.pyx +++ b/python/pylibcudf/pylibcudf/strings/repeat.pyx @@ -12,6 +12,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from ..utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["repeat_strings"] @@ -19,7 +23,7 @@ __all__ = ["repeat_strings"] cpdef Column repeat_strings( Column input, ColumnorSizeType repeat_times, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/replace.pyx b/python/pylibcudf/pylibcudf/strings/replace.pyx index 5603ac849f0c..feb0e1c1d195 100644 --- a/python/pylibcudf/pylibcudf/strings/replace.pyx +++ b/python/pylibcudf/pylibcudf/strings/replace.pyx @@ -18,6 +18,10 @@ from pylibcudf.libcudf.strings.replace cimport ( from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -29,7 +33,7 @@ cpdef Column replace( Scalar target, Scalar repl, size_type maxrepl=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replaces target string within each string with the specified replacement string. @@ -84,7 +88,7 @@ cpdef Column replace_multiple( Column target, Column repl, size_type maxrepl=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replaces target string within each string with the specified replacement string. @@ -137,7 +141,7 @@ cpdef Column replace_slice( Scalar repl=None, size_type start=0, size_type stop=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Replaces each string in the column with the provided repl string diff --git a/python/pylibcudf/pylibcudf/strings/replace_re.pyx b/python/pylibcudf/pylibcudf/strings/replace_re.pyx index 7112379911d5..05f2cb2153ab 100644 --- a/python/pylibcudf/pylibcudf/strings/replace_re.pyx +++ b/python/pylibcudf/pylibcudf/strings/replace_re.pyx @@ -16,6 +16,10 @@ from pylibcudf.libcudf.types cimport size_type from pylibcudf.scalar cimport Scalar from pylibcudf.strings.regex_program cimport RegexProgram from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -27,7 +31,7 @@ cpdef Column replace_re( RegexProgram pattern, Scalar replacement=None, size_type max_replace_count=-1, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -83,7 +87,7 @@ cpdef Column replace_with_backrefs( Column input, RegexProgram prog, str replacement, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/reverse.pyx b/python/pylibcudf/pylibcudf/strings/reverse.pyx index ba31d98f20b4..16d2dec30897 100644 --- a/python/pylibcudf/pylibcudf/strings/reverse.pyx +++ b/python/pylibcudf/pylibcudf/strings/reverse.pyx @@ -8,13 +8,17 @@ from pylibcudf.libcudf.column.column cimport column from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport reverse as cpp_reverse from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["reverse"] -cpdef Column reverse(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column reverse(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Reverses the characters within each string. Any null string entries return corresponding null output column entries. diff --git a/python/pylibcudf/pylibcudf/strings/slice.pyx b/python/pylibcudf/pylibcudf/strings/slice.pyx index dc86db07ce93..de72105e0c03 100644 --- a/python/pylibcudf/pylibcudf/strings/slice.pyx +++ b/python/pylibcudf/pylibcudf/strings/slice.pyx @@ -19,6 +19,10 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from ..utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["slice_strings"] @@ -28,7 +32,7 @@ cpdef Column slice_strings( ColumnOrScalar start=None, ColumnOrScalar stop=None, Scalar step=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Perform a slice operation on a strings column. diff --git a/python/pylibcudf/pylibcudf/strings/split/partition.pyx b/python/pylibcudf/pylibcudf/strings/split/partition.pyx index ea8735da729e..3e58c326e95a 100644 --- a/python/pylibcudf/pylibcudf/strings/split/partition.pyx +++ b/python/pylibcudf/pylibcudf/strings/split/partition.pyx @@ -13,6 +13,10 @@ from pylibcudf.libcudf.table.table cimport table from pylibcudf.scalar cimport Scalar from pylibcudf.table cimport Table from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -24,7 +28,7 @@ __all__ = ["partition", "rpartition"] cpdef Table partition( Column input, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -75,7 +79,7 @@ cpdef Table partition( cpdef Table rpartition( Column input, Scalar delimiter=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/split/split.pyx b/python/pylibcudf/pylibcudf/strings/split/split.pyx index 0a6c71ee77c6..3656e22af7e0 100644 --- a/python/pylibcudf/pylibcudf/strings/split/split.pyx +++ b/python/pylibcudf/pylibcudf/strings/split/split.pyx @@ -12,6 +12,10 @@ from pylibcudf.scalar cimport Scalar from pylibcudf.strings.regex_program cimport RegexProgram from pylibcudf.table cimport Table from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -33,7 +37,7 @@ cpdef Table split( Column strings_column, Scalar delimiter, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -87,7 +91,7 @@ cpdef Table rsplit( Column strings_column, Scalar delimiter, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -140,7 +144,7 @@ cpdef Column split_record( Column strings, Scalar delimiter, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -190,7 +194,7 @@ cpdef Column rsplit_record( Column strings, Scalar delimiter, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -241,7 +245,7 @@ cpdef Table split_re( Column input, RegexProgram prog, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -288,7 +292,7 @@ cpdef Table rsplit_re( Column input, RegexProgram prog, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -336,7 +340,7 @@ cpdef Column split_record_re( Column input, RegexProgram prog, size_type maxsplit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -380,7 +384,7 @@ cpdef Column split_record_re( return Column.from_libcudf(move(c_result), _stream, mr) cpdef Column rsplit_record_re( - Column input, RegexProgram prog, size_type maxsplit, object stream=None, + Column input, RegexProgram prog, size_type maxsplit, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ @@ -425,7 +429,7 @@ cpdef Column rsplit_record_re( cpdef Column split_part( - Column input, Scalar delimiter, size_type index, object stream=None, + Column input, Scalar delimiter, size_type index, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): cdef unique_ptr[column] c_result diff --git a/python/pylibcudf/pylibcudf/strings/strip.pyx b/python/pylibcudf/pylibcudf/strings/strip.pyx index 0e4dd400d53d..801b237a062e 100644 --- a/python/pylibcudf/pylibcudf/strings/strip.pyx +++ b/python/pylibcudf/pylibcudf/strings/strip.pyx @@ -15,6 +15,10 @@ from pylibcudf.libcudf.strings cimport strip as cpp_strip from pylibcudf.scalar cimport Scalar from pylibcudf.strings.side_type cimport side_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -25,7 +29,7 @@ cpdef Column strip( Column input, side_type side=side_type.BOTH, Scalar to_strip=None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Removes the specified characters from the beginning diff --git a/python/pylibcudf/pylibcudf/strings/translate.pyx b/python/pylibcudf/pylibcudf/strings/translate.pyx index 4c48fdd72ca3..fd911e362145 100644 --- a/python/pylibcudf/pylibcudf/strings/translate.pyx +++ b/python/pylibcudf/pylibcudf/strings/translate.pyx @@ -12,6 +12,10 @@ from pylibcudf.libcudf.strings cimport translate as cpp_translate from pylibcudf.libcudf.types cimport char_utf8 from pylibcudf.scalar cimport Scalar from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream @@ -45,7 +49,10 @@ cdef vector[pair[char_utf8, char_utf8]] _table_to_c_table(dict table): cpdef Column translate( - Column input, dict chars_table, object stream=None, DeviceMemoryResource mr=None + Column input, + dict chars_table: dict[int | str, int | str], + object stream: CudaStreamLike | None = None, + DeviceMemoryResource mr=None, ): """ Translates individual characters within each string. @@ -87,10 +94,10 @@ cpdef Column translate( cpdef Column filter_characters( Column input, - dict characters_to_filter, + dict characters_to_filter: dict[int | str, int | str], filter_type keep_characters, Scalar replacement, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """ diff --git a/python/pylibcudf/pylibcudf/strings/wrap.pyx b/python/pylibcudf/pylibcudf/strings/wrap.pyx index f6180c7f56dd..23b20d8ac82c 100644 --- a/python/pylibcudf/pylibcudf/strings/wrap.pyx +++ b/python/pylibcudf/pylibcudf/strings/wrap.pyx @@ -9,6 +9,10 @@ from pylibcudf.libcudf.column.column_view cimport column_view from pylibcudf.libcudf.strings cimport wrap as cpp_wrap from pylibcudf.libcudf.types cimport size_type from pylibcudf.utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream from cuda.bindings.cyruntime cimport cudaStream_t @@ -16,7 +20,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["wrap"] cpdef Column wrap( - Column input, size_type width, object stream=None, DeviceMemoryResource mr=None + Column input, size_type width, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Wraps strings onto multiple lines shorter than `width` by diff --git a/python/pylibcudf/pylibcudf/table.pyx b/python/pylibcudf/pylibcudf/table.pyx index 3bbf318b6834..7c6d5b06c65b 100644 --- a/python/pylibcudf/pylibcudf/table.pyx +++ b/python/pylibcudf/pylibcudf/table.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from cython.operator cimport dereference from cpython.pycapsule cimport ( @@ -34,6 +35,10 @@ from pylibcudf.libcudf.types cimport size_type from .column cimport Column from .types cimport DataType from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from pylibcudf._interop_helpers cimport ( _release_schema, _release_array, @@ -76,7 +81,7 @@ cdef class Table: """ __hash__ = None - def __init__(self, columns, num_rows=None): + def __init__(self, columns: Sequence[Column], num_rows=None): columns = tuple(columns) if not all(isinstance(c, Column) for c in columns): raise ValueError("All columns must be pylibcudf Column objects") @@ -127,7 +132,7 @@ cdef class Table: def from_arrow( obj: ArrowLike, dtype: DataType | None = None, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ) -> Table: """ @@ -359,7 +364,7 @@ cdef class Table: """The shape of this table""" return (self.num_rows(), self.num_columns()) - cpdef Table copy(self, object stream=None, DeviceMemoryResource mr=None): + cpdef Table copy(self, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Create a deep copy of the table. Parameters @@ -404,7 +409,7 @@ cdef class Table: return PyCapsule_New(raw_schema_ptr, "arrow_schema", _release_schema) - def _to_host_array(self, object stream): + def _to_host_array(self, object stream: CudaStreamLike): cdef ArrowArray* raw_host_array_ptr cdef Stream _stream = _get_stream(stream) cdef cudaStream_t _cs = _stream.view().value() diff --git a/python/pylibcudf/pylibcudf/table_equality.pyx b/python/pylibcudf/pylibcudf/table_equality.pyx index e57ae491b76c..975bb90115c0 100644 --- a/python/pylibcudf/pylibcudf/table_equality.pyx +++ b/python/pylibcudf/pylibcudf/table_equality.pyx @@ -10,6 +10,10 @@ from rmm.pylibrmm.stream cimport Stream from .table cimport Table from .utils cimport _get_stream +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["tables_equal"] @@ -19,7 +23,7 @@ cpdef bool tables_equal( Table left, Table right, null_equality nulls_equal=null_equality.EQUAL, - object stream=None, + object stream: CudaStreamLike | None = None, ): """Check if two tables are equal. diff --git a/python/pylibcudf/pylibcudf/transform.pyx b/python/pylibcudf/pylibcudf/transform.pyx index 66c392f524cd..2d29ecd4eef8 100644 --- a/python/pylibcudf/pylibcudf/transform.pyx +++ b/python/pylibcudf/pylibcudf/transform.pyx @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from collections.abc import Sequence from cython.operator cimport dereference from libcpp.memory cimport unique_ptr @@ -31,6 +32,10 @@ from .expressions cimport Expression from .gpumemoryview cimport gpumemoryview from .types cimport DataType, null_aware, output_nullability from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t ctypedef const cpp_transform.transform_input const_transform_input @@ -50,7 +55,7 @@ __all__ = [ cpdef tuple[gpumemoryview, int] nans_to_nulls( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a null mask preserving existing nulls and converting nans to null. @@ -92,7 +97,7 @@ cpdef tuple[gpumemoryview, int] nans_to_nulls( cpdef Column column_nans_to_nulls( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a column with nans converted to nulls. @@ -129,7 +134,7 @@ cpdef Column column_nans_to_nulls( cpdef Column compute_column( - Table input, Expression expr, object stream=None, DeviceMemoryResource mr=None + Table input, Expression expr, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Create a column by evaluating an expression on a table. @@ -166,7 +171,7 @@ cpdef Column compute_column( cpdef Column compute_column_jit( - Table input, Expression expr, object stream=None, DeviceMemoryResource mr=None + Table input, Expression expr, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """ Create a column by evaluating an expression on a table @@ -206,7 +211,7 @@ cpdef Column compute_column_jit( cpdef tuple[gpumemoryview, int] bools_to_mask( Column input, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a bitmask from a column of boolean elements @@ -249,7 +254,7 @@ cpdef Column mask_to_bools( Py_ssize_t bitmask, int begin_bit, int end_bit, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Creates a boolean column from given bitmask. @@ -292,13 +297,13 @@ cpdef Column mask_to_bools( cpdef Column transform( - inputs, + inputs: Sequence[Column], str transform_udf, DataType output_type, bool is_ptx, null_aware is_null_aware, output_nullability null_policy, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Create a new column by applying a transform function against @@ -400,7 +405,7 @@ cpdef Column transform( return Column.from_libcudf(move(c_result), _stream, mr) cpdef tuple[Table, Column] encode( - Table input, object stream=None, DeviceMemoryResource mr=None + Table input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Encode the rows of the given table as integers. @@ -437,7 +442,7 @@ cpdef tuple[Table, Column] encode( cpdef Table one_hot_encode( Column input, Column categories, - object stream=None, + object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None, ): """Encodes `input` by generating a new column diff --git a/python/pylibcudf/pylibcudf/transpose.pyx b/python/pylibcudf/pylibcudf/transpose.pyx index 38ad6e68e17f..eda2a5bc0f03 100644 --- a/python/pylibcudf/pylibcudf/transpose.pyx +++ b/python/pylibcudf/pylibcudf/transpose.pyx @@ -13,12 +13,16 @@ from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from .column cimport Column from .table cimport Table from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = ["transpose"] cpdef Table transpose( - Table input_table, object stream=None, DeviceMemoryResource mr=None + Table input_table, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Transpose a Table. diff --git a/python/pylibcudf/pylibcudf/types.pyx b/python/pylibcudf/pylibcudf/types.pyx index a5b6a19a0494..cbb369938a0b 100644 --- a/python/pylibcudf/pylibcudf/types.pyx +++ b/python/pylibcudf/pylibcudf/types.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, 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 @@ -26,13 +26,10 @@ from pylibcudf.libcudf.types import order as Order # no-cython-lint, isort:skip from pylibcudf.libcudf.types import sorted as Sorted # no-cython-lint, isort:skip from functools import cache +from typing import TYPE_CHECKING -try: +if TYPE_CHECKING: import pyarrow as pa - pa_err = None -except ImportError as e: - pa = None - pa_err = e try: import pyarrow as pa @@ -197,7 +194,7 @@ cdef class DataType: ret.c_obj = dt return ret - def to_arrow(self, **kwargs): + def to_arrow(self, **kwargs) -> pa.DataType: """ Convert a datatype to arrow. @@ -252,7 +249,7 @@ cdef class DataType: ) @staticmethod - def from_arrow(pa_typ) -> DataType: + def from_arrow(pa_typ: pa.DataType) -> DataType: """ Construct a DataType from a Python type. diff --git a/python/pylibcudf/pylibcudf/typing.py b/python/pylibcudf/pylibcudf/typing.py new file mode 100644 index 000000000000..965bcd054054 --- /dev/null +++ b/python/pylibcudf/pylibcudf/typing.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Protocol, TypeAlias, TypedDict + +from rmm.pylibrmm.stream import Stream + + +class HasCudaStream(Protocol): + def __cuda_stream__(self) -> object: ... + + +CudaStreamLike: TypeAlias = Stream | HasCudaStream + + +class ArrayInterfaceBase(TypedDict): + data: tuple[int, bool] + shape: tuple[int, ...] + typestr: str + version: int + + +class SupportsCudaArrayInterface(Protocol): + @property + def __cuda_array_interface__(self) -> ArrayInterfaceBase: ... + + +class SupportsArrayInterface(Protocol): + @property + def __array_interface__(self) -> ArrayInterfaceBase: ... + + +__all__ = [ + "ArrayInterfaceBase", + "CudaStreamLike", + "HasCudaStream", + "SupportsArrayInterface", + "SupportsCudaArrayInterface", +] diff --git a/python/pylibcudf/pylibcudf/unary.pyx b/python/pylibcudf/pylibcudf/unary.pyx index 37fda4c86291..3a14404c858b 100644 --- a/python/pylibcudf/pylibcudf/unary.pyx +++ b/python/pylibcudf/pylibcudf/unary.pyx @@ -18,6 +18,10 @@ from pylibcudf.libcudf.unary import \ from .column cimport Column from .types cimport DataType from .utils cimport _get_stream, _get_memory_resource +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike from cuda.bindings.cyruntime cimport cudaStream_t __all__ = [ @@ -33,7 +37,7 @@ __all__ = [ ] cpdef Column unary_operation( - Column input, unary_operator op, object stream=None, DeviceMemoryResource mr=None + Column input, unary_operator op, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Perform a unary operation on a column. @@ -70,7 +74,7 @@ cpdef Column unary_operation( return Column.from_libcudf(move(result), _stream, mr) -cpdef Column is_null(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_null(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Check whether elements of a column are null. For details, see :cpp:func:`is_null`. @@ -102,7 +106,7 @@ cpdef Column is_null(Column input, object stream=None, DeviceMemoryResource mr=N return Column.from_libcudf(move(result), _stream, mr) -cpdef Column is_valid(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_valid(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Check whether elements of a column are valid. For details, see :cpp:func:`is_valid`. @@ -135,7 +139,7 @@ cpdef Column is_valid(Column input, object stream=None, DeviceMemoryResource mr= cpdef Column cast( - Column input, DataType data_type, object stream=None, DeviceMemoryResource mr=None + Column input, DataType data_type, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Cast a column to a different data type. @@ -173,7 +177,7 @@ cpdef Column cast( cpdef Column bit_cast( - Column input, DataType data_type, object stream=None, DeviceMemoryResource mr=None + Column input, DataType data_type, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None ): """Bit-cast a column to a different data type. @@ -210,7 +214,7 @@ cpdef Column bit_cast( return Column.from_libcudf(move(result), _stream, mr) -cpdef Column is_nan(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_nan(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Check whether elements of a column are nan. For details, see :cpp:func:`is_nan`. @@ -242,7 +246,7 @@ cpdef Column is_nan(Column input, object stream=None, DeviceMemoryResource mr=No return Column.from_libcudf(move(result), _stream, mr) -cpdef Column is_not_nan(Column input, object stream=None, DeviceMemoryResource mr=None): +cpdef Column is_not_nan(Column input, object stream: CudaStreamLike | None = None, DeviceMemoryResource mr=None): """Check whether elements of a column are not nan. For details, see :cpp:func:`is_not_nan`. diff --git a/python/pylibcudf/pylibcudf/utils.pyx b/python/pylibcudf/pylibcudf/utils.pyx index 314e62f77606..576351dc798e 100644 --- a/python/pylibcudf/pylibcudf/utils.pyx +++ b/python/pylibcudf/pylibcudf/utils.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cython.operator import dereference @@ -16,10 +16,14 @@ from rmm.pylibrmm.memory_resource cimport ( get_current_device_resource, ) -from rmm.pylibrmm.stream import DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM +from rmm.pylibrmm.stream import DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, Stream import os +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pylibcudf.typing import CudaStreamLike # Check the environment for the variable CUDF_PER_THREAD_STREAM. If it is set, @@ -47,7 +51,7 @@ cdef vector[reference_wrapper[const scalar]] _as_vector(list source): return c_scalars -cpdef Stream _get_stream(object stream = None): +cpdef Stream _get_stream(object stream: CudaStreamLike | None = None): if stream is None: return CUDF_DEFAULT_STREAM if isinstance(stream, Stream):