Add a bulk ingest example

Author: sethmlarson

docs/helpers.rst

@@ -91,7 +91,7 @@ document is like ``{"word": "<myword>"}``.
 
 
 For a more complete and complex example please take a look at
-https://github.com/elastic/elasticsearch-py/blob/master/example/load.py#L76-L130
+https://github.com/elastic/elasticsearch-py/blob/master/examples/bulk-ingest
 
 The :meth:`~elasticsearch.Elasticsearch.parallel_bulk` api is a wrapper around the :meth:`~elasticsearch.Elasticsearch.bulk` api to provide threading. :meth:`~elasticsearch.Elasticsearch.parallel_bulk` returns a generator which must be consumed to produce results.
 

elasticsearch/helpers/test.py

@@ -50,7 +50,9 @@ def setUpClass(cls):
 
     def tearDown(self):
         super(ElasticsearchTestCase, self).tearDown()
-        self.client.indices.delete(index="*", ignore=404, expand_wildcards="open,closed,hidden")
+        self.client.indices.delete(
+            index="*", ignore=404, expand_wildcards="open,closed,hidden"
+        )
         self.client.indices.delete_template(name="*", ignore=404)
 
     @property

examples/bulk-ingest/README.md

@@ -0,0 +1,48 @@
+# Bulk Ingest
+
+A simple script that shows how to ingest a dataset from a file into Elasticsearch.
+The file that is used for this example is a `.csv` so each row is turned into a document.
+
+To run this example install the dependencies with `pip`:
+
+```console
+python -m pip install -r requirements.txt
+```
+
+and then run the script with Python:
+
+```console
+python bulk-ingest.py
+```
+
+You should see the script downloading the dataset into `nyc-restaurants.csv`
+
+Once all the data is loaded into Elasticsearch you can do queries on the dataset
+or create visualizations within Kibana.
+
+```python
+import elasticsearch
+
+client = elasticsearch.Elasticsearch()
+resp = client.search(
+    index="nyc-restaurants",
+    size=0,
+    body={
+        "aggs": {
+            "borough": {
+                "terms": {
+                    "field": "borough"
+                },
+                "aggs": {
+                    "grades": {
+                        "terms": {
+                            "field": "grade"
+                        }
+                    }
+                }
+            }
+        }
+    }
+)
+print(resp)
+```

examples/bulk-ingest/bulk-ingest.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+"""Script that downloads a public dataset and streams it to an Elasticsearch cluster"""
+
+import csv
+from os.path import abspath, join, dirname, exists
+import tqdm
+import urllib3
+from elasticsearch import Elasticsearch
+from elasticsearch.helpers import streaming_bulk
+
+
+NYC_RESTAURANTS = (
+    "https://data.cityofnewyork.us/api/views/43nn-pn8j/rows.csv?accessType=DOWNLOAD"
+)
+DATASET_PATH = join(dirname(abspath(__file__)), "nyc-restaurants.csv")
+CHUNK_SIZE = 16384
+
+
+def download_dataset():
+    """Downloads the public dataset if not locally downlaoded
+    and returns the number of rows are in the .csv file.
+    """
+    if not exists(DATASET_PATH):
+        http = urllib3.PoolManager()
+        resp = http.request("GET", NYC_RESTAURANTS, preload_content=False)
+
+        if resp.status != 200:
+            raise RuntimeError("Could not download dataset")
+
+        with open(DATASET_PATH, mode="wb") as f:
+            chunk = resp.read(CHUNK_SIZE)
+            while chunk:
+                f.write(chunk)
+                chunk = resp.read(CHUNK_SIZE)
+
+    with open(DATASET_PATH) as f:
+        return sum([1 for _ in f]) - 1
+
+
+def create_index(client):
+    """Creates an index in Elasticsearch if one isn't already there."""
+    client.indices.create(
+        index="nyc-restaurants",
+        body={
+            "settings": {"number_of_shards": 1},
+            "mappings": {
+                "properties": {
+                    "name": {"type": "text"},
+                    "borough": {"type": "keyword"},
+                    "cuisine": {"type": "keyword"},
+                    "grade": {"type": "keyword"},
+                    "location": {"type": "geo_point"},
+                }
+            },
+        },
+        ignore=400,
+    )
+
+
+def generate_actions():
+    """Reads the file through csv.DictReader() and for each row
+    yields a single document. This function is passed into the bulk()
+    helper to create many documents in sequence.
+    """
+    with open(DATASET_PATH, mode="r") as f:
+        reader = csv.DictReader(f)
+
+        for row in reader:
+            doc = {
+                "_id": row["CAMIS"],
+                "name": row["DBA"],
+                "borough": row["BORO"],
+                "cuisine": row["CUISINE DESCRIPTION"],
+                "grade": row["GRADE"] or None,
+            }
+
+            lat = row["Latitude"]
+            lon = row["Longitude"]
+            if lat not in ("", "0") and lon not in ("", "0"):
+                doc["location"] = {"lat": float(lat), "lon": float(lon)}
+            yield doc
+
+
+def main():
+    print("Loading dataset...")
+    number_of_docs = download_dataset()
+
+    client = Elasticsearch(
+        # Add your cluster configuration here!
+    )
+    print("Creating an index...")
+    create_index(client)
+
+    print("Indexing documents...")
+    progress = tqdm.tqdm(unit="docs", total=number_of_docs)
+    successes = 0
+    for ok, action in streaming_bulk(
+        client=client, index="nyc-restaurants", actions=generate_actions(),
+    ):
+        progress.update(1)
+        successes += ok
+    print("Indexed %d/%d documents" % (successes, number_of_docs))
+
+
+if __name__ == "__main__":
+    main()

examples/bulk-ingest/requirements.txt

@@ -0,0 +1,3 @@
+elasticsearch
+urllib3
+tqdm

tox.ini

@@ -7,22 +7,30 @@ setenv =
 commands =
     python setup.py test
 
-[testenv:lint]
+[testenv:blacken]
 deps =
-    flake8
     black
 commands =
     black --target-version=py27 \
         elasticsearch/ \
         test_elasticsearch/ \
+        examples/ \
         setup.py
+
+[testenv:lint]
+deps =
+    flake8
+    black
+commands =
     black --target-version=py27 --check \
         elasticsearch/ \
         test_elasticsearch/ \
+        examples/ \
         setup.py
     flake8 \
         elasticsearch/ \
-        test_elasticsearch/
+        test_elasticsearch/ \
+        examples/
 
 [testenv:docs]
 deps =