




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 		snake_case_   (UpperCamelCase	:  List[str]			,						UpperCamelCase	:  Tuple				):





							'''simple docstring'''


							_a     					=    []
							for part_id in partition_order:
														_a     					=    df.where(f'SPARK_PARTITION_ID() = {part_id}'				).collect()
														for row_idx, row in enumerate(UpperCamelCase				):
																					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 		snake_case_   ():





							'''simple docstring'''


							_a     					=    pyspark.sql.SparkSession.builder.master('''local[*]'''				).appName('''pyspark'''				).getOrCreate()
							_a     					=    spark.range(100				).repartition(1				)
							_a     					=    Spark(UpperCamelCase				)
							# 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 		snake_case_   ():





							'''simple docstring'''


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

							for i, (row_id, row_dict) in enumerate(generate_fn()				):
														_a					,							_a     					=    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 		snake_case_   ():





							'''simple docstring'''


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





							'''simple docstring'''


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

														_a     					=    SparkExamplesIterable(UpperCamelCase				).shuffle_data_sources(UpperCamelCase				)
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(UpperCamelCase				):
																					_a					,							_a     					=    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 		snake_case_   ():





							'''simple docstring'''


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

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

							# Partitions 1 and 3
							_a     					=    SparkExamplesIterable(UpperCamelCase				).shard_data_sources(worker_id=1			,						num_workers=2				)
							assert shard_it_a.n_shards == 2
							_a     					=    _get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase			,						[1, 3]				)
							for i, (row_id, row_dict) in enumerate(UpperCamelCase				):
														_a					,							_a     					=    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 		snake_case_   ():





							'''simple docstring'''


							_a     					=    pyspark.sql.SparkSession.builder.master('''local[*]'''				).appName('''pyspark'''				).getOrCreate()
							_a     					=    spark.range(100				).repartition(1				)
							_a     					=    Spark(UpperCamelCase				)
							# 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'''


from __future__ import annotations
def 		snake_case_   (UpperCamelCase	:  list[int]				):





							'''simple docstring'''


							if not nums:
														return 0
							_a     					=    nums[0]
							_a     					=    0
							for num in nums[1:]:
														_a					,							_a     					=    (
														    max_excluding + num,
														    max(UpperCamelCase			,						UpperCamelCase				),
														)
							return max(UpperCamelCase			,						UpperCamelCase				)


if __name__ == "__main__":
				import doctest

				doctest.testmod()

