





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 				lowercase		(				self :		List[str]     )							->						Optional[int]:
     __lowerCAmelCase								=	pa.array(TypedSequence([1, 2, 3]     )     )
     self.assertEqual(arr.type					,  pa.intaa()     )






  def 				lowercase		(				self :		Dict     )							->						Tuple:
     with self.assertRaises(lowerCAmelCase_     ):
        __lowerCAmelCase								=	pa.array(TypedSequence([1, 2, 3]     )					,  type=pa.intaa()     )






  def 				lowercase		(				self :		str     )							->						Optional[Any]:
     with self.assertRaises(lowerCAmelCase_     ):
        __lowerCAmelCase								=	pa.array(TypedSequence([1, 2, 3]					,  try_type=Value('bool'     )					,  type=Value('int64'     )     )     )






  def 				lowercase		(				self :		Dict     )							->						List[str]:
     __lowerCAmelCase								=	pa.array(TypedSequence([1, 2, 3]					,  type=Value('int32'     )     )     )
     self.assertEqual(arr.type					,  pa.intaa()     )






  def 				lowercase		(				self :		Union[str, Any]     )							->						Dict:
     with self.assertRaises((TypeError, pa.lib.ArrowInvalid)     ):
        __lowerCAmelCase								=	pa.array(TypedSequence(['foo', 'bar']					,  type=Value('int64'     )     )     )






  def 				lowercase		(				self :		List[str]     )							->						Dict:
     __lowerCAmelCase								=	pa.array(TypedSequence([1, 2, 3]					,  try_type=Value('int32'     )     )     )
     self.assertEqual(arr.type					,  pa.intaa()     )






  def 				lowercase		(				self :		List[str]     )							->						int:
     __lowerCAmelCase								=	pa.array(TypedSequence(['foo', 'bar']					,  try_type=Value('int64'     )     )     )
     self.assertEqual(arr.type					,  pa.string()     )






  def 				lowercase		(				self :		Dict     )							->						Optional[int]:
     __lowerCAmelCase								=	pa.array(TypedSequence([[[1, 2, 3]]]					,  type=ArrayaD((1, 3)					,  'int64'     )     )     )
     self.assertEqual(arr.type					,  ArrayaDExtensionType((1, 3)					,  'int64'     )     )






  def 				lowercase		(				self :		Tuple     )							->						Union[str, Any]:
     with self.assertRaises((TypeError, pa.lib.ArrowInvalid)     ):
        __lowerCAmelCase								=	pa.array(TypedSequence(['foo', 'bar']					,  type=ArrayaD((1, 3)					,  'int64'     )     )     )






  def 				lowercase		(				self :		Dict     )							->						Any:
     __lowerCAmelCase								=	pa.array(TypedSequence([[[1, 2, 3]]]					,  try_type=ArrayaD((1, 3)					,  'int64'     )     )     )
     self.assertEqual(arr.type					,  ArrayaDExtensionType((1, 3)					,  'int64'     )     )






  def 				lowercase		(				self :		Any     )							->						List[Any]:
     __lowerCAmelCase								=	pa.array(TypedSequence(['foo', 'bar']					,  try_type=ArrayaD((1, 3)					,  'int64'     )     )     )
     self.assertEqual(arr.type					,  pa.string()     )






  @require_pil
  def 				lowercase		(				self :		Optional[Any]     )							->						Any:
     import PIL.Image

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





