Feature pandas-dev#186 and Bugfix pandas-dev#189 (pandas-dev#190)

* Feature pandas-dev#186 and Bugfix pandas-dev#189

Author: bmoscon

CHANGES.md

@@ -5,6 +5,8 @@
   * Bugfix: #187 Compatibility with latest version of pytest-dbfixtures
   * Feature: #182 Improve ChunkStore read/write performance
   * Feature: #162 Rename API for ChunkStore
+  * Feature: #186 chunk_range on update
+  * Bugfix: #189 range delete does not update symbol metadata
 
 ### 1.26 (2016-07-20)
 

arctic/chunkstore/chunkstore.py

@@ -98,6 +98,7 @@ def delete(self, symbol, chunk_range=None):
             # read out chunks that fall within the range and filter out
             # data within the range
             df = self.read(symbol, chunk_range=chunk_range, filter_data=False)
+            row_adjust = len(df)
             df = self.chunker.exclude(df, chunk_range)
 
             # remove chunks, and update any remaining data
@@ -106,6 +107,12 @@ def delete(self, symbol, chunk_range=None):
             self._collection.delete_many(query)
             self.update(symbol, df)
 
+            # update symbol metadata (rows and chunk count)
+            sym = self._get_symbol_info(symbol)
+            sym[ROWS] -= row_adjust
+            sym[CHUNK_COUNT] = self._collection.count({SYMBOL: symbol})
+            self._symbols.replace_one({SYMBOL: symbol}, sym)
+
         else:
             query = {SYMBOL: symbol}
             self._collection.delete_many(query)
@@ -264,7 +271,7 @@ def __concat(self, a, b):
     def __take_new(self, a, b):
         return a
 
-    def __update(self, symbol, item, combine_method=None):
+    def __update(self, symbol, item, combine_method=None, chunk_range=None):
         if not isinstance(item, (DataFrame, Series)):
             raise Exception("Can only chunk DataFrames and Series")
 
@@ -277,11 +284,16 @@ def __update(self, symbol, item, combine_method=None):
         if sym[TYPE] == 'dataframe' and not isinstance(item, DataFrame):
             raise Exception("Cannot combine DataFrame and Series")
 
+        if chunk_range:
+            self.delete(symbol, chunk_range)
+            sym = self._get_symbol_info(symbol)
+
         bulk = self._collection.initialize_unordered_bulk_op()
         op = False
         for start, end, record in self.chunker.to_chunks(item, sym[CHUNK_SIZE]):
             # read out matching chunks
             df = self.read(symbol, chunk_range=self.chunker.to_range(start, end), filter_data=False)
+
             # assuming they exist, update them and store the original chunk
             # range for later use
             if not df.empty:
@@ -334,7 +346,7 @@ def append(self, symbol, item):
         """
         self.__update(symbol, item, combine_method=self.__concat)
 
-    def update(self, symbol, item):
+    def update(self, symbol, item, chunk_range=None):
         """
         Overwrites data in DB with data in item for the given symbol.
 
@@ -346,9 +358,19 @@ def update(self, symbol, item):
             the symbol for the given item in the DB
         item: DataFrame or Series
             the data to update
+        chunk_range: None, or a range object
+            If a range is specified, it will clear/delete the data within the
+            range and overwrite it with the data in item. This allows the user
+            to update with data that might only be a subset of the
+            original data.
         """
 
-        self.__update(symbol, item, combine_method=self.__take_new)
+        if chunk_range:
+            if self.chunker.filter(item, chunk_range).empty:
+                raise Exception('Range must be inclusive of data')
+            self.__update(symbol, item, combine_method=self.__concat, chunk_range=chunk_range)
+        else:
+            self.__update(symbol, item, combine_method=self.__take_new, chunk_range=chunk_range)
 
     def get_info(self, symbol):
         sym = self._get_symbol_info(symbol)

tests/integration/chunkstore/test_chunkstore.py

@@ -342,6 +342,8 @@ def test_update(chunkstore_lib):
     chunkstore_lib.write('chunkstore_test', df, chunk_size='D')
     chunkstore_lib.update('chunkstore_test', df2)
     assert_frame_equal(chunkstore_lib.read('chunkstore_test'), equals)
+    assert(chunkstore_lib.get_info('chunkstore_test')['rows'] == len(equals))
+    assert(chunkstore_lib.get_info('chunkstore_test')['chunk_count'] == len(equals))
 
 
 def test_update_no_overlap(chunkstore_lib):
@@ -364,8 +366,25 @@ def test_update_no_overlap(chunkstore_lib):
 
     chunkstore_lib.write('chunkstore_test', df, chunk_size='D')
     chunkstore_lib.update('chunkstore_test', df2)
-    assert_frame_equal(chunkstore_lib.read('chunkstore_test') , equals)
+    assert_frame_equal(chunkstore_lib.read('chunkstore_test'), equals)
+
+
+def test_update_chunk_range(chunkstore_lib):
+    df = DataFrame(data={'data': [1, 2, 3]},
+                   index=pd.Index(data=[dt(2015, 1, 1),
+                                        dt(2015, 1, 2),
+                                        dt(2015, 1, 3)], name='date'))
+    df2 = DataFrame(data={'data': [30]},
+                    index=pd.Index(data=[dt(2015, 1, 2)],
+                                   name='date'))
+    equals = DataFrame(data={'data': [30, 3]},
+                       index=pd.Index(data=[dt(2015, 1, 2),
+                                            dt(2015, 1, 3)],
+                                      name='date'))
 
+    chunkstore_lib.write('chunkstore_test', df, chunk_size='M')
+    chunkstore_lib.update('chunkstore_test', df2, chunk_range=DateRange(dt(2015, 1, 1), dt(2015, 1, 2)))
+    assert_frame_equal(chunkstore_lib.read('chunkstore_test'), equals)
 
 def test_append_before(chunkstore_lib):
     df = DataFrame(data={'data': [1, 2, 3]},
@@ -664,6 +683,8 @@ def test_delete_range(chunkstore_lib):
     chunkstore_lib.write('test', df, 'M')
     chunkstore_lib.delete('test', chunk_range=DateRange(dt(2016, 1, 2), dt(2016, 3, 1)))
     assert_frame_equal(chunkstore_lib.read('test'), df_result)
+    assert(chunkstore_lib.get_info('test')['rows'] == len(df_result))
+    assert(chunkstore_lib.get_info('test')['chunk_count'] == 2)
 
 
 def test_delete_range_noindex(chunkstore_lib):