mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-17 11:25:55 +08:00
ENH: Updates the blaze loader and adds more tests
This commit is contained in:
llllllllll
+326
-22
@@ -3,31 +3,31 @@ Tests for the blaze interface to the pipeline api.
|
||||
"""
|
||||
from __future__ import division
|
||||
|
||||
from collections import OrderedDict
|
||||
from datetime import timedelta
|
||||
from itertools import chain
|
||||
from unittest import TestCase
|
||||
import warnings
|
||||
|
||||
import blaze as bz
|
||||
from datashape import dshape
|
||||
from datashape import dshape, var, Record
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from pandas.util.testing import assert_frame_equal
|
||||
from toolz import flip
|
||||
from toolz import keymap
|
||||
from toolz.curried import operator as op
|
||||
|
||||
from zipline.pipeline import Pipeline
|
||||
from zipline.pipeline import Pipeline, CustomFactor
|
||||
from zipline.pipeline.data import DataSet, BoundColumn
|
||||
from zipline.pipeline.engine import SimplePipelineEngine
|
||||
from zipline.pipeline.loaders.blaze import (
|
||||
pipeline_api_from_blaze,
|
||||
BlazeLoader,
|
||||
NoDeltasWarning,
|
||||
NonNumpyField,
|
||||
NonPipelineField,
|
||||
NotPipelineCompatible,
|
||||
)
|
||||
from zipline.utils.test_utils import (
|
||||
make_simple_asset_info,
|
||||
tmp_asset_finder,
|
||||
)
|
||||
from zipline.utils.test_utils import tmp_asset_finder
|
||||
|
||||
|
||||
nameof = op.attrgetter('name')
|
||||
@@ -43,15 +43,16 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
[pd.Timestamp('2014-01-02')] * 3 +
|
||||
[pd.Timestamp('2014-01-03')] * 3
|
||||
)
|
||||
cls.df = pd.DataFrame({
|
||||
'sid': [ord('A'), ord('B'), ord('C')] * 3,
|
||||
cls.sids = sids = ord('A'), ord('B'), ord('C')
|
||||
cls.df = df = pd.DataFrame({
|
||||
'sid': sids * 3,
|
||||
'value': tuple(
|
||||
chain.from_iterable((a, a + 1, a + 2) for a in range(3)),
|
||||
),
|
||||
'asof_date': dates,
|
||||
'timestamp': dates,
|
||||
})
|
||||
cls.dshape = dshape("""
|
||||
cls.dshape = dshape_ = dshape("""
|
||||
var * {
|
||||
sid: ?int64,
|
||||
value: ?float64,
|
||||
@@ -59,12 +60,21 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
timestamp: datetime
|
||||
}
|
||||
""")
|
||||
cls.macro_df = df[df.sid == 65].drop('sid', axis=1)
|
||||
dshape_ = OrderedDict(dshape_.measure.fields)
|
||||
del dshape_['sid']
|
||||
cls.macro_dshape = var * Record(dshape_)
|
||||
|
||||
cls.garbage_loader = BlazeLoader()
|
||||
|
||||
def test_tabular(self):
|
||||
name = 'expr'
|
||||
expr = bz.Data(self.df, name=name, dshape=self.dshape)
|
||||
ds = pipeline_api_from_blaze(expr, loader=self.garbage_loader)
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
self.assertEqual(ds.__name__, name)
|
||||
self.assertTrue(issubclass(ds, DataSet))
|
||||
self.assertEqual(
|
||||
@@ -79,31 +89,51 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
self.assertIn("'datetime'", str(e.exception))
|
||||
|
||||
self.assertIs(
|
||||
pipeline_api_from_blaze(expr, loader=self.garbage_loader),
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
),
|
||||
ds,
|
||||
)
|
||||
|
||||
def test_column(self):
|
||||
exprname = 'expr'
|
||||
expr = bz.Data(self.df, name=exprname, dshape=self.dshape)
|
||||
value = pipeline_api_from_blaze(expr.value, loader=self.garbage_loader)
|
||||
value = pipeline_api_from_blaze(
|
||||
expr.value,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
self.assertEqual(value.name, 'value')
|
||||
self.assertIsInstance(value, BoundColumn)
|
||||
self.assertEqual(value.dtype, np.float64)
|
||||
|
||||
# test memoization
|
||||
self.assertIs(
|
||||
pipeline_api_from_blaze(expr.value, loader=self.garbage_loader),
|
||||
pipeline_api_from_blaze(
|
||||
expr.value,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
),
|
||||
value,
|
||||
)
|
||||
self.assertIs(
|
||||
pipeline_api_from_blaze(expr, loader=self.garbage_loader).value,
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
).value,
|
||||
value,
|
||||
)
|
||||
|
||||
# test the walk back up the tree
|
||||
self.assertIs(
|
||||
pipeline_api_from_blaze(expr, loader=self.garbage_loader),
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
),
|
||||
value.dataset,
|
||||
)
|
||||
self.assertEqual(value.dataset.__name__, exprname)
|
||||
@@ -121,15 +151,144 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
)
|
||||
|
||||
with self.assertRaises(TypeError) as e:
|
||||
pipeline_api_from_blaze(expr, loader=self.garbage_loader)
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
self.assertIn("'asof_date'", str(e.exception))
|
||||
self.assertIn(repr(str(expr.dshape.measure)), str(e.exception))
|
||||
|
||||
def test_id(self):
|
||||
name = 'expr'
|
||||
expr = bz.Data(self.df, name=name, dshape=self.dshape)
|
||||
def test_auto_deltas(self):
|
||||
expr = bz.Data(
|
||||
{'ds': self.df,
|
||||
'ds_deltas': pd.DataFrame(columns=self.df.columns)},
|
||||
dshape=var * Record((
|
||||
('ds', self.dshape.measure),
|
||||
('ds_deltas', self.dshape.measure),
|
||||
)),
|
||||
)
|
||||
loader = BlazeLoader()
|
||||
ds = pipeline_api_from_blaze(expr, loader=loader)
|
||||
ds = pipeline_api_from_blaze(expr.ds, loader=loader)
|
||||
self.assertEqual(len(loader), 1)
|
||||
exprdata = loader[ds]
|
||||
self.assertTrue(exprdata.expr.isidentical(expr.ds))
|
||||
self.assertTrue(exprdata.deltas.isidentical(expr.ds_deltas))
|
||||
|
||||
def test_auto_deltas_fail_warn(self):
|
||||
with warnings.catch_warnings(record=True) as ws:
|
||||
warnings.simplefilter('always')
|
||||
loader = BlazeLoader()
|
||||
expr = bz.Data(self.df, dshape=self.dshape)
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=loader,
|
||||
no_deltas_rule='warn',
|
||||
)
|
||||
self.assertEqual(len(ws), 1)
|
||||
w = ws[0].message
|
||||
self.assertIsInstance(w, NoDeltasWarning)
|
||||
self.assertIn(str(expr), str(w))
|
||||
|
||||
def test_auto_deltas_fail_raise(self):
|
||||
loader = BlazeLoader()
|
||||
expr = bz.Data(self.df, dshape=self.dshape)
|
||||
with self.assertRaises(ValueError) as e:
|
||||
pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=loader,
|
||||
no_deltas_rule='raise',
|
||||
)
|
||||
self.assertIn(str(expr), str(e.exception))
|
||||
|
||||
def test_non_numpy_field(self):
|
||||
expr = bz.Data(
|
||||
[],
|
||||
dshape="""
|
||||
var * {
|
||||
a: datetime,
|
||||
asof_date: datetime,
|
||||
timestamp: datetime,
|
||||
}""",
|
||||
)
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
with self.assertRaises(AttributeError):
|
||||
ds.a
|
||||
self.assertIsInstance(object.__getattribute__(ds, 'a'), NonNumpyField)
|
||||
|
||||
def test_non_pipeline_field(self):
|
||||
# NOTE: This test will fail if we ever allow string types in
|
||||
# the pipeline api. If this happens, change the dtype of the `a` field
|
||||
# of expr to another type we don't allow.
|
||||
expr = bz.Data(
|
||||
[],
|
||||
dshape="""
|
||||
var * {
|
||||
a: string,
|
||||
asof_date: datetime,
|
||||
timestamp: datetime,
|
||||
}""",
|
||||
)
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=self.garbage_loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
with self.assertRaises(AttributeError):
|
||||
ds.a
|
||||
self.assertIsInstance(
|
||||
object.__getattribute__(ds, 'a'),
|
||||
NonPipelineField,
|
||||
)
|
||||
|
||||
def test_complex_expr(self):
|
||||
expr = bz.Data(self.df, dshape=self.dshape)
|
||||
# put an Add in the table
|
||||
expr_with_add = bz.transform(expr, value=expr.value + 1)
|
||||
|
||||
# Test that we can have complex expressions with no deltas
|
||||
pipeline_api_from_blaze(
|
||||
expr_with_add,
|
||||
deltas=None,
|
||||
loader=self.garbage_loader,
|
||||
)
|
||||
|
||||
pipeline_api_from_blaze(
|
||||
expr.value + 1, # put an Add in the column
|
||||
deltas=None,
|
||||
loader=self.garbage_loader,
|
||||
)
|
||||
|
||||
deltas = bz.Data(
|
||||
pd.DataFrame(columns=self.df.columns),
|
||||
dshape=self.dshape,
|
||||
)
|
||||
with self.assertRaises(TypeError):
|
||||
pipeline_api_from_blaze(
|
||||
expr_with_add,
|
||||
deltas=deltas,
|
||||
loader=self.garbage_loader,
|
||||
)
|
||||
|
||||
with self.assertRaises(TypeError):
|
||||
pipeline_api_from_blaze(
|
||||
expr.value + 1,
|
||||
deltas=deltas,
|
||||
loader=self.garbage_loader,
|
||||
)
|
||||
|
||||
def test_id(self):
|
||||
expr = bz.Data(self.df, name='expr', dshape=self.dshape)
|
||||
loader = BlazeLoader()
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
p = Pipeline('p')
|
||||
p.add(ds.value.latest, 'value')
|
||||
dates = self.dates
|
||||
@@ -149,3 +308,148 @@ class BlazeToPipelineTestCase(TestCase):
|
||||
expected.index.levels[1].map(finder.retrieve_asset),
|
||||
))
|
||||
assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_id_macro_dataset(self):
|
||||
expr = bz.Data(self.macro_df, name='expr', dshape=self.macro_dshape)
|
||||
loader = BlazeLoader()
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
loader=loader,
|
||||
no_deltas_rule='ignore',
|
||||
)
|
||||
p = Pipeline('p')
|
||||
p.add(ds.value.latest, 'value')
|
||||
dates = self.dates
|
||||
|
||||
with tmp_asset_finder() as finder:
|
||||
result = SimplePipelineEngine(
|
||||
loader,
|
||||
dates,
|
||||
finder,
|
||||
).run_pipeline(p, dates[0], dates[-1])
|
||||
|
||||
expected = pd.DataFrame(
|
||||
[0, 0, 0, 1, 1, 1, 2, 2, 2],
|
||||
index=pd.MultiIndex.from_product((
|
||||
self.macro_df.timestamp,
|
||||
tuple(map(finder.retrieve_asset, self.sids)),
|
||||
)),
|
||||
columns=('value',),
|
||||
)
|
||||
assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_deltas(self):
|
||||
expr = bz.Data(self.df, name='expr', dshape=self.dshape)
|
||||
deltas = bz.Data(self.df.iloc[:-3], name='deltas', dshape=self.dshape)
|
||||
deltas = bz.transform(
|
||||
deltas,
|
||||
value=deltas.value + 10,
|
||||
timestamp=deltas.timestamp + timedelta(days=1),
|
||||
)
|
||||
loader = BlazeLoader()
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
deltas,
|
||||
loader=loader,
|
||||
no_deltas_rule='raise',
|
||||
)
|
||||
p = Pipeline('p')
|
||||
|
||||
expected_views = keymap(pd.Timestamp, {
|
||||
'2014-01-02': np.array([[10.0, 11.0, 12.0],
|
||||
[1.0, 2.0, 3.0]]),
|
||||
'2014-01-03': np.array([[11.0, 12.0, 13.0],
|
||||
[2.0, 3.0, 4.0]]),
|
||||
})
|
||||
assertTrue = self.assertTrue
|
||||
|
||||
class TestFactor(CustomFactor):
|
||||
inputs = ds.value,
|
||||
window_length = 2
|
||||
|
||||
def compute(self, today, assets, out, data):
|
||||
assertTrue((data == expected_views[today]).all())
|
||||
out[:] = np.max(data)
|
||||
|
||||
p.add(TestFactor(), 'value')
|
||||
dates = self.dates
|
||||
|
||||
with tmp_asset_finder() as finder:
|
||||
result = SimplePipelineEngine(
|
||||
loader,
|
||||
dates,
|
||||
finder,
|
||||
).run_pipeline(p, dates[1], dates[-1])
|
||||
|
||||
assert_frame_equal(
|
||||
result,
|
||||
pd.DataFrame(
|
||||
[12, 12, 12, 13, 13, 13],
|
||||
index=pd.MultiIndex.from_product((
|
||||
sorted(expected_views.keys()),
|
||||
tuple(map(finder.retrieve_asset, self.sids)),
|
||||
)),
|
||||
columns=('value',),
|
||||
),
|
||||
check_dtype=False,
|
||||
)
|
||||
|
||||
def test_deltas_macro_dataset(self):
|
||||
expr = bz.Data(self.macro_df, name='expr', dshape=self.macro_dshape)
|
||||
deltas = bz.Data(
|
||||
self.macro_df.iloc[:-1],
|
||||
name='deltas',
|
||||
dshape=self.macro_dshape,
|
||||
)
|
||||
deltas = bz.transform(
|
||||
deltas,
|
||||
value=deltas.value + 10,
|
||||
timestamp=deltas.timestamp + timedelta(days=1),
|
||||
)
|
||||
loader = BlazeLoader()
|
||||
ds = pipeline_api_from_blaze(
|
||||
expr,
|
||||
deltas,
|
||||
loader=loader,
|
||||
no_deltas_rule='raise',
|
||||
)
|
||||
p = Pipeline('p')
|
||||
|
||||
expected_views = keymap(pd.Timestamp, {
|
||||
'2014-01-02': np.array([[10.0, 10.0, 10.0],
|
||||
[1.0, 1.0, 1.0]]),
|
||||
'2014-01-03': np.array([[11.0, 11.0, 11.0],
|
||||
[2.0, 2.0, 2.0]]),
|
||||
})
|
||||
assertTrue = self.assertTrue
|
||||
|
||||
class TestFactor(CustomFactor):
|
||||
inputs = ds.value,
|
||||
window_length = 2
|
||||
|
||||
def compute(self, today, assets, out, data):
|
||||
assertTrue((data == expected_views[today]).all())
|
||||
out[:] = np.max(data)
|
||||
|
||||
p.add(TestFactor(), 'value')
|
||||
dates = self.dates
|
||||
|
||||
with tmp_asset_finder() as finder:
|
||||
result = SimplePipelineEngine(
|
||||
loader,
|
||||
dates,
|
||||
finder,
|
||||
).run_pipeline(p, dates[1], dates[-1])
|
||||
|
||||
assert_frame_equal(
|
||||
result,
|
||||
pd.DataFrame(
|
||||
[10, 10, 10, 11, 11, 11],
|
||||
index=pd.MultiIndex.from_product((
|
||||
sorted(expected_views.keys()),
|
||||
tuple(map(finder.retrieve_asset, self.sids)),
|
||||
)),
|
||||
columns=('value',),
|
||||
),
|
||||
check_dtype=False,
|
||||
)
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
"""Blaze integration with the pipeline API.
|
||||
"""
|
||||
from __future__ import division
|
||||
|
||||
from abc import ABCMeta, abstractproperty
|
||||
from collections import namedtuple
|
||||
from collections import namedtuple, defaultdict
|
||||
from itertools import count
|
||||
from operator import attrgetter
|
||||
import warnings
|
||||
from weakref import WeakKeyDictionary
|
||||
|
||||
import blaze as bz
|
||||
@@ -15,7 +19,6 @@ from datashape import (
|
||||
isscalar,
|
||||
promote,
|
||||
)
|
||||
from logbook import Logger
|
||||
from numpy.lib.stride_tricks import as_strided
|
||||
from odo import odo
|
||||
import pandas as pd
|
||||
@@ -37,7 +40,24 @@ valid_deltas_node_types = (
|
||||
bz.expr.Symbol,
|
||||
)
|
||||
getname = attrgetter('__name__')
|
||||
log = Logger(__name__)
|
||||
|
||||
|
||||
class _ExprRepr(object):
|
||||
"""Box for repring expressions with the str of the expression.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expr : Expr
|
||||
The expression to box for repring.
|
||||
"""
|
||||
__slots__ = 'expr',
|
||||
|
||||
def __init__(self, expr):
|
||||
self.expr = expr
|
||||
|
||||
def __repr__(self):
|
||||
return str(self.expr)
|
||||
__str__ = __repr__
|
||||
|
||||
|
||||
class ExprData(namedtuple('ExprData', 'expr deltas resources')):
|
||||
@@ -57,11 +77,11 @@ class ExprData(namedtuple('ExprData', 'expr deltas resources')):
|
||||
|
||||
def __repr__(self):
|
||||
# If the expressions have _resources() then the repr will
|
||||
# drive computation so we str them.
|
||||
# drive computation so we box them.
|
||||
cls = type(self)
|
||||
return super(ExprData, cls).__repr__(cls(
|
||||
str(self.expr),
|
||||
str(self.deltas),
|
||||
_ExprRepr(self.expr),
|
||||
_ExprRepr(self.deltas),
|
||||
self.resources,
|
||||
))
|
||||
|
||||
@@ -110,6 +130,9 @@ class NotPipelineCompatible(TypeError):
|
||||
return "'%s' is a non pipleine API compatible type'" % self.args
|
||||
|
||||
|
||||
_new_names = ('_%d' % n for n in count())
|
||||
|
||||
|
||||
@memoize
|
||||
def new_dataset(expr, deltas):
|
||||
"""Creates or returns a dataset from a pair of blaze expressions.
|
||||
@@ -134,19 +157,22 @@ def new_dataset(expr, deltas):
|
||||
for name, type_ in expr.dshape.measure.fields:
|
||||
try:
|
||||
if promote(type_, float64, promote_option=False) != float64:
|
||||
raise NotPipelineCompatible
|
||||
raise NotPipelineCompatible()
|
||||
if isinstance(type_, Option):
|
||||
type_ = type_.ty
|
||||
except TypeError:
|
||||
col = NonNumpyField(name, type_)
|
||||
except NotPipelineCompatible:
|
||||
col = NonPipelineField(name, type_)
|
||||
except TypeError:
|
||||
col = NonNumpyField(name, type_)
|
||||
else:
|
||||
col = Column(type_.to_numpy_dtype().type)
|
||||
|
||||
columns[name] = col
|
||||
|
||||
return type(expr._name, (DataSet,), columns)
|
||||
name = expr._name
|
||||
if expr._name is None:
|
||||
name = next(_new_names)
|
||||
return type(name, (DataSet,), columns)
|
||||
|
||||
|
||||
def _check_resources(name, expr, resources):
|
||||
@@ -202,6 +228,25 @@ def _check_datetime_field(name, measure):
|
||||
)
|
||||
|
||||
|
||||
class NoDeltasWarning(UserWarning):
|
||||
"""Warning used to signal that no deltas could be found and none
|
||||
were provided.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
expr : Expr
|
||||
The expression that was searched.
|
||||
"""
|
||||
def __init__(self, expr):
|
||||
self._expr = expr
|
||||
|
||||
def __str__(self):
|
||||
return 'No deltas could be infered from expr: %s' % self._expr
|
||||
|
||||
|
||||
_valid_no_deltas_rules = 'warn', 'raise', 'ignore'
|
||||
|
||||
|
||||
def _get_deltas(expr, deltas, no_deltas_rule):
|
||||
"""Find the correct deltas for the expression.
|
||||
|
||||
@@ -213,7 +258,7 @@ def _get_deltas(expr, deltas, no_deltas_rule):
|
||||
The deltas argument. If this is 'auto', then the deltas table will
|
||||
be searched for by walking up the expression tree. If this can not be
|
||||
reflected, then an action will be taken based on the 'no_deltas_rule'.
|
||||
no_deltas_rule : {'log', 'raise', 'ignore'}
|
||||
no_deltas_rule : {'warn', 'raise', 'ignore'}
|
||||
How to handle the case where deltas='auto' but no deltas could be
|
||||
found.
|
||||
|
||||
@@ -222,34 +267,32 @@ def _get_deltas(expr, deltas, no_deltas_rule):
|
||||
deltas : Expr or None
|
||||
The deltas table to use.
|
||||
"""
|
||||
if no_deltas_rule not in _get_deltas.valid_no_deltas_rules:
|
||||
if no_deltas_rule not in _valid_no_deltas_rules:
|
||||
raise ValueError(
|
||||
'no_deltas_rule must be one of: %s' %
|
||||
_get_deltas.valid_no_deltas_rules
|
||||
_valid_no_deltas_rules
|
||||
)
|
||||
|
||||
if deltas != 'auto':
|
||||
if isinstance(deltas, bz.Expr) or deltas != 'auto':
|
||||
return deltas
|
||||
|
||||
try:
|
||||
return expr._child[expr._name + '_deltas']
|
||||
except (AttributeError, KeyError):
|
||||
return expr._child[(expr._name or '') + '_deltas']
|
||||
except (ValueError, AttributeError):
|
||||
if no_deltas_rule == 'raise':
|
||||
raise ValueError(
|
||||
"no deltas table could be reflected for '%s'" % expr
|
||||
"no deltas table could be reflected for %s" % expr
|
||||
)
|
||||
elif no_deltas_rule == 'log':
|
||||
log.warn("no deltas table found for '%s'" % expr)
|
||||
elif no_deltas_rule == 'warn':
|
||||
warnings.warn(NoDeltasWarning(expr))
|
||||
return None
|
||||
|
||||
_get_deltas.valid_no_deltas_rules = 'log', 'raise', 'ignore'
|
||||
|
||||
|
||||
def pipeline_api_from_blaze(expr,
|
||||
deltas='auto',
|
||||
loader=None,
|
||||
resources=None,
|
||||
no_deltas_rule='log'):
|
||||
no_deltas_rule=_valid_no_deltas_rules[0]):
|
||||
"""Create a pipeline api object from a blaze expression.
|
||||
|
||||
Parameters
|
||||
@@ -268,7 +311,7 @@ def pipeline_api_from_blaze(expr,
|
||||
resources : dict or any, optional
|
||||
The data to execute the blaze expressions against. This is used as the
|
||||
scope for ``bz.compute``.
|
||||
no_deltas_rule : {'log', 'raise', 'ignore'}
|
||||
no_deltas_rule : {'warn', 'raise', 'ignore'}
|
||||
What should happen if ``deltas='auto'`` but no deltas can be found.
|
||||
'log' says to log a message but continue.
|
||||
'raise' says to raise an exception if no deltas can be found.
|
||||
@@ -283,19 +326,6 @@ def pipeline_api_from_blaze(expr,
|
||||
is passed, a ``BoundColumn`` on the dataset that would be constructed
|
||||
from passing the parent is returned.
|
||||
"""
|
||||
# Check if this is a single column out of a dataset.
|
||||
single_column = None
|
||||
if isscalar(expr.dshape.measure):
|
||||
# This is a single column, record which column we are to return
|
||||
# but create the entire dataset.
|
||||
single_column = expr._name
|
||||
col = expr
|
||||
for expr in expr._subterms():
|
||||
if isrecord(expr.dshape.measure):
|
||||
break
|
||||
else:
|
||||
expr = bz.Data(col, name=single_column)
|
||||
|
||||
deltas = _get_deltas(expr, deltas, no_deltas_rule)
|
||||
if deltas is not None:
|
||||
invalid_nodes = tuple(filter(
|
||||
@@ -311,6 +341,19 @@ def pipeline_api_from_blaze(expr,
|
||||
),
|
||||
)
|
||||
|
||||
# Check if this is a single column out of a dataset.
|
||||
single_column = None
|
||||
if isscalar(expr.dshape.measure):
|
||||
# This is a single column, record which column we are to return
|
||||
# but create the entire dataset.
|
||||
single_column = expr._name
|
||||
col = expr
|
||||
for expr in expr._subterms():
|
||||
if isrecord(expr.dshape.measure):
|
||||
break
|
||||
else:
|
||||
expr = bz.Data(col, name=single_column)
|
||||
|
||||
measure = expr.dshape.measure
|
||||
if not isrecord(measure) or AD_FIELD_NAME not in measure.names:
|
||||
raise TypeError(
|
||||
@@ -333,10 +376,12 @@ def pipeline_api_from_blaze(expr,
|
||||
else:
|
||||
_check_datetime_field(TS_FIELD_NAME, measure)
|
||||
|
||||
if deltas is not None and deltas.dshape.measure != measure:
|
||||
if deltas is not None and (sorted(deltas.dshape.measure.fields) !=
|
||||
sorted(measure.fields)):
|
||||
raise TypeError(
|
||||
"base measure != deltas measure ('%s' != '%s')" % (
|
||||
measure, deltas.dshape.measure,
|
||||
'base measure != deltas measure:\n%s != %s' % (
|
||||
measure,
|
||||
deltas.dshape.measure,
|
||||
),
|
||||
)
|
||||
|
||||
@@ -391,10 +436,10 @@ def inline_novel_deltas(base, deltas, dates):
|
||||
)
|
||||
|
||||
|
||||
def overwrite_from_dates(asof, dates, sparse_dates, asset_idx, value):
|
||||
def overwrite_from_dates(asof, dates, dense_dates, asset_idx, value):
|
||||
"""Construct a `Float64Overwrite` with the correct
|
||||
start and end date based on the asof date of the delta,
|
||||
the dense_dates, and the sparse_dates.
|
||||
the dense_dates, and the dense_dates.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -402,7 +447,7 @@ def overwrite_from_dates(asof, dates, sparse_dates, asset_idx, value):
|
||||
The asof date of the delta.
|
||||
dates : pd.DatetimeIndex
|
||||
The dates requested by the loader.
|
||||
sparse_dates : pd.DatetimeIndex
|
||||
dense_dates : pd.DatetimeIndex
|
||||
The dates that appeared in the dataset.
|
||||
asset_idx : int
|
||||
The index of the asset in the block.
|
||||
@@ -416,17 +461,18 @@ def overwrite_from_dates(asof, dates, sparse_dates, asset_idx, value):
|
||||
"""
|
||||
return Float64Overwrite(
|
||||
dates.searchsorted(asof),
|
||||
dates.get_loc(sparse_dates[sparse_dates.searchsorted(asof) + 1]) - 1,
|
||||
dates.get_loc(dense_dates[dense_dates.searchsorted(asof) + 1]) - 1,
|
||||
asset_idx,
|
||||
value,
|
||||
)
|
||||
|
||||
|
||||
def adjustments_from_deltas(dates,
|
||||
sparse_dates,
|
||||
column_idx,
|
||||
assets,
|
||||
deltas):
|
||||
def adjustments_from_deltas_no_sids(dates,
|
||||
dense_dates,
|
||||
column_idx,
|
||||
column_name,
|
||||
assets,
|
||||
deltas):
|
||||
"""Collect all the adjustments that occur in a dataset that does not
|
||||
have a sid column.
|
||||
|
||||
@@ -434,10 +480,12 @@ def adjustments_from_deltas(dates,
|
||||
----------
|
||||
dates : pd.DatetimeIndex
|
||||
The dates requested by the loader.
|
||||
sparse_dates : pd.DatetimeIndex
|
||||
The dates that were in the sparse data.
|
||||
dense_dates : pd.DatetimeIndex
|
||||
The dates that were in the dense data.
|
||||
column_idx : int
|
||||
The index of the column in the dataset.
|
||||
column_name : str
|
||||
The name of the column to compute deltas for.
|
||||
deltas : pd.DataFrame
|
||||
The overwrites that should be applied to the dataset.
|
||||
|
||||
@@ -451,14 +499,56 @@ def adjustments_from_deltas(dates,
|
||||
overwrite_from_dates(
|
||||
deltas.loc[kd, AD_FIELD_NAME],
|
||||
dates,
|
||||
sparse_dates,
|
||||
dense_dates,
|
||||
n,
|
||||
v,
|
||||
) for n in range(len(assets))
|
||||
) for kd, v in deltas.icol(column_idx).iteritems()
|
||||
) for kd, v in deltas[column_name].iteritems()
|
||||
}
|
||||
|
||||
|
||||
def adjustments_from_deltas_with_sids(dates,
|
||||
dense_dates,
|
||||
column_idx,
|
||||
column_name,
|
||||
assets,
|
||||
deltas):
|
||||
"""Collect all the adjustments that occur in a dataset that does not
|
||||
have a sid column.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dates : pd.DatetimeIndex
|
||||
The dates requested by the loader.
|
||||
dense_dates : pd.DatetimeIndex
|
||||
The dates that were in the dense data.
|
||||
column_idx : int
|
||||
The index of the column in the dataset.
|
||||
column_name : str
|
||||
The name of the column to compute deltas for.
|
||||
deltas : pd.DataFrame
|
||||
The overwrites that should be applied to the dataset.
|
||||
|
||||
Returns
|
||||
-------
|
||||
adjustments : dict[idx -> Float64Overwrite]
|
||||
The adjustments dictionary to feed to the adjusted array.
|
||||
"""
|
||||
adjustments = defaultdict(list)
|
||||
for sid_idx, (sid, per_sid) in enumerate(deltas[column_name].iteritems()):
|
||||
for kd, v in per_sid.iteritems():
|
||||
adjustments[dates.get_loc(kd)].append(
|
||||
overwrite_from_dates(
|
||||
deltas[AD_FIELD_NAME].loc[kd, sid],
|
||||
dates,
|
||||
dense_dates,
|
||||
sid_idx,
|
||||
v,
|
||||
),
|
||||
)
|
||||
return dict(adjustments) # no subclasses of dict
|
||||
|
||||
|
||||
class BlazeLoader(dict):
|
||||
def __init__(self, colmap=None):
|
||||
self.update(colmap or {})
|
||||
@@ -501,7 +591,7 @@ class BlazeLoader(dict):
|
||||
# be removed from scope too early otherwise.
|
||||
lower = odo(ts[ts <= dates[0]].max(), pd.Timestamp)
|
||||
return e[
|
||||
e[SID_FIELD_NAME].isin(assets) &
|
||||
(e[SID_FIELD_NAME].isin(assets) if have_sids else True) &
|
||||
((ts >= lower) if lower is not pd.NaT else True) &
|
||||
(ts <= dates[-1])
|
||||
][query_fields]
|
||||
@@ -510,17 +600,16 @@ class BlazeLoader(dict):
|
||||
bz.compute(where(expr), resources),
|
||||
pd.DataFrame,
|
||||
)
|
||||
|
||||
materialized_deltas = (
|
||||
odo(bz.compute(where(deltas), resources), pd.DataFrame)
|
||||
if deltas is not None else
|
||||
pd.DataFrame(columns=query_fields)
|
||||
)
|
||||
# Capture the original (sparse) dates that came from the resource.
|
||||
sparse_dates = pd.DatetimeIndex(materialized_expr[TS_FIELD_NAME])
|
||||
# Inline the deltas that changed our most recently known value.
|
||||
# Also, we reindex by the dates to create a dense representation of
|
||||
# the data.
|
||||
dense_output = inline_novel_deltas(
|
||||
sparse_output = inline_novel_deltas(
|
||||
materialized_expr,
|
||||
materialized_deltas,
|
||||
dates,
|
||||
@@ -529,18 +618,21 @@ class BlazeLoader(dict):
|
||||
if have_sids:
|
||||
# Unstack by the sid so that we get a multi-index on the columns
|
||||
# of datacolumn, sid.
|
||||
dense_output = dense_output.set_index(
|
||||
sparse_output = sparse_output.set_index(
|
||||
SID_FIELD_NAME,
|
||||
append=True,
|
||||
).unstack()
|
||||
sparse_deltas = materialized_deltas.set_index(
|
||||
[TS_FIELD_NAME, SID_FIELD_NAME],
|
||||
).unstack()
|
||||
|
||||
# Allocate the whole output dataframe at once instead of
|
||||
# reindexing.
|
||||
sparse_output = pd.DataFrame(
|
||||
dense_output = pd.DataFrame(
|
||||
columns=pd.MultiIndex.from_product(
|
||||
(dense_output.columns.levels[0], assets),
|
||||
(sparse_output.columns.levels[0], assets),
|
||||
names=(
|
||||
dense_output.columns.levels[0].name,
|
||||
sparse_output.columns.levels[0].name,
|
||||
SID_FIELD_NAME,
|
||||
),
|
||||
),
|
||||
@@ -548,12 +640,14 @@ class BlazeLoader(dict):
|
||||
)
|
||||
|
||||
# In place update the output based on the base.
|
||||
sparse_output.update(dense_output)
|
||||
|
||||
dense_output.update(sparse_output)
|
||||
adjustments_from_deltas = adjustments_from_deltas_with_sids
|
||||
column_view = identity
|
||||
else:
|
||||
# We use the column view to make an array per asset.
|
||||
sparse_output = dense_output.reindex(dates)
|
||||
dense_output = sparse_output.reindex(dates)
|
||||
sparse_deltas = materialized_deltas.set_index(TS_FIELD_NAME)
|
||||
adjustments_from_deltas = adjustments_from_deltas_no_sids
|
||||
|
||||
def column_view(arr, _shape=(len(dates), len(assets))):
|
||||
"""Return a virtual matrix where we make a view that
|
||||
@@ -578,17 +672,19 @@ class BlazeLoader(dict):
|
||||
sparse_output = sparse_output.ffill()
|
||||
|
||||
for column_idx, column in enumerate(columns):
|
||||
column_name = column.name
|
||||
yield adjusted_array(
|
||||
column_view(
|
||||
sparse_output[column.name].values.astype(column.dtype),
|
||||
dense_output[column_name].values.astype(column.dtype),
|
||||
),
|
||||
mask,
|
||||
adjustments_from_deltas(
|
||||
dates,
|
||||
sparse_dates,
|
||||
sparse_output.index,
|
||||
column_idx,
|
||||
column_name,
|
||||
assets,
|
||||
materialized_deltas,
|
||||
sparse_deltas,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user