code: def _rechunk_if_needed(df, num_shards: int = None):
      try:
          if num_shards:
              assert isinstance(num_shards, int) and num_shards > 0
              df = df.rebalance(axis=0, num_partitions=num_shards)
          df = df.rechunk({1: df.shape[1]})
          df = df.reset_index(drop=True)
          return df.execute()
      except Exception as e:  # pragma: no cover
          raise Exception(f"rechunk failed df.shape {df.shape}") from e
  
file_path: mars/mars/dataframe/contrib/raydataset/mldataset.py
imports: import numpy
  import pandas
  from collections import defaultdict
  from typing import Dict, Iterable, List, Tuple
  from ....utils import lazy_import
label_desc: Use image node_id 2 for calling a global function with example usage: _rechunk_if_needed(df, num_shards) and returns: df