


import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


if TYPE_CHECKING:
						import sqlitea

						import sqlalchemy


_UpperCAmelCase      :	Tuple  =datasets.utils.logging.get_logger(__name__)







@dataclass
class 					snake_case__(				datasets.BuilderConfig				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	Union[str, "sqlalchemy.sql.Selectable"]       =				None
						SCREAMING_SNAKE_CASE__   :	Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]       =				None
						SCREAMING_SNAKE_CASE__   :	Optional[Union[str, List[str]]]       =				None
						SCREAMING_SNAKE_CASE__   :	bool       =				True
						SCREAMING_SNAKE_CASE__   :	Optional[Union[List, Tuple, Dict]]       =				None
						SCREAMING_SNAKE_CASE__   :	Optional[Union[List, Dict]]       =				None
						SCREAMING_SNAKE_CASE__   :	Optional[List[str]]       =				None
						SCREAMING_SNAKE_CASE__   :	Optional[int]       =				10_000
						SCREAMING_SNAKE_CASE__   :	Optional[datasets.Features]       =				None




						def     lowercase_			(					self				)       ->     str:


										if self.sql is None:
														raise ValueError('''sql must be specified'''				)
										if self.con is None:
														raise ValueError('''con must be specified'''				)




						def     lowercase_			(					self ,    __lowercase ,    __lowercase = None ,    )       ->     List[Any]:
										lowerCAmelCase_			:   Any     					=					config_kwargs.copy()
										# We need to stringify the Selectable object to make its hash deterministic

										# The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
										lowerCAmelCase_			:   Tuple     					=					config_kwargs['''sql''']
										if not isinstance(__UpperCamelCase ,    __UpperCamelCase				):
														if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
																		import sqlalchemy

																		if isinstance(__UpperCamelCase ,    sqlalchemy.sql.Selectable				):
																						lowerCAmelCase_			:   Optional[Any]     					=					sqlalchemy.create_engine(config_kwargs['''con'''].split('''://'''				)[0] + '''://'''				)
																						lowerCAmelCase_			:   Tuple     					=					str(sql.compile(dialect=engine.dialect				)				)
																						lowerCAmelCase_			:   List[str]     					=					sql_str
																		else:
																						raise TypeError(
																						    f"""Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(__UpperCamelCase				)}: {sql}"""				)
														else:
																		raise TypeError(
																		    f"""Supported types for \'sql\' are string and sqlalchemy.sql.Selectable but got {type(__UpperCamelCase				)}: {sql}"""				)
										lowerCAmelCase_			:   int     					=					config_kwargs['''con''']
										if not isinstance(__UpperCamelCase ,    __UpperCamelCase				):
														lowerCAmelCase_			:   List[str]     					=					id(__UpperCamelCase				)
														logger.info(
														    f"""SQL connection \'con\' of type {type(__UpperCamelCase				)} couldn\'t be hashed properly. To enable hashing, specify \'con\' as URI string instead."""				)

										return super().create_config_id(__UpperCamelCase ,    custom_features=__UpperCamelCase				)




						@property
						def     lowercase_			(					self				)       ->     Any:
										lowerCAmelCase_			:   int     					=					{
										    '''index_col''': self.index_col,
										    '''columns''': self.columns,
										    '''params''': self.params,
										    '''coerce_float''': self.coerce_float,
										    '''parse_dates''': self.parse_dates,
										}
										return pd_read_sql_kwargs







