greatly simplify parallel_bulk by using Pool.imap

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()