



"""simple docstring"""







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


lowercase__		:    int				   =     datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
     import pyspark







@dataclass
class 				UpperCamelCase__       (     datasets.BuilderConfig  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					None




def     UpperCamelCase_   (   lowerCAmelCase__				:    "pyspark.sql.DataFrame"   ,       lowerCAmelCase__				:    List[int]   ,       ):

  """simple docstring"""




  import pyspark

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

  return generate_fn







class 				UpperCamelCase__       (     _BaseExamplesIterable  ):




      """simple docstring"""




      def __init__(				self    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=None   ,      ):
        lowerCAmelCase_       :    Union[str, Any]								=			df
        lowerCAmelCase_       :    Tuple								=			partition_order or range(self.df.rdd.getNumPartitions()	)
        lowerCAmelCase_       :    str								=			_generate_iterable_examples(self.df   ,      self.partition_order	)



      def __iter__(				self    :				Dict	):
        yield from self.generate_examples_fn()



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				int	):
        lowerCAmelCase_       :    int								=			list(range(self.df.rdd.getNumPartitions()	)	)
        generator.shuffle(__lowerCAmelCase	)
        return SparkExamplesIterable(self.df   ,      partition_order=__lowerCAmelCase	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Any   ,      SCREAMING_SNAKE_CASE_    :				List[Any]	):
        lowerCAmelCase_       :    Optional[Any]								=			self.split_shard_indices_by_worker(__lowerCAmelCase   ,      __lowerCAmelCase	)
        return SparkExamplesIterable(self.df   ,      partition_order=__lowerCAmelCase	)





      @property
      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        return len(self.partition_order	)







class 				UpperCamelCase__       (     datasets.DatasetBuilder  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					SparkConfig



      def __init__(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				int = None   ,      SCREAMING_SNAKE_CASE_    :				Tuple = None   ,      **SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      ):
        import pyspark

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

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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        # Returns the path of the created file.
        def create_cache_and_write_probe(SCREAMING_SNAKE_CASE_    :				Union[str, Any]	):
          # 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=__lowerCAmelCase	)
          lowerCAmelCase_       :    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(__lowerCAmelCase   ,      '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:
          lowerCAmelCase_       :    int								=			(
              self._spark.sparkContext.parallelize(range(1	)   ,      1	).mapPartitions(__lowerCAmelCase	).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 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        return datasets.DatasetInfo(features=self.config.features	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Tuple	):
        return [datasets.SplitGenerator(name=datasets.Split.TRAIN	)]



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				List[Any]	):
        import pyspark

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

        lowerCAmelCase_       :    Optional[Any]								=			self.df.count()
        lowerCAmelCase_       :    Union[str, Any]								=			df_num_rows if df_num_rows <= 1_0_0 else 1_0_0
        # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
        lowerCAmelCase_       :    Optional[int]								=			(
            self.df.limit(__lowerCAmelCase	)
            .repartition(1	)
            .mapInArrow(__lowerCAmelCase   ,      'batch_bytes: long'	)
            .agg(pyspark.sql.functions.sum('batch_bytes'	).alias('sample_bytes'	)	)
            .collect()[0]
            .sample_bytes
            / sample_num_rows
        )
        lowerCAmelCase_       :    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.
          lowerCAmelCase_       :    List[str]								=			min(__lowerCAmelCase   ,      int(approx_total_size / max_shard_size	)	)
          lowerCAmelCase_       :    List[Any]								=			self.df.repartition(__lowerCAmelCase	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      ):
        import pyspark

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

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

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

        lowerCAmelCase_       :    int								=			(
            self.df.mapInArrow(__lowerCAmelCase   ,      '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 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]   ,      SCREAMING_SNAKE_CASE_    :				Dict = "arrow"   ,      SCREAMING_SNAKE_CASE_    :				List[Any] = None   ,      SCREAMING_SNAKE_CASE_    :				str = None   ,      **SCREAMING_SNAKE_CASE_    :				Tuple   ,      ):
        self._validate_cache_dir()

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

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

        lowerCAmelCase_       :    Optional[int]								=			0
        lowerCAmelCase_       :    Optional[int]								=			0
        lowerCAmelCase_       :    int								=			0
        lowerCAmelCase_       :    List[Any]								=			[]
        lowerCAmelCase_       :    List[str]								=			[]

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

        lowerCAmelCase_       :    Optional[int]								=			total_num_examples
        lowerCAmelCase_       :    Any								=			total_num_bytes

        # should rename everything at the end
        logger.debug(F"Renaming {total_shards} shards."	)
        if total_shards > 1:
          lowerCAmelCase_       :    Dict								=			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.
          lowerCAmelCase_       :    str								=			self._fs

          # use the -SSSSS-of-NNNNN pattern
          def _rename_shard(
              SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]   ,      ):
            rename(
                __lowerCAmelCase   ,      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}"	)   ,      )

          lowerCAmelCase_       :    List[str]								=			[]
          lowerCAmelCase_       :    int								=			0
          for i in range(len(__lowerCAmelCase	)	):
            lowerCAmelCase_						,lowerCAmelCase_       :    Union[str, Any]								=			task_id_and_num_shards[i]
            for shard_id in range(__lowerCAmelCase	):
              args.append([task_id, shard_id, global_shard_id]	)
              global_shard_id += 1
          self._spark.sparkContext.parallelize(__lowerCAmelCase   ,      len(__lowerCAmelCase	)	).map(lambda SCREAMING_SNAKE_CASE_	: _rename_shard(*__lowerCAmelCase	)	).collect()
        else:
          # don't use any pattern
          lowerCAmelCase_       :    Union[str, Any]								=			0
          lowerCAmelCase_       :    Optional[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(__lowerCAmelCase   ,      ''	)   ,      )





      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				str   ,      ):
        return SparkExamplesIterable(self.df	)










"""simple docstring"""







from collections.abc import Iterable
from typing import Generic, TypeVar

lowercase__		:    Any				   =     TypeVar("""_T""")







class 				UpperCamelCase__       (     Generic[_T]  ):




      """simple docstring"""




      def __init__(				self    :				Any   ,      SCREAMING_SNAKE_CASE_    :				Iterable[_T] | None = None	):
        lowerCAmelCase_       :    list[_T]								=			list(iterable or []	)
        lowerCAmelCase_       :    list[_T]								=			[]



      def __len__(				self    :				List[str]	):
        return len(self._stacka	) + len(self._stacka	)



      def __repr__(				self    :				int	):
        return F"Queue({tuple(self._stacka[::-1] + self._stacka	)})"



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple   ,      SCREAMING_SNAKE_CASE_    :				_T	):
        self._stacka.append(SCREAMING_SNAKE_CASE_	)





      def 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        lowerCAmelCase_       :    List[str]								=			self._stacka.pop
        lowerCAmelCase_       :    int								=			self._stacka.append

        if not self._stacka:
          while self._stacka:
            stacka_append(stacka_pop()	)

        if not self._stacka:
          raise IndexError('Queue is empty'	)
        return self._stacka.pop()


if __name__ == "__main__":
     from doctest import testmod

     testmod()





