From 3ae44db90c29c9b1c135cc737e30e8b19f94f8e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Honza=20Kr=C3=A1l?= Date: Fri, 22 Nov 2013 14:29:08 +0100 Subject: [PATCH] Bulk helpers overhaul streaming_bulk for iterative consumption of the input bulk_index -> bulk callback expand_action to transform input into action/data pairs support for non-index actions through _op_type key --- Changelog.rst | 10 ++ elasticsearch/helpers.py | 159 +++++++++++++----- .../test_server/test_helpers.py | 54 +++++- 3 files changed, 176 insertions(+), 47 deletions(-) diff --git a/Changelog.rst b/Changelog.rst index b972d24c..60f4a0b8 100644 --- a/Changelog.rst +++ b/Changelog.rst @@ -3,6 +3,16 @@ Changelog ========= +0.4.4 +----- + + * `helpers.bulk_index` renamed to `helpers.bulk` (alias put in place for + backwards compatibility, to be removed in future versions) + * Added `helpers.streaming_bulk` to consume an iterator and yield results per + operation + * `helpers.bulk` and `helpers.streaming_bulk` are no longer limitted to just + index operations. + 0.4.3 ----- diff --git a/elasticsearch/helpers.py b/elasticsearch/helpers.py index 42f0ab7e..325e7050 100644 --- a/elasticsearch/helpers.py +++ b/elasticsearch/helpers.py @@ -1,16 +1,42 @@ -from itertools import islice +from itertools import islice, chain +from operator import methodcaller +try: + from itertools import imap as map +except ImportError: + pass # python 3, use builtin map from elasticsearch.exceptions import ElasticsearchException -class BulkIndexError(ElasticsearchException): pass +class BulkIndexError(ElasticsearchException): + @property + def errors(self): + """ List of errors from execution of the last chunk. """ + return self.args[1] -def bulk_index(client, docs, chunk_size=500, stats_only=False, raise_on_error=False, **kwargs): + +def expand_action(data): """ - Helper for the :meth:`~elasticsearch.Elasticsearch.bulk` api that provides - a more human friendly interface - it consumes an iterator of documents and - sends them to elasticsearch in chunks. + From one document or action definition passed in by the user extract the + action/data lines needed for elasticsearch's + :meth:`~elasticsearch.Elasticsearch.bulk` api. + """ + op_type = data.pop('_op_type', 'index') + action = {op_type: {}} + for key in ('_index', '_parent', '_percolate', '_routing', '_timestamp', + '_ttl', '_type', '_version', '_id', '_retry_on_conflict'): + if key in data: + action[op_type][key] = data.pop(key) - This function expects the doc to be in the format as returned by + # no data payload for delete + if op_type == 'delete': + return action, None + + return action, data.get('_source', data) + + +def streaming_bulk(client, actions, chunk_size=500, raise_on_error=False, expand_action_callback=expand_action, **kwargs): + """ + This function expects the action to be in the format as returned by :meth:`~elasticsearch.Elasticsearch.search`, for example:: { @@ -24,56 +50,107 @@ def bulk_index(client, docs, chunk_size=500, stats_only=False, raise_on_error=Fa } } - alternatively, if `_source` is not present, it will pop all metadata fields - from the doc and use the rest as the document data. + Alternatively, if `_source` is not present, it will pop all metadata fields + from the doc and use the rest as the document data. The dict passed in will + be modified - metadata fields will be popped out. + + Alternative actions (`_op_type` field defaults to `index`) can be sent as + well:: + + { + '_op_type': 'delete', + '_index': 'index-name', + '_type': 'document', + '_id': 42, + } + { + '_op_type': 'update', + '_index': 'index-name', + '_type': 'document', + '_id': 42, + 'doc': {'question': 'The life, universe and everything.'} + } :arg client: instance of :class:`~elasticsearch.Elasticsearch` to use - :arg docs: iterator containing the docs + :arg actions: iterable containing the actions to be executed :arg chunk_size: number of docs in one chunk sent to es (default: 500) - :arg stats_only: if `True` only report number of successful/failed - operations instead of just number of successful and a list of error responses - :arg raise_on_error: raise `BulkIndexError` if some documents failed to - index (and stop sending chunks to the server) - - Any additional keyword arguments will be passed to the bulk API itself. + :arg raise_on_error: raise `BulkIndexError` containing errors (as `.errors` + from the execution of the last chunk) + :arg expand_action_callback: callback executed on each action passed in, + should return a tuple containing the action line and the data line + (`None` if data line should be omitted). """ - success, failed = 0, 0 + actions = map(expand_action_callback, actions) - # list of errors to be collected when + # if raise on error is set, we need to collect errors per chunk before raising them errors = [] - docs = iter(docs) while True: - chunk = islice(docs, chunk_size) + chunk = islice(actions, chunk_size) bulk_actions = [] - for d in chunk: - action = {'index': {}} - for key in ('_index', '_parent', '_percolate', '_routing', - '_timestamp', '_ttl', '_type', '_version', '_id'): - if key in d: - action['index'][key] = d.pop(key) - + for action, data in chunk: bulk_actions.append(action) - bulk_actions.append(d.get('_source', d)) + if data is not None: + bulk_actions.append(data) if not bulk_actions: - return success, failed if stats_only else errors + return # send the actual request resp = client.bulk(bulk_actions, **kwargs) # go through request-reponse pairs and detect failures - for req, item in zip(bulk_actions[::2], resp['items']): - ok = item['index' if '_id' in req['index'] else 'create'].get('ok') - if not ok: - if not stats_only: - errors.append(item) - failed += 1 - else: - success += 1 + for op_type, item in chain.from_iterable(map(methodcaller('items'), resp['items'])): + ok = item.get('ok') + if not ok and raise_on_error: + errors.append({op_type: item}) - if failed and raise_on_error: - raise BulkIndexError('%i document(s) failed to index.' % failed, errors) + if not errors: + # if we are not just recording all errors to be able to raise + # them all at once, yield items individually + yield ok, {op_type: item} + + if errors: + raise BulkIndexError('%i document(s) failed to index.' % len(errors), errors) + +def bulk(client, actions, stats_only=False, **kwargs): + """ + Helper for the :meth:`~elasticsearch.Elasticsearch.bulk` api that provides + a more human friendly interface - it consumes an iterator of actions and + sends them to elasticsearch in chunks. It returns a tuple with summary + information - number of successfully executed actions and either list of + errors or number of errors if `stats_only` is set to `True`. + + See :function:`~elasticsearch.helpers.streaming_bulk` for more information + and accepted formats. + + :arg client: instance of :class:`~elasticsearch.Elasticsearch` to use + :arg actions: iterator containing the actions + :arg stats_only: if `True` only report number of successful/failed + operations instead of just number of successful and a list of error responses + + Any additional keyword arguments will be passed to + :function:`~elasticsearch.helpers.streaming_bulk` which is used to execute + the operation. + """ + success, failed = 0, 0 + + # list of errors to be collected is not stats_only + errors = [] + + for ok, item in streaming_bulk(client, actions, **kwargs): + # go through request-reponse pairs and detect failures + if not ok: + if not stats_only: + errors.append(item) + failed += 1 + else: + success += 1 + + return success, failed if stats_only else errors + +# preserve the name for backwards compatibility +bulk_index = bulk def scan(client, query=None, scroll='5m', **kwargs): """ @@ -89,7 +166,7 @@ def scan(client, query=None, scroll='5m', **kwargs): Any additional keyword arguments will be passed to the initial :meth:`~elasticsearch.Elasticsearch.search` call. """ - # initial search to + # initial search to resp = client.search(body=query, search_type='scan', scroll=scroll, **kwargs) scroll_id = resp['_scroll_id'] @@ -129,5 +206,5 @@ def reindex(client, source_index, target_index, target_client=None, chunk_size=5 h['_index'] = index yield h - return bulk_index(target_client, _change_doc_index(docs, target_index), + return bulk(target_client, _change_doc_index(docs, target_index), chunk_size=chunk_size, stats_only=True) diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py index 30b77001..8986e913 100644 --- a/test_elasticsearch/test_server/test_helpers.py +++ b/test_elasticsearch/test_server/test_helpers.py @@ -2,10 +2,52 @@ from elasticsearch import helpers from . import ElasticTestCase -class TestBulkIndex(ElasticTestCase): +class TestStreamingBulk(ElasticTestCase): def test_all_documents_get_inserted(self): docs = [{"answer": x, '_id': x} for x in range(100)] - success, failed = helpers.bulk_index(self.client, docs, index='test-index', doc_type='answers', refresh=True) + for ok, item in helpers.streaming_bulk(self.client, docs, index='test-index', doc_type='answers', refresh=True): + self.assertTrue(ok) + + self.assertEquals(100, self.client.count(index='test-index', doc_type='answers')['count']) + self.assertEquals({"answer": 42}, self.client.get(index='test-index', doc_type='answers', id=42)['_source']) + + def test_all_errors_from_chunk_are_raised_on_failure(self): + self.client.indices.create("i", + { + "mappings": {"t": {"properties": {"a": {"type": "integer"}}}}, + "settings": {"number_of_shards": 1, "number_of_replicas": 0} + }) + self.client.cluster.health(wait_for_status="yellow") + + try: + for ok, item in helpers.streaming_bulk(self.client, [{"a": "b"}, + {"a": "c"}], index="i", doc_type="t", raise_on_error=True): + self.assertTrue(ok) + except helpers.BulkIndexError as e: + self.assertEquals(2, len(e.errors)) + else: + assert False, "exception should have been raised" + + def test_different_op_types(self): + self.client.index(index='i', doc_type='t', id=45, body={}) + self.client.index(index='i', doc_type='t', id=42, body={}) + docs = [ + {'_index': 'i', '_type': 't', '_id': 47, 'f': 'v'}, + {'_op_type': 'delete', '_index': 'i', '_type': 't', '_id': 45}, + {'_op_type': 'update', '_index': 'i', '_type': 't', '_id': 42, 'doc': {'answer': 42}}, + ] + for ok, item in helpers.streaming_bulk(self.client, docs): + self.assertTrue(ok) + + self.assertFalse(self.client.exists(index='i', id=45)) + self.assertEquals({'answer': 42}, self.client.get_source(index='i', id=42)) + self.assertEquals({'f': 'v'}, self.client.get_source(index='i', id=47)) + + +class TestBulk(ElasticTestCase): + def test_all_documents_get_inserted(self): + docs = [{"answer": x, '_id': x} for x in range(100)] + success, failed = helpers.bulk(self.client, docs, index='test-index', doc_type='answers', refresh=True) self.assertEquals(100, success) self.assertFalse(failed) @@ -14,7 +56,7 @@ class TestBulkIndex(ElasticTestCase): def test_stats_only_reports_numbers(self): docs = [{"answer": x} for x in range(100)] - success, failed = helpers.bulk_index(self.client, docs, index='test-index', doc_type='answers', refresh=True, stats_only=True) + success, failed = helpers.bulk(self.client, docs, index='test-index', doc_type='answers', refresh=True, stats_only=True) self.assertEquals(100, success) self.assertEquals(0, failed) @@ -28,7 +70,7 @@ class TestBulkIndex(ElasticTestCase): }) self.client.cluster.health(wait_for_status="yellow") - success, failed = helpers.bulk_index( + success, failed = helpers.bulk( self.client, [{"a": 42}, {"a": "c", '_id': 42}], index="i", @@ -50,7 +92,7 @@ class TestBulkIndex(ElasticTestCase): }) self.client.cluster.health(wait_for_status="yellow") - self.assertRaises(helpers.BulkIndexError, helpers.bulk_index, + self.assertRaises(helpers.BulkIndexError, helpers.bulk, self.client, [{"a": 42}, {"a": "c"}], index="i", @@ -66,7 +108,7 @@ class TestBulkIndex(ElasticTestCase): }) self.client.cluster.health(wait_for_status="yellow") - success, failed = helpers.bulk_index( + success, failed = helpers.bulk( self.client, [{"a": 42}, {"a": "c"}], index="i",