



'''simple docstring'''



from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)



def   __lowerCAmelCase					(							lowerCamelCase   :						Dict				,		lowerCamelCase   :						Tuple					):






 '''simple docstring'''




 __lowerCAmelCase												=  []
 for part_id in partition_order:
  __lowerCAmelCase												=  df.where(f'''SPARK_PARTITION_ID() = {part_id}'''					).collect()
  for row_idx, row in enumerate(lowerCamelCase					):
   expected_row_ids_and_row_dicts.append((f'''{part_id}_{row_idx}''', row.asDict())					)
 return expected_row_ids_and_row_dicts



@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(1_00					).repartition(1					)
 __lowerCAmelCase												=  Spark(lowerCamelCase					)
 # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
 # that each partition can hold 2 rows.
 spark_builder._repartition_df_if_needed(max_shard_size=16					)
 # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
 assert spark_builder.df.rdd.getNumPartitions() == 50



@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(10					).repartition(2					)
 __lowerCAmelCase												=  [1, 0]
 __lowerCAmelCase												=  _generate_iterable_examples(lowerCamelCase				,		lowerCamelCase					)  # Reverse the partitions.
 __lowerCAmelCase												=  _get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase				,		lowerCamelCase					)

 for i, (row_id, row_dict) in enumerate(generate_fn()					):
  __lowerCAmelCase							,       __lowerCAmelCase												=  expected_row_ids_and_row_dicts[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(10					).repartition(1					)
 __lowerCAmelCase												=  SparkExamplesIterable(lowerCamelCase					)
 assert it.n_shards == 1
 for i, (row_id, row_dict) in enumerate(lowerCamelCase					):
  assert row_id == f'''0_{i}'''
  assert row_dict == {"id": i}



@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(30					).repartition(3					)
 # Mock the generator so that shuffle reverses the partition indices.
 with patch("numpy.random.Generator"					) as generator_mock:
  __lowerCAmelCase												=  lambda lowerCamelCase					: x.reverse()
  __lowerCAmelCase												=  _get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase				,		[2, 1, 0]					)

  __lowerCAmelCase												=  SparkExamplesIterable(lowerCamelCase					).shuffle_data_sources(lowerCamelCase					)
  assert shuffled_it.n_shards == 3
  for i, (row_id, row_dict) in enumerate(lowerCamelCase					):
   __lowerCAmelCase							,       __lowerCAmelCase												=  expected_row_ids_and_row_dicts[i]
   assert row_id == expected_row_id
   assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(20					).repartition(4					)

 # Partitions 0 and 2
 __lowerCAmelCase												=  SparkExamplesIterable(lowerCamelCase					).shard_data_sources(worker_id=0				,		num_workers=2					)
 assert shard_it_a.n_shards == 2
 __lowerCAmelCase												=  _get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase				,		[0, 2]					)
 for i, (row_id, row_dict) in enumerate(lowerCamelCase					):
  __lowerCAmelCase							,       __lowerCAmelCase												=  expected_row_ids_and_row_dicts_a[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict

 # Partitions 1 and 3
 __lowerCAmelCase												=  SparkExamplesIterable(lowerCamelCase					).shard_data_sources(worker_id=1				,		num_workers=2					)
 assert shard_it_a.n_shards == 2
 __lowerCAmelCase												=  _get_expected_row_ids_and_row_dicts_for_partition_order(lowerCamelCase				,		[1, 3]					)
 for i, (row_id, row_dict) in enumerate(lowerCamelCase					):
  __lowerCAmelCase							,       __lowerCAmelCase												=  expected_row_ids_and_row_dicts_a[i]
  assert row_id == expected_row_id
  assert row_dict == expected_row_dict




@require_not_windows
@require_dill_gt_0_3_2
def   __lowerCAmelCase					(							):






 '''simple docstring'''




 __lowerCAmelCase												=  pyspark.sql.SparkSession.builder.master("local[*]"					).appName("pyspark"					).getOrCreate()
 __lowerCAmelCase												=  spark.range(1_00					).repartition(1					)
 __lowerCAmelCase												=  Spark(lowerCamelCase					)
 # Choose a small max_shard_size for maximum partitioning.
 spark_builder._repartition_df_if_needed(max_shard_size=1					)
 # The new number of partitions should not be greater than the number of rows.
 assert spark_builder.df.rdd.getNumPartitions() == 1_00




'''simple docstring'''



# flake8: noqa
# Lint as: python3

from typing import Dict, List, Optional, Type

from .. import config
from ..utils import logging
from .formatting import (
    ArrowFormatter,
    CustomFormatter,
    Formatter,
    PandasFormatter,
    PythonFormatter,
    TensorFormatter,
    format_table,
    query_table,
)
from .np_formatter import NumpyFormatter


lowerCAmelCase	:		List[str]				=   logging.get_logger(__name__)

lowerCAmelCase	:		Dict[Optional[str], Type[Formatter]]				=   {}
lowerCAmelCase	:		Dict[Optional[str], str]				=   {}
lowerCAmelCase	:		Dict[Optional[str], Exception]				=   {}



def   __lowerCAmelCase					(							lowerCamelCase   :						type				,		lowerCamelCase   :						Optional[str]				,		lowerCamelCase   :						Optional[List[str]] = None				,		):






	'''simple docstring'''




	__lowerCAmelCase												=  aliases if aliases is not None else []
	if format_type in _FORMAT_TYPES:
		logger.warning(
		    f'''Overwriting format type \'{format_type}\' ({_FORMAT_TYPES[format_type].__name__} -> {formatter_cls.__name__})'''					)
	__lowerCAmelCase												=  formatter_cls
	for alias in set(aliases + [format_type]					):
		if alias in _FORMAT_TYPES_ALIASES:
			logger.warning(
			    f'''Overwriting format type alias \'{alias}\' ({_FORMAT_TYPES_ALIASES[alias]} -> {format_type})'''					)
		__lowerCAmelCase												=  format_type



def   __lowerCAmelCase					(							lowerCamelCase   :						Exception				,		lowerCamelCase   :						Optional[str]				,		lowerCamelCase   :						Optional[List[str]] = None					):






	'''simple docstring'''




	__lowerCAmelCase												=  aliases if aliases is not None else []
	for alias in set(aliases + [format_type]					):
		__lowerCAmelCase												=  unavailable_error


# Here we define all the available formatting functions that can be used by `Dataset.set_format`
_register_formatter(PythonFormatter, None, aliases=['''python'''])
_register_formatter(ArrowFormatter, '''arrow''', aliases=['''pa''', '''pyarrow'''])
_register_formatter(NumpyFormatter, '''numpy''', aliases=['''np'''])
_register_formatter(PandasFormatter, '''pandas''', aliases=['''pd'''])
_register_formatter(CustomFormatter, '''custom''')

if config.TORCH_AVAILABLE:
	from .torch_formatter import TorchFormatter

	_register_formatter(TorchFormatter, '''torch''', aliases=['''pt''', '''pytorch'''])
else:
	lowerCAmelCase	:		Optional[int]				=   ValueError('''PyTorch needs to be installed to be able to return PyTorch tensors.''')
	_register_unavailable_formatter(_torch_error, '''torch''', aliases=['''pt''', '''pytorch'''])

if config.TF_AVAILABLE:
	from .tf_formatter import TFFormatter

	_register_formatter(TFFormatter, '''tensorflow''', aliases=['''tf'''])
else:
	lowerCAmelCase	:		str				=   ValueError('''Tensorflow needs to be installed to be able to return Tensorflow tensors.''')
	_register_unavailable_formatter(_tf_error, '''tensorflow''', aliases=['''tf'''])

if config.JAX_AVAILABLE:
	from .jax_formatter import JaxFormatter

	_register_formatter(JaxFormatter, '''jax''', aliases=[])
else:
	lowerCAmelCase	:		Any				=   ValueError('''JAX needs to be installed to be able to return JAX arrays.''')
	_register_unavailable_formatter(_jax_error, '''jax''', aliases=[])



def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[str]					):






	'''simple docstring'''




	if format_type in _FORMAT_TYPES_ALIASES:
		return _FORMAT_TYPES_ALIASES[format_type]
	else:
		return format_type




def   __lowerCAmelCase					(							lowerCamelCase   :						Optional[str]				,		**lowerCamelCase   :						Tuple					):






	'''simple docstring'''




	__lowerCAmelCase												=  get_format_type_from_alias(lowerCamelCase					)
	if format_type in _FORMAT_TYPES:
		return _FORMAT_TYPES[format_type](**lowerCamelCase					)
	if format_type in _FORMAT_TYPES_ALIASES_UNAVAILABLE:
		raise _FORMAT_TYPES_ALIASES_UNAVAILABLE[format_type]
	else:
		raise ValueError(
		    f'''Return type should be None or selected in {list(type for type in _FORMAT_TYPES.keys() if type != None					)}, but got \'{format_type}\''''					)