





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 (							__a       :int   ,					__a       :Union[str, Any]   )		->    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(__a   ):
						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 (							)		->    Union[str, Any]:






		"""simple docstring"""

		A__  				=       pyspark.sql.SparkSession.builder.master("""local[*]"""   ).appName("""pyspark"""   ).getOrCreate()
		A__  				=       spark.range(1_0_0   ).repartition(1   )
		A__  				=       Spark(__a   )
		# 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=1_6   )
		# Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
		assert spark_builder.df.rdd.getNumPartitions() == 5_0







@require_not_windows
@require_dill_gt_0_3_2
def        __lowerCamelCase (							)		->    Optional[Any]:






		"""simple docstring"""

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

		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        __lowerCamelCase (							)		->    Dict:






		"""simple docstring"""

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







@require_not_windows
@require_dill_gt_0_3_2
def        __lowerCamelCase (							)		->    List[str]:






		"""simple docstring"""

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

				A__  				=       SparkExamplesIterable(__a   ).shuffle_data_sources(__a   )
				assert shuffled_it.n_shards == 3
				for i, (row_id, row_dict) in enumerate(__a   ):
						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        __lowerCamelCase (							)		->    Optional[int]:






		"""simple docstring"""

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

		# Partitions 0 and 2
		A__  				=       SparkExamplesIterable(__a   ).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(__a   ,					[0, 2]   )
		for i, (row_id, row_dict) in enumerate(__a   ):
				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(__a   ).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(__a   ,					[1, 3]   )
		for i, (row_id, row_dict) in enumerate(__a   ):
				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        __lowerCamelCase (							)		->    str:






		"""simple docstring"""

		A__  				=       pyspark.sql.SparkSession.builder.master("""local[*]"""   ).appName("""pyspark"""   ).getOrCreate()
		A__  				=       spark.range(1_0_0   ).repartition(1   )
		A__  				=       Spark(__a   )
		# 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_0_0










# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from accelerate import PartialState
from accelerate.utils.operations import broadcast, gather, gather_object, pad_across_processes, reduce







def        __lowerCamelCase (							__a       :List[str]   )		->    Tuple:






		"""simple docstring"""

		return (torch.arange(state.num_processes   ) + 1.0 + (state.num_processes * state.process_index)).to(state.device   )







def        __lowerCamelCase (							__a       :int   )		->    Optional[int]:






		"""simple docstring"""

		A__  				=       create_tensor(__a   )
		A__  				=       gather(__a   )
		assert gathered_tensor.tolist() == list(range(1   ,					state.num_processes**2 + 1   )   )







def        __lowerCamelCase (							__a       :Any   )		->    Any:






		"""simple docstring"""

		A__  				=       [state.process_index]
		A__  				=       gather_object(__a   )
		assert len(__a   ) == state.num_processes, F'{gathered_obj}, {len(__a   )} != {state.num_processes}'
		assert gathered_obj == list(range(state.num_processes   )   ), F'{gathered_obj} != {list(range(state.num_processes   )   )}'







def        __lowerCamelCase (							__a       :List[str]   )		->    List[str]:






		"""simple docstring"""

		A__  				=       create_tensor(__a   )
		A__  				=       broadcast(__a   )
		assert broadcasted_tensor.shape == torch.Size([state.num_processes]   )
		assert broadcasted_tensor.tolist() == list(range(1   ,					state.num_processes + 1   )   )







def        __lowerCamelCase (							__a       :Any   )		->    Any:






		"""simple docstring"""

		if state.is_main_process:
				A__  				=       torch.arange(state.num_processes + 1   ).to(state.device   )
		else:
				A__  				=       torch.arange(state.num_processes   ).to(state.device   )
		A__  				=       pad_across_processes(__a   )
		assert padded_tensor.shape == torch.Size([state.num_processes + 1]   )
		if not state.is_main_process:
				assert padded_tensor.tolist() == list(range(0   ,					state.num_processes   )   ) + [0]







def        __lowerCamelCase (							__a       :Union[str, Any]   )		->    str:






		"""simple docstring"""

		if state.num_processes != 2:
				return
		A__  				=       create_tensor(__a   )
		A__  				=       reduce(__a   ,					"""sum"""   )
		A__  				=       torch.tensor([4.0, 6]   ).to(state.device   )
		assert torch.allclose(__a   ,					__a   ), F'{reduced_tensor} != {truth_tensor}'







def        __lowerCamelCase (							__a       :List[Any]   )		->    List[str]:






		"""simple docstring"""

		if state.num_processes != 2:
				return
		A__  				=       create_tensor(__a   )
		A__  				=       reduce(__a   ,					"""mean"""   )
		A__  				=       torch.tensor([2.0, 3]   ).to(state.device   )
		assert torch.allclose(__a   ,					__a   ), F'{reduced_tensor} != {truth_tensor}'







def        __lowerCamelCase (							__a       :str   )		->    Optional[int]:






		"""simple docstring"""

		main()







def        __lowerCamelCase (							)		->    str:






		"""simple docstring"""

		A__  				=       PartialState()
		state.print(F'State: {state}'   )
		state.print("""testing gather"""   )
		test_gather(__a   )
		state.print("""testing gather_object"""   )
		test_gather_object(__a   )
		state.print("""testing broadcast"""   )
		test_broadcast(__a   )
		state.print("""testing pad_across_processes"""   )
		test_pad_across_processes(__a   )
		state.print("""testing reduce_sum"""   )
		test_reduce_sum(__a   )
		state.print("""testing reduce_mean"""   )
		test_reduce_mean(__a   )


if __name__ == "__main__":
					main()