def   a_    ( lowerCAmelCase_      :			Any,	lowerCAmelCase_      :			int	):
   __lowerCAmelCase								=	pa.BufferReader(lowerCAmelCase_	) if isinstance(lowerCAmelCase_,	pa.Buffer	) else pa.memory_map(lowerCAmelCase_	)
   __lowerCAmelCase								=	pa.ipc.open_stream(lowerCAmelCase_	)
   __lowerCAmelCase								=	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   a_    ( lowerCAmelCase_      :			str,	lowerCAmelCase_      :			int	):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   __lowerCAmelCase								=	pa.schema(lowerCAmelCase_	) if fields else None
   with ArrowWriter(stream=lowerCAmelCase_,	schema=lowerCAmelCase_,	writer_batch_size=lowerCAmelCase_	) as writer:
      writer.write({'col_1': 'foo', 'col_2': 1}	)
      writer.write({'col_1': 'bar', 'col_2': 2}	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	writer.finalize()
   assert num_examples == 2
   assert num_bytes > 0
   if not fields:
      __lowerCAmelCase								=	{'col_1': pa.string(), 'col_2': pa.intaa()}
   assert writer._schema == pa.schema(lowerCAmelCase_,	metadata=writer._schema.metadata	)
   _check_output(output.getvalue(),	expected_num_chunks=num_examples if writer_batch_size == 1 else 1	)

def   a_    ( ):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   __lowerCAmelCase								=	Features({'labels': ClassLabel(names=['neg', 'pos']	)}	)
   with ArrowWriter(stream=lowerCAmelCase_,	features=lowerCAmelCase_	) as writer:
      writer.write({'labels': 0}	)
      writer.write({'labels': 1}	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	writer.finalize()
   assert num_examples == 2
   assert num_bytes > 0
   assert writer._schema == features.arrow_schema
   assert writer._schema.metadata == features.arrow_schema.metadata
   __lowerCAmelCase								=	pa.BufferReader(output.getvalue()	)
   __lowerCAmelCase								=	pa.ipc.open_stream(lowerCAmelCase_	)
   __lowerCAmelCase								=	f.read_all()
   __lowerCAmelCase								=	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(lowerCAmelCase_	)

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

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

@pytest.mark.parametrize('writer_batch_size',	[None, 2, 10]	)
def   a_    ( lowerCAmelCase_      :			int	):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   with ArrowWriter(
       stream=lowerCAmelCase_,	writer_batch_size=lowerCAmelCase_,	hash_salt='split_name',	check_duplicates=lowerCAmelCase_,	) as writer:
      writer.write({'col_1': 'foo', 'col_2': 1},	key=1	)
      writer.write({'col_1': 'bar', 'col_2': 2},	key=2	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	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   a_    ( lowerCAmelCase_      :			str,	lowerCAmelCase_      :			Any	):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   __lowerCAmelCase								=	pa.schema(lowerCAmelCase_	) if fields else None
   with ArrowWriter(stream=lowerCAmelCase_,	schema=lowerCAmelCase_,	writer_batch_size=lowerCAmelCase_	) as writer:
      writer.write_batch({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}	)
      writer.write_batch({'col_1': [], 'col_2': []}	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	writer.finalize()
   assert num_examples == 2
   assert num_bytes > 0
   if not fields:
      __lowerCAmelCase								=	{'col_1': pa.string(), 'col_2': pa.intaa()}
   assert writer._schema == pa.schema(lowerCAmelCase_,	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   a_    ( lowerCAmelCase_      :			Union[str, Any],	lowerCAmelCase_      :			Union[str, Any]	):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   __lowerCAmelCase								=	pa.schema(lowerCAmelCase_	) if fields else None
   with ArrowWriter(stream=lowerCAmelCase_,	schema=lowerCAmelCase_,	writer_batch_size=lowerCAmelCase_	) as writer:
      writer.write_table(pa.Table.from_pydict({'col_1': ['foo', 'bar'], 'col_2': [1, 2]}	)	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	writer.finalize()
   assert num_examples == 2
   assert num_bytes > 0
   if not fields:
      __lowerCAmelCase								=	{'col_1': pa.string(), 'col_2': pa.intaa()}
   assert writer._schema == pa.schema(lowerCAmelCase_,	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   a_    ( lowerCAmelCase_      :			int,	lowerCAmelCase_      :			Any	):
   __lowerCAmelCase								=	pa.BufferOutputStream()
   __lowerCAmelCase								=	pa.schema(lowerCAmelCase_	) if fields else None
   with ArrowWriter(stream=lowerCAmelCase_,	schema=lowerCAmelCase_,	writer_batch_size=lowerCAmelCase_	) 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]}	)	)
      __lowerCAmelCase	,	__lowerCAmelCase								=	writer.finalize()
   assert num_examples == 2
   assert num_bytes > 0
   if not fields:
      __lowerCAmelCase								=	{'col_1': pa.string(), 'col_2': pa.intaa()}
   assert writer._schema == pa.schema(lowerCAmelCase_,	metadata=writer._schema.metadata	)
   _check_output(output.getvalue(),	expected_num_chunks=num_examples if writer_batch_size == 1 else 1	)

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

def   a_    ( lowerCAmelCase_      :			List[str]	):
   if pa.types.is_list(lowerCAmelCase_	):
      return get_base_dtype(arr_type.value_type	)
   else:
      return arr_type

def   a_    ( lowerCAmelCase_      :			int,	lowerCAmelCase_      :			Optional[Any]	):
   if isinstance(lst[0],	lowerCAmelCase_	):
      change_first_primitive_element_in_list(lst[0],	lowerCAmelCase_	)
   else:
      __lowerCAmelCase								=	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   a_    ( lowerCAmelCase_      :			Optional[int],	lowerCAmelCase_      :			Any,	lowerCAmelCase_      :			Dict	):
   __lowerCAmelCase								=	pa.array(TypedSequence(lowerCAmelCase_,	optimized_int_type=lowerCAmelCase_	)	)
   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   a_    ( lowerCAmelCase_      :			List[Any],	lowerCAmelCase_      :			int,	lowerCAmelCase_      :			int	):
   # in range
   __lowerCAmelCase								=	pa.array(OptimizedTypedSequence(lowerCAmelCase_,	col=lowerCAmelCase_	)	)
   assert get_base_dtype(arr.type	) == expected_dtype

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

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

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

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

@require_pil
@pytest.mark.parametrize('embed_local_files',	[False, True]	)
def   a_    ( lowerCAmelCase_      :			Optional[Any],	lowerCAmelCase_      :			Optional[Any]	):
   import PIL.Image

   __lowerCAmelCase								=	str(tmp_path / 'test_image_rgb.jpg'	)
   PIL.Image.fromarray(np.zeros((5, 5),	dtype=np.uinta	)	).save(lowerCAmelCase_,	format='png'	)
   __lowerCAmelCase								=	pa.BufferOutputStream()
   with ParquetWriter(
       stream=lowerCAmelCase_,	features=Features({'image': Image()}	),	embed_local_files=lowerCAmelCase_	) as writer:
      writer.write({'image': image_path}	)
      writer.finalize()
   __lowerCAmelCase								=	pa.BufferReader(output.getvalue()	)
   __lowerCAmelCase								=	pq.read_table(lowerCAmelCase_	)
   __lowerCAmelCase								=	pa_table.to_pydict()
   if embed_local_files:
      assert isinstance(out['image'][0]['path'],	lowerCAmelCase_	)
      with open(lowerCAmelCase_,	'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   a_    ( ):
   __lowerCAmelCase								=	pa.schema([pa.field('col_1',	pa.string(),	nullable=lowerCAmelCase_	)]	)
   __lowerCAmelCase								=	pa.BufferOutputStream()
   with ArrowWriter(stream=lowerCAmelCase_	) as writer:
      writer._build_writer(inferred_schema=lowerCAmelCase_	)
      assert writer._schema == pa.schema([pa.field('col_1',	pa.string()	)]	)








import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs

@pytest.mark.parametrize(
    'kwargs, expected',	[
        ({'num_shards': 0, 'max_num_jobs': 1}, []),
        ({'num_shards': 10, 'max_num_jobs': 1}, [range(10	)]),
        ({'num_shards': 10, 'max_num_jobs': 10}, [range(lowerCAmelCase_,	i + 1	) for i in range(10	)]),
        ({'num_shards': 1, 'max_num_jobs': 10}, [range(1	)]),
        ({'num_shards': 10, 'max_num_jobs': 3}, [range(0,	4	), range(4,	7	), range(7,	10	)]),
        ({'num_shards': 3, 'max_num_jobs': 10}, [range(0,	1	), range(1,	2	), range(2,	3	)]),
    ],	)
def   a_    ( lowerCAmelCase_      :			Optional[Any],	lowerCAmelCase_      :			Any	):
   __lowerCAmelCase								=	_distribute_shards(**lowerCAmelCase_	)
   assert out == expected

@pytest.mark.parametrize(
    'gen_kwargs, max_num_jobs, expected',	[
        ({'foo': 0}, 10, [{'foo': 0}]),
        ({'shards': [0, 1, 2, 3]}, 1, [{'shards': [0, 1, 2, 3]}]),
        ({'shards': [0, 1, 2, 3]}, 4, [{'shards': [0]}, {'shards': [1]}, {'shards': [2]}, {'shards': [3]}]),
        ({'shards': [0, 1]}, 4, [{'shards': [0]}, {'shards': [1]}]),
        ({'shards': [0, 1, 2, 3]}, 2, [{'shards': [0, 1]}, {'shards': [2, 3]}]),
    ],	)
def   a_    ( lowerCAmelCase_      :			Optional[Any],	lowerCAmelCase_      :			List[str],	lowerCAmelCase_      :			Optional[int]	):
   __lowerCAmelCase								=	_split_gen_kwargs(lowerCAmelCase_,	lowerCAmelCase_	)
   assert out == expected




@pytest.mark.parametrize(
    'gen_kwargs, expected',	[
        ({'foo': 0}, 1),
        ({'shards': [0]}, 1),
        ({'shards': [0, 1, 2, 3]}, 4),
        ({'shards': [0, 1, 2, 3], 'foo': 0}, 4),
        ({'shards': [0, 1, 2, 3], 'other': (0, 1)}, 4),
        ({'shards': [0, 1, 2, 3], 'shards2': [0, 1]}, RuntimeError),
    ],	)
def   a_    ( lowerCAmelCase_      :			int,	lowerCAmelCase_      :			Any	):
   if expected is RuntimeError:
      with pytest.raises(lowerCAmelCase_	):
         _number_of_shards_in_gen_kwargs(lowerCAmelCase_	)
   else:
      __lowerCAmelCase								=	_number_of_shards_in_gen_kwargs(lowerCAmelCase_	)
      assert out == expected

