






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


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

if TYPE_CHECKING:
     import pyspark






@dataclass
class     a_						(      datasets.BuilderConfig    ):
  lowercase							=				None







def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						)->			List[Any]:
      import pyspark

      def generate_fn():
            UpperCamelCase      = df.select("""*"""		,						pyspark.sql.functions.spark_partition_id().alias("""part_id"""   )   )
            for partition_id in partition_order:
                  UpperCamelCase      = df_with_partition_id.select("""*"""   ).where(F"part_id = {partition_id}"   ).drop("""part_id"""   )
                  UpperCamelCase      = partition_df.collect()
                  UpperCamelCase      = 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      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE=None      ,		)  ->				List[str]:



        """simple docstring"""






        UpperCamelCase      = df
        UpperCamelCase      = partition_order or range(self.df.rdd.getNumPartitions()				)
        UpperCamelCase      = _generate_iterable_examples(self.df      ,		self.partition_order				)

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



        """simple docstring"""






        yield from self.generate_examples_fn()

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				"SparkExamplesIterable":



        """simple docstring"""






        UpperCamelCase      = list(range(self.df.rdd.getNumPartitions()				)				)
        generator.shuffle(_SCREAMING_SNAKE_CASE				)
        return SparkExamplesIterable(self.df      ,		partition_order=_SCREAMING_SNAKE_CASE				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)  ->				"SparkExamplesIterable":



        """simple docstring"""






        UpperCamelCase      = self.split_shard_indices_by_worker(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
        return SparkExamplesIterable(self.df      ,		partition_order=_SCREAMING_SNAKE_CASE				)


  @property
  def   A__						(				self				)  ->				int:



        """simple docstring"""






        return len(self.partition_order				)






class     a_						(      datasets.DatasetBuilder    ):
  lowercase							=				SparkConfig

  def __init__(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				Optional[Any]:



        """simple docstring"""






        import pyspark

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

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

  def   A__						(				self				)  ->				Optional[int]:



        """simple docstring"""






        def create_cache_and_write_probe(_SCREAMING_SNAKE_CASE				):
              # 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=_SCREAMING_SNAKE_CASE				)
              UpperCamelCase      = 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(_SCREAMING_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:
              UpperCamelCase      = (
                  self._spark.sparkContext.parallelize(range(1				)      ,		1				).mapPartitions(_SCREAMING_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   A__						(				self				)  ->				str:



        """simple docstring"""






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

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				Union[str, Any]:



        """simple docstring"""






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

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE				)  ->				List[Any]:



        """simple docstring"""






        import pyspark

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

        UpperCamelCase      = self.df.count()
        UpperCamelCase      = 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      = (
            self.df.limit(_SCREAMING_SNAKE_CASE				)
            .repartition(1				)
            .mapInArrow(_SCREAMING_SNAKE_CASE      ,		"""batch_bytes: long"""				)
            .agg(pyspark.sql.functions.sum("""batch_bytes"""				).alias("""sample_bytes"""				)				)
            .collect()[0]
            .sample_bytes
            / sample_num_rows
        )
        UpperCamelCase      = 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      = min(_SCREAMING_SNAKE_CASE      ,		int(approx_total_size / max_shard_size				)				)
              UpperCamelCase      = self.df.repartition(_SCREAMING_SNAKE_CASE				)

  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		)  ->				Iterable[Tuple[int, bool, Union[int, tuple]]]:



        """simple docstring"""






        import pyspark

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

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

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

        UpperCamelCase      = (
            self.df.mapInArrow(_SCREAMING_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   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE = "arrow"      ,		_SCREAMING_SNAKE_CASE = None      ,		_SCREAMING_SNAKE_CASE = None      ,		**_SCREAMING_SNAKE_CASE      ,		)  ->				List[str]:



        """simple docstring"""






        self._validate_cache_dir()

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

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

        UpperCamelCase      = 0
        UpperCamelCase      = 0
        UpperCamelCase      = 0
        UpperCamelCase      = []
        UpperCamelCase      = []

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

        UpperCamelCase      = total_num_examples
        UpperCamelCase      = total_num_bytes

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

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

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


  def   A__						(				self      ,		_SCREAMING_SNAKE_CASE      ,		)  ->				SparkExamplesIterable:



        """simple docstring"""






        return SparkExamplesIterable(self.df				)








'''simple docstring'''

import time
from dataclasses import dataclass
from multiprocessing import Pool
from unittest import TestCase
from unittest.mock import patch

import multiprocess
import numpy as np
import pytest

from datasets.utils.py_utils import (
    NestedDataStructure,
    asdict,
    iflatmap_unordered,
    map_nested,
    temp_seed,
    temporary_assignment,
    zip_dict,
)

from .utils import require_tf, require_torch







def    lowercase__			(				__UpperCamelCase   )->			Dict:  # picklable for multiprocessing
      return x.sum()







def    lowercase__			(				__UpperCamelCase   )->			Tuple:  # picklable for multiprocessing
      return i + 1






@dataclass
class     a_						:
  lowercase							=				42
  lowercase							=				42






class     a_						(      lowerCamelCase    ):

  def   A__						(				self				)  ->				Union[str, Any]:



        """simple docstring"""






        UpperCamelCase      = {}
        UpperCamelCase      = []
        UpperCamelCase      = 1
        UpperCamelCase      = [1, 2]
        UpperCamelCase      = {"""a""": 1, """b""": 2}
        UpperCamelCase      = {"""a""": [1, 2], """b""": [3, 4]}
        UpperCamelCase      = {"""a""": {"""1""": 1}, """b""": 2}
        UpperCamelCase      = {"""a""": 1, """b""": 2, """c""": 3, """d""": 4}
        UpperCamelCase      = {}
        UpperCamelCase      = []
        UpperCamelCase      = 2
        UpperCamelCase      = [2, 3]
        UpperCamelCase      = {"""a""": 2, """b""": 3}
        UpperCamelCase      = {"""a""": [2, 3], """b""": [4, 5]}
        UpperCamelCase      = {"""a""": {"""1""": 2}, """b""": 3}
        UpperCamelCase      = {"""a""": 2, """b""": 3, """c""": 4, """d""": 5}
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = 2
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)

        UpperCamelCase      = {"""a""": np.eye(2				), """b""": np.zeros(3				), """c""": np.ones(2				)}
        UpperCamelCase      = {"""a""": 2, """b""": 0, """c""": 2}
        UpperCamelCase      = {
            """a""": np.eye(2				).astype(_SCREAMING_SNAKE_CASE				),
            """b""": np.zeros(3				).astype(_SCREAMING_SNAKE_CASE				),
            """c""": np.ones(2				).astype(_SCREAMING_SNAKE_CASE				),
        }
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		map_numpy=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		map_numpy=_SCREAMING_SNAKE_CASE				).items()}      ,		{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,		)
        self.assertEqual(map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		map_numpy=_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)      ,		_SCREAMING_SNAKE_CASE				)
        self.assertEqual(
            {k: v.tolist() for k, v in map_nested(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		map_numpy=_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				).items()}      ,		{k: v.tolist() for k, v in expected_map_nested_sna_int.items()}      ,		)
        with self.assertRaises(_SCREAMING_SNAKE_CASE				):  # can't pickle a local lambda
              map_nested(lambda _SCREAMING_SNAKE_CASE				: x + 1      ,		_SCREAMING_SNAKE_CASE      ,		num_proc=_SCREAMING_SNAKE_CASE				)

  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






        UpperCamelCase      = {"""a""": 1, """b""": 2}
        UpperCamelCase      = {"""a""": 3, """b""": 4}
        UpperCamelCase      = {"""a""": 5, """b""": 6}
        UpperCamelCase      = sorted([("""a""", (1, 3, 5)), ("""b""", (2, 4, 6))]				)
        self.assertEqual(sorted(zip_dict(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)				)      ,		_SCREAMING_SNAKE_CASE				)


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



        """simple docstring"""






        class     a_						:
          lowercase							=				"""bar"""


        UpperCamelCase      = Foo()
        self.assertEqual(foo.my_attr      ,		"""bar"""				)
        with temporary_assignment(_SCREAMING_SNAKE_CASE      ,		"""my_attr"""      ,		"""BAR"""				):
              self.assertEqual(foo.my_attr      ,		"""BAR"""				)
        self.assertEqual(foo.my_attr      ,		"""bar"""				)







