
import unittest

import numpy as np

from transformers.testing_utils import require_pytesseract, require_torch
from transformers.utils import is_pytesseract_available, is_torch_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_pytesseract_available():
      from PIL import Image

      from transformers import LayoutLMvaImageProcessor





class 				A__						(      unittest.TestCase):
      def __init__( self       ,    __magic_name__       ,    __magic_name__=7       ,    __magic_name__=3       ,    __magic_name__=1_8       ,    __magic_name__=3_0       ,    __magic_name__=4_0_0       ,    __magic_name__=True       ,    __magic_name__=None       ,    __magic_name__=True       ,    ):
       lowerCamelCase						:      Optional[Any]							     = size if size is not None else {'''height''': 1_8, '''width''': 1_8}
       lowerCamelCase						:      Dict							     = parent
       lowerCamelCase						:      Any							     = batch_size
       lowerCamelCase						:      Dict							     = num_channels
       lowerCamelCase						:      Optional[Any]							     = image_size
       lowerCamelCase						:      Dict							     = min_resolution
       lowerCamelCase						:      Optional[int]							     = max_resolution
       lowerCamelCase						:      Any							     = do_resize
       lowerCamelCase						:      List[Any]							     = size
       lowerCamelCase						:      List[Any]							     = apply_ocr






      def 		UpperCamelCase__				( self ):
       return {"do_resize": self.do_resize, "size": self.size, "apply_ocr": self.apply_ocr}





