




'''simple docstring'''

import contextlib
import os
import sqlitea

import pytest

from datasets import Dataset, Features, Value
from datasets.io.sql import SqlDatasetReader, SqlDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases, require_sqlalchemy


def __UpperCAmelCase	(   _UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Optional[int]						)      ->      List[str]:
			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


@require_sqlalchemy
@pytest.mark.parametrize("keep_in_memory"		,				[False, True]						)
def __UpperCAmelCase	(   _UpperCAmelCase		:   int		,				_UpperCAmelCase		:   Any		,				_UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Optional[Any]						)      ->      Optional[int]:
			__snake_case			=		tmp_path / "cache"
			__snake_case			=		{"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			=		SqlDatasetReader(
						    "dataset"		,				"sqlite:///" + sqlite_path		,				cache_dir=_UpperCAmelCase		,				keep_in_memory=_UpperCAmelCase						).read()
			_check_sql_dataset(_UpperCAmelCase		,				_UpperCAmelCase						)


@require_sqlalchemy
@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	(   _UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Any		,				_UpperCAmelCase		:   Tuple		,				_UpperCAmelCase		:   List[Any]						)      ->      Optional[Any]:
			__snake_case			=		tmp_path / "cache"
			__snake_case			=		{"col_1": "string", "col_2": "int64", "col_3": "float64"}
			__snake_case			=		features.copy() if features else default_expected_features
			__snake_case			=		(
			    Features({feature: Value(_UpperCAmelCase						) for feature, dtype in features.items()}						) if features is not None else None
			)
			__snake_case			=		SqlDatasetReader("dataset"		,				"sqlite:///" + sqlite_path		,				features=_UpperCAmelCase		,				cache_dir=_UpperCAmelCase						).read()
			_check_sql_dataset(_UpperCAmelCase		,				_UpperCAmelCase						)


def __UpperCAmelCase	(   _UpperCAmelCase		:   str						)      ->      Tuple:
			with contextlib.closing(sqlitea.connect(_UpperCAmelCase						)						) as con:
						__snake_case			=		con.cursor()
						cur.execute("SELECT * FROM dataset"						)
						for row in cur:
									yield row


@require_sqlalchemy
def __UpperCAmelCase	(   _UpperCAmelCase		:   Dict		,				_UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Optional[Any]						)      ->      List[str]:
			__snake_case			=		tmp_path / "cache"
			__snake_case			=		os.path.join(_UpperCAmelCase		,				"tmp.sql"						)
			__snake_case			=		SqlDatasetReader("dataset"		,				"sqlite:///" + sqlite_path		,				cache_dir=_UpperCAmelCase						).read()
			SqlDatasetWriter(_UpperCAmelCase		,				"dataset"		,				"sqlite:///" + output_sqlite_path		,				num_proc=1						).write()

			__snake_case			=		iter_sql_file(_UpperCAmelCase						)
			__snake_case			=		iter_sql_file(_UpperCAmelCase						)

			for rowa, rowa in zip(_UpperCAmelCase		,				_UpperCAmelCase						):
						assert rowa == rowa


@require_sqlalchemy
def __UpperCAmelCase	(   _UpperCAmelCase		:   Dict		,				_UpperCAmelCase		:   Dict		,				_UpperCAmelCase		:   List[Any]						)      ->      str:
			__snake_case			=		tmp_path / "cache"
			__snake_case			=		os.path.join(_UpperCAmelCase		,				"tmp.sql"						)
			__snake_case			=		SqlDatasetReader("dataset"		,				"sqlite:///" + sqlite_path		,				cache_dir=_UpperCAmelCase						).read()
			SqlDatasetWriter(_UpperCAmelCase		,				"dataset"		,				"sqlite:///" + output_sqlite_path		,				num_proc=2						).write()

			__snake_case			=		iter_sql_file(_UpperCAmelCase						)
			__snake_case			=		iter_sql_file(_UpperCAmelCase						)

			for rowa, rowa in zip(_UpperCAmelCase		,				_UpperCAmelCase						):
						assert rowa == rowa






@require_sqlalchemy
def __UpperCAmelCase	(   _UpperCAmelCase		:   List[str]		,				_UpperCAmelCase		:   Union[str, Any]		,				_UpperCAmelCase		:   Union[str, Any]						)      ->      Any:
			__snake_case			=		tmp_path / "cache"
			__snake_case			=		os.path.join(_UpperCAmelCase		,				"tmp.sql"						)
			__snake_case			=		SqlDatasetReader("dataset"		,				"sqlite:///" + sqlite_path		,				cache_dir=_UpperCAmelCase						).read()
			with pytest.raises(_UpperCAmelCase						):
						SqlDatasetWriter(_UpperCAmelCase		,				"dataset"		,				"sqlite:///" + output_sqlite_path		,				num_proc=0						).write()











'''simple docstring'''


def 						lowerCamelCase__   (				a							):
						if number < 0:
												raise ValueError('number must not be negative'							)
						return number & (number - 1) == 0


if __name__ == "__main__":
	import doctest

	doctest.testmod()

