





'''simple docstring'''


import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Value
from datasets.packaged_modules.json.json import Json
@pytest.fixture
def 							UpperCAmelCase_						(		__lowerCamelCase      :      int				):
   lowercase_    :int					    =				tmp_path / "file.jsonl"
   lowercase_    :Optional[int]					    =				textwrap.dedent(
       "\\n        {\"col_1\": -1}\n        {\"col_1\": 1, \"col_2\": 2}\n        {\"col_1\": 10, \"col_2\": 20}\n        "				)
   with open(__lowerCamelCase			,"w"				) as f:
      f.write(__lowerCamelCase				)
   return str(__lowerCamelCase				)
@pytest.fixture
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Union[str, Any]				):
   lowercase_    :Optional[Any]					    =				tmp_path / "file_utf16_encoded.jsonl"
   lowercase_    :Optional[int]					    =				textwrap.dedent(
       "\\n        {\"col_1\": -1}\n        {\"col_1\": 1, \"col_2\": 2}\n        {\"col_1\": 10, \"col_2\": 20}\n        "				)
   with open(__lowerCamelCase			,"w"			,encoding="utf-16"				) as f:
      f.write(__lowerCamelCase				)
   return str(__lowerCamelCase				)
@pytest.fixture
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Optional[int]				):
   lowercase_    :List[Any]					    =				tmp_path / "file_with_list_of_dicts.json"
   lowercase_    :Any					    =				textwrap.dedent(
       "\\n        [\n            {\"col_1\": -1},\n            {\"col_1\": 1, \"col_2\": 2},\n            {\"col_1\": 10, \"col_2\": 20}\n        ]\n        "				)
   with open(__lowerCamelCase			,"w"				) as f:
      f.write(__lowerCamelCase				)
   return str(__lowerCamelCase				)
@pytest.fixture
def 							UpperCAmelCase_						(		__lowerCamelCase      :      List[str]				):
   lowercase_    :int					    =				tmp_path / "file_with_list_of_dicts_field.json"
   lowercase_    :Any					    =				textwrap.dedent(
       "\\n        {\n            \"field1\": 1,\n            \"field2\": \"aabb\",\n            \"field3\": [\n                {\"col_1\": -1},\n                {\"col_1\": 1, \"col_2\": 2},\n                {\"col_1\": 10, \"col_2\": 20}\n            ]\n        }\n        "				)
   with open(__lowerCamelCase			,"w"				) as f:
      f.write(__lowerCamelCase				)
   return str(__lowerCamelCase				)
@pytest.mark.parametrize(
    "file_fixture, config_kwargs"			,[
        ("jsonl_file", {}),
        ("jsonl_file_utf16_encoded", {"encoding": "utf-16"}),
        ("json_file_with_list_of_dicts", {}),
        ("json_file_with_list_of_dicts_field", {"field": "field3"}),
    ]			,)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      List[Any]			,__lowerCamelCase      :      Union[str, Any]			,__lowerCamelCase      :      Dict				):
   lowercase_    :Union[str, Any]					    =				Json(**__lowerCamelCase				)
   lowercase_    :Tuple					    =				json._generate_tables([[request.getfixturevalue(__lowerCamelCase				)]]				)
   lowercase_    :Union[str, Any]					    =				pa.concat_tables([table for _, table in generator]				)
   assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20]}


@pytest.mark.parametrize(
    "file_fixture, config_kwargs"			,[
        (
            "jsonl_file",
            {"features": Features({"col_1": Value("int64"				), "col_2": Value("int64"				), "missing_col": Value("string"				)}				)},
        ),
        (
            "json_file_with_list_of_dicts",
            {"features": Features({"col_1": Value("int64"				), "col_2": Value("int64"				), "missing_col": Value("string"				)}				)},
        ),
        (
            "json_file_with_list_of_dicts_field",
            {
                "field": "field3",
                "features": Features(
                    {"col_1": Value("int64"				), "col_2": Value("int64"				), "missing_col": Value("string"				)}				),
            },
        ),
    ]			,)
def 							UpperCAmelCase_						(		__lowerCamelCase      :      Dict			,__lowerCamelCase      :      str			,__lowerCamelCase      :      List[Any]				):
   lowercase_    :Tuple					    =				Json(**__lowerCamelCase				)
   lowercase_    :str					    =				json._generate_tables([[request.getfixturevalue(__lowerCamelCase				)]]				)
   lowercase_    :List[Any]					    =				pa.concat_tables([table for _, table in generator]				)
   assert pa_table.to_pydict() == {"col_1": [-1, 1, 10], "col_2": [None, 2, 20], "missing_col": [None, None, None]}









'''simple docstring'''


