






"""simple docstring"""



def 			_snake_case    (					lowerCamelCase__     :   str  )    ->     str:
							return "".join(chr(ord(lowerCamelCase__  ) - 32  ) if "a" <= char <= "z" else char for char in word  )


if __name__ == "__main__":
						from doctest import testmod

						testmod()











"""simple docstring"""



import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


A__	:    List[str]			 = datasets.utils.logging.get_logger(__name__)







@dataclass
class        lowercase__      (							datasets.BuilderConfig       ):
			_UpperCAmelCase		:int             =			10000
			_UpperCAmelCase		:Optional[List[str]]             =			None
			_UpperCAmelCase		:Optional[datasets.Features]             =			None







class        lowercase__      (							datasets.ArrowBasedBuilder       ):
			_UpperCAmelCase		:Optional[int]             =			ParquetConfig


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


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


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  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_	:							List[str]									=table_cast(snake_case__       ,	self.info.features.arrow_schema )
										return pa_table


			def      UpperCAmelCase__						(				self					:  int       ,	snake_case__					:  Optional[Any] ):
										lowerCamelCase_	:							Tuple									=self.info.features.arrow_schema if self.info.features is not None else None
										if self.info.features is not None and self.config.columns is not None:
																	if sorted(field.name for field in schema ) != sorted(self.config.columns ):
																								raise ValueError(
																								    F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'""" )
										for file_idx, file in enumerate(itertools.chain.from_iterable(snake_case__ ) ):
																	with open(snake_case__       ,	"rb" ) as f:
																								lowerCamelCase_	:							List[str]									=pq.ParquetFile(snake_case__ )
																								try:
																															for batch_idx, record_batch in enumerate(
																															    parquet_file.iter_batches(batch_size=self.config.batch_size       ,	columns=self.config.columns ) ):
																																						lowerCamelCase_	:							Union[str, 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(snake_case__ )
																								except ValueError as e:
																															logger.error(F"""Failed to read file '{file}' with error {type(snake_case__ )}: {e}""" )
																															raise



