code: def load_dataset(client, data_dir, s3_bucket, nbytes, npartitions):
      num_bytes_per_partition = nbytes // npartitions
      filenames = []
  
      @dask.delayed
      def generate_s3_file(i, data_dir, s3_bucket):
          s3 = boto3.client("s3")
          key = "df-{}-{}.parquet.gzip".format(
              num_bytes_per_partition, i
          )
          contents = s3.list_objects(Bucket=s3_bucket, Prefix=key)
          for obj in contents.get("Contents", []):
              if obj["Key"] == key:
                  print(f"S3 partition {i} exists")
                  return
  
          filename = os.path.join(data_dir, key)
          if not os.path.exists(filename):
              print("Generating partition", filename)
              nrows = num_bytes_per_partition // 8
              dataset = pd.DataFrame(
                  np.random.randint(
                      0,
                      np.iinfo(np.int64).max,
                      size=(nrows, 1),
                      dtype=np.int64,
                  ),
                  columns=["a"],
              )
              dataset.to_parquet(filename, compression="gzip")
          print("Writing partition to S3", filename)
          with open(filename, "rb") as f:
              s3.put_object(Bucket=s3_bucket, Key=key, Body=f)
  
      x = []
      for i in range(npartitions):
          x.append(generate_s3_file(i, data_dir, s3_bucket))
      # from ray.util.dask import ProgressBarCallback
      # with ProgressBarCallback():
      # dask.compute(x, _ray_enable_progress_bar=True)
      dask.compute(x)
  
      filenames = [
          f"s3://{s3_bucket}/df-{num_bytes_per_partition}-{i}.parquet.gzip"
          for i in range(npartitions)
      ]
  
      df = None
      max_retry = 3
      retry = 0
      while retry < max_retry:
          try:
              df = dd.read_parquet(filenames)
              break
          except FileNotFoundError as e:
              print(traceback.format_exc())
              print(f"Failed to load a file. {e}")
              # Wait a little bit before retrying.
              time.sleep(30)
              retry += 1
  
      return df
  
file_path: ray/release/nightly_tests/dask_on_ray/dask_on_ray_sort.py
imports: import csv
  import json
  import os.path
  import time
  import boto3
  import dask
  import dask.dataframe
  import numpy
  import pandas
  from dask.distributed import Client
  import traceback
label_desc: Use image node_id 1 for calling a global function with example usage: load_dataset(client, data_dir, s3_bucket, nbytes, npartitions) and returns: df