






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 							snake_case_       (    datasets.BuilderConfig      ):

						"""simple docstring"""



						snake_case__        =      None







def     a_   (__A				,       __A				,       )				->						List[str]:



					"""simple docstring"""
					import pyspark

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

					return generate_fn




class 							snake_case_       (    _BaseExamplesIterable      ):

						"""simple docstring"""



						def __init__(self:					Tuple ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					Dict=None ,					)    ->     List[str]:




											'''simple docstring'''



											__a					:						List[Any]      =					df
											__a					:						List[str]      =					partition_order or range(self.df.rdd.getNumPartitions()   )
											__a					:						Optional[int]      =					_generate_iterable_examples(self.df ,					self.partition_order   )



						def __iter__(self:					List[str]   )    ->     str:




											'''simple docstring'''



											yield from self.generate_examples_fn()



						def        UpperCAmelCase__					(self:					Optional[Any] ,					__UpperCAmelCase:					List[str]   )    ->     "SparkExamplesIterable":




											'''simple docstring'''



											__a					:						Tuple      =					list(range(self.df.rdd.getNumPartitions()   )   )
											generator.shuffle(__UpperCAmelCase   )
											return SparkExamplesIterable(self.df ,					partition_order=__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					str ,					__UpperCAmelCase:					Optional[Any] ,					__UpperCAmelCase:					Optional[Any]   )    ->     "SparkExamplesIterable":




											'''simple docstring'''



											__a					:						Tuple      =					self.split_shard_indices_by_worker(__UpperCAmelCase ,					__UpperCAmelCase   )
											return SparkExamplesIterable(self.df ,					partition_order=__UpperCAmelCase   )



						@property
						def        UpperCAmelCase__					(self:					str   )    ->     int:




											'''simple docstring'''



											return len(self.partition_order   )




class 							snake_case_       (    datasets.DatasetBuilder      ):

						"""simple docstring"""



						snake_case__        =      SparkConfig



						def __init__(self:					Union[str, Any] ,					__UpperCAmelCase:					Optional[Any] ,					__UpperCAmelCase:					str = None ,					__UpperCAmelCase:					int = None ,					**__UpperCAmelCase:					Optional[int] ,					)    ->     Optional[Any]:




											'''simple docstring'''



											import pyspark

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

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



						def        UpperCAmelCase__					(self:					Tuple   )    ->     Optional[Any]:




											'''simple docstring'''



											def create_cache_and_write_probe(__UpperCAmelCase:					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=__UpperCAmelCase   )
																__a					:						Optional[Any]      =					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(__UpperCAmelCase ,					"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:
																__a					:						Optional[int]      =					(
																    self._spark.sparkContext.parallelize(range(1   ) ,					1   ).mapPartitions(__UpperCAmelCase   ).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[int]   )    ->     List[Any]:




											'''simple docstring'''



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



						def        UpperCAmelCase__					(self:					List[Any] ,					__UpperCAmelCase:					Tuple   )    ->     str:




											'''simple docstring'''



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



						def        UpperCAmelCase__					(self:					int ,					__UpperCAmelCase:					str   )    ->     List[Any]:




											'''simple docstring'''



											import pyspark

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

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



						def        UpperCAmelCase__					(self:					List[str] ,					__UpperCAmelCase:					int ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					str ,					)    ->     Iterable[Tuple[int, bool, Union[int, tuple]]]:




											'''simple docstring'''



											import pyspark

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

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

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

											__a					:						Union[str, Any]      =					(
											    self.df.mapInArrow(__UpperCAmelCase ,					"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:					Optional[Any] ,					__UpperCAmelCase:					Any ,					__UpperCAmelCase:					Any = "arrow" ,					__UpperCAmelCase:					Any = None ,					__UpperCAmelCase:					Optional[int] = None ,					**__UpperCAmelCase:					Dict ,					)    ->     List[str]:




											'''simple docstring'''



											self._validate_cache_dir()

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

											__a					:						Any      =					'''-TTTTT-SSSSS-of-NNNNN'''
											__a					:						Any      =					f'{self.name}-{split_generator.name}{SUFFIX}.{file_format}'
											__a					:						List[Any]      =					path_join(self._output_dir ,					__UpperCAmelCase   )

											__a					:						Union[str, Any]      =					0
											__a					:						Any      =					0
											__a					:						Tuple      =					0
											__a					:						Union[str, Any]      =					[]
											__a					:						List[str]      =					[]

											for task_id, content in self._prepare_split_single(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase   ):
																(
																    __a		
																)					:						List[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(__UpperCAmelCase   )

											__a					:						List[str]      =					total_num_examples
											__a					:						List[str]      =					total_num_bytes

											# should rename everything at the end
											logger.debug(f'Renaming {total_shards} shards.'   )
											if total_shards > 1:
																__a					:						Tuple      =					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.
																__a					:						Union[str, Any]      =					self._fs

																# use the -SSSSS-of-NNNNN pattern
																def _rename_shard(
																    __UpperCAmelCase:					Optional[Any] ,					__UpperCAmelCase:					List[Any] ,					__UpperCAmelCase:					Optional[Any] ,					):
																					rename(
																					    __UpperCAmelCase ,					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}'   ) ,					)

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



						def        UpperCAmelCase__					(self:					Optional[Any] ,					__UpperCAmelCase:					Any ,					)    ->     SparkExamplesIterable:




											'''simple docstring'''



											return SparkExamplesIterable(self.df   )








def        snake_case__		(    lowerCamelCase_ = 1000		):
							return sum(e for e in range(3     ,				lowerCamelCase_		) if e % 3 == 0 or e % 5 == 0		)


if __name__ == "__main__":
		print(F"{solution() = }")






