
import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


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





@dataclass
class      _lowerCAmelCase (     datasets.BuilderConfig	):
				_lowercase      		=None





class      _lowerCAmelCase (     datasets.ArrowBasedBuilder	):
				_lowercase      		=ArrowConfig




				def 					__a						( self					)    ->				Optional[int]:
							return datasets.DatasetInfo(features=self.config.features					)




				def 					__a						( self    ,    _UpperCamelCase					)    ->				Tuple:


							if not self.config.data_files:
										raise ValueError(f"""At least one data file must be specified, but got data_files={self.config.data_files}"""					)
							lowerCAmelCase_      =     dl_manager.download_and_extract(self.config.data_files					)
							if isinstance(__lowerCAmelCase    ,    (str, list, tuple)					):
										lowerCAmelCase_      =     data_files
										if isinstance(__lowerCAmelCase    ,    __lowerCAmelCase					):
													lowerCAmelCase_      =     [files]
										# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
										lowerCAmelCase_      =     [dl_manager.iter_files(__lowerCAmelCase					) for file in files]
										return [datasets.SplitGenerator(name=datasets.Split.TRAIN    ,    gen_kwargs={"files": files}					)]
							lowerCAmelCase_      =     []
							for split_name, files in data_files.items():
										if isinstance(__lowerCAmelCase    ,    __lowerCAmelCase					):
													lowerCAmelCase_      =     [files]
										# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
										lowerCAmelCase_      =     [dl_manager.iter_files(__lowerCAmelCase					) for file in files]
										# Infer features is they are stoed in the arrow schema
										if self.info.features is None:
													for file in itertools.chain.from_iterable(__lowerCAmelCase					):
																with open(__lowerCAmelCase    ,    "rb"					) as f:
																			lowerCAmelCase_      =     datasets.Features.from_arrow_schema(pa.ipc.open_stream(__lowerCAmelCase					).schema					)
																break
										splits.append(datasets.SplitGenerator(name=__lowerCAmelCase    ,    gen_kwargs={"files": files}					)					)
							return splits




				def 					__a						( self    ,    _UpperCamelCase					)    ->				int:


							if self.info.features is not None:
										# more expensive cast to support nested features with keys in a different order
										# allows str <-> int/float or str to Audio for example
										lowerCAmelCase_      =     table_cast(__lowerCAmelCase    ,    self.info.features.arrow_schema					)
							return pa_table




				def 					__a						( self    ,    _UpperCamelCase					)    ->				int:


							for file_idx, file in enumerate(itertools.chain.from_iterable(__lowerCAmelCase					)					):
										with open(__lowerCAmelCase    ,    "rb"					) as f:
													try:
																for batch_idx, record_batch in enumerate(pa.ipc.open_stream(__lowerCAmelCase					)					):
																			lowerCAmelCase_      =     pa.Table.from_batches([record_batch]					)
																			# Uncomment for debugging (will print the Arrow table size and elements)
																			# logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
																			# logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
																			yield f"""{file_idx}_{batch_idx}""", self._cast_table(__lowerCAmelCase					)
													except ValueError as e:
																logger.error(f"""Failed to read file '{file}' with error {type(__lowerCAmelCase					)}: {e}"""					)
																raise








import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


_A       =					10_00_00_00_00_00
_A       =					1_00_00

_A,				_A       =					os.path.split(__file__)
_A       =					os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))






def        lowerCamelCase__    (    __lowerCAmelCase  :  int   ,							__lowerCAmelCase  :  int    ):




			"""simple docstring"""


			lowerCAmelCase_      =     pa.Table.from_pydict({"col": [0] * chunk_size}    )
			lowerCAmelCase_      =     pa.concat_tables([table] * (num_examples // chunk_size)    )
			return datasets.Dataset(__lowerCAmelCase   ,							fingerprint="table_100B"    )





@dataclass
class      _lowerCAmelCase :
				_lowercase      		=42
				_lowercase      		=42
				_lowercase      		=42
				_lowercase      		=42




				def 					__a						( self					)    ->				List[Any]:
							lowerCAmelCase_      =     np.random.default_rng(self.seed					)
							lowerCAmelCase_      =     rng.integers(low=self.low    ,    high=self.high    ,    size=self.size					).tolist()




				def __iter__( self					)    ->				List[str]:
							return iter(self._sampled_values					)




				def __len__( self					)    ->				Optional[int]:
							return self.size






@get_duration
def        lowerCamelCase__    (    __lowerCAmelCase  :  datasets.Dataset    ):




			"""simple docstring"""


			lowerCAmelCase_      =     dataset[0]






@get_duration
def        lowerCamelCase__    (    __lowerCAmelCase  :  datasets.Dataset    ):




			"""simple docstring"""


			lowerCAmelCase_      =     dataset[-1]






@get_duration
def        lowerCamelCase__    (    __lowerCAmelCase  :  datasets.Dataset    ):




			"""simple docstring"""


			lowerCAmelCase_      =     dataset[range(len(__lowerCAmelCase    ) // 2   ,							len(__lowerCAmelCase    ) // 2 + 1024    )]






@get_duration
def        lowerCamelCase__    (    __lowerCAmelCase  :  datasets.Dataset    ):




			"""simple docstring"""


			lowerCAmelCase_      =     dataset[RandIter(0   ,							len(__lowerCAmelCase    )   ,							1024   ,							seed=42    )]






def        lowerCamelCase__    (    ):




			"""simple docstring"""


			lowerCAmelCase_      =     {"num examples": SPEED_TEST_N_EXAMPLES}
			lowerCAmelCase_      =     (get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
			print("generating dataset"    )
			lowerCAmelCase_      =     generate_100B_dataset(num_examples=__lowerCAmelCase   ,							chunk_size=__lowerCAmelCase    )
			print("Functions"    )
			for func in functions:
						print(func.__name__    )
						lowerCAmelCase_      =     func(__lowerCAmelCase    )

			with open(__lowerCAmelCase   ,							"wb"    ) as f:
						f.write(json.dumps(__lowerCAmelCase    ).encode("utf-8"    )    )


if __name__ == "__main__":  # useful to run the profiler
			benchmark_table_100B()



