Allos the size of the bulk request to be defined in bytes

Fixes #199

Author: honzakral

elasticsearch/helpers/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import unicode_literals
 
 import logging
-from itertools import islice
 from operator import methodcaller
 
 from ..exceptions import ElasticsearchException, TransportError
@@ -40,22 +39,35 @@ def expand_action(data):
 
     return action, data.get('_source', data)
 
-def _chunk_actions(actions, chunk_size, serializer):
-    while True:
-        bulk_actions = []
-        for action, data in islice(actions, chunk_size):
-            bulk_actions.append(serializer.dumps(action))
-            if data is not None:
-                bulk_actions.append(serializer.dumps(data))
-
-        if not bulk_actions:
-            return
-
+def _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
+    bulk_actions = []
+    size, action_count = 0, 0
+    for action, data in actions:
+        action = serializer.dumps(action)
+        cur_size = len(action) + 1
+
+        if data is not None:
+            data = serializer.dumps(data)
+            cur_size += len(data) + 1
+
+        # full chunk, send it and start a new one
+        if bulk_actions and (size + cur_size > max_chunk_bytes or action_count == chunk_size):
+            yield bulk_actions
+            bulk_actions = []
+            size, action_count = 0, 0
+
+        bulk_actions.append(action)
+        if data is not None:
+            bulk_actions.append(data)
+        size += cur_size
+        action_count += 1
+
+    if bulk_actions:
         yield bulk_actions
 
-def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
-        expand_action_callback=expand_action, raise_on_exception=True, 
-        **kwargs):
+def streaming_bulk(client, actions, chunk_size=500, max_chunk_bytes=100 * 1014 * 1024,
+        raise_on_error=True, expand_action_callback=expand_action,
+        raise_on_exception=True, **kwargs):
     """
     Streaming bulk consumes actions from the iterable passed in and yields
     results per action. For non-streaming usecases use
@@ -101,6 +113,7 @@ def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
     :arg client: instance of :class:`~elasticsearch.Elasticsearch` to use
     :arg actions: iterable containing the actions to be executed
     :arg chunk_size: number of docs in one chunk sent to es (default: 500)
+    :arg max_chunk_bytes: the maximum size of the request in bytes (default: 100MB)
     :arg raise_on_error: raise ``BulkIndexError`` containing errors (as `.errors`)
         from the execution of the last chunk when some occur. By default we raise.
     :arg raise_on_exception: if ``False`` then don't propagate exceptions from
@@ -115,7 +128,7 @@ def streaming_bulk(client, actions, chunk_size=500, raise_on_error=True,
     # if raise on error is set, we need to collect errors per chunk before raising them
     errors = []
 
-    for bulk_actions in _chunk_actions(actions, chunk_size, serializer):
+    for bulk_actions in _chunk_actions(actions, chunk_size, max_chunk_bytes, serializer):
         try:
             # send the actual request
             resp = client.bulk('\n'.join(bulk_actions) + '\n', **kwargs)

test_elasticsearch/test_server/test_helpers.py

@@ -1,9 +1,21 @@
-from datetime import datetime
-
 from elasticsearch import helpers, TransportError
+from elasticsearch.serializer import JSONSerializer
 
 from . import ElasticsearchTestCase
-from ..test_cases import SkipTest
+from ..test_cases import SkipTest, TestCase
+
+
+class TestChunkActions(TestCase):
+    def setUp(self):
+        super(TestChunkActions, self).setUp()
+        self.actions = [({'index': {}}, {'some': 'data', 'i': i}) for i in range(100)]
+
+    def test_chunks_are_chopped_by_byte_size(self):
+        self.assertEquals(100, len(list(helpers._chunk_actions(self.actions, 100000, 1, JSONSerializer()))))
+
+    def test_chunks_are_chopped_by_chunk_size(self):
+        self.assertEquals(10, len(list(helpers._chunk_actions(self.actions, 10, 99999999, JSONSerializer()))))
+
 
 class FailingBulkClient(object):
     def __init__(self, client, fail_at=1):