






"""simple docstring"""




import itertools
from dataclasses import dataclass
from io import StringIO
from typing import Optional

import pyarrow as pa

import datasets
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


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






@dataclass
class  UpperCamelCase    (			datasets.BuilderConfig   ):
 UpperCamelCase      : Optional[datasets.Features]    =      None
 UpperCamelCase      : str    =      "utf-8"
 UpperCamelCase      : Optional[str]    =      None
 UpperCamelCase      : int    =      10 << 20  # 10MB
 UpperCamelCase      : bool    =      False
 UpperCamelCase      : str    =      "line"







class  UpperCamelCase    (			datasets.ArrowBasedBuilder   ):
 UpperCamelCase      : str    =      TextConfig


 def _lowercase			(  self	:		Optional[Any]							)      -> List[Any]:
   return datasets.DatasetInfo(features=self.config.features							)


 def _lowercase			(  self	:		int							,		UpperCAmelCase__	:		Tuple							)      -> Any:
   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}"""							)
   _a						: Any    =			dl_manager.download_and_extract(self.config.data_files							)
   if isinstance(UpperCAmelCase__							,		(str, list, tuple)							):
     _a						: List[Any]    =			data_files
     if isinstance(UpperCAmelCase__							,		UpperCAmelCase__							):
       _a						: Optional[int]    =			[files]
     _a						: List[str]    =			[dl_manager.iter_files(UpperCAmelCase__							) for file in files]
     return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,		gen_kwargs={"""files""": files}							)]
   _a						: List[str]    =			[]
   for split_name, files in data_files.items():
     if isinstance(UpperCAmelCase__							,		UpperCAmelCase__							):
       _a						: Any    =			[files]
     _a						: Any    =			[dl_manager.iter_files(UpperCAmelCase__							) for file in files]
     splits.append(datasets.SplitGenerator(name=UpperCAmelCase__							,		gen_kwargs={"""files""": files}							)							)
   return splits


 def _lowercase			(  self	:		Optional[int]							,		UpperCAmelCase__	:		pa.Table							)      -> pa.Table:
   if self.config.features is not None:
     _a						: int    =			self.config.features.arrow_schema
     if all(not require_storage_cast(UpperCAmelCase__							) for feature in self.config.features.values()							):
       # cheaper cast
       _a						: str    =			pa_table.cast(UpperCAmelCase__							)
     else:
       # more expensive cast; allows str <-> int/float or str to Audio for example
       _a						: Any    =			table_cast(UpperCAmelCase__							,		UpperCAmelCase__							)
     return pa_table
   else:
     return pa_table.cast(pa.schema({"""text""": pa.string()}							)							)







 def _lowercase			(  self	:		List[Any]							,		UpperCAmelCase__	:		int							)      -> Optional[Any]:
   _a						: int    =			list(self.config.features							) if self.config.features is not None else ["""text"""]
   for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase__							)							):
     # open in text mode, by default translates universal newlines ("\n", "\r\n" and "\r") into "\n"
     with open(UpperCAmelCase__							,		encoding=self.config.encoding							,		errors=self.config.errors							) as f:
       if self.config.sample_by == "line":
         _a						: Optional[int]    =			0
         while True:
           _a						: Any    =			f.read(self.config.chunksize							)
           if not batch:
             break
           batch += f.readline()  # finish current line
           # StringIO.readlines, by default splits only on "\n" (and keeps line breaks)
           _a						: List[str]    =			StringIO(UpperCAmelCase__							).readlines()
           if not self.config.keep_linebreaks:
             _a						: Optional[Any]    =			[line.rstrip("""\n"""							) for line in batch]
           _a						: List[str]    =			pa.Table.from_arrays([pa.array(UpperCAmelCase__							)]							,		names=UpperCAmelCase__							)
           # 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 (file_idx, batch_idx), self._cast_table(UpperCAmelCase__							)
           batch_idx += 1
       elif self.config.sample_by == "paragraph":
         _a						: Optional[Any]    =			0
         _a						: Union[str, Any]    =			""""""
         while True:
           _a						: str    =			f.read(self.config.chunksize							)
           if not new_batch:
             break
           batch += new_batch
           batch += f.readline()  # finish current line
           _a						: str    =			batch.split("""\n\n"""							)
           _a						: Any    =			pa.Table.from_arrays(
               [pa.array([example for example in batch[:-1] if example]							)]							,		names=UpperCAmelCase__							)
           # 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 (file_idx, batch_idx), self._cast_table(UpperCAmelCase__							)
           batch_idx += 1
           _a						: Dict    =			batch[-1]
         if batch:
           _a						: Dict    =			pa.Table.from_arrays([pa.array([batch]							)]							,		names=UpperCAmelCase__							)
           yield (file_idx, batch_idx), self._cast_table(UpperCAmelCase__							)
       elif self.config.sample_by == "document":
         _a						: List[Any]    =			f.read()
         _a						: int    =			pa.Table.from_arrays([pa.array([text]							)]							,		names=UpperCAmelCase__							)
         yield file_idx, self._cast_table(UpperCAmelCase__							)









"""simple docstring"""




import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


_snake_case					   =     1000_0000_0000
_snake_case					   =     1_0000

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

def     lowerCAmelCase__  (  UpperCamelCase__  ,     UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Union[str, Any]    =			pa.Table.from_pydict({"""col""": [0] * chunk_size}       )
  _a						: Dict    =			pa.concat_tables([table] * (num_examples // chunk_size)       )
  return datasets.Dataset(UpperCamelCase__  ,     fingerprint="""table_100B"""       )






@dataclass
class  UpperCamelCase    :
 UpperCamelCase      : int
 UpperCamelCase      : int
 UpperCamelCase      : int
 UpperCamelCase      : int


 def _lowercase			(  self	:		Optional[int]							)      -> Dict:
   _a						: Tuple    =			np.random.default_rng(self.seed							)
   _a						: Dict    =			rng.integers(low=self.low							,		high=self.high							,		size=self.size							).tolist()


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







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







@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Union[str, Any]    =			dataset[0]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Optional[Any]    =			dataset[-1]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: str    =			dataset[range(len(UpperCamelCase__       ) // 2  ,     len(UpperCamelCase__       ) // 2 + 1_0_2_4       )]

@get_duration
def     lowerCAmelCase__  (  UpperCamelCase__       ):



  '''simple docstring'''
  _a						: Tuple    =			dataset[RandIter(0  ,     len(UpperCamelCase__       )  ,     1_0_2_4  ,     seed=4_2       )]

def     lowerCAmelCase__  (  ):



  '''simple docstring'''
  _a						: str    =			{"""num examples""": SPEED_TEST_N_EXAMPLES}
  _a						: Dict    =			(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
  print("""generating dataset"""       )
  _a						: List[str]    =			generate_100B_dataset(num_examples=UpperCamelCase__  ,     chunk_size=UpperCamelCase__       )
  print("""Functions"""       )
  for func in functions:
    print(func.__name__       )
    _a						: Optional[Any]    =			func(UpperCamelCase__       )

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


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