class 					snake_case__(				datasets.ArrowBasedBuilder				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	List[Any]       =				SqlConfig




						def     lowercase_			(					self				)       ->     Any:
										return datasets.DatasetInfo(features=self.config.features				)




						def     lowercase_			(					self ,    __lowercase				)       ->     Dict:
										return [datasets.SplitGenerator(name=datasets.Split.TRAIN ,    gen_kwargs={}				)]




						def     lowercase_			(					self ,    __lowercase				)       ->     List[str]:


										if self.config.features is not None:
														lowerCAmelCase_			:   str     					=					self.config.features.arrow_schema
														if all(not require_storage_cast(__UpperCamelCase				) for feature in self.config.features.values()				):
																		# cheaper cast
																		lowerCAmelCase_			:   List[str]     					=					pa.Table.from_arrays([pa_table[field.name] for field in schema] ,    schema=__UpperCamelCase				)
														else:
																		# more expensive cast; allows str <-> int/float or str to Audio for example
																		lowerCAmelCase_			:   Optional[int]     					=					table_cast(__UpperCamelCase ,    __UpperCamelCase				)
										return pa_table




						def     lowercase_			(					self				)       ->     Tuple:
										lowerCAmelCase_			:   Tuple     					=					self.config.chunksize
										lowerCAmelCase_			:   Union[str, Any]     					=					pd.read_sql(
										    self.config.sql ,    self.config.con ,    chunksize=__UpperCamelCase ,    **self.config.pd_read_sql_kwargs				)
										lowerCAmelCase_			:   Optional[int]     					=					[sql_reader] if chunksize is None else sql_reader
										for chunk_idx, df in enumerate(__UpperCamelCase				):
														lowerCAmelCase_			:   Optional[Any]     					=					pa.Table.from_pandas(__UpperCamelCase				)
														yield chunk_idx, self._cast_table(__UpperCamelCase				)

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 		lowercase_							(							__SCREAMING_SNAKE_CASE					):

 def __init__(   self  ,						__UpperCamelCase  ,						__UpperCamelCase = None  ,						__UpperCamelCase = None  ,						__UpperCamelCase = None  ,						__UpperCamelCase = False  ,						__UpperCamelCase = False  ,						__UpperCamelCase = None  ,						**__UpperCamelCase  ,						):




        """simple docstring"""


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

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        if self.streaming:
               UpperCamelCase_											=							self.builder.as_streaming_dataset(split=self.split  )
        # Build regular (map-style) dataset
        else:
               UpperCamelCase_											=							None
               UpperCamelCase_											=							None
               UpperCamelCase_											=							None
               UpperCamelCase_											=							None

               self.builder.download_and_prepare(
                   download_config=__UpperCamelCase  ,						download_mode=__UpperCamelCase  ,						verification_mode=__UpperCamelCase  ,						base_path=__UpperCamelCase  ,						num_proc=self.num_proc  ,						)
               UpperCamelCase_											=							self.builder.as_dataset(
                   split=self.split  ,						verification_mode=__UpperCamelCase  ,						in_memory=self.keep_in_memory  )
        return dataset



class 		lowercase_							:

 def __init__(   self  ,						__UpperCamelCase  ,						__UpperCamelCase  ,						__UpperCamelCase = None  ,						__UpperCamelCase = None  ,						**__UpperCamelCase  ,						):




        """simple docstring"""


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

        UpperCamelCase_											=							dataset
        UpperCamelCase_											=							path_or_buf
        UpperCamelCase_											=							batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
        UpperCamelCase_											=							num_proc
        UpperCamelCase_											=							"""utf-8"""
        UpperCamelCase_											=							to_csv_kwargs

 def lowerCamelCase_					(   self  ):




        """simple docstring"""


        UpperCamelCase_											=							self.to_csv_kwargs.pop("""path_or_buf"""  ,						__UpperCamelCase  )
        UpperCamelCase_											=							self.to_csv_kwargs.pop("""header"""  ,						__UpperCamelCase  )
        UpperCamelCase_											=							self.to_csv_kwargs.pop("""index"""  ,						__UpperCamelCase  )

        if isinstance(self.path_or_buf  ,						(str, bytes, os.PathLike)  ):
               with open(self.path_or_buf  ,						"""wb+"""  ) as buffer:
                      UpperCamelCase_											=							self._write(file_obj=__UpperCamelCase  ,						header=__UpperCamelCase  ,						index=__UpperCamelCase  ,						**self.to_csv_kwargs  )
        else:
               UpperCamelCase_											=							self._write(file_obj=self.path_or_buf  ,						header=__UpperCamelCase  ,						index=__UpperCamelCase  ,						**self.to_csv_kwargs  )
        return written

 def lowerCamelCase_					(   self  ,						__UpperCamelCase  ):




        """simple docstring"""


        UpperCamelCase_							,					UpperCamelCase_							,					UpperCamelCase_							,					UpperCamelCase_											=							args

        UpperCamelCase_											=							query_table(
            table=self.dataset.data  ,						key=slice(__UpperCamelCase  ,						offset + self.batch_size  )  ,						indices=self.dataset._indices  ,						)
        UpperCamelCase_											=							batch.to_pandas().to_csv(
            path_or_buf=__UpperCamelCase  ,						header=header if (offset == 0) else False  ,						index=__UpperCamelCase  ,						**__UpperCamelCase  )
        return csv_str.encode(self.encoding  )

 def lowerCamelCase_					(   self  ,						__UpperCamelCase  ,						__UpperCamelCase  ,						__UpperCamelCase  ,						**__UpperCamelCase  ):




        """simple docstring"""


        UpperCamelCase_											=							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"""  ,						):
                      UpperCamelCase_											=							self._batch_csv((offset, header, index, to_csv_kwargs)  )
                      written += file_obj.write(__UpperCamelCase  )

        else:
               UpperCamelCase_							,					UpperCamelCase_											=							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  ,						__UpperCamelCase  ,						__UpperCamelCase  )]  ,						)  ,						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(__UpperCamelCase  )

        return written