@require_torch
@require_pytesseract
class 				A__						(      _SCREAMING_SNAKE_CASE    ,  unittest.TestCase):
      _UpperCAmelCase							:       Union[str, Any]       			=     LayoutLMvaImageProcessor if is_pytesseract_available() else None
      def 		UpperCamelCase__				( self ):
       lowerCamelCase						:      Optional[Any]							     = LayoutLMvaImageProcessingTester(self )
      @property
      def 		UpperCamelCase__				( self ):
       return self.image_processor_tester.prepare_image_processor_dict()
      def 		UpperCamelCase__				( self ):
       lowerCamelCase						:      Tuple							     = self.image_processing_class(**self.image_processor_dict )
       self.assertTrue(hasattr(__UpperCAmelCase       ,    """do_resize""" ) )
       self.assertTrue(hasattr(__UpperCAmelCase       ,    """size""" ) )
       self.assertTrue(hasattr(__UpperCAmelCase       ,    """apply_ocr""" ) )
      def 		UpperCamelCase__				( self ):
       lowerCamelCase						:      Any							     = self.image_processing_class.from_dict(self.image_processor_dict )
       self.assertEqual(image_processor.size       ,    {"""height""": 1_8, """width""": 1_8} )

       lowerCamelCase						:      int							     = self.image_processing_class.from_dict(self.image_processor_dict       ,    size=4_2 )
       self.assertEqual(image_processor.size       ,    {"""height""": 4_2, """width""": 4_2} )
      def 		UpperCamelCase__				( self ):
       pass
      def 		UpperCamelCase__				( self ):
       # Initialize image_processing
       lowerCamelCase						:      Optional[Any]							     = self.image_processing_class(**self.image_processor_dict )
       # create random PIL images
       lowerCamelCase						:      Union[str, Any]							     = prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__UpperCAmelCase )
       for image in image_inputs:
        self.assertIsInstance(__UpperCAmelCase       ,    Image.Image )

       # Test not batched input
       lowerCamelCase						:      str							     = image_processing(image_inputs[0]       ,    return_tensors="""pt""" )
       self.assertEqual(
           encoding.pixel_values.shape       ,    (
               1,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )

       self.assertIsInstance(encoding.words       ,    __UpperCAmelCase )
       self.assertIsInstance(encoding.boxes       ,    __UpperCAmelCase )

       # Test batched
       lowerCamelCase						:      int							     = image_processing(__UpperCAmelCase       ,    return_tensors="""pt""" ).pixel_values
       self.assertEqual(
           encoded_images.shape       ,    (
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )
      def 		UpperCamelCase__				( self ):
       # Initialize image_processing
       lowerCamelCase						:      Any							     = self.image_processing_class(**self.image_processor_dict )
       # create random numpy tensors
       lowerCamelCase						:      Optional[Any]							     = prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__UpperCAmelCase       ,    numpify=__UpperCAmelCase )
       for image in image_inputs:
        self.assertIsInstance(__UpperCAmelCase       ,    np.ndarray )

       # Test not batched input
       lowerCamelCase						:      Dict							     = image_processing(image_inputs[0]       ,    return_tensors="""pt""" ).pixel_values
       self.assertEqual(
           encoded_images.shape       ,    (
               1,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )

       # Test batched
       lowerCamelCase						:      int							     = image_processing(__UpperCAmelCase       ,    return_tensors="""pt""" ).pixel_values
       self.assertEqual(
           encoded_images.shape       ,    (
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )
      def 		UpperCamelCase__				( self ):
       # Initialize image_processing
       lowerCamelCase						:      Optional[Any]							     = self.image_processing_class(**self.image_processor_dict )
       # create random PyTorch tensors
       lowerCamelCase						:      Optional[Any]							     = prepare_image_inputs(self.image_processor_tester       ,    equal_resolution=__UpperCAmelCase       ,    torchify=__UpperCAmelCase )
       for image in image_inputs:
        self.assertIsInstance(__UpperCAmelCase       ,    torch.Tensor )

       # Test not batched input
       lowerCamelCase						:      List[str]							     = image_processing(image_inputs[0]       ,    return_tensors="""pt""" ).pixel_values
       self.assertEqual(
           encoded_images.shape       ,    (
               1,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )

       # Test batched
       lowerCamelCase						:      Dict							     = image_processing(__UpperCAmelCase       ,    return_tensors="""pt""" ).pixel_values
       self.assertEqual(
           encoded_images.shape       ,    (
               self.image_processor_tester.batch_size,
               self.image_processor_tester.num_channels,
               self.image_processor_tester.size["""height"""],
               self.image_processor_tester.size["""width"""],
           )       ,    )






      def 		UpperCamelCase__				( self ):
       # with apply_OCR = True
       lowerCamelCase						:      Dict							     = LayoutLMvaImageProcessor()

       from datasets import load_dataset

       lowerCamelCase						:      Dict							     = load_dataset("""hf-internal-testing/fixtures_docvqa"""       ,    split="""test""" )

       lowerCamelCase						:      Any							     = Image.open(ds[0]["""file"""] ).convert("""RGB""" )

       lowerCamelCase						:      List[Any]							     = image_processing(__UpperCAmelCase       ,    return_tensors="""pt""" )

       self.assertEqual(encoding.pixel_values.shape       ,    (1, 3, 2_2_4, 2_2_4) )
       self.assertEqual(len(encoding.words )       ,    len(encoding.boxes ) )

       # fmt: off
       # the words and boxes were obtained with Tesseract 4.1.1
       lowerCamelCase						:      Tuple							     = [['''11:14''', '''to''', '''11:39''', '''a.m''', '''11:39''', '''to''', '''11:44''', '''a.m.''', '''11:44''', '''a.m.''', '''to''', '''12:25''', '''p.m.''', '''12:25''', '''to''', '''12:58''', '''p.m.''', '''12:58''', '''to''', '''4:00''', '''p.m.''', '''2:00''', '''to''', '''5:00''', '''p.m.''', '''Coffee''', '''Break''', '''Coffee''', '''will''', '''be''', '''served''', '''for''', '''men''', '''and''', '''women''', '''in''', '''the''', '''lobby''', '''adjacent''', '''to''', '''exhibit''', '''area.''', '''Please''', '''move''', '''into''', '''exhibit''', '''area.''', '''(Exhibits''', '''Open)''', '''TRRF''', '''GENERAL''', '''SESSION''', '''(PART''', '''|)''', '''Presiding:''', '''Lee''', '''A.''', '''Waller''', '''TRRF''', '''Vice''', '''President''', '''“Introductory''', '''Remarks”''', '''Lee''', '''A.''', '''Waller,''', '''TRRF''', '''Vice''', '''Presi-''', '''dent''', '''Individual''', '''Interviews''', '''with''', '''TRRF''', '''Public''', '''Board''', '''Members''', '''and''', '''Sci-''', '''entific''', '''Advisory''', '''Council''', '''Mem-''', '''bers''', '''Conducted''', '''by''', '''TRRF''', '''Treasurer''', '''Philip''', '''G.''', '''Kuehn''', '''to''', '''get''', '''answers''', '''which''', '''the''', '''public''', '''refrigerated''', '''warehousing''', '''industry''', '''is''', '''looking''', '''for.''', '''Plus''', '''questions''', '''from''', '''the''', '''floor.''', '''Dr.''', '''Emil''', '''M.''', '''Mrak,''', '''University''', '''of''', '''Cal-''', '''ifornia,''', '''Chairman,''', '''TRRF''', '''Board;''', '''Sam''', '''R.''', '''Cecil,''', '''University''', '''of''', '''Georgia''', '''College''', '''of''', '''Agriculture;''', '''Dr.''', '''Stanley''', '''Charm,''', '''Tufts''', '''University''', '''School''', '''of''', '''Medicine;''', '''Dr.''', '''Robert''', '''H.''', '''Cotton,''', '''ITT''', '''Continental''', '''Baking''', '''Company;''', '''Dr.''', '''Owen''', '''Fennema,''', '''University''', '''of''', '''Wis-''', '''consin;''', '''Dr.''', '''Robert''', '''E.''', '''Hardenburg,''', '''USDA.''', '''Questions''', '''and''', '''Answers''', '''Exhibits''', '''Open''', '''Capt.''', '''Jack''', '''Stoney''', '''Room''', '''TRRF''', '''Scientific''', '''Advisory''', '''Council''', '''Meeting''', '''Ballroom''', '''Foyer''']]  # noqa: E231
       lowerCamelCase						:      int							     = [[[1_4_1, 5_7, 2_1_4, 6_9], [2_2_8, 5_8, 2_5_2, 6_9], [1_4_1, 7_5, 2_1_6, 8_8], [2_3_0, 7_9, 2_8_0, 8_8], [1_4_2, 2_6_0, 2_1_8, 2_7_3], [2_3_0, 2_6_1, 2_5_5, 2_7_3], [1_4_3, 2_7_9, 2_1_8, 2_9_0], [2_3_1, 2_8_2, 2_9_0, 2_9_1], [1_4_3, 3_4_2, 2_1_8, 3_5_4], [2_3_1, 3_4_5, 2_8_9, 3_5_5], [2_0_2, 3_6_2, 2_2_7, 3_7_3], [1_4_3, 3_7_9, 2_2_0, 3_9_2], [2_3_1, 3_8_2, 2_9_1, 3_9_4], [1_4_4, 7_1_4, 2_2_0, 7_2_6], [2_3_1, 7_1_5, 2_5_6, 7_2_6], [1_4_4, 7_3_2, 2_2_0, 7_4_5], [2_3_2, 7_3_6, 2_9_1, 7_4_7], [1_4_4, 7_6_9, 2_1_8, 7_8_2], [2_3_1, 7_7_0, 2_5_6, 7_8_2], [1_4_1, 7_8_8, 2_0_2, 8_0_1], [2_1_5, 7_9_1, 2_7_4, 8_0_4], [1_4_3, 8_2_6, 2_0_4, 8_3_8], [2_1_5, 8_2_6, 2_4_0, 8_3_8], [1_4_2, 8_4_4, 2_0_2, 8_5_7], [2_1_5, 8_4_7, 2_7_4, 8_5_9], [3_3_4, 5_7, 4_2_7, 6_9], [4_4_0, 5_7, 5_2_2, 6_9], [3_6_9, 7_5, 4_6_1, 8_8], [4_6_9, 7_5, 5_1_6, 8_8], [5_2_8, 7_6, 5_6_2, 8_8], [5_7_0, 7_6, 6_6_7, 8_8], [6_7_5, 7_5, 7_1_1, 8_7], [7_2_1, 7_9, 7_7_8, 8_8], [7_8_9, 7_5, 8_4_0, 8_8], [3_6_9, 9_7, 4_7_0, 1_0_7], [4_8_4, 9_4, 5_0_7, 1_0_6], [5_1_8, 9_4, 5_6_2, 1_0_7], [5_7_6, 9_4, 6_5_5, 1_1_0], [6_6_8, 9_4, 7_9_2, 1_0_9], [8_0_4, 9_5, 8_2_9, 1_0_7], [3_6_9, 1_1_3, 4_6_5, 1_2_5], [4_7_7, 1_1_6, 5_4_7, 1_2_5], [5_6_2, 1_1_3, 6_5_8, 1_2_5], [6_7_1, 1_1_6, 7_4_8, 1_2_5], [7_6_1, 1_1_3, 8_1_1, 1_2_5], [3_6_9, 1_3_1, 4_6_5, 1_4_3], [4_7_7, 1_3_3, 5_4_8, 1_4_3], [5_6_3, 1_3_0, 6_9_8, 1_4_5], [7_1_0, 1_3_0, 8_0_2, 1_4_6], [3_3_6, 1_7_1, 4_1_2, 1_8_3], [4_2_3, 1_7_1, 5_7_2, 1_8_3], [5_8_2, 1_7_0, 7_1_6, 1_8_4], [7_2_8, 1_7_1, 8_1_7, 1_8_7], [8_2_9, 1_7_1, 8_4_4, 1_8_6], [3_3_8, 1_9_7, 4_8_2, 2_1_2], [5_0_7, 1_9_6, 5_5_7, 2_0_9], [5_6_9, 1_9_6, 5_9_5, 2_0_8], [6_1_0, 1_9_6, 7_0_2, 2_0_9], [5_0_5, 2_1_4, 5_8_3, 2_2_6], [5_9_5, 2_1_4, 6_5_6, 2_2_7], [6_7_0, 2_1_5, 8_0_7, 2_2_7], [3_3_5, 2_5_9, 5_4_3, 2_7_4], [5_5_6, 2_5_9, 7_0_8, 2_7_2], [3_7_2, 2_7_9, 4_2_2, 2_9_1], [4_3_5, 2_7_9, 4_6_0, 2_9_1], [4_7_4, 2_7_9, 5_7_4, 2_9_2], [5_8_7, 2_7_8, 6_6_4, 2_9_1], [6_7_6, 2_7_8, 7_3_8, 2_9_1], [7_5_1, 2_7_9, 8_3_4, 2_9_1], [3_7_2, 2_9_8, 4_3_4, 3_1_0], [3_3_5, 3_4_1, 4_8_3, 3_5_4], [4_9_7, 3_4_1, 6_5_5, 3_5_4], [6_6_7, 3_4_1, 7_2_8, 3_5_4], [7_4_0, 3_4_1, 8_2_5, 3_5_4], [3_3_5, 3_6_0, 4_3_0, 3_7_2], [4_4_2, 3_6_0, 5_3_4, 3_7_2], [5_4_5, 3_5_9, 6_8_7, 3_7_2], [6_9_7, 3_6_0, 7_5_4, 3_7_2], [7_6_5, 3_6_0, 8_2_3, 3_7_3], [3_3_4, 3_7_8, 4_2_8, 3_9_1], [4_4_0, 3_7_8, 5_7_7, 3_9_4], [5_9_0, 3_7_8, 7_0_5, 3_9_1], [7_2_0, 3_7_8, 8_0_1, 3_9_1], [3_3_4, 3_9_7, 4_0_0, 4_0_9], [3_7_0, 4_1_6, 5_2_9, 4_2_9], [5_4_4, 4_1_6, 5_7_6, 4_3_2], [5_8_7, 4_1_6, 6_6_5, 4_2_8], [6_7_7, 4_1_6, 8_1_4, 4_2_9], [3_7_2, 4_3_5, 4_5_2, 4_5_0], [4_6_5, 4_3_4, 4_9_5, 4_4_7], [5_1_1, 4_3_4, 6_0_0, 4_4_7], [6_1_1, 4_3_6, 6_3_7, 4_4_7], [6_4_9, 4_3_6, 6_9_4, 4_5_1], [7_0_5, 4_3_8, 8_2_4, 4_4_7], [3_6_9, 4_5_3, 4_5_2, 4_6_6], [4_6_4, 4_5_4, 5_0_9, 4_6_6], [5_2_2, 4_5_3, 6_1_1, 4_6_9], [6_2_5, 4_5_3, 7_9_2, 4_6_9], [3_7_0, 4_7_2, 5_5_6, 4_8_8], [5_7_0, 4_7_2, 6_8_4, 4_8_7], [6_9_7, 4_7_2, 7_1_8, 4_8_5], [7_3_2, 4_7_2, 8_3_5, 4_8_8], [3_6_9, 4_9_0, 4_1_1, 5_0_3], [4_2_5, 4_9_0, 4_8_4, 5_0_3], [4_9_6, 4_9_0, 6_3_5, 5_0_6], [6_4_5, 4_9_0, 7_0_7, 5_0_3], [7_1_8, 4_9_1, 7_6_1, 5_0_3], [7_7_1, 4_9_0, 8_4_0, 5_0_3], [3_3_6, 5_1_0, 3_7_4, 5_2_1], [3_8_8, 5_1_0, 4_4_7, 5_2_2], [4_6_0, 5_1_0, 4_8_9, 5_2_1], [5_0_3, 5_1_0, 5_8_0, 5_2_2], [5_9_2, 5_0_9, 7_3_6, 5_2_5], [7_4_5, 5_0_9, 7_7_0, 5_2_2], [7_8_1, 5_0_9, 8_4_0, 5_2_2], [3_3_8, 5_2_8, 4_3_4, 5_4_1], [4_4_8, 5_2_8, 5_9_6, 5_4_1], [6_0_9, 5_2_7, 6_8_7, 5_4_0], [7_0_0, 5_2_8, 7_9_2, 5_4_1], [3_3_6, 5_4_6, 3_9_7, 5_5_9], [4_0_7, 5_4_6, 4_3_1, 5_5_9], [4_4_3, 5_4_6, 5_2_5, 5_6_0], [5_3_7, 5_4_6, 6_8_0, 5_6_2], [6_8_8, 5_4_6, 7_1_4, 5_5_9], [7_2_2, 5_4_6, 8_3_7, 5_6_2], [3_3_6, 5_6_5, 4_4_9, 5_8_1], [4_6_1, 5_6_5, 4_8_5, 5_7_7], [4_9_7, 5_6_5, 6_6_5, 5_8_1], [6_8_1, 5_6_5, 7_1_8, 5_7_7], [7_3_2, 5_6_5, 8_3_7, 5_8_0], [3_3_7, 5_8_4, 4_3_8, 5_9_7], [4_5_2, 5_8_3, 5_2_1, 5_9_6], [5_3_5, 5_8_4, 6_7_7, 5_9_9], [6_9_0, 5_8_3, 7_8_7, 5_9_6], [8_0_1, 5_8_3, 8_2_5, 5_9_6], [3_3_8, 6_0_2, 4_7_8, 6_1_5], [4_9_2, 6_0_2, 5_3_0, 6_1_4], [5_4_3, 6_0_2, 6_3_8, 6_1_5], [6_5_0, 6_0_2, 6_7_6, 6_1_4], [6_8_8, 6_0_2, 7_8_8, 6_1_5], [8_0_2, 6_0_2, 8_4_3, 6_1_4], [3_3_7, 6_2_1, 5_0_2, 6_3_3], [5_1_6, 6_2_1, 6_1_5, 6_3_7], [6_2_9, 6_2_1, 7_7_4, 6_3_6], [7_8_9, 6_2_1, 8_2_7, 6_3_3], [3_3_7, 6_3_9, 4_1_8, 6_5_2], [4_3_2, 6_4_0, 5_7_1, 6_5_3], [5_8_7, 6_3_9, 7_3_1, 6_5_5], [7_4_3, 6_3_9, 7_6_9, 6_5_2], [7_8_0, 6_3_9, 8_4_1, 6_5_2], [3_3_8, 6_5_8, 4_4_0, 6_7_3], [4_5_5, 6_5_8, 4_9_1, 6_7_0], [5_0_8, 6_5_8, 6_0_2, 6_7_1], [6_1_6, 6_5_8, 6_3_8, 6_7_0], [6_5_4, 6_5_8, 8_3_5, 6_7_4], [3_3_7, 6_7_7, 4_2_9, 6_8_9], [3_3_7, 7_1_4, 4_8_2, 7_2_6], [4_9_5, 7_1_4, 5_4_8, 7_2_6], [5_6_1, 7_1_4, 6_8_3, 7_2_6], [3_3_8, 7_7_0, 4_6_1, 7_8_2], [4_7_4, 7_6_9, 5_5_4, 7_8_5], [4_8_9, 7_8_8, 5_6_2, 8_0_3], [5_7_6, 7_8_8, 6_4_3, 8_0_1], [6_5_6, 7_8_7, 7_5_1, 8_0_4], [7_6_4, 7_8_8, 8_4_4, 8_0_1], [3_3_4, 8_2_5, 4_2_1, 8_3_8], [4_3_0, 8_2_4, 5_7_4, 8_3_8], [5_8_4, 8_2_4, 7_2_3, 8_4_1], [3_3_5, 8_4_4, 4_5_0, 8_5_7], [4_6_4, 8_4_3, 5_8_3, 8_6_0], [6_2_8, 8_6_2, 7_5_5, 8_7_5], [7_6_9, 8_6_1, 8_4_8, 8_7_8]]]  # noqa: E231
       # fmt: on

       self.assertListEqual(encoding.words       ,    __UpperCAmelCase )
       self.assertListEqual(encoding.boxes       ,    __UpperCAmelCase )

       # with apply_OCR = False
       lowerCamelCase						:      Union[str, Any]							     = LayoutLMvaImageProcessor(apply_ocr=__UpperCAmelCase )

       lowerCamelCase						:      List[str]							     = image_processing(__UpperCAmelCase       ,    return_tensors="""pt""" )

       self.assertEqual(encoding.pixel_values.shape       ,    (1, 3, 2_2_4, 2_2_4) )







import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


lowercase			:				Dict   =     datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
		import pyspark







@dataclass
class 					__lowercase  (						datasets.BuilderConfig				):





							"""simple docstring"""

							UpperCAmelCase_						:		Optional[datasets.Features]								=	None






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				):
							import pyspark

							def generate_fn():
														A      :					Dict      		=					df.select('''*'''     ,				pyspark.sql.functions.spark_partition_id().alias('''part_id'''		)		)
														for partition_id in partition_order:
																					A      :					Optional[Any]      		=					df_with_partition_id.select('''*'''		).where(F'part_id = {partition_id}'		).drop('''part_id'''		)
																					A      :					Dict      		=					partition_df.collect()
																					A      :					List[str]      		=					0
																					for row in rows:
																												yield F'{partition_id}_{row_id}', row.asDict()
																												row_id += 1

							return generate_fn







class 					__lowercase  (						_BaseExamplesIterable				):





							"""simple docstring"""

							def __init__(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase=None    ,						)		->       List[str]:
														A      :					List[Any]      		=					df
														A      :					List[str]      		=					partition_order or range(self.df.rdd.getNumPartitions()						)
														A      :					Optional[int]      		=					_generate_iterable_examples(self.df    ,						self.partition_order						)




							def __iter__(					self						)		->       str:
														yield from self.generate_examples_fn()




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       "SparkExamplesIterable":
														A      :					Tuple      		=					list(range(self.df.rdd.getNumPartitions()						)						)
														generator.shuffle(__UpperCAmelCase						)
														return SparkExamplesIterable(self.df    ,						partition_order=__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase						)		->       "SparkExamplesIterable":
														A      :					Tuple      		=					self.split_shard_indices_by_worker(__UpperCAmelCase    ,						__UpperCAmelCase						)
														return SparkExamplesIterable(self.df    ,						partition_order=__UpperCAmelCase						)






							@property
							def       snake_case			(					self						)		->       int:
														return len(self.partition_order						)







class 					__lowercase  (						datasets.DatasetBuilder				):





							"""simple docstring"""

							UpperCAmelCase_						:		List[str]								=	SparkConfig




							def __init__(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       Optional[Any]:
														import pyspark

														A      :					Tuple      		=					pyspark.sql.SparkSession.builder.getOrCreate()
														A      :					List[Any]      		=					df
														A      :					List[Any]      		=					working_dir

														super().__init__(
														    cache_dir=__UpperCAmelCase    ,						config_name=str(self.df.semanticHash()						)    ,						**__UpperCAmelCase    ,						)




							def       snake_case			(					self						)		->       Optional[Any]:
														# Returns the path of the created file.
														def create_cache_and_write_probe(__UpperCAmelCase						):
																					# makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
																					# already exist.
																					os.makedirs(self._cache_dir    ,						exist_ok=__UpperCAmelCase						)
																					A      :					Optional[Any]      		=					os.path.join(self._cache_dir    ,						'''fs_test''' + uuid.uuida().hex						)
																					# Opening the file in append mode will create a new file unless it already exists, in which case it will not
																					# change the file contents.
																					open(__UpperCAmelCase    ,						'''a'''						)
																					return [probe_file]

														if self._spark.conf.get('''spark.master'''    ,						''''''						).startswith('''local'''						):
																					return

														# If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
														# accessible to the driver.
														# TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
														if self._cache_dir:
																					A      :					Optional[int]      		=					(
																					    self._spark.sparkContext.parallelize(range(1						)    ,						1						).mapPartitions(__UpperCAmelCase						).collect()
																					)
																					if os.path.isfile(probe[0]						):
																												return

														raise ValueError(
														    '''When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'''						)




							def       snake_case			(					self						)		->       List[Any]:
														return datasets.DatasetInfo(features=self.config.features						)




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       str:
														return [datasets.SplitGenerator(name=datasets.Split.TRAIN						)]




							def       snake_case			(					self    ,						__UpperCAmelCase						)		->       List[Any]:
														import pyspark

														def get_arrow_batch_size(__UpperCAmelCase						):
																					for batch in it:
																												yield pa.RecordBatch.from_pydict({'''batch_bytes''': [batch.nbytes]}						)

														A      :					Any      		=					self.df.count()
														A      :					Tuple      		=					df_num_rows if df_num_rows <= 1_00 else 1_00
														# Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
														A      :					Dict      		=					(
														    self.df.limit(__UpperCAmelCase						)
														    .repartition(1						)
														    .mapInArrow(__UpperCAmelCase    ,						'''batch_bytes: long'''						)
														    .agg(pyspark.sql.functions.sum('''batch_bytes'''						).alias('''sample_bytes'''						)						)
														    .collect()[0]
														    .sample_bytes
														    / sample_num_rows
														)
														A      :					str      		=					approx_bytes_per_row * df_num_rows
														if approx_total_size > max_shard_size:
																					# Make sure there is at least one row per partition.
																					A      :					Union[str, Any]      		=					min(__UpperCAmelCase    ,						int(approx_total_size / max_shard_size						)						)
																					A      :					List[Any]      		=					self.df.repartition(__UpperCAmelCase						)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						)		->       Iterable[Tuple[int, bool, Union[int, tuple]]]:
														import pyspark

														A      :					Dict      		=					ParquetWriter if file_format == '''parquet''' else ArrowWriter
														A      :					Optional[Any]      		=					os.path.join(self._working_dir    ,						os.path.basename(__UpperCAmelCase						)						) if self._working_dir else fpath
														A      :					int      		=					file_format == '''parquet'''

														# Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
														# pickling the SparkContext.
														A      :					Any      		=					self.config.features
														A      :					Any      		=					self._writer_batch_size
														A      :					int      		=					self._fs.storage_options

														def write_arrow(__UpperCAmelCase						):
																					# Within the same SparkContext, no two task attempts will share the same attempt ID.
																					A      :					List[Any]      		=					pyspark.TaskContext().taskAttemptId()
																					A      :					Union[str, Any]      		=					next(__UpperCAmelCase    ,						__UpperCAmelCase						)
																					if first_batch is None:
																												# Some partitions might not receive any data.
																												return pa.RecordBatch.from_arrays(
																												    [[task_id], [0], [0]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)
																					A      :					Optional[int]      		=					0
																					A      :					Optional[int]      		=					writer_class(
																					    features=__UpperCAmelCase    ,						path=working_fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						writer_batch_size=__UpperCAmelCase    ,						storage_options=__UpperCAmelCase    ,						embed_local_files=__UpperCAmelCase    ,						)
																					A      :					Tuple      		=					pa.Table.from_batches([first_batch]						)
																					writer.write_table(__UpperCAmelCase						)
																					for batch in it:
																												if max_shard_size is not None and writer._num_bytes >= max_shard_size:
																																			A		,		A      :					str      		=					writer.finalize()
																																			writer.close()
																																			yield pa.RecordBatch.from_arrays(
																																			    [[task_id], [num_examples], [num_bytes]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)
																																			shard_id += 1
																																			A      :					List[str]      		=					writer_class(
																																			    features=writer._features    ,						path=working_fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						writer_batch_size=__UpperCAmelCase    ,						storage_options=__UpperCAmelCase    ,						embed_local_files=__UpperCAmelCase    ,						)
																												A      :					str      		=					pa.Table.from_batches([batch]						)
																												writer.write_table(__UpperCAmelCase						)

																					if writer._num_bytes > 0:
																												A		,		A      :					Tuple      		=					writer.finalize()
																												writer.close()
																												yield pa.RecordBatch.from_arrays(
																												    [[task_id], [num_examples], [num_bytes]]    ,						names=['''task_id''', '''num_examples''', '''num_bytes''']    ,						)

																					if working_fpath != fpath:
																												for file in os.listdir(os.path.dirname(__UpperCAmelCase						)						):
																																			A      :					Union[str, Any]      		=					os.path.join(os.path.dirname(__UpperCAmelCase						)    ,						os.path.basename(__UpperCAmelCase						)						)
																																			shutil.move(__UpperCAmelCase    ,						__UpperCAmelCase						)

														A      :					Union[str, Any]      		=					(
														    self.df.mapInArrow(__UpperCAmelCase    ,						'''task_id: long, num_examples: long, num_bytes: long'''						)
														    .groupBy('''task_id'''						)
														    .agg(
														        pyspark.sql.functions.sum('''num_examples'''						).alias('''total_num_examples'''						)    ,						pyspark.sql.functions.sum('''num_bytes'''						).alias('''total_num_bytes'''						)    ,						pyspark.sql.functions.count('''num_bytes'''						).alias('''num_shards'''						)    ,						pyspark.sql.functions.collect_list('''num_examples'''						).alias('''shard_lengths'''						)    ,						)
														    .collect()
														)
														for row in stats:
																					yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase = "arrow"    ,						__UpperCAmelCase = None    ,						__UpperCAmelCase = None    ,						**__UpperCAmelCase    ,						)		->       List[str]:
														self._validate_cache_dir()

														A      :					int      		=					convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE						)
														self._repartition_df_if_needed(__UpperCAmelCase						)
														A      :					Any      		=					not is_remote_filesystem(self._fs						)
														A      :					Union[str, Any]      		=					os.path.join if is_local else posixpath.join

														A      :					Any      		=					'''-TTTTT-SSSSS-of-NNNNN'''
														A      :					Any      		=					f'{self.name}-{split_generator.name}{SUFFIX}.{file_format}'
														A      :					List[Any]      		=					path_join(self._output_dir    ,						__UpperCAmelCase						)

														A      :					Union[str, Any]      		=					0
														A      :					Any      		=					0
														A      :					Tuple      		=					0
														A      :					Union[str, Any]      		=					[]
														A      :					List[str]      		=					[]

														for task_id, content in self._prepare_split_single(__UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase						):
																					(
																					    (
																					    A
																					)		,		(
																					    A
																					)		,		(
																					    A
																					)		,		(
																					    A
																					)		,		
																					)      :					List[Any]      		=					content
																					if num_bytes > 0:
																												total_num_examples += num_examples
																												total_num_bytes += num_bytes
																												total_shards += num_shards
																												task_id_and_num_shards.append((task_id, num_shards)						)
																												all_shard_lengths.extend(__UpperCAmelCase						)

														A      :					List[str]      		=					total_num_examples
														A      :					List[str]      		=					total_num_bytes

														# should rename everything at the end
														logger.debug(f'Renaming {total_shards} shards.'						)
														if total_shards > 1:
																					A      :					Tuple      		=					all_shard_lengths

																					# Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
																					# pickling error due to pickling the SparkContext.
																					A      :					Union[str, Any]      		=					self._fs

																					# use the -SSSSS-of-NNNNN pattern
																					def _rename_shard(
																					    __UpperCAmelCase    ,						__UpperCAmelCase    ,						__UpperCAmelCase    ,						):
																												rename(
																												    __UpperCAmelCase    ,						fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						fpath.replace('''TTTTT-SSSSS'''    ,						f'{global_shard_id:05d}'						).replace('''NNNNN'''    ,						f'{total_shards:05d}'						)    ,						)

																					A      :					Tuple      		=					[]
																					A      :					Optional[int]      		=					0
																					for i in range(len(__UpperCAmelCase						)						):
																												A		,		A      :					Tuple      		=					task_id_and_num_shards[i]
																												for shard_id in range(__UpperCAmelCase						):
																																			args.append([task_id, shard_id, global_shard_id]						)
																																			global_shard_id += 1
																					self._spark.sparkContext.parallelize(__UpperCAmelCase    ,						len(__UpperCAmelCase						)						).map(lambda __UpperCAmelCase						: _rename_shard(*__UpperCAmelCase						)						).collect()
														else:
																					# don't use any pattern
																					A      :					int      		=					0
																					A      :					Tuple      		=					task_id_and_num_shards[0][0]
																					self._rename(
																					    fpath.replace('''SSSSS'''    ,						f'{shard_id:05d}'						).replace('''TTTTT'''    ,						f'{task_id:05d}'						)    ,						fpath.replace(__UpperCAmelCase    ,						''''''						)    ,						)






							def       snake_case			(					self    ,						__UpperCAmelCase    ,						)		->       SparkExamplesIterable:
														return SparkExamplesIterable(self.df						)