import multiprocessing
import os
from typing import BinaryIO, Optional, Union

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.csv.csv import Csv
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader
class    a_     (       _lowerCAmelCase							):





     def __init__(				self	:      Dict			,   lowercase	:      NestedDataStructureLike[PathLike]			,   lowercase	:      Optional[NamedSplit] = None			,   lowercase	:      Optional[Features] = None			,   lowercase	:      str = None			,   lowercase	:      bool = False			,   lowercase	:      bool = False			,   lowercase	:      Optional[int] = None			,   **lowercase	:      Any			,   ):





        """simple docstring"""




        super().__init__(
            lowercase			,   split=lowercase			,   features=lowercase			,   cache_dir=lowercase			,   keep_in_memory=lowercase			,   streaming=lowercase			,   num_proc=lowercase			,   **lowercase			,   )
        lowercase_    :Any					    =				path_or_paths if isinstance(lowercase			,   lowercase		) else {self.split: path_or_paths}
        lowercase_    :Any					    =				Csv(
            cache_dir=lowercase			,   data_files=lowercase			,   features=lowercase			,   **lowercase			,   )







     def 							lowercase__   (				self	:      Any		):





        """simple docstring"""




        if self.streaming:
           lowercase_    :Union[str, Any]					    =				self.builder.as_streaming_dataset(split=self.split		)
        # Build regular (map-style) dataset
        else:
           lowercase_    :Dict					    =				None
           lowercase_    :List[str]					    =				None
           lowercase_    :List[Any]					    =				None
           lowercase_    :Union[str, Any]					    =				None

           self.builder.download_and_prepare(
               download_config=lowercase			,   download_mode=lowercase			,   verification_mode=lowercase			,   base_path=lowercase			,   num_proc=self.num_proc			,   )
           lowercase_    :int					    =				self.builder.as_dataset(
               split=self.split			,   verification_mode=lowercase			,   in_memory=self.keep_in_memory		)
        return dataset







class    a_     :





     def __init__(				self	:      Dict			,   lowercase	:      Dataset			,   lowercase	:      Union[PathLike, BinaryIO]			,   lowercase	:      Optional[int] = None			,   lowercase	:      Optional[int] = None			,   **lowercase	:      Any			,   ):





        """simple docstring"""




        if num_proc is not None and num_proc <= 0:
           raise ValueError(F'num_proc {num_proc} must be an integer > 0.'		)

        lowercase_    :str					    =				dataset
        lowercase_    :List[Any]					    =				path_or_buf
        lowercase_    :Any					    =				batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        lowercase_    :int					    =				num_proc
        lowercase_    :Optional[Any]					    =				"utf-8"
        lowercase_    :Dict					    =				to_csv_kwargs





     def 							lowercase__   (				self	:      Union[str, Any]		):





        """simple docstring"""




        lowercase_    :int					    =				self.to_csv_kwargs.pop("path_or_buf"			,   lowercase		)
        lowercase_    :List[str]					    =				self.to_csv_kwargs.pop("header"			,   lowercase		)
        lowercase_    :Optional[Any]					    =				self.to_csv_kwargs.pop("index"			,   lowercase		)

        if isinstance(self.path_or_buf			,   (str, bytes, os.PathLike)		):
           with open(self.path_or_buf			,   "wb+"		) as buffer:
              lowercase_    :List[Any]					    =				self._write(file_obj=lowercase			,   header=lowercase			,   index=lowercase			,   **self.to_csv_kwargs		)
        else:
           lowercase_    :Any					    =				self._write(file_obj=self.path_or_buf			,   header=lowercase			,   index=lowercase			,   **self.to_csv_kwargs		)
        return written





     def 							lowercase__   (				self	:      List[Any]			,   lowercase	:      Union[str, Any]		):





        """simple docstring"""




        lowercase_    ,							lowercase_    ,							lowercase_    ,							lowercase_    :Optional[Any]					    =				args

        lowercase_    :Dict					    =				query_table(
            table=self.dataset.data			,   key=slice(lowercase			,   offset + self.batch_size		)			,   indices=self.dataset._indices			,   )
        lowercase_    :List[str]					    =				batch.to_pandas().to_csv(
            path_or_buf=lowercase			,   header=header if (offset == 0) else False			,   index=lowercase			,   **lowercase		)
        return csv_str.encode(self.encoding		)







     def 							lowercase__   (				self	:      Dict			,   lowercase	:      BinaryIO			,   lowercase	:      Dict			,   lowercase	:      int			,   **lowercase	:      Any		):





        """simple docstring"""




        lowercase_    :Dict					    =				0

        if self.num_proc is None or self.num_proc == 1:
           for offset in logging.tqdm(
               range(0			,   len(self.dataset		)			,   self.batch_size		)			,   unit="ba"			,   disable=not logging.is_progress_bar_enabled()			,   desc="Creating CSV from Arrow format"			,   ):
              lowercase_    :Optional[Any]					    =				self._batch_csv((offset, header, index, to_csv_kwargs)		)
              written += file_obj.write(lowercase		)

        else:
           lowercase_    ,							lowercase_    :Optional[int]					    =				len(self.dataset		), self.batch_size
           with multiprocessing.Pool(self.num_proc		) as pool:
              for csv_str in logging.tqdm(
                  pool.imap(
                      self._batch_csv			,   [(offset, header, index, to_csv_kwargs) for offset in range(0			,   lowercase			,   lowercase		)]			,   )			,   total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size			,   unit="ba"			,   disable=not logging.is_progress_bar_enabled()			,   desc="Creating CSV from Arrow format"			,   ):
                 written += file_obj.write(lowercase		)

        return written


