Add an option to use scan helper without search_type=scan, use with caution

Fixes #160

elasticsearch/helpers/__init__.py

@@ -156,7 +156,7 @@ def bulk(client, actions, stats_only=False, **kwargs):
 # preserve the name for backwards compatibility
 bulk_index = bulk
 
-def scan(client, query=None, scroll='5m', **kwargs):
+def scan(client, query=None, scroll='5m', preserve_order=False, **kwargs):
     """
     Simple abstraction on top of the
     :meth:`~elasticsearch.Elasticsearch.scroll` api - a simple iterator that
@@ -166,19 +166,30 @@ def scan(client, query=None, scroll='5m', **kwargs):
     :arg query: body for the :meth:`~elasticsearch.Elasticsearch.search` api
     :arg scroll: Specify how long a consistent view of the index should be
         maintained for scrolled search
+    :arg preserve_order: don't set the ``search_type`` to ``scan`` - this will
+        cause the scroll to paginate with preserving the order. Note that this
+        can be an extremely expensive operation and can easily lead to
+        unpredictable results, use with caution.
 
     Any additional keyword arguments will be passed to the initial
     :meth:`~elasticsearch.Elasticsearch.search` call.
     """
-    # initial search to
-    resp = client.search(body=query, search_type='scan', scroll=scroll, **kwargs)
+    if not preserve_order:
+        kwargs['search_type'] = 'scan'
+    # initial search
+    resp = client.search(body=query, scroll=scroll, **kwargs)
 
     scroll_id = resp.get('_scroll_id')
     if scroll_id is None:
         return
 
+    first_run = True
     while True:
-        resp = client.scroll(scroll_id, scroll=scroll)
+        # if we didn't set search_type to scan initial search contains data
+        if preserve_order and first_run:
+            first_run = False
+        else:
+            resp = client.scroll(scroll_id, scroll=scroll)
         if not resp['hits']['hits']:
             break
         for hit in resp['hits']['hits']:

test_elasticsearch/test_server/test_helpers.py

@@ -138,6 +138,19 @@ class TestBulk(ElasticsearchTestCase):
 
 
 class TestScan(ElasticsearchTestCase):
+    def test_order_can_be_preserved(self):
+        bulk = []
+        for x in range(100):
+            bulk.append({"index": {"_index": "test_index", "_type": "answers", "_id": x}})
+            bulk.append({"answer": x, "correct": x == 42})
+        self.client.bulk(bulk, refresh=True)
+
+        docs = list(helpers.scan(self.client, index="test_index", doc_type="answers", size=2, query={"sort": ["answer"]}, preserve_order=True))
+
+        self.assertEquals(100, len(docs))
+        self.assertEquals(list(map(str, range(100))), list(d['_id'] for d in docs))
+        self.assertEquals(list(range(100)), list(d['_source']['answer'] for d in docs))
+
     def test_all_documents_are_read(self):
         bulk = []
         for x in range(100):