Feature pandas-dev#162 (pandas-dev#188)

* Feature pandas-dev#162 - rename method for ChunkStore to allow symbol renaming

Author: bmoscon

CHANGES.md

@@ -4,6 +4,7 @@
 
   * Bugfix: #187 Compatibility with latest version of pytest-dbfixtures
   * Feature: #182 Improve ChunkStore read/write performance
+  * Feature: #162 Rename API for ChunkStore
 
 ### 1.26 (2016-07-20)
 

arctic/chunkstore/chunkstore.py

@@ -24,7 +24,6 @@
 ROWS = 'r'
 
 
-
 class ChunkStore(object):
     @classmethod
     def initialize_library(cls, arctic_lib, **kwargs):
@@ -73,12 +72,11 @@ def __str__(self):
     def __repr__(self):
         return str(self)
 
-    def _checksum(self, symbol, doc):
+    def _checksum(self, doc):
         """
         Checksum the passed in dictionary
         """
         sha = hashlib.sha1()
-        sha.update(symbol.encode('ascii'))
         sha.update(self.chunker.chunk_to_str(doc[START]).encode('ascii'))
         sha.update(self.chunker.chunk_to_str(doc[END]).encode('ascii'))
         for k in doc[DATA][COLUMNS]:
@@ -126,6 +124,31 @@ def list_symbols(self):
     def _get_symbol_info(self, symbol):
         return self._symbols.find_one({SYMBOL: symbol})
 
