[data] Defer hash-shuffle schema broadcast to first row-bearing block

[curtis-anyscale]

Why are these changes needed?

The hash shuffle used by Dataset.join() broadcasts the input schema to aggregators. When the upstream step is itself a hash shuffle (e.g. groupby.map_groups) it can emit empty / zero-column blocks. If an empty block is the first to reach the join's hash shuffle, the empty schema is broadcast to the aggregators. Aggregators end up with schema-less blocks, and pa.Table.join fails resolving the join key:

pyarrow.lib.ArrowInvalid: No match or multiple matches for key field reference
FieldRef.Name(<key>) on left side of the join

This is reproducible on the latest stable (Ray 2.55.1) with a from_pandas → groupby.map_groups → join where default_hash_shuffle_parallelism is high enough relative to the distinct key cardinality to leave partitions empty.

The fix defers the broadcast to the first row-bearing block. Zero-row marker blocks fall through with send_empty_blocks=False and are filtered as before; a row-bearing block with a real schema will arrive whenever the dataset is non-empty, and the broadcast then propagates the correct schema to every aggregator.

Minimal repro (fails on master, passes with this PR)

import pandas as pd
import ray

ray.data.DataContext.get_current().default_hash_shuffle_parallelism = 20

def noop(df: pd.DataFrame) -> pd.DataFrame:
    return df

left = (
    ray.data.from_pandas(pd.DataFrame({"key": list("abcde"), "lval": [1, 2, 3, 4, 5]}))
    .groupby("key").map_groups(noop, batch_format="pandas")
)
right = (
    ray.data.from_pandas(pd.DataFrame({"key": list("abcde"), "rval": [10, 20, 30, 40, 50]}))
    .groupby("key").map_groups(noop, batch_format="pandas")
)

left.join(right, on=("key",), join_type="inner", num_partitions=5).show()

Related issue number

N/A.

Checks

• I've signed off every commit (by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a method in Tune, I've added it in doc/source/tune/api/ under the corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gemini-code-assist]

Code Review

This pull request modifies the hash shuffle operator to defer schema broadcasting until the first row-bearing block is received, preventing issues where zero-row marker blocks cause downstream join failures. A regression test was added to verify this fix. Feedback suggests refining the broadcast logic to handle cases where a dataset is entirely empty but contains a valid schema by also checking for the presence of columns in the block metadata.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit 6d27c8f. Configure here.}

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.