greatly simplify parallel_bulk by using Pool.imap

Author: honzakral

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
-
-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))
+ pool = Pool(thread_count)
 
- 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()