
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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     Optional[int]				)			->						Any:
       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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     Tuple			,							UpperCAmelCase_	:     Union[str, Any]				)			->						List[str]:
       SCREAMING_SNAKE_CASE_							:				List[Any]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				Tuple		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     List[str]			,							UpperCAmelCase_	:     Optional[int]			,							UpperCAmelCase_	:     List[Any]				)			->						Any:
       SCREAMING_SNAKE_CASE_							:				List[Any]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				List[Any]		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       SCREAMING_SNAKE_CASE_							:				str		=features.copy() if features else default_expected_features
       SCREAMING_SNAKE_CASE_							:				Any		=(
           Features({feature: Value(UpperCAmelCase_				) for feature, dtype in features.items()}				) if features is not None else None
       )
       SCREAMING_SNAKE_CASE_							:				int		=ParquetDatasetReader(UpperCAmelCase_			,							features=UpperCAmelCase_			,							cache_dir=UpperCAmelCase_				).read()
       _check_parquet_dataset(UpperCAmelCase_			,							UpperCAmelCase_				)
@pytest.mark.parametrize('''split'''			,							[None, NamedSplit('''train'''				), '''train''', '''test''']				)
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     List[str]			,							UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     List[str]				)			->						str:
       SCREAMING_SNAKE_CASE_							:				Dict		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       SCREAMING_SNAKE_CASE_							:				str		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     int			,							UpperCAmelCase_	:     int				)			->						List[Any]:
       if issubclass(UpperCAmelCase_			,							UpperCAmelCase_				):
              SCREAMING_SNAKE_CASE_							:				int		=parquet_path
       elif issubclass(UpperCAmelCase_			,							UpperCAmelCase_				):
              SCREAMING_SNAKE_CASE_							:				int		=[parquet_path]
       SCREAMING_SNAKE_CASE_							:				List[str]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				Dict		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       SCREAMING_SNAKE_CASE_							:				List[Any]		=ParquetDatasetReader(UpperCAmelCase_			,							cache_dir=UpperCAmelCase_				).read()
       _check_parquet_dataset(UpperCAmelCase_			,							UpperCAmelCase_				)
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Optional[int]			,							UpperCAmelCase_	:     List[Any]			,							UpperCAmelCase_	:     List[Any]=("train",)				)			->						List[str]:
       assert isinstance(UpperCAmelCase_			,							UpperCAmelCase_				)
       for split in splits:
              SCREAMING_SNAKE_CASE_							:				Optional[int]		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Dict			,							UpperCAmelCase_	:     Any			,							UpperCAmelCase_	:     Tuple				)			->						Union[str, Any]:
       SCREAMING_SNAKE_CASE_							:				List[Any]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				str		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
              SCREAMING_SNAKE_CASE_							:				List[Any]		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Optional[Any]			,							UpperCAmelCase_	:     Optional[Any]			,							UpperCAmelCase_	:     List[str]				)			->						List[str]:
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				Dict		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       SCREAMING_SNAKE_CASE_							:				Optional[int]		=features.copy() if features else default_expected_features
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=(
           Features({feature: Value(UpperCAmelCase_				) for feature, dtype in features.items()}				) if features is not None else None
       )
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     int			,							UpperCAmelCase_	:     int			,							UpperCAmelCase_	:     List[str]				)			->						str:
       if split:
              SCREAMING_SNAKE_CASE_							:				int		={split: parquet_path}
       else:
              SCREAMING_SNAKE_CASE_							:				List[Any]		='''train'''
              SCREAMING_SNAKE_CASE_							:				Union[str, Any]		={'''train''': parquet_path, '''test''': parquet_path}
       SCREAMING_SNAKE_CASE_							:				Optional[int]		=tmp_path / '''cache'''
       SCREAMING_SNAKE_CASE_							:				int		={'''col_1''': '''string''', '''col_2''': '''int64''', '''col_3''': '''float64'''}
       SCREAMING_SNAKE_CASE_							:				List[str]		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Optional[int]			,							UpperCAmelCase_	:     Any				)			->						Dict:
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=ParquetDatasetWriter(UpperCAmelCase_			,							tmp_path / '''foo.parquet'''				)
       assert writer.write() > 0
       SCREAMING_SNAKE_CASE_							:				Optional[int]		=pq.ParquetFile(tmp_path / '''foo.parquet'''				)
       SCREAMING_SNAKE_CASE_							:				List[Any]		=pf.read()
       assert dataset.data.table == output_table
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Any			,							UpperCAmelCase_	:     List[str]				)			->						Optional[int]:
       SCREAMING_SNAKE_CASE_							:				List[str]		=str(shared_datadir / '''test_image_rgb.jpg'''				)
       SCREAMING_SNAKE_CASE_							:				Optional[int]		={'''image''': [image_path]}
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=Features({'''image''': Image()}				)
       SCREAMING_SNAKE_CASE_							:				str		=Dataset.from_dict(UpperCAmelCase_			,							features=UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				int		=ParquetDatasetWriter(UpperCAmelCase_			,							tmp_path / '''foo.parquet'''				)
       assert writer.write() > 0

       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=Dataset.from_parquet(str(tmp_path / '''foo.parquet'''				)				)
       assert dataset.features == reloaded_dataset.features

       SCREAMING_SNAKE_CASE_							:				Dict		=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 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Any			,							UpperCAmelCase_	:     Tuple				)			->						Dict:
       assert get_writer_batch_size(UpperCAmelCase_				) == expected


import importlib
import os

import fsspec
import pytest
from fsspec import register_implementation
from fsspec.registry import _registry as _fsspec_registry

from datasets.filesystems import COMPRESSION_FILESYSTEMS, HfFileSystem, extract_path_from_uri, is_remote_filesystem

from .utils import require_lza, require_zstandard
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     List[Any]				)			->						Dict:
       assert "mock" in _fsspec_registry
       assert "bz2" in _fsspec_registry
def 							SCREAMING_SNAKE_CASE_	(					)			->						Union[str, Any]:
       assert "mock" not in _fsspec_registry
       assert "bz2" in _fsspec_registry
def 							SCREAMING_SNAKE_CASE_	(					)			->						Dict:
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		='''mock-s3-bucket'''
       SCREAMING_SNAKE_CASE_							:				Any		=f's3://{mock_bucket}'
       SCREAMING_SNAKE_CASE_							:				Dict		=extract_path_from_uri(UpperCAmelCase_				)
       assert dataset_path.startswith('''s3://'''				) is False

       SCREAMING_SNAKE_CASE_							:				Dict		='''./local/path'''
       SCREAMING_SNAKE_CASE_							:				str		=extract_path_from_uri(UpperCAmelCase_				)
       assert dataset_path == new_dataset_path
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Union[str, Any]				)			->						List[str]:
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=is_remote_filesystem(UpperCAmelCase_				)
       assert is_remote is True

       SCREAMING_SNAKE_CASE_							:				Dict		=fsspec.filesystem('''file'''				)

       SCREAMING_SNAKE_CASE_							:				Optional[int]		=is_remote_filesystem(UpperCAmelCase_				)
       assert is_remote is False
