





import pyarrow.parquet as pq
import pytest

from datasets import Audio, Dataset, DatasetDict, Features, NamedSplit, Sequence, Value, config
from datasets.features.image import Image
from datasets.io.parquet import ParquetDatasetReader, ParquetDatasetWriter, get_writer_batch_size

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases

def   a_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Optional[Any]     ) ->    str:
   assert isinstance(lowerCAmelCase_,     lowerCAmelCase_     )
   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   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      List[str]     ) ->    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							     =     ParquetDatasetReader(lowerCAmelCase_,     cache_dir=lowerCAmelCase_,     keep_in_memory=lowerCAmelCase_     ).read()
   _check_parquet_dataset(lowerCAmelCase_,     lowerCAmelCase_     )

@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   a_   (			lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      Optional[int],     lowerCAmelCase_				:      Tuple     ) ->    List[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(lowerCAmelCase_     ) for feature, dtype in features.items()}     ) if features is not None else None
   )
   __lowerCAmelCase							     =     ParquetDatasetReader(lowerCAmelCase_,     features=lowerCAmelCase_,     cache_dir=lowerCAmelCase_     ).read()
   _check_parquet_dataset(lowerCAmelCase_,     lowerCAmelCase_     )

@pytest.mark.parametrize('split',     [None, NamedSplit('train'     ), 'train', 'test']     )
def   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      int,     lowerCAmelCase_				:      Dict     ) ->    Dict:
   __lowerCAmelCase							     =     tmp_path / """cache"""
   __lowerCAmelCase							     =     {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   __lowerCAmelCase							     =     ParquetDatasetReader(lowerCAmelCase_,     cache_dir=lowerCAmelCase_,     split=lowerCAmelCase_     ).read()
   _check_parquet_dataset(lowerCAmelCase_,     lowerCAmelCase_     )
   assert dataset.split == split if split else "train"

@pytest.mark.parametrize('path_type',     [str, list]     )
def   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      Optional[Any]     ) ->    Optional[Any]:
   if issubclass(lowerCAmelCase_,     lowerCAmelCase_     ):
      __lowerCAmelCase							     =     parquet_path
   elif issubclass(lowerCAmelCase_,     lowerCAmelCase_     ):
      __lowerCAmelCase							     =     [parquet_path]
   __lowerCAmelCase							     =     tmp_path / """cache"""
   __lowerCAmelCase							     =     {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   __lowerCAmelCase							     =     ParquetDatasetReader(lowerCAmelCase_,     cache_dir=lowerCAmelCase_     ).read()
   _check_parquet_dataset(lowerCAmelCase_,     lowerCAmelCase_     )

def   a_   (			lowerCAmelCase_				:      int,     lowerCAmelCase_				:      str,     lowerCAmelCase_				:      Union[str, Any]=("train",)     ) ->    Tuple:
   assert isinstance(lowerCAmelCase_,     lowerCAmelCase_     )
   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   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      Any     ) ->    Tuple:
   __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							     =     ParquetDatasetReader(
          {'train': parquet_path},     cache_dir=lowerCAmelCase_,     keep_in_memory=lowerCAmelCase_     ).read()
   _check_parquet_datasetdict(lowerCAmelCase_,     lowerCAmelCase_     )

@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   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      Dict     ) ->    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(lowerCAmelCase_     ) for feature, dtype in features.items()}     ) if features is not None else None
   )
   __lowerCAmelCase							     =     ParquetDatasetReader({'train': parquet_path},     features=lowerCAmelCase_,     cache_dir=lowerCAmelCase_     ).read()
   _check_parquet_datasetdict(lowerCAmelCase_,     lowerCAmelCase_     )

@pytest.mark.parametrize('split',     [None, NamedSplit('train'     ), 'train', 'test']     )
def   a_   (			lowerCAmelCase_				:      Optional[Any],     lowerCAmelCase_				:      List[str],     lowerCAmelCase_				:      Optional[int]     ) ->    Any:
   if split:
      __lowerCAmelCase							     =     {split: parquet_path}
   else:
      __lowerCAmelCase							     =     """train"""
      __lowerCAmelCase							     =     {"""train""": parquet_path, """test""": parquet_path}
   __lowerCAmelCase							     =     tmp_path / """cache"""
   __lowerCAmelCase							     =     {"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
   __lowerCAmelCase							     =     ParquetDatasetReader(lowerCAmelCase_,     cache_dir=lowerCAmelCase_     ).read()
   _check_parquet_datasetdict(lowerCAmelCase_,     lowerCAmelCase_,     splits=list(path.keys()     )     )
   assert all(dataset[split].split == split for split in path.keys()     )

def   a_   (			lowerCAmelCase_				:      Dict,     lowerCAmelCase_				:      Any     ) ->    Dict:
   __lowerCAmelCase							     =     ParquetDatasetWriter(lowerCAmelCase_,     tmp_path / 'foo.parquet'     )
   assert writer.write() > 0
   __lowerCAmelCase							     =     pq.ParquetFile(tmp_path / 'foo.parquet'     )
   __lowerCAmelCase							     =     pf.read()
   assert dataset.data.table == output_table

def   a_   (			lowerCAmelCase_				:      Tuple,     lowerCAmelCase_				:      Tuple     ) ->    Dict:
   __lowerCAmelCase							     =     str(shared_datadir / 'test_image_rgb.jpg'     )
   __lowerCAmelCase							     =     {"""image""": [image_path]}
   __lowerCAmelCase							     =     Features({'image': Image()}     )
   __lowerCAmelCase							     =     Dataset.from_dict(lowerCAmelCase_,     features=lowerCAmelCase_     )
   __lowerCAmelCase							     =     ParquetDatasetWriter(lowerCAmelCase_,     tmp_path / 'foo.parquet'     )
   assert writer.write() > 0

   __lowerCAmelCase							     =     Dataset.from_parquet(str(tmp_path / 'foo.parquet'     )     )
   assert dataset.features == reloaded_dataset.features

   __lowerCAmelCase							     =     ParquetDatasetReader(str(tmp_path / 'foo.parquet'     ),     streaming=lowerCAmelCase_     ).read()
   assert dataset.features == reloaded_iterable_dataset.features




@pytest.mark.parametrize(
    'feature, expected',     [
        (Features({'foo': Value('int32'     )}     ), None),
        (Features({'image': Image(), 'foo': Value('int32'     )}     ), config.PARQUET_ROW_GROUP_SIZE_FOR_IMAGE_DATASETS),
        (Features({'nested': Sequence(Audio()     )}     ), config.PARQUET_ROW_GROUP_SIZE_FOR_AUDIO_DATASETS),
    ],     )
def   a_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Optional[Any]     ) ->    List[Any]:
   assert get_writer_batch_size(lowerCAmelCase_     ) == expected








