import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

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


UpperCAmelCase_			    =							datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
  import pyspark




@dataclass
class  __magic_name__  (		datasets.BuilderConfig    ):





   """simple docstring"""




   lowerCAmelCase      :    Optional[datasets.Features]  			=   None




def       lowerCAmelCase_     (		lowercase:							"pyspark.sql.DataFrame"  ,					lowercase:							List[int]  ,					)    ->				Union[str, Any]:







 '''simple docstring'''
 import pyspark

 def generate_fn():
  _UpperCamelCase:     Dict							= df.select('''*'''  ,					pyspark.sql.functions.spark_partition_id().alias('''part_id'''		)		)
  for partition_id in partition_order:
   _UpperCamelCase:     Dict							= df_with_partition_id.select('''*'''		).where(F"""part_id = {partition_id}"""		).drop('''part_id'''		)
   _UpperCamelCase:     Tuple							= partition_df.collect()
   _UpperCamelCase:     Tuple							= 0
   for row in rows:
    yield F"""{partition_id}_{row_id}""", row.asDict()
    row_id += 1

 return generate_fn




class  __magic_name__  (		_BaseExamplesIterable    ):





   """simple docstring"""




   def __init__(	self					:  List[Any]		,  _lowercase					:  "pyspark.sql.DataFrame"		,  _lowercase					:  Optional[Any]=None		,  ):



    """simple docstring"""



    _UpperCamelCase:     Union[str, Any]							= df
    _UpperCamelCase:     Dict							= partition_order or range(self.df.rdd.getNumPartitions()							)
    _UpperCamelCase:     List[Any]							= _generate_iterable_examples(self.df		,  self.partition_order							)



   def __iter__(	self					:  List[str]							):



    """simple docstring"""



    yield from self.generate_examples_fn()



   def   lowerCAmelCase						(	self					:  Tuple		,  _lowercase					:  np.random.Generator							):



    """simple docstring"""



    _UpperCamelCase:     int							= list(range(self.df.rdd.getNumPartitions()							)							)
    generator.shuffle(_lowercase							)
    return SparkExamplesIterable(self.df		,  partition_order=_lowercase							)



   def   lowerCAmelCase						(	self					:  List[str]		,  _lowercase					:  int		,  _lowercase					:  int							):



    """simple docstring"""



    _UpperCamelCase:     str							= self.split_shard_indices_by_worker(_lowercase		,  _lowercase							)
    return SparkExamplesIterable(self.df		,  partition_order=_lowercase							)



   @property
   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    return len(self.partition_order							)