@pytest.mark.parametrize('''compression_fs_class'''			,							UpperCAmelCase_				)
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Union[str, Any]			,							UpperCAmelCase_	:     Any			,							UpperCAmelCase_	:     Optional[Any]			,							UpperCAmelCase_	:     Optional[Any]			,							UpperCAmelCase_	:     int			,							UpperCAmelCase_	:     str			,							UpperCAmelCase_	:     Any				)			->						List[Any]:
       SCREAMING_SNAKE_CASE_							:				str		={'''gzip''': gz_file, '''xz''': xz_file, '''zstd''': zstd_file, '''bz2''': bza_file, '''lz4''': lza_file}
       SCREAMING_SNAKE_CASE_							:				List[Any]		=input_paths[compression_fs_class.protocol]
       if input_path is None:
              SCREAMING_SNAKE_CASE_							:				List[str]		=f'for \'{compression_fs_class.protocol}\' compression protocol, '
              if compression_fs_class.protocol == "lz4":
                     reason += require_lza.kwargs["reason"]
              elif compression_fs_class.protocol == "zstd":
                     reason += require_zstandard.kwargs["reason"]
              pytest.skip(UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				Union[str, Any]		=fsspec.filesystem(compression_fs_class.protocol			,							fo=UpperCAmelCase_				)
       assert isinstance(UpperCAmelCase_			,							UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				str		=os.path.basename(UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				Optional[Any]		=expected_filename[: expected_filename.rindex('''.'''				)]
       assert fs.glob('''*'''				) == [expected_filename]
       with fs.open(UpperCAmelCase_			,							'''r'''			,							encoding='''utf-8'''				) as f, open(UpperCAmelCase_			,							encoding='''utf-8'''				) as expected_file:
              assert f.read() == expected_file.read()
@pytest.mark.parametrize('''protocol'''			,							['''zip''', '''gzip''']				)
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     Any			,							UpperCAmelCase_	:     Union[str, Any]			,							UpperCAmelCase_	:     str				)			->						Any:
       SCREAMING_SNAKE_CASE_							:				Tuple		={'''zip''': zip_jsonl_path, '''gzip''': jsonl_gz_path}
       SCREAMING_SNAKE_CASE_							:				List[Any]		=compressed_file_paths[protocol]
       SCREAMING_SNAKE_CASE_							:				str		='''dataset.jsonl'''
       SCREAMING_SNAKE_CASE_							:				Dict		=f'{protocol}://{member_file_path}::{compressed_file_path}'
       SCREAMING_SNAKE_CASE_	,     *SCREAMING_SNAKE_CASE_							:				Dict		=fsspec.get_fs_token_paths(UpperCAmelCase_				)
       assert fs.isfile(UpperCAmelCase_				)
       assert not fs.isfile('''non_existing_''' + member_file_path				)
@pytest.mark.integration
def 							SCREAMING_SNAKE_CASE_	(					UpperCAmelCase_	:     List[str]			,							UpperCAmelCase_	:     Optional[int]			,							UpperCAmelCase_	:     Optional[int]			,							UpperCAmelCase_	:     int				)			->						int:
       SCREAMING_SNAKE_CASE_							:				Dict		=hf_api.dataset_info(UpperCAmelCase_			,							token=UpperCAmelCase_				)
       SCREAMING_SNAKE_CASE_							:				Optional[int]		=HfFileSystem(repo_info=UpperCAmelCase_			,							token=UpperCAmelCase_				)
       assert sorted(hffs.glob('''*'''				)				) == [".gitattributes", "data"]
       assert hffs.isdir('''data'''				)
       assert hffs.isfile('''.gitattributes'''				) and hffs.isfile('''data/text_data.txt'''				)
       with open(UpperCAmelCase_				) as f:
              assert hffs.open('''data/text_data.txt'''			,							'''r'''				).read() == f.read()

def 							SCREAMING_SNAKE_CASE_	(					)			->						Union[str, Any]:
       SCREAMING_SNAKE_CASE_							:				Tuple		='''bz2'''

       # Import module
       import datasets.filesystems

       # Overwrite protocol and reload
       register_implementation(UpperCAmelCase_			,							UpperCAmelCase_			,							clobber=UpperCAmelCase_				)
       with pytest.warns(UpperCAmelCase_				) as warning_info:
              importlib.reload(datasets.filesystems				)

       assert len(UpperCAmelCase_				) == 1
       assert (
           str(warning_info[0].message				)
           == f'A filesystem protocol was already set for {protocol} and will be overwritten.'
       )
