



import copy
import os
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from datasets.arrow_writer import ArrowWriter, OptimizedTypedSequence, ParquetWriter, TypedSequence
from datasets.features import ArrayaD, ClassLabel, Features, Image, Value
from datasets.features.features import ArrayaDExtensionType, cast_to_python_objects
from datasets.keyhash import DuplicatedKeysError, InvalidKeyError

from .utils import require_pil

class        UpperCAmelCase_  (       UpperCamelCase_	):

						'''simple docstring'''




						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([1, 2, 3]		)		)
													self.assertEqual(arr.type				,							pa.intaa()		)






						def 	_A     (   self		):







													'''simple docstring'''

													with self.assertRaises(_A		):
																				__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([1, 2, 3]		)				,							type=pa.intaa()		)






						def 	_A     (   self		):







													'''simple docstring'''

													with self.assertRaises(_A		):
																				__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([1, 2, 3]				,							try_type=Value('bool'		)				,							type=Value('int64'		)		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([1, 2, 3]				,							type=Value('int32'		)		)		)
													self.assertEqual(arr.type				,							pa.intaa()		)






						def 	_A     (   self		):







													'''simple docstring'''

													with self.assertRaises((TypeError, pa.lib.ArrowInvalid)		):
																				__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence(['foo', 'bar']				,							type=Value('int64'		)		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([1, 2, 3]				,							try_type=Value('int32'		)		)		)
													self.assertEqual(arr.type				,							pa.intaa()		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence(['foo', 'bar']				,							try_type=Value('int64'		)		)		)
													self.assertEqual(arr.type				,							pa.string()		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([[[1, 2, 3]]]				,							type=ArrayaD((1, 3)				,							'int64'		)		)		)
													self.assertEqual(arr.type				,							ArrayaDExtensionType((1, 3)				,							'int64'		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													with self.assertRaises((TypeError, pa.lib.ArrowInvalid)		):
																				__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence(['foo', 'bar']				,							type=ArrayaD((1, 3)				,							'int64'		)		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([[[1, 2, 3]]]				,							try_type=ArrayaD((1, 3)				,							'int64'		)		)		)
													self.assertEqual(arr.type				,							ArrayaDExtensionType((1, 3)				,							'int64'		)		)






						def 	_A     (   self		):







													'''simple docstring'''

													__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence(['foo', 'bar']				,							try_type=ArrayaD((1, 3)				,							'int64'		)		)		)
													self.assertEqual(arr.type				,							pa.string()		)






						@require_pil
						def 	_A     (   self		):







													'''simple docstring'''

													import PIL.Image

													__SCREAMING_SNAKE_CASE									= PIL.Image.fromarray(np.arange(10				,							dtype=np.uinta		).reshape(2				,							5		)		)
													with patch(
													    'datasets.arrow_writer.cast_to_python_objects'				,							side_effect=_A		) as mock_cast_to_python_objects:
																				__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence([{'path': None, 'bytes': b'image_bytes'}, pil_image]				,							type=Image()		)		)
																				__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= mock_cast_to_python_objects.call_args_list[-1]
																				self.assertIn('optimize_list_casting'				,							_A		)
																				self.assertFalse(kwargs['optimize_list_casting']		)







def      __lowercase       (			a__							,       a__       )    ->			List[str]:
							__SCREAMING_SNAKE_CASE									= pa.BufferReader(a__       ) if isinstance(a__							,       pa.Buffer       ) else pa.memory_map(a__       )
							__SCREAMING_SNAKE_CASE									= pa.ipc.open_stream(a__       )
							__SCREAMING_SNAKE_CASE									= f.read_all()
							assert len(pa_table.to_batches()       ) == expected_num_chunks
							assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}
							del pa_table






@pytest.mark.parametrize('writer_batch_size'							,       [None, 1, 10]       )
@pytest.mark.parametrize(
    'fields'							,       [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]       )
