[data][tests] Fix sort.py test ray-project#42264

Fix test script for 100TB sort. Verified fix locally. Closes ray-project#39527. Signed-off-by: Stephanie Wang <[email protected]>
c21 · Jan 9, 2024 · 601709e · 601709e
1 parent 08c0202
commit 601709e
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/release/nightly_tests/dataset/sort.py b/release/nightly_tests/dataset/sort.py
@@ -109,6 +109,11 @@ def make_block(count: int, row_size_bytes: int) -> Block:
  print("Using polars for sort")
  ctx = DataContext.get_current()
  ctx.use_polars = True
+ ctx = DataContext.get_current()
+ if args.limit_num_blocks is not None:
+ DataContext.get_current().set_config(
+ "debug_limit_shuffle_execution_to_num_blocks", args.limit_num_blocks
+ )
 
  num_partitions = int(args.num_partitions)
  partition_size = int(float(args.partition_size))
@@ -130,14 +135,9 @@ def run_benchmark(args):
  )
 
  if args.shuffle:
- ds = ds.random_shuffle(
- _debug_limit_shuffle_execution_to_num_blocks=args.limit_num_blocks
- )
+ ds = ds.random_shuffle()
  else:
- ds = ds.sort(
- key="c_0",
- _debug_limit_shuffle_execution_to_num_blocks=args.limit_num_blocks,
- )
+ ds = ds.sort(key="c_0")
  exc = None
  try:
  ds = ds.materialize()