[Data] - Add shuffle aggregation type to JoinOperator

[goutamvenkat-anyscale]

Why are these changes needed?

Add shuffle aggregation type to join operator.

Related issue number

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 introduces a shuffle_aggregation_type parameter to the JoinOperator, allowing for custom shuffle aggregation logic. The implementation is clean and correctly integrates the new parameter. My review includes a suggestion to improve the type validation for this new parameter to provide more informative error messages for invalid inputs, enhancing the operator's robustness and user experience.

[gemini-code-assist]

The issubclass() built-in will raise a TypeError if shuffle_aggregation_type is not a class (e.g., if an instance of a class is passed). While the default error message is acceptable, we can provide a more user-friendly and consistent error message by handling this case explicitly. This improves error handling and makes debugging easier for users of this operator.

        if shuffle_aggregation_type is not None:
            try:
                is_sub = issubclass(shuffle_aggregation_type, StatefulShuffleAggregation)
            except TypeError:
                is_sub = False
            if not is_sub:
                raise TypeError(
                    f"shuffle_aggregation_type must be a class and a subclass of StatefulShuffleAggregation, got {shuffle_aggregation_type}"
                )