<commit_before># 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.

from bisect import bisect_left
from collections import defaultdict
from concurrent.futures import Future
from functools import lru_cache
from itertools import chain
from math import ceil
from operator import attrgetter
from statistics import mean
import gc
import logging
import os

from cytoolz.itertoolz import pluck

from bndl.compute.dataset import Dataset, Partition
from bndl.compute.scheduler import DependenciesFailed
from bndl.compute.storage import ContainerFactory, InMemoryContainer, SharedMemoryContainer
from bndl.compute.tasks import TaskCancelled
from bndl.compute.tasks import task_context
from bndl.net import rmi
from bndl.net.connection import NotConnected
from bndl.net.rmi import InvocationException
from bndl.util.collection import batch as batch_data, ensure_collection, flatten
from bndl.util.conf import Float
from bndl.util.funcs import star_prefetch, identity
from bndl.util.hash import portable_hash
from cyheapq import merge


PARTITIONED_SORT_SPLITS = 128


logger = logging.getLogger(__name__)


block_size_mb_net = Float(8, desc='Target (maximum) size (in megabytes) of blocks created in the '
                                  'shuffle write (during spill and for transmission to shuffle '
                                  'read).')
block_size_mb_merge = Float(100, desc='Target (maximum) size (in megabytes) of blocks created by '
                                      'spilling for merge sort in shuffle read. Determines size '
                                      'of files to merge.')
chunk_size_mb_merge = Float(1, desc='Target (maximum) size (in megabytes) of chunks created by '
                                    'spilling for merge sort in shuffle read. Determines size '
                                    'of reads during merge.')


