






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

			for i, (row_id, row_dict) in enumerate(generate_fn()		):
						UpperCAmelCase_    ,      UpperCAmelCase_      :       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 __UpperCAmelCase   (					)    ->   Union[str, Any]:
			UpperCAmelCase_      :       str          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       List[Any]          =       spark.range(1_0		).repartition(1		)
			UpperCAmelCase_      :       Optional[int]          =       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 __UpperCAmelCase   (					)    ->   Union[str, Any]:
			UpperCAmelCase_      :       Dict          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       Dict          =       spark.range(3_0		).repartition(3		)
			# Mock the generator so that shuffle reverses the partition indices.
			with patch('''numpy.random.Generator'''		) as generator_mock:
						UpperCAmelCase_      :       Any          =       lambda A		: x.reverse()
						UpperCAmelCase_      :       Dict          =       _get_expected_row_ids_and_row_dicts_for_partition_order(A			,       [2, 1, 0]		)

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

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

			# Partitions 1 and 3
			UpperCAmelCase_      :       Tuple          =       SparkExamplesIterable(A		).shard_data_sources(worker_id=1			,       num_workers=2		)
			assert shard_it_a.n_shards == 2
			UpperCAmelCase_      :       List[str]          =       _get_expected_row_ids_and_row_dicts_for_partition_order(A			,       [1, 3]		)
			for i, (row_id, row_dict) in enumerate(A		):
						UpperCAmelCase_    ,      UpperCAmelCase_      :       Union[str, 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 __UpperCAmelCase   (					)    ->   Any:
			UpperCAmelCase_      :       List[Any]          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       List[str]          =       spark.range(1_0_0		).repartition(1		)
			UpperCAmelCase_      :       int          =       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










'''simple docstring'''





_UpperCamelCase						:						Dict  =					range(2, 20 + 1)
_UpperCamelCase						:						str  =					[10**k for k in range(ks[-1] + 1)]
_UpperCamelCase						:						dict[int, dict[int, list[list[int]]]]  =					{}
def __UpperCAmelCase   (					A					:			List[Any]			,       A					:			str			,       A					:			Union[str, Any]			,       A					:			Union[str, Any]		)    ->   Optional[Any]:
			UpperCAmelCase_      :       Optional[Any]          =       sum(a_i[j] for j in range(A			,       len(A		)		)		)
			UpperCAmelCase_      :       Union[str, Any]          =       sum(a_i[j] * base[j] for j in range(min(len(A		)			,       A		)		)		)

			UpperCAmelCase_    ,      UpperCAmelCase_      :       Optional[int]          =       0, 0
			UpperCAmelCase_      :       str          =       n - i

			UpperCAmelCase_      :       Union[str, Any]          =       memo.get(A		)

			if sub_memo is not None:
						UpperCAmelCase_      :       Dict          =       sub_memo.get(A		)

						if jumps is not None and len(A		) > 0:
									# find and make the largest jump without going over
									UpperCAmelCase_      :       str          =       -1
									for _k in range(len(A		) - 1			,       -1			,       -1		):
												if jumps[_k][2] <= k and jumps[_k][1] <= max_dn:
															UpperCAmelCase_      :       Any          =       _k
															break

									if max_jump >= 0:
												UpperCAmelCase_    ,      UpperCAmelCase_    ,      UpperCAmelCase_      :       Any          =       jumps[max_jump]
												# since the difference between jumps is cached, add c
												UpperCAmelCase_      :       Union[str, Any]          =       diff + c
												for j in range(min(A			,       len(A		)		)		):
															UpperCAmelCase_    ,      UpperCAmelCase_      :       Tuple          =       divmod(A			,       1_0		)
												if new_c > 0:
															add(A			,       A			,       A		)

						else:
									UpperCAmelCase_      :       Tuple          =       []
			else:
						UpperCAmelCase_      :       List[str]          =       {c: []}
						UpperCAmelCase_      :       List[str]          =       sub_memo

			if dn >= max_dn or c + diff >= base[k]:
						return diff, dn

			if k > ks[0]:
						while True:
									# keep doing smaller jumps
									UpperCAmelCase_    ,      UpperCAmelCase_      :       Optional[Any]          =       next_term(A			,       k - 1			,       i + dn			,       A		)
									diff += _diff
									dn += terms_jumped

									if dn >= max_dn or c + diff >= base[k]:
												break
			else:
						# would be too small a jump, just compute sequential terms instead
						UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       compute(A			,       A			,       i + dn			,       A		)
						diff += _diff
						dn += terms_jumped

			UpperCAmelCase_      :       List[str]          =       sub_memo[c]

			# keep jumps sorted by # of terms skipped
			UpperCAmelCase_      :       List[Any]          =       0
			while j < len(A		):
						if jumps[j][1] > dn:
									break
						j += 1

			# cache the jump for this value digitsum(b) and c
			sub_memo[c].insert(A			,       (diff, dn, k)		)
			return (diff, dn)
def __UpperCAmelCase   (					A					:			List[str]			,       A					:			Dict			,       A					:			List[str]			,       A					:			Union[str, Any]		)    ->   str:
			if i >= n:
						return 0, i
			if k > len(A		):
						a_i.extend([0 for _ in range(k - len(A		)		)]		)

			# note: a_i -> b * 10^k + c
			# ds_b -> digitsum(b)
			# ds_c -> digitsum(c)
			UpperCAmelCase_      :       int          =       i
			UpperCAmelCase_    ,      UpperCAmelCase_    ,      UpperCAmelCase_      :       List[Any]          =       0, 0, 0
			for j in range(len(A		)		):
						if j >= k:
									ds_b += a_i[j]
						else:
									ds_c += a_i[j]

			while i < n:
						i += 1
						UpperCAmelCase_      :       Optional[int]          =       ds_c + ds_b
						diff += addend
						UpperCAmelCase_      :       Any          =       0
						for j in range(A		):
									UpperCAmelCase_      :       Dict          =       a_i[j] + addend
									UpperCAmelCase_    ,      UpperCAmelCase_      :       Tuple          =       divmod(A			,       1_0		)

									ds_c += a_i[j]

						if addend > 0:
									break

			if addend > 0:
						add(A			,       A			,       A		)
			return diff, i - start_i
def __UpperCAmelCase   (					A					:			Dict			,       A					:			Union[str, Any]			,       A					:			str		)    ->   List[Any]:
			for j in range(A			,       len(A		)		):
						UpperCAmelCase_      :       List[str]          =       digits[j] + addend
						if s >= 1_0:
									UpperCAmelCase_    ,      UpperCAmelCase_      :       Any          =       divmod(A			,       1_0		)
									UpperCAmelCase_      :       str          =       addend // 1_0 + quotient
						else:
									UpperCAmelCase_      :       Any          =       s
									UpperCAmelCase_      :       Any          =       addend // 1_0

						if addend == 0:
									break

			while addend > 0:
						UpperCAmelCase_    ,      UpperCAmelCase_      :       Union[str, Any]          =       divmod(A			,       1_0		)
						digits.append(A		)
def __UpperCAmelCase   (					A					:			int = 1_0**1_5		)    ->   int:
			UpperCAmelCase_      :       Any          =       [1]
			UpperCAmelCase_      :       Union[str, Any]          =       1
			UpperCAmelCase_      :       List[Any]          =       0
			while True:
						UpperCAmelCase_    ,      UpperCAmelCase_      :       str          =       next_term(A			,       2_0			,       i + dn			,       A		)
						dn += terms_jumped
						if dn == n - i:
									break

			UpperCAmelCase_      :       Any          =       0
			for j in range(len(A		)		):
						a_n += digits[j] * 1_0**j
			return a_n


if __name__ == "__main__":
						print(f'''{solution() = }''')


