Repository: elastic/rally

Issue:
Allow to ingest only a subset of the document corpus
For benchmarking smaller configurations (e.g. with very little memory), ingesting a full document corpus with tens of GB may be too much. For these cases we should allow users to provide a track parameter which allows them to define a percentage of the document corpus to ingest.

Notes: 

* Actually, this is already possible at the moment by forcing a specific number of iterations for the bulk operation. However, this is a bit hard to understand because the number of iterations is also tied to the bulk size.
* This is not meant to be accurate down to document level but only on a granularity of individual bulk requests (e.g. suppose we have a corpus with 1 million documents and a bulk size of 10.000 documents. If a user specifies 2.5% (= 25.000 documents) this will correspond to three bulk requests (i.e. 30.000 documents)).

Discussion:

Patch:
diff --git a/docs/track.rst b/docs/track.rst
index ec849cb1..4fcfefb1 100644
--- a/docs/track.rst
+++ b/docs/track.rst
@@ -308,9 +308,10 @@ bulk
 With the operation type ``bulk`` you can execute `bulk requests <http://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html>`_. It supports the following properties:
 
 * ``bulk-size`` (mandatory): Defines the bulk size in number of documents.
+* ``ingest-percentage`` (optional, defaults to 100): A number between (0, 100] that defines how much of the document corpus will be bulk-indexed.
 * ``corpora`` (optional): A list of document corpus names that should be targeted by this bulk-index operation. Only needed if the ``corpora`` section contains more than one document corpus and you don't want to index all of them with this operation.
 * ``indices`` (optional): A list of index names that defines which indices should be used by this bulk-index operation. Rally will then only select the documents files that have a matching ``target-index`` specified.
-* ``batch-size`` (optional): Defines how many documents Rally will read at once. This is an expert setting and only meant to avoid accidental bottlenecks for very small bulk sizes (e.g. if you want to benchmark with a bulk-size of 1, you should set batch-size higher).
+* ``batch-size`` (optional): Defines how many documents Rally will read at once. This is an expert setting and only meant to avoid accidental bottlenecks for very small bulk sizes (e.g. if you want to benchmark with a bulk-size of 1, you should set ``batch-size`` higher).
 * ``pipeline`` (optional): Defines the name of an (existing) ingest pipeline that should be used (only supported from Elasticsearch 5.0).
 * ``conflicts`` (optional): Type of index conflicts to simulate. If not specified, no conflicts will be simulated. Valid values are: 'sequential' (A document id is replaced with a document id with a sequentially increasing id), 'random' (A document id is replaced with a document id with a random other id).
 
diff --git a/esrally/track/params.py b/esrally/track/params.py
index e9838253..48ef30eb 100644
--- a/esrally/track/params.py
+++ b/esrally/track/params.py
@@ -1,6 +1,7 @@
 import logging
 import random
 import time
+import math
 import types
 import inspect
 from enum import Enum
@@ -482,6 +483,14 @@ class BulkIndexParamSource(ParamSource):
         except ValueError:
             raise exceptions.InvalidSyntax("'batch-size' must be numeric")
 
+        try:
+            self.ingest_percentage = float(params.get("ingest-percentage", 100.0))
+            if self.ingest_percentage <= 0 or self.ingest_percentage > 100.0:
+                raise exceptions.InvalidSyntax(
+                    "'ingest-percentage' must be in the range (0.0, 100.0] but was {:.1f}".format(self.ingest_percentage))
+        except ValueError:
+            raise exceptions.InvalidSyntax("'ingest-percentage' must be numeric")
+
     def used_corpora(self, t, params):
         corpora = []
         track_corpora_names = [corpus.name for corpus in t.corpora]
@@ -504,7 +513,7 @@ class BulkIndexParamSource(ParamSource):
 
     def partition(self, partition_index, total_partitions):
         return PartitionBulkIndexParamSource(self.corpora, partition_index, total_partitions, self.batch_size, self.bulk_size,
-                                             self.id_conflicts, self.pipeline, self._params)
+                                             self.ingest_percentage, self.id_conflicts, self.pipeline, self._params)
 
     def params(self):
         raise exceptions.RallyError("Do not use a BulkIndexParamSource without partitioning")
@@ -514,7 +523,7 @@ class BulkIndexParamSource(ParamSource):
 
 
 class PartitionBulkIndexParamSource:
