[Python] Add OrderScheme partitioning metadata

[rjzamora]

Depends on #993 (adds the C++ API)
Adds mechanism to track ordering within ChannelMetadata. Needed in cudf-polars to keep track of sorted data.

Proposed API

@dataclass(frozen=True, slots=True)
class OrderKey:
    """Sort key: column index, direction, and null ordering (``pylibcudf`` enums)."""

    column_index: int
    order: plc.types.Order
    null_order: plc.types.NullOrder

class OrderScheme:
    def __init__(
        self,
        keys: Sequence[OrderKey],
        boundaries: TableChunk,
        *,
        strict_boundaries: bool = False,
    ) -> None: ...
    @property
    def keys(self) -> tuple[OrderKey, ...]: ...
    @property
    def strict_boundaries(self) -> bool: ...
    @property
    def num_boundaries(self) -> int: ...
    def replace_keys(self, new_keys: Sequence[OrderKey]) -> OrderScheme: ...
    def boundaries_aligned_with(self, other: OrderScheme) -> bool: ...
    # Shallow equality: keys + strict_boundaries + boundary shape only.
    def __eq__(self, other: object) -> bool: ...
    def __repr__(self) -> str: ...

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[rjzamora]

\okay to test

[rjzamora]

@TomAugspurger - I'd be interested to get your thoughts on the specific contract proposed here.

[TomAugspurger]

High-level, I think this makes sense as a way to support what we need out of cudf-polars. Left some comments and questions on the implementation.

[TomAugspurger]

Looks like a todo here. Solving that here, or file a separate issue to track it?

[rjzamora]

I updated this comment a bit. I think its "okay" for us to do a "shallow" equality and document that this is the case. We can implement a separate API for strict equality later (or just let cudf-polars do this pylibcudf).

[TomAugspurger]

This seems thorny.

Do we want to avoid overloading == here entirely? I suspect that if we do ever need to compare two OrderScheme structs in a way that includes a comparison of the boundary values, then we'll need an API that provides a stream and mr.

[rjzamora]

Agree that this is thorny. I don't think we should ever require == to enforce "deep" equality here. We need a separate API for that.

[TomAugspurger]

Do we need this? Is it just order_scheme.boundaries is not None?

[TomAugspurger]

I don't see this import elsewhere. Do we need it?

In

rapidsmpf/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx

Line 14 in c7ecfe5

from rmm.pylibrmm.stream cimport Stream

, we use Stream._from_cudaStream_t. But I'm not sure about the type.

[wence-]

I feel like the cython bindings want a bit of a tidy up now we've changed how the C++ works.

[wence-]

Oh, how about:

Suggested change

	cdef vector[cpp_OrderKey] cpp_keys
	for key in keys:
	if not isinstance(key, OrderKey):
	raise TypeError(
	f"keys must contain OrderKey objects, got {type(key).__name__}"
	)
	cpp_keys.push_back((<OrderKey>key)._key)
	cdef vector[cpp_OrderKey] cpp_keys
	for key in keys:
	cpp_keys.push_back((<OrderKey?>key)._key)

[rjzamora]

I feel like the cython bindings want a bit of a tidy up now we've changed how the C++ works.

Thanks for the quick review @wence- ! I addressed your suggestions, but let me know if you're looking for a larger overhaul here.

[rjzamora]

Thanks (again) for the review @wence- !

[rjzamora]

/merge