



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     lowerCAmelCase_		(  datasets.BuilderConfig							):
       UpperCAmelCase__    : Optional[datasets.Features]       					=  None






def 		UpperCamelCase     (       snake_case__	:   List[Any]	,					snake_case__	:   str	,					)     ->					Union[str, Any]:
   import pyspark

   def generate_fn():
      UpperCamelCase   :  List[Any]						     =     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   :  List[Any]						     =     partition_df.collect()
         UpperCamelCase   :  List[Any]						     =     0
         for row in rows:
            yield F"""{partition_id}_{row_id}""", row.asDict()
            row_id += 1

   return generate_fn



class     lowerCAmelCase_		(  _BaseExamplesIterable							):
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_=None,      )					->		Any:
          UpperCamelCase   :  str						     =     df
          UpperCamelCase   :  Tuple						     =     partition_order or range(self.df.rdd.getNumPartitions()	)
          UpperCamelCase   :  Optional[int]						     =     _generate_iterable_examples(self.df,      self.partition_order	)
       def __iter__(      self	)					->		Optional[int]:
          yield from self.generate_examples_fn()
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		"SparkExamplesIterable":
          UpperCamelCase   :  Optional[int]						     =     list(range(self.df.rdd.getNumPartitions()	)	)
          generator.shuffle(SCREAMING_SNAKE_CASE_	)
          return SparkExamplesIterable(self.df,      partition_order=SCREAMING_SNAKE_CASE_	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		"SparkExamplesIterable":
          UpperCamelCase   :  List[str]						     =     self.split_shard_indices_by_worker(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          return SparkExamplesIterable(self.df,      partition_order=SCREAMING_SNAKE_CASE_	)






       @property
       def  snake_case_	(      self	)					->		int:
          return len(self.partition_order	)






class     lowerCAmelCase_		(  datasets.DatasetBuilder							):
       UpperCAmelCase__    : str       					=  SparkConfig
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = None,      **SCREAMING_SNAKE_CASE_,      )					->		str:
          import pyspark

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

          super().__init__(
              cache_dir=SCREAMING_SNAKE_CASE_,      config_name=str(self.df.semanticHash()	),      **SCREAMING_SNAKE_CASE_,      )
       def  snake_case_	(      self	)					->		Tuple:

          def create_cache_and_write_probe(SCREAMING_SNAKE_CASE_	):
             # 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=SCREAMING_SNAKE_CASE_	)
             UpperCamelCase   :  Tuple						     =     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(SCREAMING_SNAKE_CASE_,      '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(SCREAMING_SNAKE_CASE_	).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  snake_case_	(      self	)					->		Dict:
          return datasets.DatasetInfo(features=self.config.features	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Dict:
          return [datasets.SplitGenerator(name=datasets.Split.TRAIN	)]
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Optional[int]:
          import pyspark

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

          UpperCamelCase   :  Optional[Any]						     =     self.df.count()
          UpperCamelCase   :  int						     =     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   :  int						     =     (
              self.df.limit(SCREAMING_SNAKE_CASE_	)
              .repartition(1	)
              .mapInArrow(SCREAMING_SNAKE_CASE_,      'batch_bytes: long'	)
              .agg(pyspark.sql.functions.sum('batch_bytes'	).alias('sample_bytes'	)	)
              .collect()[0]
              .sample_bytes
              / sample_num_rows
          )
          UpperCamelCase   :  Dict						     =     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   :  List[Any]						     =     min(SCREAMING_SNAKE_CASE_,      int(approx_total_size / max_shard_size	)	)
             UpperCamelCase   :  Tuple						     =     self.df.repartition(SCREAMING_SNAKE_CASE_	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      )					->		Iterable[Tuple[int, bool, Union[int, tuple]]]:
          import pyspark

          UpperCamelCase   :  Optional[int]						     =     ParquetWriter if file_format == 'parquet' else ArrowWriter
          UpperCamelCase   :  Optional[int]						     =     os.path.join(self._working_dir,      os.path.basename(SCREAMING_SNAKE_CASE_	)	) if self._working_dir else fpath
          UpperCamelCase   :  str						     =     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   :  List[str]						     =     self.config.features
          UpperCamelCase   :  str						     =     self._writer_batch_size
          UpperCamelCase   :  Tuple						     =     self._fs.storage_options

          def write_arrow(SCREAMING_SNAKE_CASE_	):
             # Within the same SparkContext, no two task attempts will share the same attempt ID.
             UpperCamelCase   :  str						     =     pyspark.TaskContext().taskAttemptId()
             UpperCamelCase   :  Dict						     =     next(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
             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   :  Union[str, Any]						     =     0
             UpperCamelCase   :  Any						     =     writer_class(
                 features=SCREAMING_SNAKE_CASE_,      path=working_fpath.replace('SSSSS',      F"""{shard_id:05d}"""	).replace('TTTTT',      F"""{task_id:05d}"""	),      writer_batch_size=SCREAMING_SNAKE_CASE_,      storage_options=SCREAMING_SNAKE_CASE_,      embed_local_files=SCREAMING_SNAKE_CASE_,      )
             UpperCamelCase   :  Union[str, Any]						     =     pa.Table.from_batches([first_batch]	)
             writer.write_table(SCREAMING_SNAKE_CASE_	)
             for batch in it:
                if max_shard_size is not None and writer._num_bytes >= max_shard_size:
                   UpperCamelCase   :  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   :  Optional[int]						     =     writer_class(
                       features=writer._features,      path=working_fpath.replace('SSSSS',      F"""{shard_id:05d}"""	).replace('TTTTT',      F"""{task_id:05d}"""	),      writer_batch_size=SCREAMING_SNAKE_CASE_,      storage_options=SCREAMING_SNAKE_CASE_,      embed_local_files=SCREAMING_SNAKE_CASE_,      )
                UpperCamelCase   :  Any						     =     pa.Table.from_batches([batch]	)
                writer.write_table(SCREAMING_SNAKE_CASE_	)

             if writer._num_bytes > 0:
                UpperCamelCase   :  Optional[Any]						     =     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(SCREAMING_SNAKE_CASE_	)	):
                   UpperCamelCase   :  Any						     =     os.path.join(os.path.dirname(SCREAMING_SNAKE_CASE_	),      os.path.basename(SCREAMING_SNAKE_CASE_	)	)
                   shutil.move(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  Union[str, Any]						     =     (
              self.df.mapInArrow(SCREAMING_SNAKE_CASE_,      '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  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = "arrow",      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = None,      **SCREAMING_SNAKE_CASE_,      )					->		Any:
          self._validate_cache_dir()

          UpperCamelCase   :  List[Any]						     =     convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE	)
          self._repartition_df_if_needed(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     not is_remote_filesystem(self._fs	)
          UpperCamelCase   :  Any						     =     os.path.join if is_local else posixpath.join

          UpperCamelCase   :  List[str]						     =     '-TTTTT-SSSSS-of-NNNNN'
          UpperCamelCase   :  List[Any]						     =     F"""{self.name}-{split_generator.name}{SUFFIX}.{file_format}"""
          UpperCamelCase   :  Optional[int]						     =     path_join(self._output_dir,      SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  List[str]						     =     0
          UpperCamelCase   :  int						     =     0
          UpperCamelCase   :  Union[str, Any]						     =     0
          UpperCamelCase   :  Optional[int]						     =     []
          UpperCamelCase   :  Tuple						     =     []

          for task_id, content in self._prepare_split_single(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	):
             (
                 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(SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  Union[str, Any]						     =     total_num_examples
          UpperCamelCase   :  List[Any]						     =     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   :  Dict						     =     self._fs

             # use the -SSSSS-of-NNNNN pattern
             def _rename_shard(
                 SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      ):
                rename(
                    SCREAMING_SNAKE_CASE_,      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   :  Union[str, Any]						     =     []
             UpperCamelCase   :  int						     =     0
             for i in range(len(SCREAMING_SNAKE_CASE_	)	):
                UpperCamelCase   :  List[str]						     =     task_id_and_num_shards[i]
                for shard_id in range(SCREAMING_SNAKE_CASE_	):
                   args.append([task_id, shard_id, global_shard_id]	)
                   global_shard_id += 1
             self._spark.sparkContext.parallelize(SCREAMING_SNAKE_CASE_,      len(SCREAMING_SNAKE_CASE_	)	).map(lambda SCREAMING_SNAKE_CASE_	: _rename_shard(*SCREAMING_SNAKE_CASE_	)	).collect()
          else:
             # don't use any pattern
             UpperCamelCase   :  List[Any]						     =     0
             UpperCamelCase   :  int						     =     task_id_and_num_shards[0][0]
             self._rename(
                 fpath.replace('SSSSS',      F"""{shard_id:05d}"""	).replace('TTTTT',      F"""{task_id:05d}"""	),      fpath.replace(SCREAMING_SNAKE_CASE_,      ''	),      )






       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      )					->		SparkExamplesIterable:
          return SparkExamplesIterable(self.df	)









from abc import ABC, abstractmethod
from argparse import ArgumentParser



class     lowerCAmelCase_		(  a__							):
       @staticmethod
       @abstractmethod
       def  snake_case_	(      SCREAMING_SNAKE_CASE_	)					->		int:
          raise NotImplementedError()






       @abstractmethod
       def  snake_case_	(      self	)					->		str:
          raise NotImplementedError()



