Skip to content

Commit

Permalink
BUG: Fix factor.top(1) with no groupby.
Browse files Browse the repository at this point in the history
  • Loading branch information
Scott Sanderson committed Jun 21, 2018
1 parent bcc31fa commit 9798372
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 0 deletions.
27 changes: 27 additions & 0 deletions tests/pipeline/test_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -1529,3 +1529,30 @@ def test_run_chunked_pipeline(self):
chunksize=22
)
self.assertTrue(chunked_result.equals(pipeline_result))


class MaximumRegressionTest(WithSeededRandomPipelineEngine,
ZiplineTestCase):
ASSET_FINDER_EQUITY_SIDS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10)

def test_no_groupby_maximum(self):
# This is a regression test for a bug where factor.top(1) would fail
# when not passed a groupby parameters

factor = TestingDataSet.float_col.latest
maximum = factor.top(1)
pipe = Pipeline({'factor': factor, 'maximum': maximum})
result = self.run_pipeline(
pipe, self.trading_days[-5], self.trading_days[-1]
)

# We should have one maximum every day.
maxes_per_day = result.groupby(level=0)['maximum'].sum()
self.assertTrue((maxes_per_day == 1).all())

# The maximum computed by pipeline should match the maximum computed by
# doing a groupby in pandas.
groupby_max = result.groupby(level=0).max()
pipeline_max = result[result.maximum].reset_index(level=1, drop=True)

assert_equal(groupby_max, pipeline_max)
6 changes: 6 additions & 0 deletions zipline/pipeline/filters/filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@
repeat_first_axis,
)

from ..sentinels import NotSpecified


def concat_tuples(*tuples):
"""
Expand Down Expand Up @@ -593,6 +595,10 @@ class MaximumFilter(Filter, StandardOutputs):
window_length = 0

def __new__(cls, factor, groupby, mask):
from zipline.pipeline.classifiers import Everything
if groupby is NotSpecified:
groupby = Everything()

return super(MaximumFilter, cls).__new__(
cls,
inputs=(factor, groupby),
Expand Down

0 comments on commit 9798372

Please sign in to comment.