-    def __init__(self, corpora, partition_index, total_partitions, batch_size, bulk_size, id_conflicts=None,
+    def __init__(self, corpora, partition_index, total_partitions, batch_size, bulk_size, ingest_percentage, id_conflicts=None,
                  pipeline=None, original_params=None):
         """
 
@@ -523,14 +532,17 @@ class PartitionBulkIndexParamSource:
         :param total_partitions: The total number of partitions (i.e. clients) for bulk index operations.
         :param batch_size: The number of documents to read in one go.
         :param bulk_size: The size of bulk index operations (number of documents per bulk).
+        :param ingest_percentage: A number between (0.0, 100.0] that defines how much of the whole corpus should be ingested.
         :param id_conflicts: The type of id conflicts.
         :param pipeline: The name of the ingest pipeline to run.
+        :param original_params: The original dict passed to the parent parameter source.
         """
         self.corpora = corpora
         self.partition_index = partition_index
         self.total_partitions = total_partitions
         self.batch_size = batch_size
         self.bulk_size = bulk_size
+        self.ingest_percentage = ingest_percentage
         self.id_conflicts = id_conflicts
         self.pipeline = pipeline
         self.internal_params = bulk_data_based(total_partitions, partition_index, corpora, batch_size,
@@ -543,7 +555,8 @@ class PartitionBulkIndexParamSource:
         return next(self.internal_params)
 
     def size(self):
-        return number_of_bulks(self.corpora, self.partition_index, self.total_partitions, self.bulk_size)
+        all_bulks = number_of_bulks(self.corpora, self.partition_index, self.total_partitions, self.bulk_size)
+        return math.ceil((all_bulks * self.ingest_percentage) / 100)
 
 
 def number_of_bulks(corpora, partition_index, total_partitions, bulk_size):


Test patch:
diff --git a/tests/track/params_test.py b/tests/track/params_test.py
index 13794565..c6b9d8f2 100644
--- a/tests/track/params_test.py
+++ b/tests/track/params_test.py
@@ -475,11 +475,39 @@ class BulkIndexParamSourceTests(TestCase):
 
         self.assertEqual("Unknown 'conflicts' setting [crazy]", ctx.exception.args[0])
 
+    def test_create_with_ingest_percentage_too_low(self):
+        with self.assertRaises(exceptions.InvalidSyntax) as ctx:
+            params.BulkIndexParamSource(track=track.Track(name="unit-test"), params={
+                "bulk-size": 5000,
+                "ingest-percentage": 0.0
+            })
+
+        self.assertEqual("'ingest-percentage' must be in the range (0.0, 100.0] but was 0.0", ctx.exception.args[0])
+
+    def test_create_with_ingest_percentage_too_high(self):
+        with self.assertRaises(exceptions.InvalidSyntax) as ctx:
+            params.BulkIndexParamSource(track=track.Track(name="unit-test"), params={
+                "bulk-size": 5000,
+                "ingest-percentage": 100.1
+            })
+
+        self.assertEqual("'ingest-percentage' must be in the range (0.0, 100.0] but was 100.1", ctx.exception.args[0])
+
+    def test_create_with_ingest_percentage_not_numeric(self):
+        with self.assertRaises(exceptions.InvalidSyntax) as ctx:
+            params.BulkIndexParamSource(track=track.Track(name="unit-test"), params={
+                "bulk-size": 5000,
+                "ingest-percentage": "100 percent"
+            })
+
+        self.assertEqual("'ingest-percentage' must be numeric", ctx.exception.args[0])
+
     def test_create_valid_param_source(self):
         self.assertIsNotNone(params.BulkIndexParamSource(track.Track(name="unit-test"), params={
             "conflicts": "random",
             "bulk-size": 5000,
             "batch-size": 20000,
+            "ingest-percentage": 20.5,
             "pipeline": "test-pipeline"
         }))
 
@@ -565,6 +593,63 @@ class BulkIndexParamSourceTests(TestCase):
 
         self.assertEqual("The provided corpus ['does_not_exist'] does not match any of the corpora ['default'].", ctx.exception.args[0])
 
+    def test_ingests_all_documents_by_default(self):
+        corpora = [
+            track.DocumentCorpus(name="default", documents=[
+                track.Documents(source_format=track.Documents.SOURCE_FORMAT_BULK,
+                                number_of_documents=300000,
+                                target_index="test-idx",
+                                target_type="test-type"
+                                )
+            ]),
+            track.DocumentCorpus(name="special", documents=[
+                track.Documents(source_format=track.Documents.SOURCE_FORMAT_BULK,
+                                number_of_documents=700000,
+                                target_index="test-idx2",
+                                target_type="type"
+                                )
+            ]),
+        ]
+
+        source = params.BulkIndexParamSource(
+            track=track.Track(name="unit-test", corpora=corpora),
+            params={
+                "bulk-size": 10000
+            })
+
+        partition = source.partition(0, 1)
+        # # no ingest-percentage specified, should issue all one hundred bulk requests
+        self.assertEqual(100, partition.size())
+
+    def test_restricts_number_of_bulks_if_required(self):
+        corpora = [
+            track.DocumentCorpus(name="default", documents=[
+                track.Documents(source_format=track.Documents.SOURCE_FORMAT_BULK,
+                                number_of_documents=300000,
+                                target_index="test-idx",
+                                target_type="test-type"
+                                )
+            ]),
+            track.DocumentCorpus(name="special", documents=[
+                track.Documents(source_format=track.Documents.SOURCE_FORMAT_BULK,
+                                number_of_documents=700000,
+                                target_index="test-idx2",
+                                target_type="type"
+                                )
+            ]),
+        ]
+
+        source = params.BulkIndexParamSource(
+            track=track.Track(name="unit-test", corpora=corpora),
+            params={
+                "bulk-size": 10000,
+                "ingest-percentage": 2.5
+            })
+
+        partition = source.partition(0, 1)
+        # should issue three bulks of size 10.000
+        self.assertEqual(3, partition.size())
+
 
 class BulkDataGeneratorTests(TestCase):
     class TestBulkReader: