



"""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_  (			_lowercase,					_lowercase	):







						'''simple docstring'''


						assert isinstance(lowercase_,					lowercase_	)
						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_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :str						=			tmp_path / """cache"""
						snake_case_ :Optional[Any]						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												snake_case_ :Union[str, Any]						=			ParquetDatasetReader(lowercase_,					cache_dir=lowercase_,					keep_in_memory=lowercase_	).read()
						_check_parquet_dataset(lowercase_,					lowercase_	)
@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_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :Any						=			tmp_path / """cache"""
						snake_case_ :List[str]						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						snake_case_ :List[Any]						=			features.copy() if features else default_expected_features
						snake_case_ :Optional[Any]						=			(
						    Features({feature: Value(lowercase_	) for feature, dtype in features.items()}	) if features is not None else None
						)
						snake_case_ :Tuple						=			ParquetDatasetReader(lowercase_,					features=lowercase_,					cache_dir=lowercase_	).read()
						_check_parquet_dataset(lowercase_,					lowercase_	)
@pytest.mark.parametrize("""split""",					[None, NamedSplit("""train"""	), """train""", """test"""]	)
def    A_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :List[str]						=			tmp_path / """cache"""
						snake_case_ :Optional[int]						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						snake_case_ :Dict						=			ParquetDatasetReader(lowercase_,					cache_dir=lowercase_,					split=lowercase_	).read()
						_check_parquet_dataset(lowercase_,					lowercase_	)
						assert dataset.split == split if split else "train"
@pytest.mark.parametrize("""path_type""",					[str, list]	)
def    A_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						if issubclass(lowercase_,					lowercase_	):
												snake_case_ :Tuple						=			parquet_path
						elif issubclass(lowercase_,					lowercase_	):
												snake_case_ :int						=			[parquet_path]
						snake_case_ :Optional[Any]						=			tmp_path / """cache"""
						snake_case_ :str						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						snake_case_ :Optional[Any]						=			ParquetDatasetReader(lowercase_,					cache_dir=lowercase_	).read()
						_check_parquet_dataset(lowercase_,					lowercase_	)
def    A_  (			_lowercase,					_lowercase,					_lowercase=("train",)	):







						'''simple docstring'''


						assert isinstance(lowercase_,					lowercase_	)
						for split in splits:
												snake_case_ :Tuple						=			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_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :Dict						=			tmp_path / """cache"""
						snake_case_ :int						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
												snake_case_ :List[Any]						=			ParquetDatasetReader(
												    {"""train""": parquet_path},					cache_dir=lowercase_,					keep_in_memory=lowercase_	).read()
						_check_parquet_datasetdict(lowercase_,					lowercase_	)
@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_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :Dict						=			tmp_path / """cache"""
						snake_case_ :List[Any]						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						snake_case_ :Optional[Any]						=			features.copy() if features else default_expected_features
						snake_case_ :List[Any]						=			(
						    Features({feature: Value(lowercase_	) for feature, dtype in features.items()}	) if features is not None else None
						)
						snake_case_ :List[Any]						=			ParquetDatasetReader({"""train""": parquet_path},					features=lowercase_,					cache_dir=lowercase_	).read()
						_check_parquet_datasetdict(lowercase_,					lowercase_	)
@pytest.mark.parametrize("""split""",					[None, NamedSplit("""train"""	), """train""", """test"""]	)
def    A_  (			_lowercase,					_lowercase,					_lowercase	):







						'''simple docstring'''


						if split:
												snake_case_ :int						=			{split: parquet_path}
						else:
												snake_case_ :List[str]						=			"""train"""
												snake_case_ :Dict						=			{"""train""": parquet_path, """test""": parquet_path}
						snake_case_ :Optional[Any]						=			tmp_path / """cache"""
						snake_case_ :List[str]						=			{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
						snake_case_ :Dict						=			ParquetDatasetReader(lowercase_,					cache_dir=lowercase_	).read()
						_check_parquet_datasetdict(lowercase_,					lowercase_,					splits=list(path.keys()	)	)
						assert all(dataset[split].split == split for split in path.keys()	)
def    A_  (			_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :Tuple						=			ParquetDatasetWriter(lowercase_,					tmp_path / """foo.parquet"""	)
						assert writer.write() > 0
						snake_case_ :int						=			pq.ParquetFile(tmp_path / """foo.parquet"""	)
						snake_case_ :Optional[Any]						=			pf.read()
						assert dataset.data.table == output_table
def    A_  (			_lowercase,					_lowercase	):







						'''simple docstring'''


						snake_case_ :Union[str, Any]						=			str(shared_datadir / """test_image_rgb.jpg"""	)
						snake_case_ :Any						=			{"""image""": [image_path]}
						snake_case_ :Optional[Any]						=			Features({"""image""": Image()}	)
						snake_case_ :str						=			Dataset.from_dict(lowercase_,					features=lowercase_	)
						snake_case_ :str						=			ParquetDatasetWriter(lowercase_,					tmp_path / """foo.parquet"""	)
						assert writer.write() > 0

						snake_case_ :int						=			Dataset.from_parquet(str(tmp_path / """foo.parquet"""	)	)
						assert dataset.features == reloaded_dataset.features

						snake_case_ :List[Any]						=			ParquetDatasetReader(str(tmp_path / """foo.parquet"""	),					streaming=lowercase_	).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_  (			_lowercase,					_lowercase	):







						'''simple docstring'''


						assert get_writer_batch_size(lowercase_	) == expected



import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict

@pytest.mark.parametrize(
    """files"""    ,      [
        ["""full:README.md""", """dataset_infos.json"""],
        ["""empty:README.md""", """dataset_infos.json"""],
        ["""dataset_infos.json"""],
        ["""full:README.md"""],
    ]    ,      )
def 						UpperCamelCase					(lowercase_: List[str]    ,      lowercase_: Optional[int]					)     ->      List[str]:
	A__							:    Tuple        =		tmp_path_factory.mktemp("""dset_infos_dir"""					)
	if "full:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write("""---\ndataset_info:\n  dataset_size: 42\n---"""					)
	if "empty:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write(""""""					)
    # we want to support dataset_infos.json for backward compatibility
	if "dataset_infos.json" in files:
		with open(dataset_infos_dir / """dataset_infos.json"""    ,      """w"""					) as f:
			f.write("""{\"default\": {\"dataset_size\": 42}}"""					)
	A__							:    List[Any]        =		DatasetInfosDict.from_directory(lowercase_					)
	assert dataset_infos
	assert dataset_infos["default"].dataset_size == 42

