Merge pull request #20 from bmoscon/elastic

Elasticsearch improvements

config.yaml

@@ -52,9 +52,12 @@ pass_through:
 
 
 elastic:
-    host: 'http://localhost:9200'
+    host: 'http://127.0.0.1:9200'
     user: null
     token: null
+    shards: 10
+    replicas: 0
+    refresh_interval: '30s'
 
 influx:
     host: 'http://127.0.0.1:8086'

cryptostore/aggregator/redis.py

@@ -39,12 +39,20 @@ def read(self, exchange, dtype, pair):
 
         LOG.info("%s: Read %d messages from Redis", key, len(data[0][1]))
         ret = []
+
         for update_id, update in data[0][1]:
             if dtype in {L2_BOOK, L3_BOOK}:
                 update = json.loads(update['data'])
                 update = book_flatten(update, update['timestamp'], update['delta'])
+                for u in update:
+                    for k in ('size', 'amount', 'price', 'timestamp'):
+                        if k in u:
+                            u[k] = float(u[k])
                 ret.extend(update)
             if dtype == TRADES:
+                for k in ('size', 'amount', 'price', 'timestamp'):
+                    if k in update:
+                        update[k] = float(update[k])
                 ret.append(update)
             self.ids[key].append(update_id)
 

cryptostore/collector.py

@@ -42,7 +42,7 @@ def run(self):
             trade_cb = TradeStream
             book_cb = BookStream
             book_up = BookDeltaStream if not depth and self.config['book_delta'] else None
-            kwargs = {'host': self.config['redis']['ip'], 'port': self.config['redis']['port']}
+            kwargs = {'host': self.config['redis']['ip'], 'port': self.config['redis']['port'], 'numeric_type': float}
         elif cache == 'kafka':
             from cryptofeed.backends.kafka import TradeKafka, BookKafka, BookDeltaKafka
             trade_cb = TradeKafka

cryptostore/data/arctic.py

@@ -26,14 +26,11 @@ def write(self, exchange, data_type, pair, timestamp):
 
         if data_type == TRADES:
             df['id'] = df['id'].astype(str)
-            df['size'] = df.amount.astype('float64')
-            df['price'] = df.price.astype('float64')
+            df['size'] = df.amount
             df['date'] = pd.to_datetime(df['timestamp'], unit='s')
             df = df.drop(['pair', 'feed', 'amount'], axis=1)
             chunk_size = 'H'
         elif data_type in { L2_BOOK, L3_BOOK }:
-            df['size'] = df['size'].astype('float64')
-            df['price'] = df.price.astype('float64')
             df['date'] = pd.to_datetime(df['timestamp'], unit='s')
             chunk_size = 'T'
 

cryptostore/data/elastic.py

@@ -6,12 +6,16 @@
 '''
 import json
 import itertools
+import logging
 
 import requests
 
 from cryptostore.data.store import Store
 
 
+LOG = logging.getLogger('cryptostore')
+
+
 def chunk(iterable, length):
     return (iterable[i : i + length] for i in range(0, len(iterable), length))
 
@@ -22,16 +26,33 @@ def __init__(self, config: dict):
         self.host = config.host
         self.user = config.user
         self.token = config.token
+        self.settings = {'settings': {
+                            "index" : {
+                                "number_of_shards" : config.shards,
+                                "number_of_replicas" : config.replicas,
+                                "refresh_interval": config.refresh_interval
+                                }
+                            }
+                        }
 
     def aggregate(self, data):
         self.data = data
 
     def write(self, exchange, data_type, pair, timestamp):
-        for c in chunk(self.data, 100000):
-            data = itertools.chain(*zip([json.dumps({ "index":{} })] * len(c), [json.dumps(d) for d in c]))
+        if requests.head(f"{self.host}/{data_type}").status_code != 200:
+            r = requests.put(f"{self.host}/{data_type}", data=json.dumps(self.settings), auth=(self.user, self.token), headers={'content-type': 'application/json'})
+            if r.status_code != 200:
+                LOG.error("Elasticsearch Index creation failed: %s", r.text)
+            r.raise_for_status()
+
+        LOG.info("Writing %d documents to Elasticsearch", len(self.data))
+        for c in chunk(self.data, 10000):
+            data = itertools.chain(*zip(['{"index": {}}'] * len(c), [json.dumps(d) for d in c]))
             data = '\n'.join(data)
             data = f"{data}\n"
             r = requests.post(f"{self.host}/{data_type}/{data_type}/_bulk", auth=(self.user, self.token), data=data, headers={'content-type': 'application/x-ndjson'})
+            if r.status_code != 200:
+                LOG.error("Elasticsearch insert failed: %s", r.text)
             r.raise_for_status()
         self.data = None
 

cryptostore/data/parquet.py

@@ -52,8 +52,6 @@ def aggregate(self, data):
         for entry in data:
             for key in entry:
                 val = entry[key]
-                if key in {'timestamp', 'amount', 'size', 'price'}:
-                    val = float(val)
                 cols[key].append(val)
         arrays = [pa.array(cols[col]) for col in cols]
         table = pa.Table.from_arrays(arrays, names=names)