class  __magic_name__  (		datasets.DatasetBuilder    ):





   """simple docstring"""




   lowerCAmelCase      :    str  			=   SparkConfig



   def __init__(	self					:  Tuple		,  _lowercase					:  "pyspark.sql.DataFrame"		,  _lowercase					:  str = None		,  _lowercase					:  str = None		,  **_lowercase					:  Optional[int]		,  ):



    """simple docstring"""



    import pyspark

    _UpperCamelCase:     List[str]							= pyspark.sql.SparkSession.builder.getOrCreate()
    _UpperCamelCase:     str							= df
    _UpperCamelCase:     Optional[Any]							= working_dir

    super().__init__(
        cache_dir=_lowercase		,  config_name=str(self.df.semanticHash()							)		,  **_lowercase		,  )



   def   lowerCAmelCase						(	self					:  List[str]							):



    """simple docstring"""



    def create_cache_and_write_probe(_lowercase					:  Optional[int]							):
     # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
     # already exist.
     os.makedirs(self._cache_dir		,  exist_ok=_lowercase							)
     _UpperCamelCase:     List[str]							= os.path.join(self._cache_dir		,  '''fs_test''' + uuid.uuida().hex							)
     # Opening the file in append mode will create a new file unless it already exists, in which case it will not
     # change the file contents.
     open(_lowercase		,  '''a'''							)
     return [probe_file]

    if self._spark.conf.get('''spark.master'''		,  ''''''							).startswith('''local'''							):
     return

    # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
    # accessible to the driver.
    # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
    if self._cache_dir:
     _UpperCamelCase:     Optional[int]							= (
         self._spark.sparkContext.parallelize(range(1							)		,  1							).mapPartitions(_lowercase							).collect()
     )
     if os.path.isfile(probe[0]							):
      return

    raise ValueError(
        '''When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'''							)



   def   lowerCAmelCase						(	self					:  Union[str, Any]							):



    """simple docstring"""



    return datasets.DatasetInfo(features=self.config.features							)



   def   lowerCAmelCase						(	self					:  Dict		,  _lowercase					:  datasets.download.download_manager.DownloadManager							):



    """simple docstring"""



    return [datasets.SplitGenerator(name=datasets.Split.TRAIN							)]



   def   lowerCAmelCase						(	self					:  int		,  _lowercase					:  int							):



    """simple docstring"""



    import pyspark

    def get_arrow_batch_size(_lowercase					:  Tuple							):
     for batch in it:
      yield pa.RecordBatch.from_pydict({'''batch_bytes''': [batch.nbytes]}							)

    _UpperCamelCase:     Optional[Any]							= self.df.count()
    _UpperCamelCase:     Union[str, Any]							= 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.
    _UpperCamelCase:     Dict							= (
        self.df.limit(_lowercase							)
        .repartition(1							)
        .mapInArrow(_lowercase		,  '''batch_bytes: long'''							)
        .agg(pyspark.sql.functions.sum('''batch_bytes'''							).alias('''sample_bytes'''							)							)
        .collect()[0]
        .sample_bytes
        / sample_num_rows
    )
    _UpperCamelCase:     Any							= approx_bytes_per_row * df_num_rows
    if approx_total_size > max_shard_size:
     # Make sure there is at least one row per partition.
     _UpperCamelCase:     Tuple							= min(_lowercase		,  int(approx_total_size / max_shard_size							)							)
     _UpperCamelCase:     Optional[int]							= self.df.repartition(_lowercase							)



   def   lowerCAmelCase						(	self					:  Any		,  _lowercase					:  str		,  _lowercase					:  str		,  _lowercase					:  int		,  ):



    """simple docstring"""



    import pyspark

    _UpperCamelCase:     Union[str, Any]							= ParquetWriter if file_format == '''parquet''' else ArrowWriter
    _UpperCamelCase:     Union[str, Any]							= os.path.join(self._working_dir		,  os.path.basename(_lowercase							)							) if self._working_dir else fpath
    _UpperCamelCase:     Union[str, Any]							= file_format == '''parquet'''

    # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
    # pickling the SparkContext.
    _UpperCamelCase:     str							= self.config.features
    _UpperCamelCase:     Optional[int]							= self._writer_batch_size
    _UpperCamelCase:     List[str]							= self._fs.storage_options

    def write_arrow(_lowercase					:  Optional[int]							):
     # Within the same SparkContext, no two task attempts will share the same attempt ID.
     _UpperCamelCase:     Union[str, Any]							= pyspark.TaskContext().taskAttemptId()
     _UpperCamelCase:     Optional[Any]							= next(_lowercase		,  _lowercase							)
     if first_batch is None:
      # Some partitions might not receive any data.
      return pa.RecordBatch.from_arrays(
          [[task_id], [0], [0]]		,  names=['''task_id''', '''num_examples''', '''num_bytes''']		,  )
     _UpperCamelCase:     Dict							= 0
     _UpperCamelCase:     List[str]							= writer_class(
         features=_lowercase		,  path=working_fpath.replace('''SSSSS'''		,  f"""{shard_id:05d}"""							).replace('''TTTTT'''		,  f"""{task_id:05d}"""							)		,  writer_batch_size=_lowercase		,  storage_options=_lowercase		,  embed_local_files=_lowercase		,  )
     _UpperCamelCase:     List[str]							= pa.Table.from_batches([first_batch]							)
     writer.write_table(_lowercase							)
     for batch in it:
      if max_shard_size is not None and writer._num_bytes >= max_shard_size:
       _UpperCamelCase	,							_UpperCamelCase:     Optional[int]							= writer.finalize()
       writer.close()
       yield pa.RecordBatch.from_arrays(
           [[task_id], [num_examples], [num_bytes]]		,  names=['''task_id''', '''num_examples''', '''num_bytes''']		,  )
       shard_id += 1
       _UpperCamelCase:     List[str]							= writer_class(
           features=writer._features		,  path=working_fpath.replace('''SSSSS'''		,  f"""{shard_id:05d}"""							).replace('''TTTTT'''		,  f"""{task_id:05d}"""							)		,  writer_batch_size=_lowercase		,  storage_options=_lowercase		,  embed_local_files=_lowercase		,  )
      _UpperCamelCase:     List[str]							= pa.Table.from_batches([batch]							)
      writer.write_table(_lowercase							)

     if writer._num_bytes > 0:
      _UpperCamelCase	,							_UpperCamelCase:     Dict							= writer.finalize()
      writer.close()
      yield pa.RecordBatch.from_arrays(
          [[task_id], [num_examples], [num_bytes]]		,  names=['''task_id''', '''num_examples''', '''num_bytes''']		,  )

     if working_fpath != fpath:
      for file in os.listdir(os.path.dirname(_lowercase							)							):
       _UpperCamelCase:     int							= os.path.join(os.path.dirname(_lowercase							)		,  os.path.basename(_lowercase							)							)
       shutil.move(_lowercase		,  _lowercase							)

    _UpperCamelCase:     List[Any]							= (
        self.df.mapInArrow(_lowercase		,  '''task_id: long, num_examples: long, num_bytes: long'''							)
        .groupBy('''task_id'''							)
        .agg(
            pyspark.sql.functions.sum('''num_examples'''							).alias('''total_num_examples'''							)		,  pyspark.sql.functions.sum('''num_bytes'''							).alias('''total_num_bytes'''							)		,  pyspark.sql.functions.count('''num_bytes'''							).alias('''num_shards'''							)		,  pyspark.sql.functions.collect_list('''num_examples'''							).alias('''shard_lengths'''							)		,  )
        .collect()
    )
    for row in stats:
     yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)



   def   lowerCAmelCase						(	self					:  Any		,  _lowercase					:  "datasets.SplitGenerator"		,  _lowercase					:  str = "arrow"		,  _lowercase					:  Optional[Union[str, int]] = None		,  _lowercase					:  Optional[int] = None		,  **_lowercase					:  Tuple		,  ):



    """simple docstring"""



    self._validate_cache_dir()

    _UpperCamelCase:     Dict							= convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE							)
    self._repartition_df_if_needed(_lowercase							)
    _UpperCamelCase:     int							= not is_remote_filesystem(self._fs							)
    _UpperCamelCase:     Union[str, Any]							= os.path.join if is_local else posixpath.join

    _UpperCamelCase:     Tuple							= '''-TTTTT-SSSSS-of-NNNNN'''
    _UpperCamelCase:     List[Any]							= f"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
    _UpperCamelCase:     Union[str, Any]							= path_join(self._output_dir		,  _lowercase							)

    _UpperCamelCase:     Dict							= 0
    _UpperCamelCase:     Optional[int]							= 0
    _UpperCamelCase:     Optional[int]							= 0
    _UpperCamelCase:     int							= []
    _UpperCamelCase:     Union[str, Any]							= []

    for task_id, content in self._prepare_split_single(_lowercase		,  _lowercase		,  _lowercase							):
     (
         (
         _UpperCamelCase
     )	,							(
         _UpperCamelCase
     )	,							(
         _UpperCamelCase
     )	,							(
         _UpperCamelCase
     )	,							
     ):     Optional[Any]							= content
     if num_bytes > 0:
      total_num_examples += num_examples
      total_num_bytes += num_bytes
      total_shards += num_shards
      task_id_and_num_shards.append((task_id, num_shards)							)
      all_shard_lengths.extend(_lowercase							)

    _UpperCamelCase:     Dict							= total_num_examples
    _UpperCamelCase:     Tuple							= total_num_bytes

    # should rename everything at the end
    logger.debug(f"""Renaming {total_shards} shards."""							)
    if total_shards > 1:
     _UpperCamelCase:     List[Any]							= all_shard_lengths

     # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
     # pickling error due to pickling the SparkContext.
     _UpperCamelCase:     Tuple							= self._fs

     # use the -SSSSS-of-NNNNN pattern
     def _rename_shard(
         _lowercase					:  int		,  _lowercase					:  int		,  _lowercase					:  int		,  ):
      rename(
          _lowercase		,  fpath.replace('''SSSSS'''		,  f"""{shard_id:05d}"""							).replace('''TTTTT'''		,  f"""{task_id:05d}"""							)		,  fpath.replace('''TTTTT-SSSSS'''		,  f"""{global_shard_id:05d}"""							).replace('''NNNNN'''		,  f"""{total_shards:05d}"""							)		,  )

     _UpperCamelCase:     List[Any]							= []
     _UpperCamelCase:     List[Any]							= 0
     for i in range(len(_lowercase							)							):
      _UpperCamelCase	,							_UpperCamelCase:     Optional[Any]							= task_id_and_num_shards[i]
      for shard_id in range(_lowercase							):
       args.append([task_id, shard_id, global_shard_id]							)
       global_shard_id += 1
     self._spark.sparkContext.parallelize(_lowercase		,  len(_lowercase							)							).map(lambda _lowercase							: _rename_shard(*_lowercase							)							).collect()
    else:
     # don't use any pattern
     _UpperCamelCase:     List[str]							= 0
     _UpperCamelCase:     Any							= task_id_and_num_shards[0][0]
     self._rename(
         fpath.replace('''SSSSS'''		,  f"""{shard_id:05d}"""							).replace('''TTTTT'''		,  f"""{task_id:05d}"""							)		,  fpath.replace(_lowercase		,  ''''''							)		,  )



   def   lowerCAmelCase						(	self					:  Dict		,  _lowercase					:  "datasets.SplitGenerator"		,  ):



    """simple docstring"""



    return SparkExamplesIterable(self.df							)
import pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict




@pytest.mark.parametrize(
    '''split_dict'''  ,					[
        SplitDict(),
        SplitDict({'''train''': SplitInfo(name='''train'''  ,					num_bytes=1_337  ,					num_examples=42  ,					dataset_name='''my_dataset'''		)}		),
        SplitDict({'''train''': SplitInfo(name='''train'''  ,					num_bytes=1_337  ,					num_examples=42		)}		),
        SplitDict({'''train''': SplitInfo()}		),
    ]  ,					)
def       lowerCAmelCase_     (		lowercase:							SplitDict		)    ->				List[str]:







 '''simple docstring'''
 _UpperCamelCase:     List[str]							= split_dict._to_yaml_list()
 assert len(lowercase		) == len(lowercase		)
 _UpperCamelCase:     Tuple							= SplitDict._from_yaml_list(lowercase		)
 for split_name, split_info in split_dict.items():
  # dataset_name field is deprecated, and is therefore not part of the YAML dump
  _UpperCamelCase:     Optional[Any]							= None
  # the split name of split_dict takes over the name of the split info object
  _UpperCamelCase:     Any							= split_name
 assert split_dict == reloaded





@pytest.mark.parametrize(
    '''split_info'''  ,					[SplitInfo(), SplitInfo(dataset_name=lowercase		), SplitInfo(dataset_name='''my_dataset'''		)]		)
def       lowerCAmelCase_     (		lowercase:							str		)    ->				Tuple:







 '''simple docstring'''
 # For backward compatibility, we need asdict(split_dict) to return split info dictrionaries with the "dataset_name"
 # field even if it's deprecated. This way old versionso of `datasets` can still reload dataset_infos.json files
 _UpperCamelCase:     Optional[Any]							= asdict(SplitDict({'''train''': split_info}		)		)
 assert "dataset_name" in split_dict_asdict["train"]
 assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name