Make sure we don't fetch chunks that don't span the start point

Author: Artjoms Iskovs

arctic/tickstore/tickstore.py

@@ -176,17 +176,28 @@ def _mongo_date_range_query(self, symbol, date_range):
                             # Throw away everything but the start of every chunk and the symbol
                             {'$project': {'_id': 0, 's': 1, 'sy': 1}},
                             # For every symbol, get the latest chunk start (that is still before
-                            # our sought start, so all of these chunks span the start point)
+                            # our sought start)
                             {'$group': {'_id': '$sy', 'start': {'$max': '$s'}}},
-                            # Take the earliest one of these chunk starts
                             {'$sort': {'start': 1}},
-                            {'$limit': 1}])
+                            ])
+            # Now we need to get the earliest start of the chunk that still spans the start point.
+            # Since we got them sorted by start, we just need to fetch their ends as well and stop
+            # when we've seen the first such chunk
             try:
-                first_dt = next(result)['start']
+                for candidate in result:
+                    chunk = self._collection.find_one({'s': candidate['start'], 'sy': candidate['_id']}, {'e': 1})
+                    if chunk['e'].replace(tzinfo=mktz('UTC')) > start:
+                        first_dt = candidate['start'].replace(tzinfo=mktz('UTC'))
+                        break
             except StopIteration:
                 pass
+            
         if first_dt:
             start_range['$gte'] = first_dt
+        else:
+            # If all chunks start inside of the range, make sure
+            # we don't fetch everything
+            start_range['$gte'] = start
 
         # Find the end bound
         if date_range.end:

tests/integration/tickstore/test_ts_read.py

@@ -409,6 +409,50 @@ def test_read_chunk_boundaries(tickstore_lib):
     assert len(tickstore_lib.read(['SYM1', 'SYM2'], columns=None, date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')), dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC'))))) == 4
 
 
+def test_read_spanning_chunks(tickstore_lib):
+    SYM1_DATA = [
+              {'a': 1.,
+               'b': 2.,
+               'index': dt(2013, 6, 1, 11, 00, tzinfo=mktz('UTC'))
+               },
+               {'a': 3.,
+               'b': 4.,
+               'index': dt(2013, 6, 1, 12, 00, tzinfo=mktz('UTC'))
+               },
+             # Chunk boundary here
+               {'a': 5.,
+               'b': 6.,
+               'index': dt(2013, 6, 1, 14, 00, tzinfo=mktz('UTC'))
+               }
+              ]
+    SYM2_DATA = [
+                  {'a': 7.,
+                   'b': 8.,
+                   'index': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC'))
+                   },
+                   {'a': 9.,
+                   'b': 10.,
+                   'index': dt(2013, 6, 1, 13, 30, tzinfo=mktz('UTC'))
+                   },
+                 # Chunk boundary here
+                   {'a': 11.,
+                   'b': 12.,
+                   'index': dt(2013, 6, 1, 14, 30, tzinfo=mktz('UTC'))
+                   }
+                  ]
+    tickstore_lib._chunk_size = 2
+    tickstore_lib.write('SYM1', SYM1_DATA)
+    tickstore_lib.write('SYM2', SYM2_DATA)
+    
+    # Even though the latest chunk that's the closest to the start point for SYM1 starts at 11:00, it ends before the start point,
+    # so we want to ignore it and start from SYM2 (12:30) instead.
+    assert tickstore_lib._mongo_date_range_query(
+                                                 ['SYM1', 'SYM2'],
+                                                 date_range=DateRange(dt(2013, 6, 1, 12, 45, tzinfo=mktz('UTC')),
+                                                                      dt(2013, 6, 1, 15, 00, tzinfo=mktz('UTC')))) == \
+        {'s': {'$gte': dt(2013, 6, 1, 12, 30, tzinfo=mktz('UTC')), '$lte': dt(2013, 6, 1, 15, 0, tzinfo=mktz('UTC'))}}
+
+
 def test_read_longs(tickstore_lib):
     DUMMY_DATA = [
                   {'a': 1,

tests/unit/tickstore/test_tickstore.py

@@ -18,18 +18,27 @@ def test_mongo_date_range_query():
     self = create_autospec(TickStore)
     self._collection = create_autospec(Collection)
     self._symbol_query.return_value = {"sy": { "$in" : [ "s1" , "s2"]}}
-    self._collection.aggregate.return_value = iter([{"_id": "s1", "start": dt(2014, 1, 1, 0, 0, tzinfo=mktz())}])
+    self._collection.aggregate.return_value = iter([{"_id": "s1", "start": dt(2014, 1, 1, 0, 0, tzinfo=mktz())},
+                                                    {"_id": "s2", "start": dt(2014, 1, 1, 12, 0, tzinfo=mktz())}])
+    
+    self._collection.find_one.side_effect = [
+        {'e': dt(2014, 1, 1, 15, 0, tzinfo=mktz())},
+        {'e': dt(2014, 1, 2, 12, 0, tzinfo=mktz())}]
 
     query = TickStore._mongo_date_range_query(self, 'sym', DateRange(dt(2014, 1, 2, 0, 0, tzinfo=mktz()),
                                                                      dt(2014, 1, 3, 0, 0, tzinfo=mktz())))
-    
+
     assert self._collection.aggregate.call_args_list == [call([
      {"$match": {"s": {"$lte": dt(2014, 1, 2, 0, 0, tzinfo=mktz())}, "sy": { "$in" : [ "s1" , "s2"]}}},
      {"$project": {"_id": 0, "s": 1, "sy": 1}},
      {"$group": {"_id": "$sy", "start": {"$max": "$s"}}},
-     {"$sort": {"start": 1}},
-     {"$limit": 1}])]
-    assert query == {'s': {'$gte': dt(2014, 1, 1, 0, 0, tzinfo=mktz()), '$lte': dt(2014, 1, 3, 0, 0, tzinfo=mktz())}}
+     {"$sort": {"start": 1}}])]
+
+    assert self._collection.find_one.call_args_list == [
+        call({'sy': 's1', 's': dt(2014, 1, 1, 0, 0, tzinfo=mktz())}, {'e': 1}),
+        call({'sy': 's2', 's': dt(2014, 1, 1, 12, 0, tzinfo=mktz())}, {'e': 1})]
+
+    assert query == {'s': {'$gte': dt(2014, 1, 1, 12, 0, tzinfo=mktz()), '$lte': dt(2014, 1, 3, 0, 0, tzinfo=mktz())}}
 
 
 def test_mongo_date_range_query_asserts():