

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 							A_			(		datasets.BuilderConfig ):
     _lowercase      :					Optional[datasets.Features]       =     None







class 							A_			(		datasets.ArrowBasedBuilder ):
     _lowercase      :					Union[str, Any]       =     ArrowConfig


     def   UpperCAmelCase  (							self      :		Dict )							->     List[str]:
      return datasets.DatasetInfo(features=self.config.features )


     def   UpperCAmelCase  (							self      :		Dict      ,  UpperCAmelCase      :		List[Any] )							->     str:
      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:       Optional[int]            =				dl_manager.download_and_extract(self.config.data_files )
      if isinstance(UpperCAmelCase      ,  (str, list, tuple) ):
       __lowerCAmelCase:       List[str]            =				data_files
       if isinstance(UpperCAmelCase      ,  UpperCAmelCase ):
        __lowerCAmelCase:       List[str]            =				[files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       __lowerCAmelCase:       Union[str, Any]            =				[dl_manager.iter_files(UpperCAmelCase ) for file in files]
       return [datasets.SplitGenerator(name=datasets.Split.TRAIN      ,  gen_kwargs={'files': files} )]
      __lowerCAmelCase:       Optional[Any]            =				[]
      for split_name, files in data_files.items():
       if isinstance(UpperCAmelCase      ,  UpperCAmelCase ):
        __lowerCAmelCase:       Dict            =				[files]
       # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
       __lowerCAmelCase:       List[str]            =				[dl_manager.iter_files(UpperCAmelCase ) 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(UpperCAmelCase ):
         with open(UpperCAmelCase      ,  'rb' ) as f:
          __lowerCAmelCase:       List[Any]            =				datasets.Features.from_arrow_schema(pa.ipc.open_stream(UpperCAmelCase ).schema )
         break
       splits.append(datasets.SplitGenerator(name=UpperCAmelCase      ,  gen_kwargs={'files': files} ) )
      return splits


     def   UpperCAmelCase  (							self      :		List[str]      ,  UpperCAmelCase      :		pa.Table )							->     pa.Table:
      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:       Dict            =				table_cast(UpperCAmelCase      ,  self.info.features.arrow_schema )
      return pa_table


     def   UpperCAmelCase  (							self      :		Optional[Any]      ,  UpperCAmelCase      :		Any )							->     List[Any]:
      for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase ) ):
       with open(UpperCAmelCase      ,  'rb' ) as f:
        try:
         for batch_idx, record_batch in enumerate(pa.ipc.open_stream(UpperCAmelCase ) ):
          __lowerCAmelCase:       List[Any]            =				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(UpperCAmelCase )
        except ValueError as e:
         logger.error(F'''Failed to read file \'{file}\' with error {type(UpperCAmelCase )}: {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			=       1_0_0_0_0_0_0_0_0_0_0_0
_a			=       1_0_0_0_0

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







def _a			(	SCREAMING_SNAKE_CASE		:      int   , SCREAMING_SNAKE_CASE		:      int		)   ->   datasets.Dataset:

 """simple docstring"""






 __lowerCAmelCase:       List[Any]            =				pa.Table.from_pydict({'col': [0] * chunk_size}		)
 __lowerCAmelCase:       Dict            =				pa.concat_tables([table] * (num_examples // chunk_size)		)
 return datasets.Dataset(SCREAMING_SNAKE_CASE   , fingerprint='table_100B'		)






@dataclass
class 							A_			:
     _lowercase      :					int
     _lowercase      :					int
     _lowercase      :					int
     _lowercase      :					int


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


     def __iter__(							self      :		Any )							->     Optional[int]:
      return iter(self._sampled_values )


     def __len__(							self      :		Union[str, Any] )							->     Dict:
      return self.size







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   List[str]:

 """simple docstring"""






 __lowerCAmelCase:       List[Any]            =				dataset[0]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Dict:

 """simple docstring"""






 __lowerCAmelCase:       Dict            =				dataset[-1]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Any:

 """simple docstring"""






 __lowerCAmelCase:       Union[str, Any]            =				dataset[range(len(SCREAMING_SNAKE_CASE		) // 2   , len(SCREAMING_SNAKE_CASE		) // 2 + 10_24		)]







@get_duration
def _a			(	SCREAMING_SNAKE_CASE		:      datasets.Dataset		)   ->   Optional[Any]:

 """simple docstring"""






 __lowerCAmelCase:       int            =				dataset[RandIter(0   , len(SCREAMING_SNAKE_CASE		)   , 10_24   , seed=42		)]







def _a			(	)   ->   Tuple:

 """simple docstring"""






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

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


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





