have _chunk_actions always return data and actions

This commit is contained in:
Honza Král
2017-07-18 14:18:09 -04:00
parent c8231cad0e
commit f9f4f0c61a
2 changed files with 32 additions and 44 deletions
+10 -29
View File
@@ -47,13 +47,12 @@ def expand_action(data):
return action, data.get('_source', data) return action, data.get('_source', data)
def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer, include_data=False): def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
""" """
Split actions into chunks by number or size, serialize them into strings in Split actions into chunks by number or size, serialize them into strings in
the process. the process.
""" """
bulk_actions = [] bulk_actions, bulk_data = [], []
bulk_data = []
size, action_count = 0, 0 size, action_count = 0, 0
for action, data in actions: for action, data in actions:
action = serializer.dumps(action) action = serializer.dumps(action)
@@ -65,20 +64,13 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer, include_dat
# full chunk, send it and start a new one # full chunk, send it and start a new one
if bulk_actions and (size + cur_size > max_chunk_bytes or action_count == chunk_size): if bulk_actions and (size + cur_size > max_chunk_bytes or action_count == chunk_size):
if include_data:
yield bulk_data, bulk_actions yield bulk_data, bulk_actions
else: bulk_actions, bulk_data = [], []
yield bulk_actions
bulk_actions = []
bulk_data = []
size, action_count = 0, 0 size, action_count = 0, 0
bulk_actions.append(action) bulk_actions.append(action)
if data is not None: if data is not None:
bulk_actions.append(data) bulk_actions.append(data)
if include_data:
if data is not None:
bulk_data.append((action, data)) bulk_data.append((action, data))
else: else:
bulk_data.append((action, )) bulk_data.append((action, ))
@@ -87,12 +79,9 @@ def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer, include_dat
action_count += 1 action_count += 1
if bulk_actions: if bulk_actions:
if include_data:
yield bulk_data, bulk_actions yield bulk_data, bulk_actions
else:
yield bulk_actions
def _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_error=True, **kwargs): def _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception=True, raise_on_error=True, **kwargs):
""" """
Send a bulk request to elasticsearch and process the output. Send a bulk request to elasticsearch and process the output.
""" """
@@ -111,22 +100,14 @@ def _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_
err_message = str(e) err_message = str(e)
exc_errors = [] exc_errors = []
# deserialize the data back, thisis expensive but only run on for data in bulk_data:
# errors if raise_on_exception is false, so shouldn't be a real
# issue
bulk_data = map(client.transport.serializer.loads, bulk_actions)
while True:
try:
# collect all the information about failed actions # collect all the information about failed actions
action = next(bulk_data) op_type, action = data[0].popitem()
op_type, action = action.popitem()
info = {"error": err_message, "status": e.status_code, "exception": e} info = {"error": err_message, "status": e.status_code, "exception": e}
if op_type != 'delete': if op_type != 'delete':
info['data'] = next(bulk_data) info['data'] = data[1]
info.update(action) info.update(action)
exc_errors.append({op_type: info}) exc_errors.append({op_type: info})
except StopIteration:
break
# emulate standard behavior for failed actions # emulate standard behavior for failed actions
if raise_on_error: if raise_on_error:
@@ -175,8 +156,8 @@ def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 *
""" """
actions = map(expand_action_callback, actions) actions = map(expand_action_callback, actions)
for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer): for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer):
for result in _process_bulk_chunk(client, bulk_actions, raise_on_exception, raise_on_error, **kwargs): for result in _process_bulk_chunk(client, bulk_actions, bulk_data, raise_on_exception, raise_on_error, **kwargs):
yield result yield result
def bulk(client, actions, stats_only=False, **kwargs): def bulk(client, actions, stats_only=False, **kwargs):
@@ -256,7 +237,7 @@ def parallel_bulk(client, actions, thread_count=4, chunk_size=500,
try: try:
for result in pool.imap( for result in pool.imap(
lambda chunk: list(_process_bulk_chunk(client, chunk, **kwargs)), lambda bulk_data, bulk_actions: list(_process_bulk_chunk(client, bulk_actions, bulk_data, **kwargs)),
_chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer) _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer)
): ):
for item in result: for item in result:
+12 -5
View File
@@ -30,16 +30,24 @@ def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1
""" """
actions = map(expand_action_callback, actions) actions = map(expand_action_callback, actions)
for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer, include_data=True): for bulk_data, bulk_actions in _chunk_actions(actions, chunk_size,
max_chunk_bytes,
client.transport.serializer):
retry = 0 retry = 0
while max_retries == -1 or retry <= max_retries: while max_retries == -1 or retry <= max_retries:
to_retry = [] to_retry, to_retry_data = [], []
to_retry_data = []
if retry: if retry:
time.sleep(min(max_backoff, initial_backoff * 2**(retry-1))) time.sleep(min(max_backoff, initial_backoff * 2**(retry-1)))
try: try:
for data, (ok, info) in zip(bulk_data, _process_bulk_chunk(client, bulk_actions, raise_on_exception=True, raise_on_error=False, **kwargs)): for data, (ok, info) in zip(bulk_data,
_process_bulk_chunk(client,
bulk_actions,
bulk_data,
raise_on_exception=True,
raise_on_error=False,
**kwargs)):
if not ok: if not ok:
action, info = info.popitem() action, info = info.popitem()
if info['status'] == 429 and (retry+1) <= max_retries: if info['status'] == 429 and (retry+1) <= max_retries:
@@ -52,7 +60,6 @@ def backoff_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1024 * 1
except TransportError as e: except TransportError as e:
if e.status_code != 429: if e.status_code != 429:
raise raise
retry += 1 retry += 1
else: else:
if not to_retry: if not to_retry: