diff --git a/elasticsearch/helpers/parallel.py b/elasticsearch/helpers/parallel.py index 7132e568..036f8da1 100644 --- a/elasticsearch/helpers/parallel.py +++ b/elasticsearch/helpers/parallel.py @@ -1,66 +1,24 @@ from multiprocessing.dummy import Pool -from queue import Empty, Queue -from threading import Event +from . import _process_bulk_chunk, _chunk_actions, expand_action -from . import streaming_bulk -def consume(queue, done): +def parallel_bulk(client, actions, thread_count=4, chunk_size=500, + max_chunk_bytes=100 * 1014 * 1024, + expand_action_callback=expand_action, **kwargs): """ - Create an iterator on top of a Queue. + Parallel version of the bulk helper. """ - while True: - try: - yield queue.get(True, .01) - except Empty: - if done.is_set(): - break + actions = map(expand_action_callback, actions) -def wrapped_bulk(client, input, output, done, **kwargs): - """ - Wrap a call to streaming_bulk by feeding it data frm a queue and writing - the outputs to another queue. - """ - try: - for result in streaming_bulk(client, consume(input, done), **kwargs): - output.put(result) - except: - done.set() - raise + pool = Pool(thread_count) -def feed_data(actions, input, done): - """ - Feed data from an iterator into a queue. - """ - for a in actions: - input.put(a, True) - - # error short-circuit - if done.is_set(): - break - done.set() - - -def parallel_bulk(client, actions, thread_count=5, **kwargs): - """ - Paralel version of the bulk helper. It runs a thread pool with a thread for - a producer and ``thread_count`` threads for. - """ - done = Event() - input, output = Queue(), Queue() - pool = Pool(thread_count + 1) - - results = [ - pool.apply_async(wrapped_bulk, (client, input, output, done), kwargs) - for _ in range(thread_count)] - pool.apply_async(feed_data, (actions, input, done)) - - while True: - try: - yield output.get(True, .01) - except Empty: - if done.is_set() and all(r.ready() for r in results): - break + for result in pool.imap( + lambda chunk: list(_process_bulk_chunk(client, chunk, **kwargs)), + _chunk_actions(actions, chunk_size, max_chunk_bytes, client.transport.serializer) + ): + for item in result: + yield item pool.close() pool.join()