91cae4db-ee39-47cb-81ee-210e6aabfa97

# coding=utf-8
# Copyright 2018-2023 EvaDB
#
# 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.
import itertools
from typing import Iterator

from evadb.database import EvaDBDatabase
from evadb.executor.abstract_executor import AbstractExecutor
from evadb.models.storage.batch import Batch
from evadb.plan_nodes.rebatch_plan import RebatchPlan
from evadb.utils.generic_utils import get_size


class RebatchExecutor(AbstractExecutor):
    """
    Rebatch the data pipeline

    Arguments:
        node (AbstractPlan): The Rebatch Plan

    """

    def __init__(self, db: EvaDBDatabase, node: RebatchPlan):
        super().__init__(db, node)
        self.batch_mem_size = node.batch_mem_size
        self.batch_size = node.batch_size

    def exec(self, *args, **kwargs) -> Iterator[Batch]:
        child_iter = self.children[0].exec(**kwargs)

        # Calculate the proper batch size, use the larger one when both batch_mem_size and batch_size are set.
        first_batch = next(child_iter)
        sample_size = get_size(first_batch[:1])
        batch_size = int(max(1, self.batch_size, self.batch_mem_size / sample_size))

        current_size = 0
        current_batches = []
        for batch in itertools.chain((first_batch for _ in (0,)), child_iter):
            current_batches.append(batch)
            current_size += len(batch)
            if batch_size <= current_size:
                one_batch = Batch.concat(current_batches, copy=False)
                pos = 0
                while pos + batch_size <= current_size:
                    new_batch = one_batch[pos : pos + batch_size]
                    new_batch.reset_index()
                    yield new_batch
                    pos += batch_size
                if pos < current_size:
                    current_batches = [one_batch[pos:]]
                    current_size = len(current_batches[0])
                else:
                    current_batches = []
                    current_size = 0
        last_batch = Batch.concat(current_batches, copy=False)
        yield last_batch

1	# coding=utf-8
2	# Copyright 2018-2023 EvaDB
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15	import itertools	×
16	from typing import Iterator	×
17
18	from evadb.database import EvaDBDatabase	×
19	from evadb.executor.abstract_executor import AbstractExecutor	×
20	from evadb.models.storage.batch import Batch	×
21	from evadb.plan_nodes.rebatch_plan import RebatchPlan	×
22	from evadb.utils.generic_utils import get_size	×
23
24
25	class RebatchExecutor(AbstractExecutor):	×
26	"""
27	Rebatch the data pipeline
28
29	Arguments:
30	node (AbstractPlan): The Rebatch Plan
31
32	"""
33
34	def __init__(self, db: EvaDBDatabase, node: RebatchPlan):	×
35	super().__init__(db, node)	×
36	self.batch_mem_size = node.batch_mem_size	×
37	self.batch_size = node.batch_size	×
38
39	def exec(self, args, *kwargs) -> Iterator[Batch]:	×
40	child_iter = self.children[0].exec(**kwargs)	×
41
42	# Calculate the proper batch size, use the larger one when both batch_mem_size and batch_size are set.
43	first_batch = next(child_iter)	×
44	sample_size = get_size(first_batch[:1])	×
45	batch_size = int(max(1, self.batch_size, self.batch_mem_size / sample_size))	×
46
47	current_size = 0	×
48	current_batches = []	×
49	for batch in itertools.chain((first_batch for _ in (0,)), child_iter):	×
50	current_batches.append(batch)	×
51	current_size += len(batch)	×
52	if batch_size <= current_size:	×
53	one_batch = Batch.concat(current_batches, copy=False)	×
54	pos = 0	×
55	while pos + batch_size <= current_size:	×
56	new_batch = one_batch[pos : pos + batch_size]	×
57	new_batch.reset_index()	×
58	yield new_batch	×
59	pos += batch_size	×
60	if pos < current_size:	×
61	current_batches = [one_batch[pos:]]	×
62	current_size = len(current_batches[0])	×
63	else:
64	current_batches = []	×
65	current_size = 0	×
66	last_batch = Batch.concat(current_batches, copy=False)	×
67	yield last_batch	×

georgia-tech-db / eva / 91cae4db-ee39-47cb-81ee-210e6aabfa97

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous