Joining data#

Note

Joins are experimental, and some behavior might not work as expected. Joins are available in Ray 2.46 and later.

Ray Data joins multiple Dataset instances on the key columns you provide, using any of the supported join types:

import ray

doubles_ds = ray.data.range(4).map(
    lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
)

squares_ds = ray.data.range(4).map(
    lambda row: {"id": row["id"], "square": int(row["id"]) ** 2}
)

doubles_and_squares_ds = doubles_ds.join(
    squares_ds,
    join_type="inner",
    num_partitions=2,
    on=("id",),
)

Ray Data supports the following join types:

  • Inner and outer joins: Inner, left outer, right outer, and full outer.

  • Semi joins: Left semi and right semi joins return all rows that have at least one matching row in the other table. They return only the columns from the requested side.

  • Anti joins: Left anti and right anti joins return rows that have no matching rows in the other table. They return only the columns from the requested side.

See Dataset.join for the current list.

Internally, joins hash-partition both datasets on the join keys and join each partition with Polars. By default, joins use shuffle v2, which is the ShuffleStrategy.SHUFFLE_V2 strategy. See Tune shuffle v2 for the memory-related settings. To fall back to the previous hash-shuffle implementation, set ray.data.DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE before you create a Dataset.

Configure joins#

Joins are generally memory-intensive and require accurate memory accounting and projection, so they’re sensitive to skew and imbalance in the dataset.

Joins accept the following two partitioning parameters. Only num_partitions affects join performance:

  • num_partitions: Required. The number of hash partitions to split both incoming datasets into. See Configure the number of partitions for tuning guidance.

  • partition_size_hint: Deprecated. A hint to the join operator about the estimated average size of an individual partition, in bytes. Ray Data ignores this parameter, and a future release removes it. Passing a value emits a DeprecationWarning. Instead of a hint, the join path sizes reduce-task memory from observed partition sizes.

Configure the number of partitions#

The number of partitions, also called blocks, sets a trade-off between the size of the batch of rows that each task handles and the memory that the operation on those rows requires.

As a rule of thumb, keep partitions large, but not so large that they cause out-of-memory (OOM) errors. Joined partitions that are too large to fit in memory cause OOM errors. Don’t create too many small partitions either, because passing a large number of smaller objects adds overhead.

Configure the number of aggregators#

Aggregators are worker actors that perform the joins, aggregations, and shuffling. They receive individual partition chunks from the incoming blocks and aggregate them as the operation requires.

Consider the following when you configure the number of aggregators in your pool:

  • The number of aggregators defaults to the smallest of num_partitions, the number of CPUs in the cluster, and DataContext.max_hash_shuffle_aggregators, which is 128 by default.

  • An individual aggregator might handle more than one partition. Ray Data splits partitions evenly among the aggregators in round-robin fashion.

  • Aggregators are stateful components that hold the partitions in memory during shuffling.

Note

As a rule of thumb, avoid setting num_partitions far higher than the number of aggregators, because doing so might create bottlenecks.

To cap the number of aggregators, set DataContext.max_hash_shuffle_aggregators. Ray Data also caps the number of aggregators at the number of CPUs in the cluster. To allocate one partition per aggregator, set max_hash_shuffle_aggregators >= num_partitions and make sure the cluster has at least num_partitions CPUs.