




"""simple docstring"""


import enum
import os
from hashlib import shaaaa
from typing import Optional

from .. import config
from .logging import get_logger


A_      			=       get_logger(__name__)






class       lowercase(  enum.Enum				):







 '''simple docstring'''





 lowercase__				      =      "all_checks"
 lowercase__				      =      "basic_checks"
 lowercase__				      =      "no_checks"







class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












def       UpperCAmelCase__ (snake_case__							:      Optional[dict]	,     snake_case__							:      dict	,     snake_case__							:      Optional[Any]=None					):
  """simple docstring"""

  if expected_checksums is None:
    logger.info("""Unable to verify checksums."""					)
    return
  if len(set(snake_case__					) - set(snake_case__					)					) > 0:
    raise ExpectedMoreDownloadedFiles(str(set(snake_case__					) - set(snake_case__					)					)					)
  if len(set(snake_case__					) - set(snake_case__					)					) > 0:
    raise UnexpectedDownloadedFile(str(set(snake_case__					) - set(snake_case__					)					)					)
  _snake_case							:	Dict			=						[url for url in expected_checksums if expected_checksums[url] != recorded_checksums[url]]
  _snake_case							:	Any			=						""" for """ + verification_name if verification_name is not None else """"""
  if len(snake_case__					) > 0:
    raise NonMatchingChecksumError(
        F"Checksums didn't match{for_verification_name}:\n"
        F"{bad_urls}\n"
        """Set `verification_mode='no_checks'` to skip checksums verification and ignore this error"""					)
  logger.info("""All the checksums matched successfully""" + for_verification_name					)






class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












class       lowercase(  __a				):







 '''simple docstring'''












def       UpperCAmelCase__ (snake_case__							:      Optional[dict]	,     snake_case__							:      dict					):
  """simple docstring"""

  if expected_splits is None:
    logger.info("""Unable to verify splits sizes."""					)
    return
  if len(set(snake_case__					) - set(snake_case__					)					) > 0:
    raise ExpectedMoreSplits(str(set(snake_case__					) - set(snake_case__					)					)					)
  if len(set(snake_case__					) - set(snake_case__					)					) > 0:
    raise UnexpectedSplits(str(set(snake_case__					) - set(snake_case__					)					)					)
  _snake_case							:	Optional[Any]			=						[
      {"""expected""": expected_splits[name], """recorded""": recorded_splits[name]}
      for name in expected_splits
      if expected_splits[name].num_examples != recorded_splits[name].num_examples
  ]
  if len(snake_case__					) > 0:
    raise NonMatchingSplitsSizesError(str(snake_case__					)					)
  logger.info("""All the splits matched successfully."""					)


def       UpperCAmelCase__ (snake_case__							:      str	,     snake_case__							:      bool = True					):
  """simple docstring"""

  if record_checksum:
    _snake_case							:	int			=						shaaaa()
    with open(snake_case__	,     """rb"""					) as f:
      for chunk in iter(lambda: f.read(1 << 20					)	,     B""""""					):
        m.update(snake_case__					)
      _snake_case							:	List[str]			=						m.hexdigest()
  else:
    _snake_case							:	List[Any]			=						None
  return {"num_bytes": os.path.getsize(snake_case__					), "checksum": checksum}






def       UpperCAmelCase__ (snake_case__							:      Optional[int]					):
  """simple docstring"""

  if dataset_size and config.IN_MEMORY_MAX_SIZE:
    return dataset_size < config.IN_MEMORY_MAX_SIZE
  else:
    return False








"""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       UpperCAmelCase__ (snake_case__							:      Optional[Any]	,     snake_case__							:      Dict					):
  """simple docstring"""

  assert isinstance(snake_case__	,     snake_case__					)
  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       UpperCAmelCase__ (snake_case__							:      str	,     snake_case__							:      Optional[int]	,     snake_case__							:      Dict					):
  """simple docstring"""

  _snake_case							:	str			=						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							:	str			=						ParquetDatasetReader(snake_case__	,     cache_dir=snake_case__	,     keep_in_memory=snake_case__					).read()
  _check_parquet_dataset(snake_case__	,     snake_case__					)


@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       UpperCAmelCase__ (snake_case__							:      Dict	,     snake_case__							:      int	,     snake_case__							:      List[Any]					):
  """simple docstring"""

  _snake_case							:	str			=						tmp_path / """cache"""
  _snake_case							:	List[Any]			=						{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
  _snake_case							:	List[Any]			=						features.copy() if features else default_expected_features
  _snake_case							:	List[Any]			=						(
      Features({feature: Value(snake_case__					) for feature, dtype in features.items()}					) if features is not None else None
  )
  _snake_case							:	Optional[Any]			=						ParquetDatasetReader(snake_case__	,     features=snake_case__	,     cache_dir=snake_case__					).read()
  _check_parquet_dataset(snake_case__	,     snake_case__					)


@pytest.mark.parametrize("""split"""	,     [None, NamedSplit("""train"""					), """train""", """test"""]					)
def       UpperCAmelCase__ (snake_case__							:      Optional[int]	,     snake_case__							:      int	,     snake_case__							:      int					):
  """simple docstring"""

  _snake_case							:	List[str]			=						tmp_path / """cache"""
  _snake_case							:	Any			=						{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
  _snake_case							:	int			=						ParquetDatasetReader(snake_case__	,     cache_dir=snake_case__	,     split=snake_case__					).read()
  _check_parquet_dataset(snake_case__	,     snake_case__					)
  assert dataset.split == split if split else "train"


@pytest.mark.parametrize("""path_type"""	,     [str, list]					)
def       UpperCAmelCase__ (snake_case__							:      Dict	,     snake_case__							:      str	,     snake_case__							:      str					):
  """simple docstring"""

  if issubclass(snake_case__	,     snake_case__					):
    _snake_case							:	Optional[Any]			=						parquet_path
  elif issubclass(snake_case__	,     snake_case__					):
    _snake_case							:	int			=						[parquet_path]
  _snake_case							:	Union[str, Any]			=						tmp_path / """cache"""
  _snake_case							:	Tuple			=						{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
  _snake_case							:	List[str]			=						ParquetDatasetReader(snake_case__	,     cache_dir=snake_case__					).read()
  _check_parquet_dataset(snake_case__	,     snake_case__					)


def       UpperCAmelCase__ (snake_case__							:      str	,     snake_case__							:      Optional[int]	,     snake_case__							:      str=("train",)					):
  """simple docstring"""

  assert isinstance(snake_case__	,     snake_case__					)
  for split in splits:
    _snake_case							:	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       UpperCAmelCase__ (snake_case__							:      str	,     snake_case__							:      str	,     snake_case__							:      List[Any]					):
  """simple docstring"""

  _snake_case							:	Tuple			=						tmp_path / """cache"""
  _snake_case							:	Optional[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							:	Tuple			=						ParquetDatasetReader(
        {"""train""": parquet_path}	,     cache_dir=snake_case__	,     keep_in_memory=snake_case__					).read()
  _check_parquet_datasetdict(snake_case__	,     snake_case__					)


@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       UpperCAmelCase__ (snake_case__							:      List[str]	,     snake_case__							:      Tuple	,     snake_case__							:      List[Any]					):
  """simple docstring"""

  _snake_case							:	Optional[int]			=						tmp_path / """cache"""
  _snake_case							:	Dict			=						{"""col_1""": """string""", """col_2""": """int64""", """col_3""": """float64"""}
  _snake_case							:	Optional[Any]			=						features.copy() if features else default_expected_features
  _snake_case							:	Dict			=						(
      Features({feature: Value(snake_case__					) for feature, dtype in features.items()}					) if features is not None else None
  )
  _snake_case							:	Optional[int]			=						ParquetDatasetReader({"""train""": parquet_path}	,     features=snake_case__	,     cache_dir=snake_case__					).read()
  _check_parquet_datasetdict(snake_case__	,     snake_case__					)


@pytest.mark.parametrize("""split"""	,     [None, NamedSplit("""train"""					), """train""", """test"""]					)
def       UpperCAmelCase__ (snake_case__							:      List[str]	,     snake_case__							:      Optional[Any]	,     snake_case__							:      Tuple					):
  """simple docstring"""

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


def       UpperCAmelCase__ (snake_case__							:      Tuple	,     snake_case__							:      Tuple					):
  """simple docstring"""

  _snake_case							:	List[Any]			=						ParquetDatasetWriter(snake_case__	,     tmp_path / """foo.parquet"""					)
  assert writer.write() > 0
  _snake_case							:	str			=						pq.ParquetFile(tmp_path / """foo.parquet"""					)
  _snake_case							:	int			=						pf.read()
  assert dataset.data.table == output_table


def       UpperCAmelCase__ (snake_case__							:      int	,     snake_case__							:      int					):
  """simple docstring"""

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

  _snake_case							:	List[str]			=						Dataset.from_parquet(str(tmp_path / """foo.parquet"""					)					)
  assert dataset.features == reloaded_dataset.features

  _snake_case							:	Optional[Any]			=						ParquetDatasetReader(str(tmp_path / """foo.parquet"""					)	,     streaming=snake_case__					).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       UpperCAmelCase__ (snake_case__							:      List[Any]	,     snake_case__							:      str					):
  """simple docstring"""

  assert get_writer_batch_size(snake_case__					) == expected


