


'''simple docstring'''






import json
import os
import tempfile

import transformers

import datasets
from utils import generate_example_dataset, get_duration


__UpperCAmelCase					       =							500_000

__UpperCAmelCase       ,    __UpperCAmelCase					       =							os.path.split(__file__)
__UpperCAmelCase					       =							os.path.join(RESULTS_BASEPATH, '''results''', RESULTS_FILENAME.replace('''.py''', '''.json'''))


@get_duration
def       _snake_case     (  A				,					**A )		->					int:
							lowerCAmelCase__      =			dataset.map(**A )


@get_duration
def       _snake_case     (  A				,					**A )		->					Union[str, Any]:
							lowerCAmelCase__      =			dataset.filter(**A )


def       _snake_case     (  )		->					Any:
							lowerCAmelCase__      =			{'''num examples''': SPEED_TEST_N_EXAMPLES}
							with tempfile.TemporaryDirectory() as tmp_dir:
														lowerCAmelCase__      =			datasets.Features({'''text''': datasets.Value('''string''' ), '''numbers''': datasets.Value('''float32''' )} )
														lowerCAmelCase__      =			generate_example_dataset(
														    os.path.join(A				,					'''dataset.arrow''' )				,					A				,					num_examples=A )

														lowerCAmelCase__      =			transformers.AutoTokenizer.from_pretrained('''bert-base-cased'''				,					use_fast=A )

														def tokenize(A ):
																					return tokenizer(examples['''text'''] )

														lowerCAmelCase__      =			map(A )

														lowerCAmelCase__      =			map(A				,					batched=A )

														lowerCAmelCase__      =			map(A				,					function=lambda A : None				,					batched=A )

														with dataset.formatted_as(type='''numpy''' ):
																					lowerCAmelCase__      =			map(A				,					function=lambda A : None				,					batched=A )

														with dataset.formatted_as(type='''pandas''' ):
																					lowerCAmelCase__      =			map(A				,					function=lambda A : None				,					batched=A )

														with dataset.formatted_as(type='''torch'''				,					columns='''numbers''' ):
																					lowerCAmelCase__      =			map(A				,					function=lambda A : None				,					batched=A )

														with dataset.formatted_as(type='''tensorflow'''				,					columns='''numbers''' ):
																					lowerCAmelCase__      =			map(A				,					function=lambda A : None				,					batched=A )

														lowerCAmelCase__      =			map(A				,					function=A				,					batched=A )

														lowerCAmelCase__      =			filter(A )

														# Activate later when tokenizer support batched inputs
														# with dataset.formatted_as(type='numpy'):
														#     times[func.__name__ + " fast-tokenizer batched numpy"] = func(dataset, function=tokenize, batched=True)

							with open(A				,					'''wb''' ) as f:
														f.write(json.dumps(A ).encode('''utf-8''' ) )


if __name__ == "__main__":  # useful to run the profiler
		benchmark_map_filter()



'''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       _snake_case     (  A				,					A )		->					List[Any]:
							lowerCAmelCase__      =			[]
							for part_id in partition_order:
														lowerCAmelCase__      =			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       _snake_case     (  )		->					Tuple:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(100 ).repartition(1 )
							lowerCAmelCase__      =			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=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     (  )		->					Optional[int]:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(10 ).repartition(2 )
							lowerCAmelCase__      =			[1, 0]
							lowerCAmelCase__      =			_generate_iterable_examples(A				,					A )  # Reverse the partitions.
							lowerCAmelCase__      =			_get_expected_row_ids_and_row_dicts_for_partition_order(A				,					A )

							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       _snake_case     (  )		->					Optional[Any]:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(10 ).repartition(1 )
							lowerCAmelCase__      =			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       _snake_case     (  )		->					Union[str, Any]:
							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 A : x.reverse()
														lowerCAmelCase__      =			_get_expected_row_ids_and_row_dicts_for_partition_order(A				,					[2, 1, 0] )

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

							# Partitions 0 and 2
							lowerCAmelCase__      =			SparkExamplesIterable(A ).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(A				,					[0, 2] )
							for i, (row_id, row_dict) in enumerate(A ):
														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(A ).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(A				,					[1, 3] )
							for i, (row_id, row_dict) in enumerate(A ):
														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       _snake_case     (  )		->					Dict:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(100 ).repartition(1 )
							lowerCAmelCase__      =			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() == 100