




'''simple docstring'''







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      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   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 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      A__       =     ParquetDatasetReader(UpperCAmelCase      ,cache_dir=UpperCAmelCase      ,keep_in_memory=UpperCAmelCase						).read()
   _check_parquet_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 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   A__       =     features.copy() if features else default_expected_features
   A__       =     (
       Features({feature: Value(UpperCAmelCase						) for feature, dtype in features.items()}						) if features is not None else None
   )
   A__       =     ParquetDatasetReader(UpperCAmelCase      ,features=UpperCAmelCase      ,cache_dir=UpperCAmelCase						).read()
   _check_parquet_dataset(UpperCAmelCase      ,UpperCAmelCase						)







@pytest.mark.parametrize('split'      ,[None, NamedSplit('train'						), 'train', 'test']						)
def 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   A__       =     ParquetDatasetReader(UpperCAmelCase      ,cache_dir=UpperCAmelCase      ,split=UpperCAmelCase						).read()
   _check_parquet_dataset(UpperCAmelCase      ,UpperCAmelCase						)
   assert dataset.split == split if split else "train"







@pytest.mark.parametrize('path_type'      ,[str, list]						)
def 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   if issubclass(UpperCAmelCase      ,UpperCAmelCase						):
      A__       =     parquet_path
   elif issubclass(UpperCAmelCase      ,UpperCAmelCase						):
      A__       =     [parquet_path]
   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   A__       =     ParquetDatasetReader(UpperCAmelCase      ,cache_dir=UpperCAmelCase						).read()
   _check_parquet_dataset(UpperCAmelCase      ,UpperCAmelCase						)







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase=("train",)						):






   '''simple docstring'''







   assert isinstance(UpperCAmelCase      ,UpperCAmelCase						)
   for split in splits:
      A__       =     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      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
      A__       =     ParquetDatasetReader(
          {'train': parquet_path}      ,cache_dir=UpperCAmelCase      ,keep_in_memory=UpperCAmelCase						).read()
   _check_parquet_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 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   A__       =     features.copy() if features else default_expected_features
   A__       =     (
       Features({feature: Value(UpperCAmelCase						) for feature, dtype in features.items()}						) if features is not None else None
   )
   A__       =     ParquetDatasetReader({'train': parquet_path}      ,features=UpperCAmelCase      ,cache_dir=UpperCAmelCase						).read()
   _check_parquet_datasetdict(UpperCAmelCase      ,UpperCAmelCase						)







@pytest.mark.parametrize('split'      ,[None, NamedSplit('train'						), 'train', 'test']						)
def 						_A      ( UpperCAmelCase      ,UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   if split:
      A__       =     {split: parquet_path}
   else:
      A__       =     'train'
      A__       =     {'train': parquet_path, 'test': parquet_path}
   A__       =     tmp_path / 'cache'
   A__       =     {'col_1': 'string', 'col_2': 'int64', 'col_3': 'float64'}
   A__       =     ParquetDatasetReader(UpperCAmelCase      ,cache_dir=UpperCAmelCase						).read()
   _check_parquet_datasetdict(UpperCAmelCase      ,UpperCAmelCase      ,splits=list(path.keys()						)						)
   assert all(dataset[split].split == split for split in path.keys()						)







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     ParquetDatasetWriter(UpperCAmelCase      ,tmp_path / 'foo.parquet'						)
   assert writer.write() > 0
   A__       =     pq.ParquetFile(tmp_path / 'foo.parquet'						)
   A__       =     pf.read()
   assert dataset.data.table == output_table







def 						_A      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   A__       =     str(shared_datadir / 'test_image_rgb.jpg'						)
   A__       =     {'image': [image_path]}
   A__       =     Features({'image': Image()}						)
   A__       =     Dataset.from_dict(UpperCAmelCase      ,features=UpperCAmelCase						)
   A__       =     ParquetDatasetWriter(UpperCAmelCase      ,tmp_path / 'foo.parquet'						)
   assert writer.write() > 0

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

   A__       =     ParquetDatasetReader(str(tmp_path / 'foo.parquet'						)      ,streaming=UpperCAmelCase						).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      ( UpperCAmelCase      ,UpperCAmelCase						):






   '''simple docstring'''







   assert get_writer_batch_size(UpperCAmelCase						) == expected








'''simple docstring'''







def 						_A      ( UpperCAmelCase = 1      ,UpperCAmelCase = 1000						):






   '''simple docstring'''







   A__       =     1
   A__       =     0

   for divide_by_number in range(UpperCAmelCase      ,digit + 1						):
      A__       =     []
      A__       =     numerator
      for _ in range(1      ,digit + 1						):
         if now_divide in has_been_divided:
            if longest_list_length < len(UpperCAmelCase						):
               A__       =     len(UpperCAmelCase						)
               A__       =     divide_by_number
         else:
            has_been_divided.append(UpperCAmelCase						)
            A__       =     now_divide * 10 % divide_by_number

   return the_digit


# Tests
if __name__ == "__main__":
    import doctest

    doctest.testmod()