+    def rename(self, from_symbol, to_symbol):
+        """
+        Rename a symbol
+
+        Parameters
+        ----------
+        from_symbol: str
+            the existing symbol that will be renamed
+        to_symbol: str
+            the new symbol name
+        """
+
+        sym = self._get_symbol_info(from_symbol)
+        if not sym:
+            raise NoDataFoundException('No data found for %s' % (from_symbol))
+
+        if self._get_symbol_info(to_symbol) is not None:
+            raise Exception('Symbol %s already exists' % (to_symbol))
+
+        mongo_retry(self._collection.update_many)({SYMBOL: from_symbol},
+                                                  {'$set': {SYMBOL: to_symbol}})
+
+        mongo_retry(self._symbols.update_one)({SYMBOL: from_symbol},
+                                              {'$set': {SYMBOL: to_symbol}})
+
     def read(self, symbol, chunk_range=None, columns=None, filter_data=True):
         """
         Reads data for a given symbol from the database.
@@ -160,7 +183,7 @@ def read(self, symbol, chunk_range=None, columns=None, filter_data=True):
             spec.update(self.chunker.to_mongo(chunk_range))
 
         segments = []
-        for _, x in enumerate(self._collection.find(spec, sort=[(START, pymongo.ASCENDING)],)):
+        for x in self._collection.find(spec, sort=[(START, pymongo.ASCENDING)],):
             segments.append(x[DATA])
 
         data = self.serializer.deserialize(segments, columns)
@@ -192,7 +215,7 @@ def write(self, symbol, item, chunk_size):
         doc[CHUNK_SIZE] = chunk_size
         doc[ROWS] = len(item)
         doc[TYPE] = 'dataframe' if isinstance(item, DataFrame) else 'series'
-        
+
         sym = self._get_symbol_info(symbol)
         if sym:
             previous_shas = set([Binary(x[SHA]) for x in self._collection.find({SYMBOL: symbol},
@@ -212,7 +235,7 @@ def write(self, symbol, item, chunk_size):
             chunk[START] = start
             chunk[END] = end
             chunk[SYMBOL] = symbol
-            chunk[SHA] = self._checksum(symbol, chunk)
+            chunk[SHA] = self._checksum(chunk)
 
             if chunk[SHA] not in previous_shas:
                 op = True
@@ -248,13 +271,12 @@ def __update(self, symbol, item, combine_method=None):
         sym = self._get_symbol_info(symbol)
         if not sym:
             raise NoDataFoundException("Symbol does not exist.")
-        
+
         if sym[TYPE] == 'series' and not isinstance(item, Series):
             raise Exception("Cannot combine Series and DataFrame")
         if sym[TYPE] == 'dataframe' and not isinstance(item, DataFrame):
             raise Exception("Cannot combine DataFrame and Series")
 
-
         bulk = self._collection.initialize_unordered_bulk_op()
         op = False
         for start, end, record in self.chunker.to_chunks(item, sym[CHUNK_SIZE]):
@@ -278,19 +300,19 @@ def __update(self, symbol, item, combine_method=None):
             data = self.serializer.serialize(record)
             op = True
 
-            segment = {DATA: data}
-            segment[TYPE] = 'dataframe' if isinstance(record, DataFrame) else 'series'
-            segment[START] = start
-            segment[END] = end
-            sha = self._checksum(symbol, segment)
-            segment[SHA] = sha
+            chunk = {DATA: data}
+            chunk[TYPE] = 'dataframe' if isinstance(record, DataFrame) else 'series'
+            chunk[START] = start
+            chunk[END] = end
+            sha = self._checksum(chunk)
+            chunk[SHA] = sha
             if new_chunk:
                 # new chunk
                 bulk.find({SYMBOL: symbol, SHA: sha}
-                          ).upsert().update_one({'$set': segment})
+                          ).upsert().update_one({'$set': chunk})
             else:
                 bulk.find({SYMBOL: symbol, START: start, END: end}
-                          ).update_one({'$set': segment})
+                          ).update_one({'$set': chunk})
 
         if op:
             bulk.execute()

tests/integration/chunkstore/test_chunkstore.py

@@ -7,6 +7,9 @@
 import numpy as np
 import random
 import pytest
+import pymongo
+
+from arctic.chunkstore.chunkstore import START, SYMBOL
 
 
 def test_write_dataframe(chunkstore_lib):
@@ -902,3 +905,43 @@ def test_read_column_subset(chunkstore_lib):
     chunkstore_lib.write('test', df, 'D')
     r = chunkstore_lib.read('test', columns=['prev_close', 'volume'])
     assert_frame_equal(r, df[['prev_close', 'volume']])
+
+
+def test_rename(chunkstore_lib):
+    df = DataFrame(data={'data': [1, 2, 3, 4, 5, 6, 7, 8, 9],
+                         'open': [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9],
+                         'close': [1.2, 2.3, 3.4, 4.5, 5.6, 6.7, 7.8, 8.9, 9.0],
+                         'prev_close': [.1, .2, .3, .4, .5, .6, .7, .8, .8],
+                         'volume': [1000, 2000, 3000, 4000, 5000, 6000, 7000, 8000, 9000]
+                         },
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 1, 3), 1),
+                                                 (dt(2016, 2, 1), 1),
+                                                 (dt(2016, 2, 2), 1),
+                                                 (dt(2016, 2, 3), 1),
+                                                 (dt(2016, 3, 1), 1),
+                                                 (dt(2016, 3, 2), 1),
+                                                 (dt(2016, 3, 3), 1)],
+                                                names=['date', 'id'])
+                   )
+
+    chunkstore_lib.write('test', df, 'D')
+    assert_frame_equal(chunkstore_lib.read('test'), df)
+    chunkstore_lib.rename('test', 'new_name')
+    assert_frame_equal(chunkstore_lib.read('new_name'), df)
+
+    with pytest.raises(Exception) as e:
+        chunkstore_lib.rename('new_name', 'new_name')
+    assert('already exists' in str(e))
+
+    assert('test' not in chunkstore_lib.list_symbols())
+
+    '''
+    read out all chunks that have symbol set to 'test'. List should be empty
+    '''
+    chunks = []
+    for x in chunkstore_lib._collection.find({SYMBOL: 'test'}, sort=[(START, pymongo.ASCENDING)],):
+        chunks.append(x)
+
+    assert(len(chunks) == 0)