



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_	       =			datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
				import pyspark






@dataclass
class 						__a (     datasets.BuilderConfig				):
	SCREAMING_SNAKE_CASE       =       None






def 					__lowerCAmelCase		(  __lowerCamelCase	:    "pyspark.sql.DataFrame"				,   __lowerCamelCase	:    List[int]				,   )  ->						List[Any]:
					import pyspark

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

					return generate_fn






class 						__a (     _BaseExamplesIterable				):



	def __init__(						self							:    List[str]				,	snake_case_							:    "pyspark.sql.DataFrame"				,	snake_case_							:    Dict=None				,	)->							int:
						__lowerCAmelCase	   =df
						__lowerCAmelCase	   =partition_order or range(self.df.rdd.getNumPartitions())
						__lowerCAmelCase	   =_generate_iterable_examples(self.df				,	self.partition_order)



	def __iter__(						self							:    List[Any])->							List[str]:
						yield from self.generate_examples_fn()



	def 		UpperCamelCase							(						self							:    List[Any]				,	snake_case_							:    np.random.Generator)->							"SparkExamplesIterable":
						__lowerCAmelCase	   =list(range(self.df.rdd.getNumPartitions()))
						generator.shuffle(snake_case_)
						return SparkExamplesIterable(self.df				,	partition_order=snake_case_)



	def 		UpperCamelCase							(						self							:    Any				,	snake_case_							:    int				,	snake_case_							:    int)->							"SparkExamplesIterable":
						__lowerCAmelCase	   =self.split_shard_indices_by_worker(snake_case_				,	snake_case_)
						return SparkExamplesIterable(self.df				,	partition_order=snake_case_)



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






class 						__a (     datasets.DatasetBuilder				):
	SCREAMING_SNAKE_CASE       =       SparkConfig



	def __init__(						self							:    str				,	snake_case_							:    "pyspark.sql.DataFrame"				,	snake_case_							:    str = None				,	snake_case_							:    str = None				,	**snake_case_							:    Optional[Any]				,	)->							Optional[int]:
						import pyspark

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

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



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



	def 		UpperCamelCase							(						self							:    Dict				,	snake_case_							:    datasets.download.download_manager.DownloadManager)->							Tuple:
						return [datasets.SplitGenerator(name=datasets.Split.TRAIN)]



	def 		UpperCamelCase							(						self							:    int				,	snake_case_							:    List[Any])->							List[Any]:
						import pyspark

						def get_arrow_batch_size(snake_case_							:    List[str]):
											for batch in it:
																yield pa.RecordBatch.from_pydict({"""batch_bytes""": [batch.nbytes]})

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



	def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    str				,	snake_case_							:    str				,	snake_case_							:    int				,	)->							Iterable[Tuple[int, bool, Union[int, tuple]]]:
						import pyspark

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

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

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

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

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

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

						__lowerCAmelCase	   =0
						__lowerCAmelCase	   =0
						__lowerCAmelCase	   =0
						__lowerCAmelCase	   =[]
						__lowerCAmelCase	   =[]

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

						__lowerCAmelCase	   =total_num_examples
						__lowerCAmelCase	   =total_num_bytes

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

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

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



	def 		UpperCamelCase							(						self							:    str				,	snake_case_							:    "datasets.SplitGenerator"				,	)->							SparkExamplesIterable:
						return SparkExamplesIterable(self.df)









import os
from typing import Dict, List, Tuple, TypeVar, Union


lowercase_	       =			TypeVar('''T''')

lowercase_	       =			Union[List[T], Tuple[T, ...]]
lowercase_	       =			Union[T, List[T], Dict[str, T]]
lowercase_	       =			Union[str, bytes, os.PathLike]




