From 13a2b1c6379dd370e9611f893208f32e86a8b8e6 Mon Sep 17 00:00:00 2001 From: Eddie Hebert Date: Tue, 1 Jan 2013 20:56:25 -0500 Subject: [PATCH] Changes date_sort to use heapq module. Uses heapq.merge to sort input from mulitple sources instead of our own sort module. From profiling heapq.merge is more efficient than our own efforts. --- zipline/gens/composites.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/zipline/gens/composites.py b/zipline/gens/composites.py index 02b40f48..04824c3a 100644 --- a/zipline/gens/composites.py +++ b/zipline/gens/composites.py @@ -13,10 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import heapq + from itertools import chain -from zipline.gens.utils import roundrobin, done_message -from zipline.gens.sort import date_sort +from zipline.gens.utils import done_message + + +def decorate_source(source): + for message in source: + yield ((message.dt, message.source_id), message) def date_sorted_sources(*sources): @@ -24,23 +30,11 @@ def date_sorted_sources(*sources): Takes an iterable of sources, generating namestrings and piping their output into date_sort. """ + sorted_stream = heapq.merge(*(decorate_source(s) for s in sources)) - for source in sources: - assert iter(source), "Source %s not iterable" % source - assert hasattr(source, 'get_hash'), "No get_hash" - - # Get name hashes to pass to date_sort. - names = [source.get_hash() for source in sources] - - # Convert the list of generators into a flat stream by pulling - # one element at a time from each. - stream_in = roundrobin(sources, names) - - # Guarantee the flat stream will be sorted by date, using - # source_id as tie-breaker, which is fully deterministic (given - # deterministic string representation for all args/kwargs) - - return date_sort(stream_in, names) + # Strip out key decoration + for _, message in sorted_stream: + yield message def sequential_transforms(stream_in, *transforms):