

'''simple docstring'''





import io
import json

import fsspec
import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.json import JsonDatasetReader, JsonDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase	)       ->			int:
					assert isinstance(__UpperCamelCase    ,__UpperCamelCase	)
					assert dataset.num_rows == 4
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_1", "col_2", "col_3"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('keep_in_memory'    ,[False, True]	)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			List[str]:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,cache_dir=__UpperCamelCase    ,keep_in_memory=__UpperCamelCase	).read()
					_check_json_dataset(__UpperCamelCase    ,__UpperCamelCase	)







@pytest.mark.parametrize(
    'features'    ,[
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ]    ,)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Dict:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					lowerCamelCase_  				=      features.copy() if features else default_expected_features
					lowerCamelCase_  				=      (
					    Features({feature: Value(__UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
					)
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,features=__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					_check_json_dataset(__UpperCamelCase    ,__UpperCamelCase	)







@pytest.mark.parametrize(
    'features'    ,[
        None,
        {'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'},
    ]    ,)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Optional[int]:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_3': 'float64', 'col_1': 'string', 'col_2': 'int64'}
					lowerCamelCase_  				=      features.copy() if features else default_expected_features
					lowerCamelCase_  				=      (
					    Features({feature: Value(__UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
					)
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,features=__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					assert isinstance(__UpperCamelCase    ,__UpperCamelCase	)
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_3", "col_1", "col_2"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase	)       ->			str:
					# jsonl_312_path features are {"col_3": "float64", "col_1": "string", "col_2": "int64"}
					lowerCamelCase_  				=      {'col_2': 'int64', 'col_3': 'float64', 'col_1': 'string'}
					lowerCamelCase_  				=      features.copy()
					lowerCamelCase_  				=      (
					    Features({feature: Value(__UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
					)
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,features=__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					assert isinstance(__UpperCamelCase    ,__UpperCamelCase	)
					assert dataset.num_rows == 2
					assert dataset.num_columns == 3
					assert dataset.column_names == ["col_2", "col_3", "col_1"]
					for feature, expected_dtype in expected_features.items():
										assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('split'    ,[None, NamedSplit('train'	), 'train', 'test']	)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			List[Any]:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,cache_dir=__UpperCamelCase    ,split=__UpperCamelCase	).read()
					_check_json_dataset(__UpperCamelCase    ,__UpperCamelCase	)
					assert dataset.split == split if split else "train"







@pytest.mark.parametrize('path_type'    ,[str, list]	)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Optional[Any]:
					if issubclass(__UpperCamelCase    ,__UpperCamelCase	):
										lowerCamelCase_  				=      jsonl_path
					elif issubclass(__UpperCamelCase    ,__UpperCamelCase	):
										lowerCamelCase_  				=      [jsonl_path]
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					_check_json_dataset(__UpperCamelCase    ,__UpperCamelCase	)







def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase=("train",)	)       ->			Optional[int]:
					assert isinstance(__UpperCamelCase    ,__UpperCamelCase	)
					for split in splits:
										lowerCamelCase_  				=      dataset_dict[split]
										assert dataset.num_rows == 4
										assert dataset.num_columns == 3
										assert dataset.column_names == ["col_1", "col_2", "col_3"]
										for feature, expected_dtype in expected_features.items():
															assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize('keep_in_memory'    ,[False, True]	)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Dict:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
										lowerCamelCase_  				=      JsonDatasetReader({'train': jsonl_path}    ,cache_dir=__UpperCamelCase    ,keep_in_memory=__UpperCamelCase	).read()
					_check_json_datasetdict(__UpperCamelCase    ,__UpperCamelCase	)







@pytest.mark.parametrize(
    'features'    ,[
        None,
        {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'},
        {'col_1': 'string', 'col_2': 'string', 'col_3': 'string'},
        {'col_1': 'int32', 'col_2': 'int32', 'col_3': 'int32'},
        {'col_1': 'float32', 'col_2': 'float32', 'col_3': 'float32'},
    ]    ,)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Any:
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					lowerCamelCase_  				=      features.copy() if features else default_expected_features
					lowerCamelCase_  				=      (
					    Features({feature: Value(__UpperCamelCase	) for feature, dtype in features.items()}	) if features is not None else None
					)
					lowerCamelCase_  				=      JsonDatasetReader({'train': jsonl_path}    ,features=__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					_check_json_datasetdict(__UpperCamelCase    ,__UpperCamelCase	)







@pytest.mark.parametrize('split'    ,[None, NamedSplit('train'	), 'train', 'test']	)
def   _UpperCamelCase	(		__UpperCamelCase    ,__UpperCamelCase    ,__UpperCamelCase	)       ->			Any:
					if split:
										lowerCamelCase_  				=      {split: jsonl_path}
					else:
										lowerCamelCase_  				=      'train'
										lowerCamelCase_  				=      {'train': jsonl_path, 'test': jsonl_path}
					lowerCamelCase_  				=      tmp_path / 'cache'
					lowerCamelCase_  				=      {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
					lowerCamelCase_  				=      JsonDatasetReader(__UpperCamelCase    ,cache_dir=__UpperCamelCase	).read()
					_check_json_datasetdict(__UpperCamelCase    ,__UpperCamelCase    ,splits=list(path.keys()	)	)
					assert all(dataset[split].split == split for split in path.keys()	)







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			Union[str, Any]:
					return json.load(__UpperCamelCase	)







def   _UpperCamelCase	(		__UpperCamelCase	)       ->			List[str]:
					return [json.loads(__UpperCamelCase	) for line in buffer]







class  UpperCAmelCase :

		'''simple docstring'''






		@pytest.mark.parametrize('lines, load_json_function'   ,			[(True, load_json_lines), (False, load_json)]		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[int]:






							'''simple docstring'''



							with io.BytesIO() as buffer:
												JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			lines=SCREAMING_SNAKE_CASE_		).write()
												buffer.seek(0		)
												lowerCamelCase_  				=      load_json_function(SCREAMING_SNAKE_CASE_		)
							assert isinstance(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							assert isinstance(exported_content[0]   ,			SCREAMING_SNAKE_CASE_		)
							assert len(SCREAMING_SNAKE_CASE_		) == 10
		@pytest.mark.parametrize(
		    'orient, container, keys, len_at'   ,			[
		        ('records', list, {'tokens', 'labels', 'answers', 'id'}, None),
		        ('split', dict, {'columns', 'data'}, 'data'),
		        ('index', dict, set('0123456789'		), None),
		        ('columns', dict, {'tokens', 'labels', 'answers', 'id'}, 'tokens'),
		        ('values', list, None, None),
		        ('table', dict, {'schema', 'data'}, 'data'),
		    ]   ,			)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     str:






							'''simple docstring'''



							with io.BytesIO() as buffer:
												JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			lines=SCREAMING_SNAKE_CASE_   ,			orient=SCREAMING_SNAKE_CASE_		).write()
												buffer.seek(0		)
												lowerCamelCase_  				=      load_json(SCREAMING_SNAKE_CASE_		)
							assert isinstance(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							if keys:
												if container is dict:
																	assert exported_content.keys() == keys
												else:
																	assert exported_content[0].keys() == keys
							else:
												assert not hasattr(SCREAMING_SNAKE_CASE_   ,			'keys'		) and not hasattr(exported_content[0]   ,			'keys'		)
							if len_at:
												assert len(exported_content[len_at]		) == 10
							else:
												assert len(SCREAMING_SNAKE_CASE_		) == 10
		@pytest.mark.parametrize('lines, load_json_function'   ,			[(True, load_json_lines), (False, load_json)]		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     int:






							'''simple docstring'''



							with io.BytesIO() as buffer:
												JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			lines=SCREAMING_SNAKE_CASE_   ,			num_proc=2		).write()
												buffer.seek(0		)
												lowerCamelCase_  				=      load_json_function(SCREAMING_SNAKE_CASE_		)
							assert isinstance(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							assert isinstance(exported_content[0]   ,			SCREAMING_SNAKE_CASE_		)
							assert len(SCREAMING_SNAKE_CASE_		) == 10
		@pytest.mark.parametrize(
		    'orient, container, keys, len_at'   ,			[
		        ('records', list, {'tokens', 'labels', 'answers', 'id'}, None),
		        ('split', dict, {'columns', 'data'}, 'data'),
		        ('index', dict, set('0123456789'		), None),
		        ('columns', dict, {'tokens', 'labels', 'answers', 'id'}, 'tokens'),
		        ('values', list, None, None),
		        ('table', dict, {'schema', 'data'}, 'data'),
		    ]   ,			)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Any:






							'''simple docstring'''



							with io.BytesIO() as buffer:
												JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			lines=SCREAMING_SNAKE_CASE_   ,			orient=SCREAMING_SNAKE_CASE_   ,			num_proc=2		).write()
												buffer.seek(0		)
												lowerCamelCase_  				=      load_json(SCREAMING_SNAKE_CASE_		)
							assert isinstance(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)
							if keys:
												if container is dict:
																	assert exported_content.keys() == keys
												else:
																	assert exported_content[0].keys() == keys
							else:
												assert not hasattr(SCREAMING_SNAKE_CASE_   ,			'keys'		) and not hasattr(exported_content[0]   ,			'keys'		)
							if len_at:
												assert len(exported_content[len_at]		) == 10
							else:
												assert len(SCREAMING_SNAKE_CASE_		) == 10
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_		)  ->     Optional[Any]:






							'''simple docstring'''



							with pytest.raises(SCREAMING_SNAKE_CASE_		):
												with io.BytesIO() as buffer:
																	JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			num_proc=0		)






		@pytest.mark.parametrize('compression, extension'   ,			[('gzip', 'gz'), ('bz2', 'bz2'), ('xz', 'xz')]		)
		def 			UpperCamelCase(					self   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_		)  ->     Tuple:






							'''simple docstring'''



							lowerCamelCase_  				=      tmp_path_factory.mktemp('data'		) / f'''test.json.{extension}'''
							lowerCamelCase_  				=      str(shared_datadir / f'''test_file.json.{extension}'''		)
							JsonDatasetWriter(SCREAMING_SNAKE_CASE_   ,			SCREAMING_SNAKE_CASE_   ,			compression=SCREAMING_SNAKE_CASE_		).write()

							with fsspec.open(SCREAMING_SNAKE_CASE_   ,			'rb'   ,			compression='infer'		) as f:
												lowerCamelCase_  				=      f.read()
							with fsspec.open(SCREAMING_SNAKE_CASE_   ,			'rb'   ,			compression='infer'		) as f:
												lowerCamelCase_  				=      f.read()
							assert exported_content == original_content










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							:			Optional[Any]        =							datasets.utils.logging.get_logger(__name__)

@dataclass
class   lowerCAmelCase				(   datasets.BuilderConfig     ):
 UpperCAmelCase__						=						1_00_00
 UpperCAmelCase__						=						None
 UpperCAmelCase__						=						None


class   lowerCAmelCase				(   datasets.ArrowBasedBuilder     ):
 UpperCAmelCase__						=						ParquetConfig
 def 	A_      ( self    : int      )     ->      Optional[int]:
     return datasets.DatasetInfo(features=self.config.features      )
 def 	A_      ( self    : int     ,	UpperCAmelCase    : str      )     ->      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__       : List[Any]   					=						dl_manager.download_and_extract(self.config.data_files      )
     if isinstance(UpperCAmelCase     ,	(str, list, tuple)      ):
         lowerCamelCase__       : int   					=						data_files
         if isinstance(UpperCAmelCase     ,	UpperCAmelCase      ):
             lowerCamelCase__       : int   					=						[files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         lowerCamelCase__       : str   					=						[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__       : Optional[Any]   					=						[files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         lowerCamelCase__       : Dict   					=						[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__       : Optional[Any]   					=						datasets.Features.from_arrow_schema(pq.read_schema(UpperCAmelCase      )      )
                 break
         splits.append(datasets.SplitGenerator(name=UpperCAmelCase     ,	gen_kwargs={'files': files}      )      )
     return splits
 def 	A_      ( self    : Optional[int]     ,	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__       : Any   					=						table_cast(UpperCAmelCase     ,	self.info.features.arrow_schema      )
     return pa_table







 def 	A_      ( self    : Any     ,	UpperCAmelCase    : Tuple      )     ->      List[str]:
     lowerCamelCase__       : Optional[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(UpperCAmelCase      )      ):
         with open(UpperCAmelCase     ,	'rb'      ) as f:
             lowerCamelCase__       : List[Any]   					=						pq.ParquetFile(UpperCAmelCase      )
             try:
                 for batch_idx, record_batch in enumerate(
                     parquet_file.iter_batches(batch_size=self.config.batch_size     ,	columns=self.config.columns      )      ):
                     lowerCamelCase__       : Tuple   					=						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