class Bucket:
    '''
    Base bucket class. Implements spilling to disk / serializing (spilling) to memory. Any
    implementation must implement add(element) and extend(sequence). The sorted property
    indicates whether the elements in the buckets are considered sorted (and thus elements
    from different buckets should be merge sorted in the shuffle read).
    '''

    sorted = False

    def __init__(self, bucket_id, key, comb, block_size_mb, chunk_size_mb, memory_container, disk_container):
        self.id = bucket_id
        self.key = key
        self.comb = comb
        self.block_size_mb = block_size_mb
        self.chunk_size_mb = chunk_size_mb
        self.memory_container = memory_container
        self.disk_container = disk_container

        # batches (lists) of blocks serialized / spilled
        self.batches = []

        # estimate of the size of an element (after combining)
        self.element_size = None


    @property
    def memory_size(self):
        '''
        The memory footprint of a bucket in bytes. This includes elements in the bucket and the
        size of serialized in memory blocks.

        When this bucket hasn't been serialized / spilled yet, the memory_size is grossly
        inaccurate; but can be used to compare with other buckets with similar contents which
        haven't been serialized / spilled either.
        '''
        return len(self) * (self.element_size or 1)


    def _estimate_element_size(self, data):
        '''
        Estimate the size of a element from data.
        '''
        try:
            test_set = data[:max(10, len(data) // 1000)]
            c = self.memory_container(self.id + ('tmp',))
            c.write(test_set)
            return c.size // len(test_set)
        finally:
            c.clear()


    def _to_blocks(self):
        # apply combiner if any
        data = ensure_collection(self.comb(self)) if self.comb else list(self)
        # and clear self
        self.clear()

        # calculate the size of an element for splitting the batch into several blocks
        if self.element_size is None:
            self.element_size = self._estimate_element_size(data)

        # split into several blocks and chunks if necessary
        block_size_recs = ceil(self.block_size_mb * 1024 * 1024 / self.element_size)

        if block_size_recs > len(data):
            blocks = [data]
        else:
            blocks = list(batch_data(data, block_size_recs))
        if self.chunk_size_mb != self.block_size_mb:
            chunk_size_recs = ceil(self.chunk_size_mb * 1024 * 1024 / self.element_size)<commit_msg>fail with the right exception if something goes wrong in the data container<commit_after># 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.

from bisect import bisect_left
from collections import defaultdict
from concurrent.futures import Future
from functools import lru_cache
from itertools import chain
from math import ceil
from operator import attrgetter
from statistics import mean
import gc
import logging
import os

from cytoolz.itertoolz import pluck

from bndl.compute.dataset import Dataset, Partition
from bndl.compute.scheduler import DependenciesFailed
from bndl.compute.storage import ContainerFactory, InMemoryContainer, SharedMemoryContainer
from bndl.compute.tasks import TaskCancelled
from bndl.compute.tasks import task_context
from bndl.net import rmi
from bndl.net.connection import NotConnected
from bndl.net.rmi import InvocationException
from bndl.util.collection import batch as batch_data, ensure_collection, flatten
from bndl.util.conf import Float
from bndl.util.funcs import star_prefetch, identity
from bndl.util.hash import portable_hash
from cyheapq import merge


PARTITIONED_SORT_SPLITS = 128


logger = logging.getLogger(__name__)


block_size_mb_net = Float(8, desc='Target (maximum) size (in megabytes) of blocks created in the '
                                  'shuffle write (during spill and for transmission to shuffle '
                                  'read).')
block_size_mb_merge = Float(100, desc='Target (maximum) size (in megabytes) of blocks created by '
                                      'spilling for merge sort in shuffle read. Determines size '
                                      'of files to merge.')
chunk_size_mb_merge = Float(1, desc='Target (maximum) size (in megabytes) of chunks created by '
                                    'spilling for merge sort in shuffle read. Determines size '
                                    'of reads during merge.')


class Bucket:
    '''
    Base bucket class. Implements spilling to disk / serializing (spilling) to memory. Any
    implementation must implement add(element) and extend(sequence). The sorted property
    indicates whether the elements in the buckets are considered sorted (and thus elements
    from different buckets should be merge sorted in the shuffle read).
    '''

    sorted = False

    def __init__(self, bucket_id, key, comb, block_size_mb, chunk_size_mb, memory_container, disk_container):
        self.id = bucket_id
        self.key = key
        self.comb = comb
        self.block_size_mb = block_size_mb
        self.chunk_size_mb = chunk_size_mb
        self.memory_container = memory_container
        self.disk_container = disk_container

        # batches (lists) of blocks serialized / spilled
        self.batches = []

        # estimate of the size of an element (after combining)
        self.element_size = None


    @property
    def memory_size(self):
        '''
        The memory footprint of a bucket in bytes. This includes elements in the bucket and the
        size of serialized in memory blocks.

        When this bucket hasn't been serialized / spilled yet, the memory_size is grossly
        inaccurate; but can be used to compare with other buckets with similar contents which
        haven't been serialized / spilled either.
        '''
        return len(self) * (self.element_size or 1)


    def _estimate_element_size(self, data):
        '''
        Estimate the size of a element from data.
        '''
        test_set = data[:max(10, len(data) // 1000)]
        c = self.memory_container(self.id + ('tmp',))
        try:
            c.write(test_set)
            return c.size // len(test_set)
        finally:
            c.clear()


    def _to_blocks(self):
        # apply combiner if any
        data = ensure_collection(self.comb(self)) if self.comb else list(self)
        # and clear self
        self.clear()

        # calculate the size of an element for splitting the batch into several blocks
        if self.element_size is None:
            self.element_size = self._estimate_element_size(data)

        # split into several blocks and chunks if necessary
        block_size_recs = ceil(self.block_size_mb * 1024 * 1024 / self.element_size)

        if block_size_recs > len(data):
            blocks = [data]
        else:
            blocks = list(batch_data(data, block_size_recs))
        if self.chunk_size_mb != self.block_size_mb:
            chunk_size_recs = ceil(self.chunk_size_mb * 1024 * 1024 / self.element_size)
{'fraction_non_alphanumeric': 0.0633901705115346, 'fraction_numerical': 0.007021063189568706, 'mean_word_length': 2.735856125889846, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 14, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3064134', 'n_tokens_mistral': 2751, 'n_tokens_neox': 2601, 'n_words': 1003}