code: def _repartition_df_if_needed(self, max_shard_size):
      import pyspark
  
      def get_arrow_batch_size(it):
          for batch in it:
              yield pa.RecordBatch.from_pydict(
                  {"batch_bytes": [batch.nbytes]}
              )
  
      df_num_rows = self.df.count()
      sample_num_rows = df_num_rows if df_num_rows <= 100 else 100
      # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
      approx_bytes_per_row = (
          self.df.limit(sample_num_rows)
          .repartition(1)
          .mapInArrow(get_arrow_batch_size, "batch_bytes: long")
          .agg(
              pyspark.sql.functions.sum("batch_bytes").alias(
                  "sample_bytes"
              )
          )
          .collect()[0]
          .sample_bytes
          / sample_num_rows
      )
      approx_total_size = approx_bytes_per_row * df_num_rows
      if approx_total_size > max_shard_size:
          # Make sure there is at least one row per partition.
          new_num_partitions = min(
              df_num_rows, int(approx_total_size / max_shard_size)
          )
          self.df = self.df.repartition(new_num_partitions)
  
imports: import os
  import posixpath
  import uuid
  from dataclasses import dataclass
  from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union
  import numpy
  import pyarrow
  import datasets
  from datasets.arrow_writer import ArrowWriter, ParquetWriter
  from datasets.config import MAX_SHARD_SIZE
  from datasets.filesystems import is_remote_filesystem, rename
  from datasets.iterable_dataset import _BaseExamplesIterable
  from datasets.utils.py_utils import convert_file_size_to_int
label_desc: Use image node_id 5 for calling the Spark obj's underlying member method code with example usage: obj._repartition_df_if_needed(max_shard_size) without return types