


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


__UpperCAmelCase : Any 			=				datasets.utils.logging.get_logger(__name__)



@dataclass
class    UpperCAmelCase_	(							datasets.BuilderConfig):
       '''simple docstring'''
       __UpperCamelCase      :     int							=					10000
       __UpperCamelCase      :     Optional[List[str]]							=					None
       __UpperCamelCase      :     Optional[datasets.Features]							=					None




class    UpperCAmelCase_	(							datasets.ArrowBasedBuilder):
       '''simple docstring'''
       __UpperCamelCase      :     str							=					ParquetConfig






       def  _lowercase   (     self						):




              """simple docstring"""




              return datasets.DatasetInfo(features=self.config.features						)






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




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






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              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
                     UpperCamelCase    :	List[Any]   	=		table_cast(__SCREAMING_SNAKE_CASE ,					self.info.features.arrow_schema						)
              return pa_table






       def  _lowercase   (     self ,					__SCREAMING_SNAKE_CASE						):




              """simple docstring"""




              UpperCamelCase    :	List[Any]   	=		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(__SCREAMING_SNAKE_CASE						)						):
                     with open(__SCREAMING_SNAKE_CASE ,					'''rb'''						) as f:
                            UpperCamelCase    :	List[Any]   	=		pq.ParquetFile(__SCREAMING_SNAKE_CASE						)
                            try:
                                   for batch_idx, record_batch in enumerate(
                                       parquet_file.iter_batches(batch_size=self.config.batch_size ,					columns=self.config.columns						)						):
                                          UpperCamelCase    :	int   	=		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(__SCREAMING_SNAKE_CASE						)
                            except ValueError as e:
                                   logger.error(f"""Failed to read file '{file}' with error {type(__SCREAMING_SNAKE_CASE						)}: {e}"""						)
                                   raise





from .imports import is_tqdm_available


if is_tqdm_available():
  from tqdm.auto import tqdm as _tqdm

from ..state import PartialState


def      a	(							SCREAMING_SNAKE_CASE_      :		bool = True							,			*SCREAMING_SNAKE_CASE_      :		List[str]							,			**SCREAMING_SNAKE_CASE_      :		Tuple       ):






       """simple docstring"""





       if not is_tqdm_available():
              raise ImportError('''Accelerate\'s `tqdm` module requires `tqdm` to be installed. Please run `pip install tqdm`.'''       )
       UpperCamelCase    :	int   	=		False
       if main_process_only:
              UpperCamelCase    :	int   	=		PartialState().local_process_index == 0
       return _tqdm(*SCREAMING_SNAKE_CASE_							,			**SCREAMING_SNAKE_CASE_							,			disable=SCREAMING_SNAKE_CASE_       )