import builtins
import sys

from ...utils.imports import _is_package_available
from . import cursor, input
from .helpers import Direction, clear_line, forceWrite, linebreak, move_cursor, reset_cursor, writeColor
from .keymap import KEYMAP


_snake_case  :   Union[str, Any]						   =  False
try:
     _snake_case  :   Tuple						   =  _is_package_available('google.colab')
except ModuleNotFoundError:
     pass



@input.register
class     _UpperCAmelCase   :

  """simple docstring"""


  def __init__(    self			:		Union[str, Any] ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		list = []					)							->						Optional[int]:
     __lowerCAmelCase							     =     0
     __lowerCAmelCase							     =     choices
     __lowerCAmelCase							     =     prompt
     if sys.platform == "win32":
        __lowerCAmelCase							     =     '*'
     else:
        __lowerCAmelCase							     =     '➔ '






  def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		Dict ,  lowerCAmelCase_			:		str = ""					)							->						Any:
     if sys.platform != "win32":
        writeColor(self.choices[index] ,  3_2 ,  lowerCAmelCase_					)
     else:
        forceWrite(self.choices[index] ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		List[Any] ,  lowerCAmelCase_			:		int					)							->						List[Any]:
     if index == self.position:
        forceWrite(f""" {self.arrow_char} """					)
        self.write_choice(lowerCAmelCase_					)
     else:
        forceWrite(f"""    {self.choices[index]}"""					)
     reset_cursor()






  def     lowercase		(    self			:		Dict ,  lowerCAmelCase_			:		Direction ,  lowerCAmelCase_			:		int = 1					)							->						Union[str, Any]:
     __lowerCAmelCase							     =     self.position
     if direction == Direction.DOWN:
        if self.position + 1 >= len(self.choices					):
           return
        self.position += num_spaces
     else:
        if self.position - 1 < 0:
           return
        self.position -= num_spaces
     clear_line()
     self.print_choice(lowerCAmelCase_					)
     move_cursor(lowerCAmelCase_ ,  direction.name					)
     self.print_choice(self.position					)






  @input.mark(KEYMAP['up']					)
  def     lowercase		(    self			:		Optional[int]					)							->						Tuple:
     self.move_direction(Direction.UP					)






  @input.mark(KEYMAP['down']					)
  def     lowercase		(    self			:		str					)							->						Optional[Any]:
     self.move_direction(Direction.DOWN					)






  @input.mark(KEYMAP['newline']					)
  def     lowercase		(    self			:		str					)							->						Any:
     move_cursor(len(self.choices					) - self.position ,  'DOWN'					)
     return self.position






  @input.mark(KEYMAP['interrupt']					)
  def     lowercase		(    self			:		Dict					)							->						Optional[Any]:
     move_cursor(len(self.choices					) - self.position ,  'DOWN'					)
     raise KeyboardInterrupt






  @input.mark_multiple(*[KEYMAP[str(lowerCAmelCase_					)] for number in range(1_0					)]					)
  def     lowercase		(    self			:		List[Any]					)							->						Optional[int]:
     __lowerCAmelCase							     =     int(chr(self.current_selection					)					)
     __lowerCAmelCase							     =     index - self.position
     if index == self.position:
        return
     if index < len(self.choices					):
        if self.position > index:
           self.move_direction(Direction.UP ,  -movement					)
        elif self.position < index:
           self.move_direction(Direction.DOWN ,  lowerCAmelCase_					)
        else:
           return
     else:
        return






  def     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		int = 0					)							->						int:
     if self.prompt:
        linebreak()
        forceWrite(self.prompt ,  '\n'					)
        if in_colab:
           forceWrite('Please input a choice index (starting from 0), and press enter' ,  '\n'					)
        else:
           forceWrite('Please select a choice using the arrow or number keys, and selecting with enter' ,  '\n'					)
     __lowerCAmelCase							     =     default_choice
     for i in range(len(self.choices					)					):
        self.print_choice(lowerCAmelCase_					)
        forceWrite('\n'					)
     move_cursor(len(self.choices					) - self.position ,  'UP'					)
     with cursor.hide():
        while True:
           if in_colab:
              try:
                 __lowerCAmelCase							     =     int(builtins.input()					)
              except ValueError:
                 __lowerCAmelCase							     =     default_choice
           else:
              __lowerCAmelCase							     =     self.handle_input()
           if choice is not None:
              reset_cursor()
              for _ in range(len(self.choices					) + 1					):
                 move_cursor(1 ,  'UP'					)
                 clear_line()
              self.write_choice(lowerCAmelCase_ ,  '\n'					)
              return choice

