


'''simple docstring'''







import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam







class A__       (						datasets.BeamBasedBuilder      ):






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					Dict:




             '''simple docstring'''
             return datasets.DatasetInfo(
                 features=datasets.Features({"""content""": datasets.Value("""string"""	)}	)     ,						supervised_keys=SCREAMING_SNAKE_CASE     ,						)






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




             '''simple docstring'''
             return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,						gen_kwargs={"""examples""": get_test_dummy_examples()}	)]






      def __UpperCAmelCase       ( self			:List[str]     ,						SCREAMING_SNAKE_CASE			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:str	)       ->					List[Any]:




             '''simple docstring'''
             import apache_beam as beam

             return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE	)







class A__       (						datasets.BeamBasedBuilder      ):






      def __UpperCAmelCase       ( self			:List[str]	)       ->					Optional[int]:




             '''simple docstring'''
             return datasets.DatasetInfo(
                 features=datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""	)}	)}	)     ,						supervised_keys=SCREAMING_SNAKE_CASE     ,						)






      def __UpperCAmelCase       ( self			:Any     ,						SCREAMING_SNAKE_CASE			:Optional[int]     ,						SCREAMING_SNAKE_CASE			:int	)       ->					Optional[int]:




             '''simple docstring'''
             return [
                 datasets.SplitGenerator(name=datasets.Split.TRAIN     ,						gen_kwargs={"""examples""": get_test_nested_examples()}	)
             ]






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




             '''simple docstring'''
             import apache_beam as beam

             return pipeline | "Load Examples" >> beam.Create(SCREAMING_SNAKE_CASE	)


def 	SCREAMING_SNAKE_CASE_					(  )  ->			int:
       return [(i, {"content": content}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]  )]


def 	SCREAMING_SNAKE_CASE_					(  )  ->			List[Any]:
       return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["""foo""", """bar""", """foobar"""]  )]







class A__       (						UpperCAmelCase__      ):






      @require_beam
      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : int     					=len(get_test_dummy_examples()	)
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _a    : Optional[int]     					=DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE     ,						beam_runner="""DirectRunner"""	)
                    builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						f"{builder.name}-train.arrow"	)	)	)
                    self.assertDictEqual(builder.info.features     ,						datasets.Features({"""content""": datasets.Value("""string"""	)}	)	)
                    _a    : int     					=builder.as_dataset()
                    self.assertEqual(dset["""train"""].num_rows     ,						SCREAMING_SNAKE_CASE	)
                    self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples     ,						SCREAMING_SNAKE_CASE	)
                    self.assertDictEqual(dset["""train"""][0]     ,						get_test_dummy_examples()[0][1]	)
                    self.assertDictEqual(
                        dset["""train"""][expected_num_examples - 1]     ,						get_test_dummy_examples()[expected_num_examples - 1][1]	)
                    self.assertTrue(
                        os.path.exists(os.path.join(SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						"""dataset_info.json"""	)	)	)
                    del dset






      @require_beam
      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					int:




             '''simple docstring'''
             import apache_beam as beam

             _a    : int     					=beam.io.parquetio.WriteToParquet

             _a    : Union[str, Any]     					=len(get_test_dummy_examples()	)
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _a    : Optional[int]     					=DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE     ,						beam_runner="""DirectRunner"""	)
                    with patch("""apache_beam.io.parquetio.WriteToParquet"""	) as write_parquet_mock:
                           _a    : Union[str, Any]     					=partial(SCREAMING_SNAKE_CASE     ,						num_shards=2	)
                           builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(
                                SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						f"{builder.name}-train-00000-of-00002.arrow"	)	)	)
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(
                                SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						f"{builder.name}-train-00000-of-00002.arrow"	)	)	)
                    self.assertDictEqual(builder.info.features     ,						datasets.Features({"""content""": datasets.Value("""string"""	)}	)	)
                    _a    : Tuple     					=builder.as_dataset()
                    self.assertEqual(dset["""train"""].num_rows     ,						SCREAMING_SNAKE_CASE	)
                    self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples     ,						SCREAMING_SNAKE_CASE	)
                    # Order is not preserved when sharding, so we just check that all the elements are there
                    self.assertListEqual(sorted(dset["""train"""]["""content"""]	)     ,						sorted(["""foo""", """bar""", """foobar"""]	)	)
                    self.assertTrue(
                        os.path.exists(os.path.join(SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						"""dataset_info.json"""	)	)	)
                    del dset






      @require_beam
      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Optional[Any]:




             '''simple docstring'''
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _a    : List[Any]     					=DummyBeamDataset(cache_dir=SCREAMING_SNAKE_CASE	)
                    self.assertRaises(datasets.builder.MissingBeamOptions     ,						builder.download_and_prepare	)






      @require_beam
      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					List[str]:




             '''simple docstring'''
             _a    : List[Any]     					=len(get_test_nested_examples()	)
             with tempfile.TemporaryDirectory() as tmp_cache_dir:
                    _a    : str     					=NestedBeamDataset(cache_dir=SCREAMING_SNAKE_CASE     ,						beam_runner="""DirectRunner"""	)
                    builder.download_and_prepare()
                    self.assertTrue(
                        os.path.exists(
                            os.path.join(SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						f"{builder.name}-train.arrow"	)	)	)
                    self.assertDictEqual(
                        builder.info.features     ,						datasets.Features({"""a""": datasets.Sequence({"""b""": datasets.Value("""string"""	)}	)}	)	)
                    _a    : str     					=builder.as_dataset()
                    self.assertEqual(dset["""train"""].num_rows     ,						SCREAMING_SNAKE_CASE	)
                    self.assertEqual(dset["""train"""].info.splits["""train"""].num_examples     ,						SCREAMING_SNAKE_CASE	)
                    self.assertDictEqual(dset["""train"""][0]     ,						get_test_nested_examples()[0][1]	)
                    self.assertDictEqual(
                        dset["""train"""][expected_num_examples - 1]     ,						get_test_nested_examples()[expected_num_examples - 1][1]	)
                    self.assertTrue(
                        os.path.exists(os.path.join(SCREAMING_SNAKE_CASE     ,						builder.name     ,						"""default"""     ,						"""0.0.0"""     ,						"""dataset_info.json"""	)	)	)
                    del dset






'''simple docstring'''







import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs


@pytest.mark.parametrize(
    """kwargs, expected"""    ,[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(_UpperCAmelCase    ,i + 1  ) for i in range(10  )]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0    ,4  ), range(4    ,7  ), range(7    ,10  )]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0    ,1  ), range(1    ,2  ), range(2    ,3  )]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Union[str, Any]    ,_UpperCAmelCase	:	Dict  )  ->			Optional[Any]:
       _a    : Tuple     					=_distribute_shards(**_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected"""    ,[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	List[Any]    ,_UpperCAmelCase	:	Union[str, Any]  )  ->			List[str]:
       _a    : List[str]     					=_split_gen_kwargs(_UpperCAmelCase    ,_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, expected"""    ,[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	List[Any]  )  ->			Union[str, Any]:
       if expected is RuntimeError:
              with pytest.raises(_UpperCAmelCase  ):
                     _number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
       else:
              _a    : Dict     					=_number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
              assert out == expected


