hunk.code_window: using rpc::ActorTableData;
   using rpc::ErrorTableData;
   using rpc::GcsNodeInfo;
   using rpc::HeartbeatTableData;
   using rpc::JobTableData;
   using rpc::ObjectLocationInfo;
   using rpc::ObjectTableData;
   using rpc::PlacementGroupTableData;
  
hunk.labels: keep keep keep replace keep keep keep keep
hunk.file_path: src/ray/gcs/gcs_server/gcs_table_storage.h
hunk.type: replace
file: import dask
  import dask.dataframe as dd
  import json
  import pandas as pd
  import numpy as np
  import os.path
  import csv
  import boto3
  
  from dask.distributed import Client
  
  import time
  
  
  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))
      dask.compute(x)
  
      filenames = [
          f"s3://{s3_bucket}/df-{num_bytes_per_partition}-{i}.parquet.gzip"
          for i in range(npartitions)
      ]
      df = dd.read_parquet(filenames)
      return df
  
  
  def load_dataset_files(client, data_dir, file_path, nbytes, npartitions):
      num_bytes_per_partition = nbytes // npartitions
      filenames = []
  
      @dask.delayed
      def generate_file(i, data_dir, file_path):
          key = "{}/df-{}-{}.parquet.gzip".format(file_path,
                                                  num_bytes_per_partition, i)
          from os import path
          if path.exists(key):
              print(f"The file {key} already exists. Do nothing")
              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 a file", filename)
  
      x = []
      for i in range(npartitions):
          x.append(generate_file(i, data_dir, file_path))
      dask.compute(x)
  
      filenames = [
          f"{file_path}/df-{num_bytes_per_partition}-{i}.parquet.gzip"
          for i in range(npartitions)
      ]
      df = dd.read_parquet(filenames)
      return df
  
  
  def trial(client,
            data_dir,
            nbytes,
            n_partitions,
            generate_only,
            s3_bucket=None,
            file_path=None):
      if s3_bucket:
          df = load_dataset(client, data_dir, s3_bucket, nbytes, n_partitions)
      elif file_path:
          df = load_dataset_files(client, data_dir, file_path, nbytes,
                                  n_partitions)
  
      if generate_only:
          return []
  
      times = []
      start = time.time()
      for i in range(10):
          print("Trial {} start".format(i))
          trial_start = time.time()
  
          print(
              df.set_index("a", shuffle="tasks", max_branch=float("inf")).head(
                  10, npartitions=-1))
  
          trial_end = time.time()
          duration = trial_end - trial_start
          times.append(duration)
          print("Trial {} done after {}".format(i, duration))
  
          if time.time() - start > 60:
              break
      return times
  
  
  if __name__ == "__main__":
      import argparse
      parser = argparse.ArgumentParser()
  
      parser.add_argument("--nbytes", type=int, default=1_000_000)
      parser.add_argument("--npartitions", type=int, default=100, required=False)
      # Max partition size is 1GB.
      parser.add_argument(
          "--max-partition-size", type=int, default=1000_000_000, required=False)
      parser.add_argument("--num-nodes", type=int, default=1)
      parser.add_argument("--dask-tasks", action="store_true")
      parser.add_argument("--generate-only", action="store_true")
      parser.add_argument("--ray", action="store_true")
      parser.add_argument("--data-dir", default="/home/ubuntu/dask-benchmarks")
      parser.add_argument("--s3-bucket")
      parser.add_argument("--file-path")
  
      parser.add_argument("--dask-nprocs", type=int, default=0)
      parser.add_argument("--dask-nthreads", type=int, default=0)
      parser.add_argument("--dask-memlimit", type=int, default=0)
      args = parser.parse_args()
      assert not (args.s3_bucket
                  and args.file_path), "Provide S3 bucket or file path."
      if args.ray:
          import ray
          ray.init(address="auto")
          from ray.util.dask import ray_dask_get, dataframe_optimize
          dask.config.set(
              scheduler=ray_dask_get, dataframe_optimize=dataframe_optimize)
          client = None
      else:
          assert args.dask_nprocs != -0
          assert args.dask_nthreads != -0
          assert args.dask_memlimit != -0
  
          if args.dask_tasks:
              print("Using task-based Dask shuffle")
              dask.config.set(shuffle="tasks")
          else:
              print("Using disk-based Dask shuffle")
  
          client = Client("localhost:8786")
  
      print(
          trial(
              client,
              args.data_dir,
              1000,
              10,
              args.generate_only,
              s3_bucket=args.s3_bucket,
              file_path=args.file_path))
      print("WARMUP DONE")
  
      npartitions = args.npartitions
      if args.nbytes // npartitions > args.max_partition_size:
          npartitions = args.nbytes // args.max_partition_size
  
      duration = []
      output = trial(
          client,
          args.data_dir,
          args.nbytes,
          npartitions,
          args.generate_only,
          s3_bucket=args.s3_bucket,
          file_path=args.file_path)
      print("mean over {} trials: {} +- {}".format(
          len(output), np.mean(output), np.std(output)))
  
      print(ray.internal.internal_api.memory_summary(stats_only=True))
      duration = np.mean(output)
  
      with open(os.environ["TEST_OUTPUT_JSON"], "w") as f:
          f.write(json.dumps({"duration": duration, "success": 1}))
  
      write_header = not os.path.exists("output.csv") or os.path.getsize(
          "output.csv") == 0
      with open("output.csv", "a+") as csvfile:
          fieldnames = [
              "num_nodes", "nbytes", "npartitions", "dask_tasks", "dask_nprocs",
              "dask_nthreads", "dask_memlimit", "duration"
          ]
          writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
          if write_header:
              writer.writeheader()
          row = {
              "num_nodes": args.num_nodes,
              "nbytes": args.nbytes,
              "npartitions": npartitions,
              "dask_tasks": args.dask_tasks,
              "dask_nprocs": args.dask_nprocs,
              "dask_nthreads": args.dask_nthreads,
              "dask_memlimit": args.dask_memlimit,
          }
          for output in output:
              row["duration"] = output
              writer.writerow(row)
  