[Release Tests] Adding configuration to test aggregations with different shuffle strategies (#51318)

## Why are these changes needed?

Added configuration to test aggregations with different shuffle
strategies.

---------

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Author: alexeykudinkin

release/nightly_tests/dataset/groupby_benchmark.py

@@ -6,6 +6,8 @@
 import ray
 
 from benchmark import Benchmark
+from python.ray.data import DataContext
+from python.ray.data.context import ShuffleStrategy
 
 
 def parse_args() -> argparse.Namespace:
@@ -24,6 +26,14 @@ def parse_args() -> argparse.Namespace:
         type=str,
         help="Which columns to group by",
     )
+    parser.add_argument(
+        "--shuffle-strategy",
+        required=False,
+        default=ShuffleStrategy.SORT_SHUFFLE_PULL_BASED,
+        nargs="?",
+        type=str,
+        help="Strategy to use when shuffling data (see ShuffleStrategy for accepted values)",
+    )
 
     consume_group = parser.add_mutually_exclusive_group()
     consume_group.add_argument("--aggregate", action="store_true")
@@ -39,6 +49,11 @@ def main(args):
     def benchmark_fn():
         path = f"s3://ray-benchmark-data/tpch/parquet/sf{args.sf}/lineitem"
 
+        # Configure appropriate shuffle-strategy
+        DataContext.get_current().shuffle_strategy = ShuffleStrategy(
+            args.shuffle_strategy
+        )
+
         grouped_ds = ray.data.read_parquet(path).groupby(args.group_by)
         consume_fn(grouped_ds)
 

release/release_data_tests.yaml

@@ -135,16 +135,17 @@
     timeout: 3600
 
   variations:
-    - __suffix__: few_groups
+    - __suffix__: few_groups (sort_shuffle_pull_based)
       run:
         script: >
-          python groupby_benchmark.py --sf 10 --aggregate
-          --group-by column08 column13 column14
-    - __suffix__: many_groups
+          python groupby_benchmark.py --sf 10 --aggregate --group-by column08 column13 column14
+          --shuffle-strategy sort_shuffle_pull_based
+
+    - __suffix__: many_groups (sort_shuffle_pull_based)
       run:
         script: >
-          python groupby_benchmark.py --sf 10 --aggregate
-          --group-by column02 column14
+          python groupby_benchmark.py --sf 10 --aggregate --group-by column02 column14
+          --shuffle-strategy sort_shuffle_pull_based
 
 - name: map_groups
 
@@ -155,16 +156,17 @@
     timeout: 3600
 
   variations:
-    - __suffix__: few_groups
+    - __suffix__: few_groups (sort_shuffle_pull_based)
       run:
         script: >
-          python groupby_benchmark.py --sf 10 --map-groups
-          --group-by column08 column13 column14
-    - __suffix__: many_groups
+          python groupby_benchmark.py --sf 10 --map-groups --group-by column08 column13 column14
+          --shuffle-strategy sort_shuffle_pull_based
+
+    - __suffix__: many_groups (sort_shuffle_pull_based)
       run:
         script: >
-          python groupby_benchmark.py --sf 10 --map-groups
-          --group-by column02 column14
+          python groupby_benchmark.py --sf 10 --map-groups --group-by column02 column14
+          --shuffle-strategy sort_shuffle_pull_based
 
 #######################
 # Streaming split tests