





"""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


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

if TYPE_CHECKING:
       import pyspark





@dataclass
class    lowerCamelCase__						(		datasets.BuilderConfig    ):
      a :     Optional[datasets.Features]											=		None






def lowerCAmelCase_      (				UpperCamelCase__	: "pyspark.sql.DataFrame"				,   UpperCamelCase__	: List[int]				,   ):






 """simple docstring"""







 import pyspark

 def generate_fn():
  __lowercase									=	df.select("""*"""				,   pyspark.sql.functions.spark_partition_id().alias("""part_id""" ) )
  for partition_id in partition_order:
   __lowercase									=	df_with_partition_id.select("""*""" ).where(f'''part_id = {partition_id}''' ).drop("""part_id""" )
   __lowercase									=	partition_df.collect()
   __lowercase									=	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	:	Dict      ,       A_	:	"pyspark.sql.DataFrame"      ,       A_	:	Union[str, Any]=None      ,       ):






       '''simple docstring'''


       __lowercase									=	df
       __lowercase									=	partition_order or range(self.df.rdd.getNumPartitions()							)
       __lowercase									=	_generate_iterable_examples(self.df      ,       self.partition_order							)
      def __iter__(		self	:	Optional[Any]							):






       '''simple docstring'''


       yield from self.generate_examples_fn()
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]      ,       A_	:	np.random.Generator							):






       '''simple docstring'''


       __lowercase									=	list(range(self.df.rdd.getNumPartitions()							)							)
       generator.shuffle(A_							)
       return SparkExamplesIterable(self.df      ,       partition_order=A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]      ,       A_	:	int      ,       A_	:	int							):






       '''simple docstring'''


       __lowercase									=	self.split_shard_indices_by_worker(A_      ,       A_							)
       return SparkExamplesIterable(self.df      ,       partition_order=A_							)






      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       return len(self.partition_order							)






class    lowerCamelCase__						(		datasets.DatasetBuilder    ):
      a :     List[Any]											=		SparkConfig
      def __init__(		self	:	Any      ,       A_	:	"pyspark.sql.DataFrame"      ,       A_	:	str = None      ,       A_	:	str = None      ,       **A_	:	Dict      ,       ):






       '''simple docstring'''


       import pyspark

       __lowercase									=	pyspark.sql.SparkSession.builder.getOrCreate()
       __lowercase									=	df
       __lowercase									=	working_dir

       super().__init__(
           cache_dir=A_      ,       config_name=str(self.df.semanticHash()							)      ,       **A_      ,       )
      def SCREAMING_SNAKE_CASE_	(		self	:	Union[str, Any]							):






       '''simple docstring'''


       def create_cache_and_write_probe(A_	:	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=A_							)
        __lowercase									=	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(A_      ,       """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:
        __lowercase									=	(
            self._spark.sparkContext.parallelize(range(1							)      ,       1							).mapPartitions(A_							).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	:	List[str]							):






       '''simple docstring'''


       return datasets.DatasetInfo(features=self.config.features							)
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict      ,       A_	:	datasets.download.download_manager.DownloadManager							):






       '''simple docstring'''


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






       '''simple docstring'''


       import pyspark

       def get_arrow_batch_size(A_	:	Optional[Any]							):
        for batch in it:
         yield pa.RecordBatch.from_pydict({"""batch_bytes""": [batch.nbytes]}							)

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






       '''simple docstring'''


       import pyspark

       __lowercase									=	ParquetWriter if file_format == """parquet""" else ArrowWriter
       __lowercase									=	os.path.join(self._working_dir      ,       os.path.basename(A_							)							) if self._working_dir else fpath
       __lowercase									=	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.
       __lowercase									=	self.config.features
       __lowercase									=	self._writer_batch_size
       __lowercase									=	self._fs.storage_options

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

        if writer._num_bytes > 0:
         __lowercase							,		__lowercase									=	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(A_							)							):
          __lowercase									=	os.path.join(os.path.dirname(A_							)      ,       os.path.basename(A_							)							)
          shutil.move(A_      ,       A_							)

       __lowercase									=	(
           self.df.mapInArrow(A_      ,       """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	:	Union[str, Any]      ,       A_	:	"datasets.SplitGenerator"      ,       A_	:	str = "arrow"      ,       A_	:	Optional[Union[str, int]] = None      ,       A_	:	Optional[int] = None      ,       **A_	:	Union[str, Any]      ,       ):






       '''simple docstring'''


       self._validate_cache_dir()

       __lowercase									=	convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE							)
       self._repartition_df_if_needed(A_							)
       __lowercase									=	not is_remote_filesystem(self._fs							)
       __lowercase									=	os.path.join if is_local else posixpath.join

       __lowercase									=	"""-TTTTT-SSSSS-of-NNNNN"""
       __lowercase									=	F'''{self.name}-{split_generator.name}{SUFFIX}.{file_format}'''
       __lowercase									=	path_join(self._output_dir      ,       A_							)

       __lowercase									=	0
       __lowercase									=	0
       __lowercase									=	0
       __lowercase									=	[]
       __lowercase									=	[]

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

       __lowercase									=	total_num_examples
       __lowercase									=	total_num_bytes

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

        # use the -SSSSS-of-NNNNN pattern
        def _rename_shard(
            A_	:	int      ,       A_	:	int      ,       A_	:	int      ,       ):
         rename(
             A_      ,       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}'''							)      ,       )

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






      def SCREAMING_SNAKE_CASE_	(		self	:	int      ,       A_	:	"datasets.SplitGenerator"      ,       ):






       '''simple docstring'''


       return SparkExamplesIterable(self.df							)




"""simple docstring"""

import random


def lowerCAmelCase_      (				UpperCamelCase__	: int				,   UpperCamelCase__	: float				,   UpperCamelCase__	: bool = False ):






 """simple docstring"""







 __lowercase									=	{i: [] for i in range(UpperCamelCase__ )}

 # if probability is greater or equal than 1, then generate a complete graph
 if probability >= 1:
  return complete_graph(UpperCamelCase__ )
 # if probability is lower or equal than 0, then return a graph without edges
 if probability <= 0:
  return graph

 # for each couple of nodes, add an edge from u to v
 # if the number randomly generated is greater than probability probability
 for i in range(UpperCamelCase__ ):
  for j in range(i + 1				,   UpperCamelCase__ ):
   if random.random() < probability:
    graph[i].append(UpperCamelCase__ )
    if not directed:
     # if the graph is undirected, add an edge in from j to i, either
     graph[j].append(UpperCamelCase__ )
 return graph


def lowerCAmelCase_      (				UpperCamelCase__	: int ):






 """simple docstring"""







 return {
     i: [j for j in range(UpperCamelCase__ ) if i != j] for i in range(UpperCamelCase__ )
 }


if __name__ == "__main__":
       import doctest

       doctest.testmod()