@pytest.mark.parametrize(
    """dataset_info"""    ,      [
        DatasetInfo(),
        DatasetInfo(
            description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      ),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: DatasetInfo					)     ->      List[Any]:
	A__							:    Union[str, Any]        =		str(lowercase_					)
	dataset_info.write_to_directory(lowercase_					)
	A__							:    List[Any]        =		DatasetInfo.from_directory(lowercase_					)
	assert dataset_info == reloaded
	assert os.path.exists(os.path.join(lowercase_    ,      """dataset_info.json"""					)					)

def 						UpperCamelCase					()     ->      List[Any]:
	A__							:    Union[str, Any]        =		DatasetInfo(
	    description="""foo"""    ,      citation="""bar"""    ,      homepage="""https://foo.bar"""    ,      license="""CC0"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      post_processed={}    ,      supervised_keys=()    ,      task_templates=[]    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train""", """num_examples""": 42}]    ,      download_checksums={}    ,      download_size=1337    ,      post_processing_size=442    ,      dataset_size=1234    ,      size_in_bytes=1337 + 442 + 1234    ,      )
	A__							:    Dict        =		dataset_info._to_yaml_dict()
	assert sorted(lowercase_					) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML					)
	for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
		assert key in dataset_info_yaml_dict
		assert isinstance(dataset_info_yaml_dict[key]    ,      (list, dict, int, str)					)
	A__							:    Union[str, Any]        =		yaml.safe_dump(lowercase_					)
	A__							:    List[Any]        =		yaml.safe_load(lowercase_					)
	assert dataset_info_yaml_dict == reloaded

def 						UpperCamelCase					()     ->      List[str]:
	A__							:    Optional[int]        =		DatasetInfo()
	A__							:    List[Any]        =		dataset_info._to_yaml_dict()
	assert dataset_info_yaml_dict == {}


@pytest.mark.parametrize(
    """dataset_infos_dict"""    ,      [
        DatasetInfosDict(),
        DatasetInfosDict({"""default""": DatasetInfo()}					),
        DatasetInfosDict({"""my_config_name""": DatasetInfo()}					),
        DatasetInfosDict(
            {
                """default""": DatasetInfo(
                    description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      )
            }					),
        DatasetInfosDict(
            {
                """v1""": DatasetInfo(dataset_size=42					),
                """v2""": DatasetInfo(dataset_size=1337					),
            }					),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: Tuple    ,      lowercase_: DatasetInfosDict					)     ->      Optional[Any]:
	A__							:    List[Any]        =		str(lowercase_					)
	dataset_infos_dict.write_to_directory(lowercase_					)
	A__							:    Dict        =		DatasetInfosDict.from_directory(lowercase_					)

	# the config_name of the dataset_infos_dict take over the attribute
	for config_name, dataset_info in dataset_infos_dict.items():
		A__							:    Optional[int]        =		config_name
		# the yaml representation doesn't include fields like description or citation
		# so we just test that we can recover what we can from the yaml
		A__							:    List[str]        =		DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()					)
	assert dataset_infos_dict == reloaded

	if dataset_infos_dict:
		assert os.path.exists(os.path.join(lowercase_    ,      """README.md"""					)					)