def      __lowercase       (			a__							,       a__       )    ->			str:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							__SCREAMING_SNAKE_CASE									= pa.schema(a__       ) if fields else None
							with ArrowWriter(stream=a__							,       schema=a__							,       writer_batch_size=a__       ) as writer:
														writer.write({'col_1': 'foo', 'col_2': 1}       )
														writer.write({'col_1': 'bar', 'col_2': 2}       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							if not fields:
														__SCREAMING_SNAKE_CASE									= {'col_1': pa.string(), 'col_2': pa.intaa()}
							assert writer._schema == pa.schema(a__							,       metadata=writer._schema.metadata       )
							_check_output(output.getvalue()							,       expected_num_chunks=num_examples if writer_batch_size == 1 else 1       )






def      __lowercase       (			)    ->			Union[str, Any]:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							__SCREAMING_SNAKE_CASE									= Features({'labels': ClassLabel(names=['neg', 'pos']       )}       )
							with ArrowWriter(stream=a__							,       features=a__       ) as writer:
														writer.write({'labels': 0}       )
														writer.write({'labels': 1}       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							assert writer._schema == features.arrow_schema
							assert writer._schema.metadata == features.arrow_schema.metadata
							__SCREAMING_SNAKE_CASE									= pa.BufferReader(output.getvalue()       )
							__SCREAMING_SNAKE_CASE									= pa.ipc.open_stream(a__       )
							__SCREAMING_SNAKE_CASE									= f.read_all()
							__SCREAMING_SNAKE_CASE									= pa_table.schema
							assert pa_table.num_rows == 2
							assert schema == features.arrow_schema
							assert schema.metadata == features.arrow_schema.metadata
							assert features == Features.from_arrow_schema(a__       )






@pytest.mark.parametrize('writer_batch_size'							,       [None, 1, 10]       )
def      __lowercase       (			a__       )    ->			Any:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ArrowWriter(
							    stream=a__							,       writer_batch_size=a__							,       hash_salt='split_name'							,       check_duplicates=a__							,       ) as writer:
														with pytest.raises(a__       ):
																					writer.write({'col_1': 'foo', 'col_2': 1}							,       key=[1, 2]       )
																					__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()






@pytest.mark.parametrize('writer_batch_size'							,       [None, 2, 10]       )
def      __lowercase       (			a__       )    ->			Dict:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ArrowWriter(
							    stream=a__							,       writer_batch_size=a__							,       hash_salt='split_name'							,       check_duplicates=a__							,       ) as writer:
														with pytest.raises(a__       ):
																					writer.write({'col_1': 'foo', 'col_2': 1}							,       key=10       )
																					writer.write({'col_1': 'bar', 'col_2': 2}							,       key=10       )
																					__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()






@pytest.mark.parametrize('writer_batch_size'							,       [None, 2, 10]       )
def      __lowercase       (			a__       )    ->			Union[str, Any]:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ArrowWriter(
							    stream=a__							,       writer_batch_size=a__							,       hash_salt='split_name'							,       check_duplicates=a__							,       ) as writer:
														writer.write({'col_1': 'foo', 'col_2': 1}							,       key=1       )
														writer.write({'col_1': 'bar', 'col_2': 2}							,       key=2       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							_check_output(output.getvalue()							,       expected_num_chunks=num_examples if writer_batch_size == 1 else 1       )






@pytest.mark.parametrize('writer_batch_size'							,       [None, 1, 10]       )
@pytest.mark.parametrize(
    'fields'							,       [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]       )
def      __lowercase       (			a__							,       a__       )    ->			Optional[Any]:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							__SCREAMING_SNAKE_CASE									= pa.schema(a__       ) if fields else None
							with ArrowWriter(stream=a__							,       schema=a__							,       writer_batch_size=a__       ) as writer:
														writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}       )
														writer.write_batch({'col_1': [], 'col_2': []}       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							if not fields:
														__SCREAMING_SNAKE_CASE									= {'col_1': pa.string(), 'col_2': pa.intaa()}
							assert writer._schema == pa.schema(a__							,       metadata=writer._schema.metadata       )
							_check_output(output.getvalue()							,       expected_num_chunks=num_examples if writer_batch_size == 1 else 1       )






@pytest.mark.parametrize('writer_batch_size'							,       [None, 1, 10]       )
@pytest.mark.parametrize(
    'fields'							,       [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]       )
def      __lowercase       (			a__							,       a__       )    ->			str:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							__SCREAMING_SNAKE_CASE									= pa.schema(a__       ) if fields else None
							with ArrowWriter(stream=a__							,       schema=a__							,       writer_batch_size=a__       ) as writer:
														writer.write_table(pa.Table.from_pydict({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}       )       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							if not fields:
														__SCREAMING_SNAKE_CASE									= {'col_1': pa.string(), 'col_2': pa.intaa()}
							assert writer._schema == pa.schema(a__							,       metadata=writer._schema.metadata       )
							_check_output(output.getvalue()							,       expected_num_chunks=num_examples if writer_batch_size == 1 else 1       )






@pytest.mark.parametrize('writer_batch_size'							,       [None, 1, 10]       )
@pytest.mark.parametrize(
    'fields'							,       [None, {'col_1': pa.string(), 'col_2': pa.intaa()}, {'col_1': pa.string(), 'col_2': pa.intaa()}]       )
def      __lowercase       (			a__							,       a__       )    ->			Optional[Any]:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							__SCREAMING_SNAKE_CASE									= pa.schema(a__       ) if fields else None
							with ArrowWriter(stream=a__							,       schema=a__							,       writer_batch_size=a__       ) as writer:
														writer.write_row(pa.Table.from_pydict({'col_1': ['foo'], 'col_2': [1]}       )       )
														writer.write_row(pa.Table.from_pydict({'col_1': ['bar'], 'col_2': [2]}       )       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							if not fields:
														__SCREAMING_SNAKE_CASE									= {'col_1': pa.string(), 'col_2': pa.intaa()}
							assert writer._schema == pa.schema(a__							,       metadata=writer._schema.metadata       )
							_check_output(output.getvalue()							,       expected_num_chunks=num_examples if writer_batch_size == 1 else 1       )






def      __lowercase       (			)    ->			str:
							with tempfile.TemporaryDirectory() as tmp_dir:
														__SCREAMING_SNAKE_CASE									= {'col_1': pa.string(), 'col_2': pa.intaa()}
														__SCREAMING_SNAKE_CASE									= os.path.join(a__							,       'test.arrow'       )
														with ArrowWriter(path=a__							,       schema=pa.schema(a__       )       ) as writer:
																					writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}       )
																					__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
														assert num_examples == 2
														assert num_bytes > 0
														assert writer._schema == pa.schema(a__							,       metadata=writer._schema.metadata       )
														_check_output(a__							,       1       )






def      __lowercase       (			a__       )    ->			Any:
							if pa.types.is_list(a__       ):
														return get_base_dtype(arr_type.value_type       )
							else:
														return arr_type






def      __lowercase       (			a__							,       a__       )    ->			str:
							if isinstance(lst[0]							,       a__       ):
														change_first_primitive_element_in_list(lst[0]							,       a__       )
							else:
														__SCREAMING_SNAKE_CASE									= value






@pytest.mark.parametrize('optimized_int_type, expected_dtype'							,       [(None, pa.intaa()), (Value('int32'       ), pa.intaa())]       )
@pytest.mark.parametrize('sequence'							,       [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]       )
def      __lowercase       (			a__							,       a__							,       a__       )    ->			Union[str, Any]:
							__SCREAMING_SNAKE_CASE									= pa.array(TypedSequence(a__							,       optimized_int_type=a__       )       )
							assert get_base_dtype(arr.type       ) == expected_dtype






@pytest.mark.parametrize(
    'col, expected_dtype'							,       [
        ('attention_mask', pa.inta()),
        ('special_tokens_mask', pa.inta()),
        ('token_type_ids', pa.inta()),
        ('input_ids', pa.intaa()),
        ('other', pa.intaa()),
    ]							,       )
@pytest.mark.parametrize('sequence'							,       [[1, 2, 3], [[1, 2, 3]], [[[1, 2, 3]]]]       )
def      __lowercase       (			a__							,       a__							,       a__       )    ->			List[str]:
							# in range
							__SCREAMING_SNAKE_CASE									= pa.array(OptimizedTypedSequence(a__							,       col=a__       )       )
							assert get_base_dtype(arr.type       ) == expected_dtype

							# not in range
							if col != "other":
														# avoids errors due to in-place modifications
														__SCREAMING_SNAKE_CASE									= copy.deepcopy(a__       )
														__SCREAMING_SNAKE_CASE									= np.iinfo(expected_dtype.to_pandas_dtype()       ).max + 1
														change_first_primitive_element_in_list(a__							,       a__       )
														__SCREAMING_SNAKE_CASE									= pa.array(OptimizedTypedSequence(a__							,       col=a__       )       )
														assert get_base_dtype(arr.type       ) == pa.intaa()






@pytest.mark.parametrize('raise_exception'							,       [False, True]       )
def      __lowercase       (			a__							,       a__       )    ->			Optional[int]:
							__SCREAMING_SNAKE_CASE									= str(tmp_path / 'dataset-train.arrow'       )
							try:
														with ArrowWriter(path=a__       ) as writer:
																					if raise_exception:
																												raise pa.lib.ArrowInvalid()
																					else:
																												writer.stream.close()
							except pa.lib.ArrowInvalid:
														pass
							finally:
														assert writer.stream.closed






def      __lowercase       (			a__       )    ->			Any:
							__SCREAMING_SNAKE_CASE									= 'mock://dataset-train.arrow'
							with ArrowWriter(path=a__							,       storage_options=mockfs.storage_options       ) as writer:
														assert isinstance(writer._fs							,       type(a__       )       )
														assert writer._fs.storage_options == mockfs.storage_options
														writer.write({'col_1': 'foo', 'col_2': 1}       )
														writer.write({'col_1': 'bar', 'col_2': 2}       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							assert mockfs.exists(a__       )






def      __lowercase       (			)    ->			Union[str, Any]:
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ParquetWriter(stream=a__       ) as writer:
														writer.write({'col_1': 'foo', 'col_2': 1}       )
														writer.write({'col_1': 'bar', 'col_2': 2}       )
														__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= writer.finalize()
							assert num_examples == 2
							assert num_bytes > 0
							__SCREAMING_SNAKE_CASE									= pa.BufferReader(output.getvalue()       )
							__SCREAMING_SNAKE_CASE									= pq.read_table(a__       )
							assert pa_table.to_pydict() == {"col_1": ["foo", "bar"], "col_2": [1, 2]}






@require_pil
@pytest.mark.parametrize('embed_local_files'							,       [False, True]       )
def      __lowercase       (			a__							,       a__       )    ->			Dict:
							import PIL.Image

							__SCREAMING_SNAKE_CASE									= str(tmp_path / 'test_image_rgb.jpg'       )
							PIL.Image.fromarray(np.zeros((5, 5)							,       dtype=np.uinta       )       ).save(a__							,       format='png'       )
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ParquetWriter(
							    stream=a__							,       features=Features({'image': Image()}       )							,       embed_local_files=a__       ) as writer:
														writer.write({'image': image_path}       )
														writer.finalize()
							__SCREAMING_SNAKE_CASE									= pa.BufferReader(output.getvalue()       )
							__SCREAMING_SNAKE_CASE									= pq.read_table(a__       )
							__SCREAMING_SNAKE_CASE									= pa_table.to_pydict()
							if embed_local_files:
														assert isinstance(out['image'][0]['path']							,       a__       )
														with open(a__							,       'rb'       ) as f:
																					assert out["image"][0]["bytes"] == f.read()
							else:
														assert out["image"][0]["path"] == image_path
														assert out["image"][0]["bytes"] is None






def      __lowercase       (			)    ->			List[str]:
							__SCREAMING_SNAKE_CASE									= pa.schema([pa.field('col_1'							,       pa.string()							,       nullable=a__       )]       )
							__SCREAMING_SNAKE_CASE									= pa.BufferOutputStream()
							with ArrowWriter(stream=a__       ) as writer:
														writer._build_writer(inferred_schema=a__       )
														assert writer._schema == pa.schema([pa.field('col_1'							,       pa.string()       )]       )











from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCAmelCase__		:					Optional[int]							={'''configuration_fnet''': ['''FNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''FNetConfig''']}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__		:					Optional[int]							=['''FNetTokenizer''']

try:
			if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__		:					Optional[Any]							=['''FNetTokenizerFast''']

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase__		:					Tuple							=[
			    '''FNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''FNetForMaskedLM''',
			    '''FNetForMultipleChoice''',
			    '''FNetForNextSentencePrediction''',
			    '''FNetForPreTraining''',
			    '''FNetForQuestionAnswering''',
			    '''FNetForSequenceClassification''',
			    '''FNetForTokenClassification''',
			    '''FNetLayer''',
			    '''FNetModel''',
			    '''FNetPreTrainedModel''',
			]


if TYPE_CHECKING:
			from .configuration_fnet import FNET_PRETRAINED_CONFIG_ARCHIVE_MAP, FNetConfig

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_fnet import FNetTokenizer

			try:
						if not is_tokenizers_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_fnet_fast import FNetTokenizerFast

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_fnet import (
						    FNET_PRETRAINED_MODEL_ARCHIVE_LIST,
						    FNetForMaskedLM,
						    FNetForMultipleChoice,
						    FNetForNextSentencePrediction,
						    FNetForPreTraining,
						    FNetForQuestionAnswering,
						    FNetForSequenceClassification,
						    FNetForTokenClassification,
						    FNetLayer,
						    FNetModel,
						    FNetPreTrainedModel,
						)


else:
			import sys

			lowerCAmelCase__		:					Any							=_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)






