Experimental helper for doing bulk requests in parallel

Author: honzakral

elasticsearch/helpers/parallel.py

@@ -0,0 +1,66 @@
+from multiprocessing.dummy import Pool
+from queue import Empty, Queue
+
+from threading import Event
+
+from . import streaming_bulk
+
+def consume(queue, done):
+ """
+ Create an iterator on top of a Queue.
+ """
+ while True:
+ try:
+ yield queue.get(True, .01)
+ except Empty:
+ if done.is_set():
+ break
+
+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))
+
+ while True:
+ try:
+ yield output.get(True, .01)
+ except Empty:
+ if done.is_set() and all(r.ready() for r in results):
+ break
+
+ pool.close()
+ pool.join()