From 9b5d2c062e77fb8b45195feb31f4536db264b38c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Honza=20Kr=C3=A1l?= <honza.kral@gmail.com>
Date: Mon, 15 Dec 2014 15:15:03 +0100
Subject: [PATCH] Add an option to use scan helper without search_type=scan,
 use with caution

Fixes #160
---
 elasticsearch/helpers/__init__.py             | 19 +++++++++++++++----
 .../test_server/test_helpers.py               | 13 +++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/elasticsearch/helpers/__init__.py b/elasticsearch/helpers/__init__.py
index a4d7f080..07fa8b25 100644
--- a/elasticsearch/helpers/__init__.py
+++ b/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']:
diff --git a/test_elasticsearch/test_server/test_helpers.py b/test_elasticsearch/test_server/test_helpers.py
index 67bafb3b..2674f324 100644
--- a/test_elasticsearch/test_server/test_helpers.py
+++ b/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):