Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 2bcc689

Browse filesBrowse files
aneepcttseaver
aneepct
authored and
tseaver
committed
BigTable: Create MutationBatcher for bigtable (googleapis#5651)
1 parent fa2a5f1 commit 2bcc689
Copy full SHA for 2bcc689

6 files changed

+393Lines changed: 393 additions & 0 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎bigtable/google/cloud/bigtable/batcher.py‎

Copy file name to clipboard
+151Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
# Copyright 2018 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""User friendly container for Google Cloud Bigtable MutationBatcher."""
16+
17+
18+
FLUSH_COUNT = 1000
19+
MAX_MUTATIONS = 100000
20+
MAX_ROW_BYTES = 5242880 # 5MB
21+
22+
23+
class MaxMutationsError(ValueError):
24+
"""The number of mutations for bulk request is too big."""
25+
26+
27+
class MutationsBatcher(object):
28+
""" A MutationsBatcher is used in batch cases where the number of mutations
29+
is large or unknown. It will store DirectRows in memory until one of the
30+
size limits is reached, or an explicit call to flush() is performed. When
31+
a flush event occurs, the DirectRows in memory will be sent to Cloud
32+
Bigtable. Batching mutations is more efficient than sending individual
33+
request.
34+
35+
This class is not suited for usage in systems where each mutation
36+
needs to guaranteed to be sent, since calling mutate may only result in an
37+
in-memory change. In a case of a system crash, any DirectRows remaining in
38+
memory will not necessarily be sent to the service, even after the
39+
completion of the mutate() method.
40+
41+
TODO: Performance would dramatically improve if this class had the
42+
capability of asynchronous, parallel RPCs.
43+
44+
:type table: class
45+
:param table: class:`~google.cloud.bigtable.table.Table`.
46+
47+
:type flush_count: int
48+
:param flush_count: (Optional) Max number of rows to flush. If it
49+
reaches the max number of rows it calls finish_batch() to mutate the
50+
current row batch. Default is FLUSH_COUNT (1000 rows).
51+
52+
:type max_row_bytes: int
53+
:param max_row_bytes: (Optional) Max number of row mutations size to
54+
flush. If it reaches the max number of row mutations size it calls
55+
finish_batch() to mutate the current row batch. Default is MAX_ROW_BYTES
56+
(5 MB).
57+
"""
58+
59+
def __init__(self, table, flush_count=FLUSH_COUNT,
60+
max_row_bytes=MAX_ROW_BYTES):
61+
self.rows = []
62+
self.total_mutation_count = 0
63+
self.total_size = 0
64+
self.table = table
65+
self.flush_count = flush_count
66+
self.max_row_bytes = max_row_bytes
67+
68+
def mutate(self, row):
69+
""" Add a row to the batch. If the current batch meets one of the size
70+
limits, the batch is sent synchronously.
71+
72+
Example:
73+
>>> # Batcher for max row bytes
74+
>>> batcher = table.mutations_batcher(max_row_bytes=1024)
75+
>>>
76+
>>> row = table.row(b'row_key')
77+
>>>
78+
>>> # In batcher mutate will flush current batch if it
79+
>>> # reaches the max_row_bytes
80+
>>> batcher.mutate(row)
81+
>>>
82+
>>> batcher.flush()
83+
84+
:type row: class
85+
:param row: class:`~google.cloud.bigtable.row.DirectRow`.
86+
87+
:raises: One of the following:
88+
* :exc:`~.table._BigtableRetryableError` if any
89+
row returned a transient error.
90+
* :exc:`RuntimeError` if the number of responses doesn't
91+
match the number of rows that were retried
92+
* :exc:`.batcher.MaxMutationsError` if any row exceeds max
93+
mutations count.
94+
"""
95+
mutation_count = len(row._get_mutations())
96+
if mutation_count > MAX_MUTATIONS:
97+
raise MaxMutationsError(
98+
'The row key {} exceeds the number of mutations {}.'.format(
99+
row.row_key, mutation_count), )
100+
101+
if (self.total_mutation_count + mutation_count) >= MAX_MUTATIONS:
102+
self.flush()
103+
104+
self.rows.append(row)
105+
self.total_mutation_count += mutation_count
106+
self.total_size += row.get_mutations_size()
107+
108+
if (self.total_size >= self.max_row_bytes or
109+
len(self.rows) >= self.flush_count):
110+
self.flush()
111+
112+
def mutate_rows(self, rows):
113+
""" Add a row to the batch. If the current batch meets one of the size
114+
limits, the batch is sent synchronously.
115+
116+
Example:
117+
>>> # Batcher for flush count
118+
>>> batcher = table.mutations_batcher(flush_count=2)
119+
>>>
120+
>>> row1 = table.row(b'row_key_1')
121+
>>> row2 = table.row(b'row_key_2')
122+
>>> row3 = table.row(b'row_key_3')
123+
>>> row4 = table.row(b'row_key_4')
124+
>>>
125+
>>> # In batcher mutate will flush current batch if it
126+
>>> # reaches the max flush_count
127+
>>> batcher.mutate_rows([row_1, row_2, row_3, row_4])
128+
>>>
129+
>>> batcher.flush()
130+
131+
:type rows: list:[`~google.cloud.bigtable.row.DirectRow`]
132+
:param rows: list:[`~google.cloud.bigtable.row.DirectRow`].
133+
134+
:raises: One of the following:
135+
* :exc:`~.table._BigtableRetryableError` if any
136+
row returned a transient error.
137+
* :exc:`RuntimeError` if the number of responses doesn't
138+
match the number of rows that were retried
139+
* :exc:`.batcher.MaxMutationsError` if any row exceeds max
140+
mutations count.
141+
"""
142+
for row in rows:
143+
self.mutate(row)
144+
145+
def flush(self):
146+
""" Sends the current. batch to Cloud Bigtable. """
147+
if len(self.rows) is not 0:
148+
self.table.mutate_rows(self.rows)
149+
self.total_mutation_count = 0
150+
self.total_size = 0
151+
self.rows = []
Collapse file

‎bigtable/google/cloud/bigtable/row.py‎

Copy file name to clipboardExpand all lines: bigtable/google/cloud/bigtable/row.py
+9Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,15 @@ def _get_mutations(self, state=None): # pylint: disable=unused-argument
286286
"""
287287
return self._pb_mutations
288288

289+
def get_mutations_size(self):
290+
""" Gets the total mutations size for current row """
291+
292+
mutation_size = 0
293+
for mutation in self._get_mutations():
294+
mutation_size += mutation.ByteSize()
295+
296+
return mutation_size
297+
289298
def set_cell(self, column_family_id, column, value, timestamp=None):
290299
"""Sets a value in this row.
291300
Collapse file

‎bigtable/google/cloud/bigtable/table.py‎

Copy file name to clipboardExpand all lines: bigtable/google/cloud/bigtable/table.py
+23Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
from google.cloud._helpers import _to_bytes
2525
from google.cloud.bigtable.column_family import _gc_rule_from_pb
2626
from google.cloud.bigtable.column_family import ColumnFamily
27+
from google.cloud.bigtable.batcher import MutationsBatcher
28+
from google.cloud.bigtable.batcher import (FLUSH_COUNT, MAX_ROW_BYTES)
2729
from google.cloud.bigtable.row import AppendRow
2830
from google.cloud.bigtable.row import ConditionalRow
2931
from google.cloud.bigtable.row import DirectRow
@@ -483,6 +485,27 @@ def drop_by_prefix(self, row_key_prefix, timeout=None):
483485
table_admin_client.drop_row_range(
484486
self.name, row_key_prefix=_to_bytes(row_key_prefix))
485487

488+
def mutations_batcher(self, flush_count=FLUSH_COUNT,
489+
max_row_bytes=MAX_ROW_BYTES):
490+
"""Factory to create a mutation batcher associated with this instance.
491+
492+
:type table: class
493+
:param table: class:`~google.cloud.bigtable.table.Table`.
494+
495+
:type flush_count: int
496+
:param flush_count: (Optional) Maximum number of rows per batch. If it
497+
reaches the max number of rows it calls finish_batch() to
498+
mutate the current row batch. Default is FLUSH_COUNT (1000
499+
rows).
500+
501+
:type max_row_bytes: int
502+
:param max_row_bytes: (Optional) Max number of row mutations size to
503+
flush. If it reaches the max number of row mutations size it
504+
calls finish_batch() to mutate the current row batch.
505+
Default is MAX_ROW_BYTES (5 MB).
506+
"""
507+
return MutationsBatcher(self, flush_count, max_row_bytes)
508+
486509

487510
class _RetryableMutateRowsWorker(object):
488511
"""A callable worker that can retry to mutate rows with transient errors.
Collapse file

‎bigtable/tests/unit/test_batcher.py‎

Copy file name to clipboard
+179Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
# Copyright 2018 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
import unittest
17+
18+
import mock
19+
20+
from ._testing import _make_credentials
21+
22+
from google.cloud.bigtable.batcher import MutationsBatcher
23+
from google.cloud.bigtable.row import DirectRow
24+
25+
26+
class TestMutationsBatcher(unittest.TestCase):
27+
from grpc import StatusCode
28+
29+
TABLE_ID = 'table-id'
30+
TABLE_NAME = '/tables/' + TABLE_ID
31+
32+
# RPC Status Codes
33+
SUCCESS = StatusCode.OK.value[0]
34+
35+
@staticmethod
36+
def _get_target_class():
37+
from google.cloud.bigtable.table import Table
38+
39+
return Table
40+
41+
def _make_table(self, *args, **kwargs):
42+
return self._get_target_class()(*args, **kwargs)
43+
44+
@staticmethod
45+
def _get_target_client_class():
46+
from google.cloud.bigtable.client import Client
47+
48+
return Client
49+
50+
def _make_client(self, *args, **kwargs):
51+
return self._get_target_client_class()(*args, **kwargs)
52+
53+
def test_constructor(self):
54+
credentials = _make_credentials()
55+
client = self._make_client(project='project-id',
56+
credentials=credentials, admin=True)
57+
58+
instance = client.instance(instance_id='instance-id')
59+
table = self._make_table(self.TABLE_ID, instance)
60+
61+
mutation_batcher = MutationsBatcher(table)
62+
self.assertEqual(table, mutation_batcher.table)
63+
64+
def test_mutate_row(self):
65+
table = _Table(self.TABLE_NAME)
66+
mutation_batcher = MutationsBatcher(table=table)
67+
68+
rows = [DirectRow(row_key=b'row_key'),
69+
DirectRow(row_key=b'row_key_2'),
70+
DirectRow(row_key=b'row_key_3'),
71+
DirectRow(row_key=b'row_key_4')]
72+
73+
mutation_batcher.mutate_rows(rows)
74+
mutation_batcher.flush()
75+
76+
self.assertEqual(table.mutation_calls, 1)
77+
78+
def test_mutate_rows(self):
79+
table = _Table(self.TABLE_NAME)
80+
mutation_batcher = MutationsBatcher(table=table)
81+
82+
row = DirectRow(row_key=b'row_key')
83+
row.set_cell('cf1', b'c1', 1)
84+
row.set_cell('cf1', b'c2', 2)
85+
row.set_cell('cf1', b'c3', 3)
86+
row.set_cell('cf1', b'c4', 4)
87+
88+
mutation_batcher.mutate(row)
89+
90+
mutation_batcher.flush()
91+
92+
self.assertEqual(table.mutation_calls, 1)
93+
94+
def test_flush_with_no_rows(self):
95+
table = _Table(self.TABLE_NAME)
96+
mutation_batcher = MutationsBatcher(table=table)
97+
mutation_batcher.flush()
98+
99+
self.assertEqual(table.mutation_calls, 0)
100+
101+
def test_add_row_with_max_flush_count(self):
102+
table = _Table(self.TABLE_NAME)
103+
mutation_batcher = MutationsBatcher(table=table, flush_count=3)
104+
105+
row_1 = DirectRow(row_key=b'row_key_1')
106+
row_2 = DirectRow(row_key=b'row_key_2')
107+
row_3 = DirectRow(row_key=b'row_key_3')
108+
109+
mutation_batcher.mutate(row_1)
110+
mutation_batcher.mutate(row_2)
111+
mutation_batcher.mutate(row_3)
112+
113+
self.assertEqual(table.mutation_calls, 1)
114+
115+
@mock.patch('google.cloud.bigtable.batcher.MAX_MUTATIONS', new=3)
116+
def test_mutate_row_with_max_mutations_failure(self):
117+
from google.cloud.bigtable.batcher import MaxMutationsError
118+
119+
table = _Table(self.TABLE_NAME)
120+
mutation_batcher = MutationsBatcher(table=table)
121+
122+
row = DirectRow(row_key=b'row_key')
123+
row.set_cell('cf1', b'c1', 1)
124+
row.set_cell('cf1', b'c2', 2)
125+
row.set_cell('cf1', b'c3', 3)
126+
row.set_cell('cf1', b'c4', 4)
127+
128+
with self.assertRaises(MaxMutationsError):
129+
mutation_batcher.mutate(row)
130+
131+
@mock.patch('google.cloud.bigtable.batcher.MAX_MUTATIONS', new=3)
132+
def test_mutate_row_with_max_mutations(self):
133+
table = _Table(self.TABLE_NAME)
134+
mutation_batcher = MutationsBatcher(table=table)
135+
136+
row = DirectRow(row_key=b'row_key')
137+
row.set_cell('cf1', b'c1', 1)
138+
row.set_cell('cf1', b'c2', 2)
139+
row.set_cell('cf1', b'c3', 3)
140+
141+
mutation_batcher.mutate(row)
142+
mutation_batcher.flush()
143+
144+
self.assertEqual(table.mutation_calls, 1)
145+
146+
def test_mutate_row_with_max_row_bytes(self):
147+
table = _Table(self.TABLE_NAME)
148+
mutation_batcher = MutationsBatcher(table=table,
149+
max_row_bytes=3 * 1024 * 1024)
150+
151+
number_of_bytes = 1 * 1024 * 1024
152+
max_value = b'1' * number_of_bytes
153+
154+
row = DirectRow(row_key=b'row_key')
155+
row.set_cell('cf1', b'c1', max_value)
156+
row.set_cell('cf1', b'c2', max_value)
157+
row.set_cell('cf1', b'c3', max_value)
158+
159+
mutation_batcher.mutate(row)
160+
161+
self.assertEqual(table.mutation_calls, 1)
162+
163+
164+
class _Instance(object):
165+
166+
def __init__(self, client=None):
167+
self._client = client
168+
169+
170+
class _Table(object):
171+
172+
def __init__(self, name, client=None):
173+
self.name = name
174+
self._instance = _Instance(client)
175+
self.mutation_calls = 0
176+
177+
def mutate_rows(self, rows):
178+
self.mutation_calls += 1
179+
return rows

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.