






import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
	import faiss



@require_faiss
class 							SCREAMING_SNAKE_CASE__		(       __UpperCAmelCase		):





					"""simple docstring"""






					def  _lowercase  (      self  :						Dict ):
						snake_case__				:   Optional[Any]     =						tempfile.mkdtemp()
						snake_case__				:   int     =						8

						# DPR tok
						snake_case__				:   Union[str, Any]     =						[
						    "[UNK]",
						    "[CLS]",
						    "[SEP]",
						    "[PAD]",
						    "[MASK]",
						    "want",
						    "##want",
						    "##ed",
						    "wa",
						    "un",
						    "runn",
						    "##ing",
						    ",",
						    "low",
						    "lowest",
						]
						snake_case__				:   str     =						os.path.join(self.tmpdirname							,	"dpr_tokenizer" )
						os.makedirs(UpperCAmelCase_							,	exist_ok=UpperCAmelCase_ )
						snake_case__				:   Any     =						os.path.join(UpperCAmelCase_							,	DPR_VOCAB_FILES_NAMES["vocab_file"] )
						with open(self.vocab_file							,	"w"							,	encoding="utf-8" ) as vocab_writer:
							vocab_writer.write("".join([x + "\n" for x in vocab_tokens] ) )

						# BART tok
						snake_case__				:   Dict     =						[
						    "l",
						    "o",
						    "w",
						    "e",
						    "r",
						    "s",
						    "t",
						    "i",
						    "d",
						    "n",
						    "\u0120",
						    "\u0120l",
						    "\u0120n",
						    "\u0120lo",
						    "\u0120low",
						    "er",
						    "\u0120lowest",
						    "\u0120newer",
						    "\u0120wider",
						    "<unk>",
						]
						snake_case__				:   Optional[int]     =						dict(zip(UpperCAmelCase_							,	range(len(UpperCAmelCase_ ) ) ) )
						snake_case__				:   List[str]     =						["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
						snake_case__				:   List[str]     =						{"unk_token": "<unk>"}

						snake_case__				:   List[str]     =						os.path.join(self.tmpdirname							,	"bart_tokenizer" )
						os.makedirs(UpperCAmelCase_							,	exist_ok=UpperCAmelCase_ )
						snake_case__				:   str     =						os.path.join(UpperCAmelCase_							,	BART_VOCAB_FILES_NAMES["vocab_file"] )
						snake_case__				:   int     =						os.path.join(UpperCAmelCase_							,	BART_VOCAB_FILES_NAMES["merges_file"] )
						with open(self.vocab_file							,	"w"							,	encoding="utf-8" ) as fp:
							fp.write(json.dumps(UpperCAmelCase_ ) + "\n" )
						with open(self.merges_file							,	"w"							,	encoding="utf-8" ) as fp:
							fp.write("\n".join(UpperCAmelCase_ ) )



					def  _lowercase  (      self  :						str ):
						return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	"dpr_tokenizer" ) )



					def  _lowercase  (      self  :						Any ):
						return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	"dpr_tokenizer" ) )



					def  _lowercase  (      self  :						List[str] ):
						return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname							,	"bart_tokenizer" ) )



					def  _lowercase  (      self  :						str ):
						shutil.rmtree(self.tmpdirname )



					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   List[Any]     =						Dataset.from_dict(
						    {
						        "id": ["0", "1"],
						        "text": ["foo", "bar"],
						        "title": ["Foo", "Bar"],
						        "embeddings": [np.ones(self.retrieval_vector_size ), 2 * np.ones(self.retrieval_vector_size )],
						    } )
						dataset.add_faiss_index("embeddings"							,	string_factory="Flat"							,	metric_type=faiss.METRIC_INNER_PRODUCT )
						return dataset



					def  _lowercase  (      self  :						Union[str, Any] ):
						snake_case__				:   List[str]     =						self.get_dummy_dataset()
						snake_case__				:   Optional[Any]     =						RagConfig(
						    retrieval_vector_size=self.retrieval_vector_size							,	question_encoder=DPRConfig().to_dict()							,	generator=BartConfig().to_dict()							,	)
						with patch("transformers.models.rag.retrieval_rag.load_dataset" ) as mock_load_dataset:
							snake_case__				:   Any     =						dataset
							snake_case__				:   str     =						RagRetriever(
							    UpperCAmelCase_							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	)
						return retriever



					def  _lowercase  (      self  :						str							,	__A  :						List[str] ):
						snake_case__				:   Dict     =						self.get_dummy_dataset()
						snake_case__				:   int     =						RagConfig(
						    retrieval_vector_size=self.retrieval_vector_size							,	question_encoder=DPRConfig().to_dict()							,	generator=BartConfig().to_dict()							,	index_name="custom"							,	)
						if from_disk:
							snake_case__				:   List[str]     =						os.path.join(self.tmpdirname							,	"dataset" )
							snake_case__				:   Union[str, Any]     =						os.path.join(self.tmpdirname							,	"index.faiss" )
							dataset.get_index("embeddings" ).save(os.path.join(self.tmpdirname							,	"index.faiss" ) )
							dataset.drop_index("embeddings" )
							dataset.save_to_disk(os.path.join(self.tmpdirname							,	"dataset" ) )
							del dataset
							snake_case__				:   Union[str, Any]     =						RagRetriever(
							    UpperCAmelCase_							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	)
						else:
							snake_case__				:   str     =						RagRetriever(
							    UpperCAmelCase_							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer()							,	index=CustomHFIndex(config.retrieval_vector_size							,	UpperCAmelCase_ )							,	)
						return retriever



					def  _lowercase  (      self  :						Any ):
						snake_case__				:   Union[str, Any]     =						Dataset.from_dict(
						    {
						        "id": ["0", "1"],
						        "text": ["foo", "bar"],
						        "title": ["Foo", "Bar"],
						        "embeddings": [np.ones(self.retrieval_vector_size + 1 ), 2 * np.ones(self.retrieval_vector_size + 1 )],
						    } )
						dataset.add_faiss_index("embeddings"							,	string_factory="Flat"							,	metric_type=faiss.METRIC_INNER_PRODUCT )

						snake_case__				:   List[Any]     =						os.path.join(self.tmpdirname							,	"hf_bert_base.hnswSQ8_correct_phi_128.c_index" )
						dataset.save_faiss_index("embeddings"							,	index_file_name + ".index.dpr" )
						pickle.dump(dataset["id"]							,	open(index_file_name + ".index_meta.dpr"							,	"wb" ) )

						snake_case__				:   List[Any]     =						os.path.join(self.tmpdirname							,	"psgs_w100.tsv.pkl" )
						snake_case__				:   Union[str, Any]     =						{sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
						pickle.dump(UpperCAmelCase_							,	open(UpperCAmelCase_							,	"wb" ) )

						snake_case__				:   Any     =						RagConfig(
						    retrieval_vector_size=self.retrieval_vector_size							,	question_encoder=DPRConfig().to_dict()							,	generator=BartConfig().to_dict()							,	index_name="legacy"							,	index_path=self.tmpdirname							,	)
						snake_case__				:   Union[str, Any]     =						RagRetriever(
						    UpperCAmelCase_							,	question_encoder_tokenizer=self.get_dpr_tokenizer()							,	generator_tokenizer=self.get_bart_tokenizer() )
						return retriever



					def  _lowercase  (      self  :						Optional[int] ):
						snake_case__				:   List[Any]     =						1
						snake_case__				:   int     =						self.get_dummy_canonical_hf_index_retriever()
						snake_case__				:   Tuple     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__,		snake_case__,		snake_case__				:   Any     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=UpperCAmelCase_ )
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertEqual(len(UpperCAmelCase_ )							,	2 )
						self.assertEqual(sorted(doc_dicts[0] )							,	["embeddings", "id", "text", "title"] )
						self.assertEqual(len(doc_dicts[0]["id"] )							,	UpperCAmelCase_ )
						self.assertEqual(doc_dicts[0]["id"][0]							,	"1" )  # max inner product is reached with second doc
						self.assertEqual(doc_dicts[1]["id"][0]							,	"0" )  # max inner product is reached with first doc
						self.assertListEqual(doc_ids.tolist()							,	[[1], [0]] )



					def  _lowercase  (      self  :						Any ):
						snake_case__				:   List[Any]     =						self.get_dummy_canonical_hf_index_retriever()
						with tempfile.TemporaryDirectory() as tmp_dirname:
							with patch("transformers.models.rag.retrieval_rag.load_dataset" ) as mock_load_dataset:
								snake_case__				:   str     =						self.get_dummy_dataset()
								retriever.save_pretrained(UpperCAmelCase_ )
								snake_case__				:   List[str]     =						RagRetriever.from_pretrained(UpperCAmelCase_ )
								self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
								snake_case__				:   Tuple     =						np.array(
								    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
								snake_case__				:   Tuple     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=1 )
								self.assertTrue(out is not None )



					def  _lowercase  (      self  :						Dict ):
						snake_case__				:   str     =						1
						snake_case__				:   Tuple     =						self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_ )
						snake_case__				:   Any     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__,		snake_case__,		snake_case__				:   Union[str, Any]     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=UpperCAmelCase_ )
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertEqual(len(UpperCAmelCase_ )							,	2 )
						self.assertEqual(sorted(doc_dicts[0] )							,	["embeddings", "id", "text", "title"] )
						self.assertEqual(len(doc_dicts[0]["id"] )							,	UpperCAmelCase_ )
						self.assertEqual(doc_dicts[0]["id"][0]							,	"1" )  # max inner product is reached with second doc
						self.assertEqual(doc_dicts[1]["id"][0]							,	"0" )  # max inner product is reached with first doc
						self.assertListEqual(doc_ids.tolist()							,	[[1], [0]] )



					def  _lowercase  (      self  :						Dict ):
						snake_case__				:   List[str]     =						self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_ )
						with tempfile.TemporaryDirectory() as tmp_dirname:
							retriever.save_pretrained(UpperCAmelCase_ )
							snake_case__				:   Union[str, Any]     =						RagRetriever.from_pretrained(UpperCAmelCase_ )
							self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
							snake_case__				:   Optional[Any]     =						np.array(
							    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
							snake_case__				:   Any     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=1 )
							self.assertTrue(out is not None )



					def  _lowercase  (      self  :						List[str] ):
						snake_case__				:   Optional[int]     =						1
						snake_case__				:   Optional[Any]     =						self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_ )
						snake_case__				:   Union[str, Any]     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__,		snake_case__,		snake_case__				:   Optional[Any]     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=UpperCAmelCase_ )
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertEqual(len(UpperCAmelCase_ )							,	2 )
						self.assertEqual(sorted(doc_dicts[0] )							,	["embeddings", "id", "text", "title"] )
						self.assertEqual(len(doc_dicts[0]["id"] )							,	UpperCAmelCase_ )
						self.assertEqual(doc_dicts[0]["id"][0]							,	"1" )  # max inner product is reached with second doc
						self.assertEqual(doc_dicts[1]["id"][0]							,	"0" )  # max inner product is reached with first doc
						self.assertListEqual(doc_ids.tolist()							,	[[1], [0]] )



					def  _lowercase  (      self  :						Dict ):
						snake_case__				:   Optional[Any]     =						self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_ )
						with tempfile.TemporaryDirectory() as tmp_dirname:
							retriever.save_pretrained(UpperCAmelCase_ )
							snake_case__				:   Optional[int]     =						RagRetriever.from_pretrained(UpperCAmelCase_ )
							self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
							snake_case__				:   Dict     =						np.array(
							    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
							snake_case__				:   List[Any]     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=1 )
							self.assertTrue(out is not None )



					def  _lowercase  (      self  :						Optional[Any] ):
						snake_case__				:   Tuple     =						1
						snake_case__				:   Optional[Any]     =						self.get_dummy_legacy_index_retriever()
						snake_case__				:   Dict     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__,		snake_case__,		snake_case__				:   Optional[int]     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=UpperCAmelCase_ )
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertEqual(len(UpperCAmelCase_ )							,	2 )
						self.assertEqual(sorted(doc_dicts[0] )							,	["text", "title"] )
						self.assertEqual(len(doc_dicts[0]["text"] )							,	UpperCAmelCase_ )
						self.assertEqual(doc_dicts[0]["text"][0]							,	"bar" )  # max inner product is reached with second doc
						self.assertEqual(doc_dicts[1]["text"][0]							,	"foo" )  # max inner product is reached with first doc
						self.assertListEqual(doc_ids.tolist()							,	[[1], [0]] )



					def  _lowercase  (      self  :						List[str] ):
						snake_case__				:   Any     =						self.get_dummy_legacy_index_retriever()
						with tempfile.TemporaryDirectory() as tmp_dirname:
							retriever.save_pretrained(UpperCAmelCase_ )
							snake_case__				:   List[Any]     =						RagRetriever.from_pretrained(UpperCAmelCase_ )
							self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
							snake_case__				:   Dict     =						np.array(
							    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
							snake_case__				:   Optional[Any]     =						retriever.retrieve(UpperCAmelCase_							,	n_docs=1 )
							self.assertTrue(out is not None )



					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def  _lowercase  (      self  :						Union[str, Any] ):
						import torch

						snake_case__				:   List[str]     =						1
						snake_case__				:   int     =						self.get_dummy_canonical_hf_index_retriever()
						snake_case__				:   List[Any]     =						[[5, 7], [1_0, 1_1]]
						snake_case__				:   int     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__				:   Dict     =						retriever(UpperCAmelCase_							,	UpperCAmelCase_							,	prefix=retriever.config.generator.prefix							,	n_docs=UpperCAmelCase_ )
						snake_case__,		snake_case__,		snake_case__				:   str     =						(
						    out["context_input_ids"],
						    out["context_attention_mask"],
						    out["retrieved_doc_embeds"],
						)
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
						self.assertIsInstance(UpperCAmelCase_							,	UpperCAmelCase_ )
						self.assertIsInstance(UpperCAmelCase_							,	np.ndarray )

						snake_case__				:   str     =						retriever(
						    UpperCAmelCase_							,	UpperCAmelCase_							,	prefix=retriever.config.generator.prefix							,	n_docs=UpperCAmelCase_							,	return_tensors="pt"							,	)
						snake_case__,		snake_case__,		snake_case__,		snake_case__				:   Dict     =						(  # noqa: F841
						    out["context_input_ids"],
						    out["context_attention_mask"],
						    out["retrieved_doc_embeds"],
						    out["doc_ids"],
						)
						self.assertEqual(retrieved_doc_embeds.shape							,	(2, n_docs, self.retrieval_vector_size) )
						self.assertIsInstance(UpperCAmelCase_							,	torch.Tensor )
						self.assertIsInstance(UpperCAmelCase_							,	torch.Tensor )
						self.assertIsInstance(UpperCAmelCase_							,	torch.Tensor )






					@require_torch
					@require_tokenizers
					@require_sentencepiece
					def  _lowercase  (      self  :						List[Any] ):
						snake_case__				:   Tuple     =						self.get_dpr_ctx_encoder_tokenizer()
						snake_case__				:   Tuple     =						1
						snake_case__				:   Union[str, Any]     =						self.get_dummy_custom_hf_index_retriever(from_disk=UpperCAmelCase_ )
						retriever.set_ctx_encoder_tokenizer(UpperCAmelCase_ )

						snake_case__				:   Optional[Any]     =						[[5, 7], [1_0, 1_1]]
						snake_case__				:   Tuple     =						np.array(
						    [np.ones(self.retrieval_vector_size ), -np.ones(self.retrieval_vector_size )]							,	dtype=np.floataa )
						snake_case__				:   Union[str, Any]     =						retriever(UpperCAmelCase_							,	UpperCAmelCase_							,	prefix=retriever.config.generator.prefix							,	n_docs=UpperCAmelCase_ )

						self.assertEqual(
						    len(UpperCAmelCase_ )							,	6 )  # check whether the retriever output consist of 6 attributes including tokenized docs
						self.assertEqual(
						    all(k in out for k in ("tokenized_doc_ids", "tokenized_doc_attention_mask") )							,	UpperCAmelCase_ )  # check for doc token related keys in dictionary.











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

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

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

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

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



