'''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			(   UpperCAmelCase__ :      Tuple   ,   UpperCAmelCase__ :      Optional[Any]				)    ->   List[str]:
							lowercase_      :    Tuple       		= []
							for part_id in partition_order:
														lowercase_      :    List[Any]       		= df.where(F'''SPARK_PARTITION_ID() = {part_id}'''				).collect()
														for row_idx, row in enumerate(snake_case__				):
																					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			(   )    ->   Tuple:
							lowercase_      :    str       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    str       		= spark.range(100				).repartition(1				)
							lowercase_      :    Dict       		= Spark(snake_case__				)
							# 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			(   )    ->   Any:
							lowercase_      :    Any       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    Any       		= spark.range(10				).repartition(2				)
							lowercase_      :    List[str]       		= [1, 0]
							lowercase_      :    str       		= _generate_iterable_examples(snake_case__   ,   snake_case__				)  # Reverse the partitions.
							lowercase_      :    List[Any]       		= _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,   snake_case__				)

							for i, (row_id, row_dict) in enumerate(generate_fn()				):
														lowercase_		,		lowercase_      :    int       		= 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			(   )    ->   Dict:
							lowercase_      :    Optional[Any]       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    List[str]       		= spark.range(10				).repartition(1				)
							lowercase_      :    Optional[int]       		= SparkExamplesIterable(snake_case__				)
							assert it.n_shards == 1
							for i, (row_id, row_dict) in enumerate(snake_case__				):
														assert row_id == F'''0_{i}'''
														assert row_dict == {"id": i}




@require_not_windows
@require_dill_gt_0_3_2
def 							lowerCamelCase			(   )    ->   Dict:
							lowercase_      :    Optional[Any]       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    List[Any]       		= spark.range(30				).repartition(3				)
							# Mock the generator so that shuffle reverses the partition indices.
							with patch("""numpy.random.Generator"""				) as generator_mock:
														lowercase_      :    Optional[Any]       		= lambda UpperCAmelCase__				: x.reverse()
														lowercase_      :    Tuple       		= _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,   [2, 1, 0]				)

														lowercase_      :    int       		= SparkExamplesIterable(snake_case__				).shuffle_data_sources(snake_case__				)
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(snake_case__				):
																					lowercase_		,		lowercase_      :    int       		= 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			(   )    ->   Any:
							lowercase_      :    int       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    int       		= spark.range(20				).repartition(4				)

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

							# Partitions 1 and 3
							lowercase_      :    Optional[Any]       		= SparkExamplesIterable(snake_case__				).shard_data_sources(worker_id=1   ,   num_workers=2				)
							assert shard_it_a.n_shards == 2
							lowercase_      :    Optional[int]       		= _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,   [1, 3]				)
							for i, (row_id, row_dict) in enumerate(snake_case__				):
														lowercase_		,		lowercase_      :    Optional[Any]       		= 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			(   )    ->   Optional[Any]:
							lowercase_      :    Dict       		= pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
							lowercase_      :    List[str]       		= spark.range(100				).repartition(1				)
							lowercase_      :    int       		= Spark(snake_case__				)
							# 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() == 100







'''simple docstring'''




_lowercase		:				int	       =		"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"




def 							lowerCamelCase			(   UpperCAmelCase__ :      bytes				)    ->   bytes:
							# Make sure the supplied data is a bytes-like object
							if not isinstance(UpperCAmelCase__   ,   UpperCAmelCase__				):
														lowercase_      :    Union[str, Any]       		= F'''a bytes-like object is required, not \'{data.__class__.__name__}\''''
														raise TypeError(UpperCAmelCase__				)

							lowercase_      :    Dict       		= """""".join(bin(UpperCAmelCase__				)[2:].zfill(8				) for byte in data				)

							lowercase_      :    Union[str, Any]       		= len(UpperCAmelCase__				) % 6 != 0

							if padding_needed:
														# The padding that will be added later
														lowercase_      :    List[Any]       		= b"""=""" * ((6 - len(UpperCAmelCase__				) % 6) // 2)

														# Append binary_stream with arbitrary binary digits (0's by default) to make its
														# length a multiple of 6.
														binary_stream += "0" * (6 - len(UpperCAmelCase__				) % 6)
							else:
														lowercase_      :    Union[str, Any]       		= b""""""

							# Encode every 6 binary digits to their corresponding Base64 character
							return (
							    "".join(
							        B64_CHARSET[int(binary_stream[index : index + 6]   ,   2				)]
							        for index in range(0   ,   len(UpperCAmelCase__				)   ,   6				)				).encode()
							    + padding
							)




def 							lowerCamelCase			(   UpperCAmelCase__ :      str				)    ->   bytes:
							# Make sure encoded_data is either a string or a bytes-like object
							if not isinstance(UpperCAmelCase__   ,   UpperCAmelCase__				) and not isinstance(UpperCAmelCase__   ,   UpperCAmelCase__				):
														lowercase_      :    List[str]       		= (
														    """argument should be a bytes-like object or ASCII string, """
														    F'''not \'{encoded_data.__class__.__name__}\''''
														)
														raise TypeError(UpperCAmelCase__				)

							# In case encoded_data is a bytes-like object, make sure it contains only
							# ASCII characters so we convert it to a string object
							if isinstance(UpperCAmelCase__   ,   UpperCAmelCase__				):
														try:
																					lowercase_      :    Optional[int]       		= encoded_data.decode("""utf-8"""				)
														except UnicodeDecodeError:
																					raise ValueError("""base64 encoded data should only contain ASCII characters"""				)

							lowercase_      :    Any       		= encoded_data.count("""="""				)

							# Check if the encoded string contains non base64 characters
							if padding:
														assert all(
														    char in B64_CHARSET for char in encoded_data[:-padding]				), "Invalid base64 character(s) found."
							else:
														assert all(
														    char in B64_CHARSET for char in encoded_data				), "Invalid base64 character(s) found."

							# Check the padding
							assert len(UpperCAmelCase__				) % 4 == 0 and padding < 3, "Incorrect padding"

							if padding:
														# Remove padding if there is one
														lowercase_      :    Optional[int]       		= encoded_data[:-padding]

														lowercase_      :    Any       		= """""".join(
														    bin(B64_CHARSET.index(UpperCAmelCase__				)				)[2:].zfill(6				) for char in encoded_data				)[: -padding * 2]
							else:
														lowercase_      :    int       		= """""".join(
														    bin(B64_CHARSET.index(UpperCAmelCase__				)				)[2:].zfill(6				) for char in encoded_data				)

							lowercase_      :    Optional[int]       		= [
							    int(binary_stream[index : index + 8]   ,   2				)
							    for index in range(0   ,   len(UpperCAmelCase__				)   ,   8				)
							]

							return bytes(UpperCAmelCase__				)


if __name__ == "__main__":
						import doctest

						doctest.testmod()





