Align on standard ETL Metadata schema

[rjzamora]

In order to make adaptive decisions about shuffling/broadcasting/etc during run time, we need to keep track of some basic metadata. In cudf-polars, this metadata is propagated through the network using a distinct Channel between every node. After #811, we will be able to use the same Channel for both data and metadata payloads. However, in order for cudf-polars to use any C++ Join/GroupBy nodes, we will (probably) still need cudf-polars and RapidsMPF to align on a common/standard metadata schema.

At the very minimum, the channel metadata should keep track of the expected local chunk count, and whether the local data coming through the channel is already broadcasted (duplicated) across all ranks. In order to avoid "re-shuffling" data unnecessarily, we also want to track the "partitioning" status of the data being sent into an actor.

Current Metadata Schema

The current Metadata definition in cudf-polars looks like:

class Metadata:
    """Metadata payload for an ETL workload."""

    local_count: int
    """Local chunk-count estimate for the current rank."""
    global_count: int | None
    """Global chunk-count estimate across all ranks."""
    partitioning: HashPartitioned | None
    """How the data is partitioned, or None if not partitioned."""
    duplicated: bool
    """Whether the data is duplicated (identical) on all workers."""

class HashPartitioned:
    """Hash-partitioning metadata.

    columns: tuple[str, ...]
    """Columns the data is hash-partitioned on."""
    scope: Literal["local", "global"]
    """Whether data is partitioned locally (within a rank) or globally (across all ranks)."""
    count: int
    """The modulus used for hash partitioning (number of partitions)."""

Partitioning metadata

An important "open" question is how to best encode partitioning status of the data being sent through a channel. For now, the partitioning attribute can either be empty (None), or it can be set to HashPartitioned. This makes it easy to keep track of data that has already been globally shuffled. However, it does not make it easy to distinguish between data that was directly shuffled into N_g global chunks, and data that was shuffled between R ranks, and then N_l local chunks.

[rjzamora]

@wence- - I'm interested in your thoughts on this. I'm not happy with the current partitioning metadata in cudf-polars, because I don't think it extends (nicely) to the multi-level partitioning we will likely use more-frequently when we move to adaptive planning.

My current thinking is that something like:

(Note: I've edited this proposal a few times. #819 demonstrates a variation on this.)

struct HashSpec {
    std::vector<std::string> columns;
    int64_t modulus;
};

struct Aligned {};  // Marker: Partitioning inherited from parent level

using PartitioningSpec = std::variant<HashSpec, Aligned>;

struct Partitioning {
    std::optional<PartitioningSpec> inter_rank;
    std::optional<PartitioningSpec> local;
};

could handle complex partitioning more easily. For example, when data is directly shuffled to N_g global chunks (as we do with the static plan), we would have:

partitioning = Partitioning(inter_rank=HashSpec(("x",), N_g), local=Aligned())

Meaning: We are both globally partitioned using N_g as the modulus, and the local partitioning is aligned with the global partitioning. Therefore, we can do chunk-wise GroupBy aggregations. We can also do chunk-wise Joins (only) with data having the same partitioning mapping.

When data is globally shuffled by rank, and then locally into N_l chunks, we would have:

partitioning = Partitioning(inter_rank=HashSpec(("x",), N_g), local=HashSpec(("x",), N_l))

If this data is then locally "repartitioned" for some reason, this may become:

partitioning = Partitioning(inter_rank=HashSpec(("x",), N_g), local=None)

Meaning: We do not need to move data between ranks to perform GroupBy aggregations. We can also avoid global data movement to join with data having the same "inter-rank" partitioning (we just need to perform a local intra-rank shuffle to a matching number of local chunks).

Note (probably not important): I think this approach also allows us to adopt more-complex partitioning levels in the future (if we want). For example, if we decide to decompose the system into multiple groups of 1+ nodes, we could use a structure like the following to avoid shuffling outside of the current "group":

partitioning = Partitioning(inter_group=..., intra_group=..., local=...)