[dataset] Support push-based shuffle in groupby operations

[stephanie-wang]

Why are these changes needed?

Allows option for push-based shuffle in groupby operations.

[stephanie-wang]

Turns out reduce completely aggregates the inputs, but we need to only partially aggregate the inputs during the merge stage, so I made some changes to the BlockAccessor interface to support this. @ericl or @clarkzinzow can you review the new changes?

[clarkzinzow]

LGTM overall, I added something similar to this for a generic tree-reduce operation in a prototype, happy to see this land!

Just a few nits, I'd like to see if the use_push_based_shuffle parameterized fixture would work since that should decrease the code impact on the tests by a good bit.

[clarkzinzow]

Note that this changes the return type to Tuple[Block[Tuple[KeyType, Union[U, AggType]]], BlockMetadata]

[clarkzinzow]

Nit: You should be able to reduce the context munging and try-finally boilerplate by parameterizing over a few fixtures that does this for you, e.g.

@pytest.fixture(params=[True, False])
def use_push_baed_shuffle(request):
    ctx = ray.data.context.DatasetContext.get_current()
    original = ctx.use_push_based_shuffle
    ctx.use_push_based_shuffle = request.param
    yield
    ctx.use_push_based_shuffle = original


def test_groupby_arrow(ray_start_regular_shared, use_push_based_shuffle):
    # Test empty dataset.
    agg_ds = (
        ray.data.range_table(10)
        .filter(lambda r: r["value"] > 10)
        .groupby("value")
        .count()
    )
    assert agg_ds.count() == 0