@pytest.mark.parametrize(
    """iterable_length, num_proc, expected_num_proc"""		,						[
        (1, None, 1),
        (1, 1, 1),
        (2, None, 1),
        (2, 1, 1),
        (2, 2, 1),
        (2, 3, 1),
        (3, 2, 1),
        (16, 16, 16),
        (16, 17, 16),
        (17, 16, 16),
    ]		,						)
def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase		,						__UpperCamelCase   )->			List[str]:
      with patch("""datasets.utils.py_utils._single_map_nested"""   ) as mock_single_map_nested, patch(
          """datasets.parallel.parallel.Pool"""   ) as mock_multiprocessing_pool:
            UpperCamelCase      = {F"{i}": i for i in range(__UpperCamelCase   )}
            UpperCamelCase      = map_nested(lambda __UpperCamelCase   : x + 10		,						__UpperCamelCase		,						num_proc=__UpperCamelCase		,						parallel_min_length=16   )
            if expected_num_proc == 1:
                  assert mock_single_map_nested.called
                  assert not mock_multiprocessing_pool.called
            else:
                  assert not mock_single_map_nested.called
                  assert mock_multiprocessing_pool.called
                  assert mock_multiprocessing_pool.call_args[0][0] == expected_num_proc






class     a_						(      lowerCamelCase    ):

  @require_tf
  def   A__						(				self				)  ->				Any:



        """simple docstring"""






        import tensorflow as tf
        from tensorflow.keras import layers

        UpperCamelCase      = layers.Dense(2				)

        def gen_random_output():
              UpperCamelCase      = tf.random.uniform((1, 3)				)
              return model(_SCREAMING_SNAKE_CASE				).numpy()

        with temp_seed(42      ,		set_tensorflow=_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = gen_random_output()
        with temp_seed(42      ,		set_tensorflow=_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = gen_random_output()
        UpperCamelCase      = gen_random_output()

        np.testing.assert_equal(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
        self.assertGreater(np.abs(outa - outa				).sum()      ,		0				)

  @require_torch
  def   A__						(				self				)  ->				int:



        """simple docstring"""






        import torch

        def gen_random_output():
              UpperCamelCase      = torch.nn.Linear(3      ,		2				)
              UpperCamelCase      = torch.rand(1      ,		3				)
              return model(_SCREAMING_SNAKE_CASE				).detach().numpy()

        with temp_seed(42      ,		set_pytorch=_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = gen_random_output()
        with temp_seed(42      ,		set_pytorch=_SCREAMING_SNAKE_CASE				):
              UpperCamelCase      = gen_random_output()
        UpperCamelCase      = gen_random_output()

        np.testing.assert_equal(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
        self.assertGreater(np.abs(outa - outa				).sum()      ,		0				)


  def   A__						(				self				)  ->				Dict:



        """simple docstring"""






        def gen_random_output():
              return np.random.rand(1      ,		3				)

        with temp_seed(42				):
              UpperCamelCase      = gen_random_output()
        with temp_seed(42				):
              UpperCamelCase      = gen_random_output()
        UpperCamelCase      = gen_random_output()

        np.testing.assert_equal(_SCREAMING_SNAKE_CASE      ,		_SCREAMING_SNAKE_CASE				)
        self.assertGreater(np.abs(outa - outa				).sum()      ,		0				)







@pytest.mark.parametrize("""input_data"""		,						[{}]   )
def    lowercase__			(				__UpperCamelCase   )->			List[str]:
      UpperCamelCase      = NestedDataStructure(__UpperCamelCase   ).data
      assert output_data == input_data







@pytest.mark.parametrize(
    """data, expected_output"""		,						[
        ({}, []),
        ([], []),
        ("""foo""", ["""foo"""]),
        (["""foo""", """bar"""], ["""foo""", """bar"""]),
        ([["""foo""", """bar"""]], ["""foo""", """bar"""]),
        ([[["""foo"""], ["""bar"""]]], ["""foo""", """bar"""]),
        ([[["""foo"""], """bar"""]], ["""foo""", """bar"""]),
        ({"""a""": 1, """b""": 2}, [1, 2]),
        ({"""a""": [1, 2], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[1, 2]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[[3], [4]]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [[3, 4]]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, 4]}, [1, 2, 3, 4]),
        ({"""a""": [[[1], [2]]], """b""": [3, [4]]}, [1, 2, 3, 4]),
        ({"""a""": {"""1""": 1}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": 2}, [1, 2]),
        ({"""a""": {"""1""": [1]}, """b""": [2]}, [1, 2]),
    ]		,						)
def    lowercase__			(				__UpperCamelCase		,						__UpperCamelCase   )->			str:
      UpperCamelCase      = NestedDataStructure(__UpperCamelCase   ).flatten()
      assert output == expected_output







def    lowercase__			(				)->			Union[str, Any]:
      UpperCamelCase      = A(x=1		,						y="""foobar"""   )
      UpperCamelCase      = {"""x""": 1, """y""": """foobar"""}
      assert asdict(__UpperCamelCase   ) == expected_output

      UpperCamelCase      = {"""a""": {"""b""": A(x=10		,						y="""foo"""   )}, """c""": [A(x=20		,						y="""bar"""   )]}
      UpperCamelCase      = {"""a""": {"""b""": {"""x""": 10, """y""": """foo"""}}, """c""": [{"""x""": 20, """y""": """bar"""}]}
      assert asdict(__UpperCamelCase   ) == expected_output

      with pytest.raises(__UpperCamelCase   ):
            asdict([1, A(x=10		,						y="""foo"""   )]   )







def    lowercase__			(				__UpperCamelCase   )->			List[Any]:
      return text.split()







def    lowercase__			(				__UpperCamelCase   )->			List[str]:
      yield (time.time(), content)
      time.sleep(2   )
      yield (time.time(), content)







def    lowercase__			(				)->			int:
      with Pool(2   ) as pool:
            UpperCamelCase      = list(iflatmap_unordered(__UpperCamelCase		,						_split_text		,						kwargs_iterable=[{"""text""": """hello there"""}] * 10   )   )
            assert out.count("""hello"""   ) == 10
            assert out.count("""there"""   ) == 10
            assert len(__UpperCamelCase   ) == 20

      # check multiprocess from pathos (uses dill for pickling)
      with multiprocess.Pool(2   ) as pool:
            UpperCamelCase      = list(iflatmap_unordered(__UpperCamelCase		,						_split_text		,						kwargs_iterable=[{"""text""": """hello there"""}] * 10   )   )
            assert out.count("""hello"""   ) == 10
            assert out.count("""there"""   ) == 10
            assert len(__UpperCamelCase   ) == 20

      # check that we get items as fast as possible
      with Pool(2   ) as pool:
            UpperCamelCase      = []
            for yield_time, content in iflatmap_unordered(
                __UpperCamelCase		,						_aseconds_generator_of_aitems_with_timing		,						kwargs_iterable=[{"""content""": """a"""}, {"""content""": """b"""}]   ):
                  assert yield_time < time.time() + 0.1, "we should each item directly after it was yielded"
                  out.append(__UpperCamelCase   )
            assert out.count("""a"""   ) == 2
            assert out.count("""b"""   ) == 2
            assert len(__UpperCamelCase   ) == 4
