

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

if TYPE_CHECKING:
 import pyspark




@dataclass
class 	A_			(    datasets.BuilderConfig							):






     '''simple docstring'''



     UpperCAmelCase_						:     Optional[datasets.Features]  		=  None




def 				UpperCamelCase(					UpperCAmelCase_  ,	UpperCAmelCase_  ,	):
 import pyspark

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

 return generate_fn




class 	A_			(    _BaseExamplesIterable							):






     '''simple docstring'''



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


     def __iter__(		self			:					Optional[int]	)						->     Tuple:
      yield from self.generate_examples_fn()


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


     def        UpperCAmelCase_	(		self			:					Union[str, Any]					,      lowercase_			:					int					,      lowercase_			:					int	)						->     "SparkExamplesIterable":
      UpperCAmelCase	:     Union[str, Any]						 =			self.split_shard_indices_by_worker(lowercase_					,      lowercase_	)
      return SparkExamplesIterable(self.df					,      partition_order=lowercase_	)




     @property
     def        UpperCAmelCase_	(		self			:					List[str]	)						->     int:
      return len(self.partition_order	)




class 	A_			(    datasets.DatasetBuilder							):






     '''simple docstring'''



     UpperCAmelCase_						:     Union[str, Any]  		=  SparkConfig


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

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

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


     def        UpperCAmelCase_	(		self			:					Union[str, Any]	)						->     Dict:
      # Returns the path of the created file.
      def create_cache_and_write_probe(lowercase_			:					Tuple	):
       # 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=lowercase_	)
       UpperCAmelCase	:     Union[str, 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(lowercase_					,      '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	:     Any						 =			(
           self._spark.sparkContext.parallelize(range(1	)					,      1	).mapPartitions(lowercase_	).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			:					List[str]	)						->     Optional[Any]:
      return datasets.DatasetInfo(features=self.config.features	)


     def        UpperCAmelCase_	(		self			:					Optional[Any]					,      lowercase_			:					datasets.download.download_manager.DownloadManager	)						->     int:
      return [datasets.SplitGenerator(name=datasets.Split.TRAIN	)]


     def        UpperCAmelCase_	(		self			:					int					,      lowercase_			:					Any	)						->     Optional[int]:
      import pyspark

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

      UpperCAmelCase	:     int						 =			self.df.count()
      UpperCAmelCase	:     Optional[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	:     Union[str, Any]						 =			(
          self.df.limit(lowercase_	)
          .repartition(1	)
          .mapInArrow(lowercase_					,      'batch_bytes: long'	)
          .agg(pyspark.sql.functions.sum('batch_bytes'	).alias('sample_bytes'	)	)
          .collect()[0]
          .sample_bytes
          / sample_num_rows
      )
      UpperCAmelCase	:     int						 =			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	:     str						 =			min(lowercase_					,      int(approx_total_size / max_shard_size	)	)
       UpperCAmelCase	:     Dict						 =			self.df.repartition(lowercase_	)


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

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

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

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

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

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

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

      UpperCAmelCase	:     Optional[int]						 =			0
      UpperCAmelCase	:     Union[str, Any]						 =			0
      UpperCAmelCase	:     int						 =			0
      UpperCAmelCase	:     Union[str, Any]						 =			[]
      UpperCAmelCase	:     Union[str, Any]						 =			[]

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

      UpperCAmelCase	:     Union[str, Any]						 =			total_num_examples
      UpperCAmelCase	:     Any						 =			total_num_bytes

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

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




     def        UpperCAmelCase_	(		self			:					Tuple					,      lowercase_			:					"datasets.SplitGenerator"					,      )						->     SparkExamplesIterable:
      return SparkExamplesIterable(self.df	)







'''simple docstring'''

from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


lowercase__         =      {"configuration_van": ["VAN_PRETRAINED_CONFIG_ARCHIVE_MAP", "VanConfig"]}


try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowercase__         =      [
     "VAN_PRETRAINED_MODEL_ARCHIVE_LIST",
     "VanForImageClassification",
     "VanModel",
     "VanPreTrainedModel",
 ]

if TYPE_CHECKING:
 from .configuration_van import VAN_PRETRAINED_CONFIG_ARCHIVE_MAP, VanConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_van import (
      VAN_PRETRAINED_MODEL_ARCHIVE_LIST,
      VanForImageClassification,
      VanModel,
      VanPreTrainedModel,
  )

else:
 import sys

 lowercase__         =      _LazyModule(__name__, globals()["__file__"], _import_structure)




