From e880fa3e3453719338a4115d3787a7f88f047dfa Mon Sep 17 00:00:00 2001 From: Richard Frank Date: Thu, 10 Sep 2015 15:34:24 -0400 Subject: [PATCH] PERF: Batch load atomic terms by dataset Added CompositeTerm and now we dispatch more generally on atomic --- tests/pipeline/test_term.py | 57 +++----- zipline/pipeline/classifier.py | 4 +- zipline/pipeline/data/dataset.py | 6 +- zipline/pipeline/engine.py | 33 +++-- zipline/pipeline/expression.py | 4 +- zipline/pipeline/factors/factor.py | 4 +- zipline/pipeline/filters/filter.py | 4 +- zipline/pipeline/graph.py | 43 +++--- zipline/pipeline/term.py | 225 +++++++++++++++++------------ zipline/pipeline/visualize.py | 2 +- 10 files changed, 213 insertions(+), 169 deletions(-) diff --git a/tests/pipeline/test_term.py b/tests/pipeline/test_term.py index 9fc9fac8..c10bf16b 100644 --- a/tests/pipeline/test_term.py +++ b/tests/pipeline/test_term.py @@ -82,11 +82,16 @@ def to_dict(l): class DependencyResolutionTestCase(TestCase): - def setup(self): - pass + def check_dependency_order(self, ordered_terms): + seen = set() - def teardown(self): - pass + for term in ordered_terms: + if not term.atomic: + for input_ in term.inputs: + self.assertIn(input_, seen) + self.assertIn(term.mask, seen) + + seen.add(term) def test_single_factor(self): """ @@ -97,12 +102,12 @@ class DependencyResolutionTestCase(TestCase): resolution_order = list(graph.ordered()) self.assertEqual(len(resolution_order), 4) - self.assertIs(resolution_order[0], AssetExists()) - self.assertEqual( - set([resolution_order[1], resolution_order[2]]), - set([SomeDataSet.foo, SomeDataSet.bar]), - ) - self.assertEqual(resolution_order[-1], SomeFactor()) + self.check_dependency_order(resolution_order) + self.assertIn(AssetExists(), resolution_order) + self.assertIn(SomeDataSet.foo, resolution_order) + self.assertIn(SomeDataSet.bar, resolution_order) + self.assertIn(SomeFactor(), resolution_order) + self.assertEqual(graph.node[SomeDataSet.foo]['extra_rows'], 4) self.assertEqual(graph.node[SomeDataSet.bar]['extra_rows'], 4) @@ -121,18 +126,14 @@ class DependencyResolutionTestCase(TestCase): # SomeFactor, its inputs, and AssetExists() self.assertEqual(len(resolution_order), 4) - - self.assertIs(resolution_order[0], AssetExists()) + self.check_dependency_order(resolution_order) + self.assertIn(AssetExists(), resolution_order) self.assertEqual(graph.extra_rows[AssetExists()], 4) - self.assertEqual( - set([resolution_order[1], resolution_order[2]]), - set([bar, buzz]), - ) - self.assertEqual( - resolution_order[-1], - SomeFactor([bar, buzz], window_length=5), - ) + self.assertIn(bar, resolution_order) + self.assertIn(buzz, resolution_order) + self.assertIn(SomeFactor([bar, buzz], window_length=5), + resolution_order) self.assertEqual(graph.extra_rows[bar], 4) self.assertEqual(graph.extra_rows[buzz], 4) @@ -148,20 +149,8 @@ class DependencyResolutionTestCase(TestCase): # bar should only appear once. self.assertEqual(len(resolution_order), 6) - indices = { - term: resolution_order.index(term) - for term in resolution_order - } - - self.assertEqual(indices[AssetExists()], 0) - - # Verify that f1's dependencies will be computed before f1. - self.assertLess(indices[SomeDataSet.foo], indices[f1]) - self.assertLess(indices[SomeDataSet.bar], indices[f1]) - - # Verify that f2's dependencies will be computed before f2. - self.assertLess(indices[SomeDataSet.bar], indices[f2]) - self.assertLess(indices[SomeDataSet.buzz], indices[f2]) + self.assertEqual(len(set(resolution_order)), 6) + self.check_dependency_order(resolution_order) def test_disallow_recursive_lookback(self): diff --git a/zipline/pipeline/classifier.py b/zipline/pipeline/classifier.py index aa1263ad..16219d58 100644 --- a/zipline/pipeline/classifier.py +++ b/zipline/pipeline/classifier.py @@ -2,8 +2,8 @@ classifier.py """ -from zipline.pipeline.term import Term +from zipline.pipeline.term import CompositeTerm -class Classifier(Term): +class Classifier(CompositeTerm): pass diff --git a/zipline/pipeline/data/dataset.py b/zipline/pipeline/data/dataset.py index 916c352d..211d03b0 100644 --- a/zipline/pipeline/data/dataset.py +++ b/zipline/pipeline/data/dataset.py @@ -6,7 +6,7 @@ from six import ( with_metaclass, ) -from zipline.pipeline.term import Term +from zipline.pipeline.term import AtomicTerm from zipline.pipeline.factors import Latest @@ -25,7 +25,7 @@ class Column(object): return BoundColumn(dtype=self.dtype, dataset=dataset, name=name) -class BoundColumn(Term): +class BoundColumn(AtomicTerm): """ A Column of data that's been concretely bound to a particular dataset. """ @@ -33,8 +33,6 @@ class BoundColumn(Term): def __new__(cls, dtype, dataset, name): return super(BoundColumn, cls).__new__( cls, - inputs=(), - window_length=0, domain=dataset.domain, dtype=dtype, dataset=dataset, diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py index d11746f0..d1d421d3 100644 --- a/zipline/pipeline/engine.py +++ b/zipline/pipeline/engine.py @@ -240,7 +240,15 @@ class SimplePipelineEngine(object): offset = graph.extra_rows[mask] - graph.extra_rows[term] return workspace[mask][offset:], dates[offset:] - def _inputs_for_term(self, term, workspace, graph): + def _mask_and_dates_for_atomic_terms(self, terms, workspace, graph, dates): + max_extra_rows = max(graph.extra_rows[term] for term in terms) + + mask = self._root_mask_term + offset = graph.extra_rows[mask] - max_extra_rows + return workspace[mask][offset:], dates[offset:] + + @staticmethod + def _inputs_for_term(term, workspace, graph): """ Compute inputs for the given term. @@ -273,6 +281,12 @@ class SimplePipelineEngine(object): out.append(input_data) return out + @staticmethod + def _atomic_dataset_terms(graph, match): + for term in graph.atomic_terms: + if term.dataset == match.dataset: + yield term + def compute_chunk(self, graph, dates, assets, initial_workspace): """ Compute the Pipeline terms in the graph for the requested start and end @@ -310,15 +324,11 @@ class SimplePipelineEngine(object): if term in workspace: continue - # Asset labels are always the same, but date labels vary by how - # many extra rows are needed. - mask, mask_dates = self._mask_and_dates_for_term( - term, workspace, graph, dates - ) if term.atomic: - # FUTURE OPTIMIZATION: Scan the resolution order for terms in - # the same dataset and load them here as well. - to_load = [term] + to_load = list(self._atomic_dataset_terms(graph, term)) + mask, mask_dates = self._mask_and_dates_for_atomic_terms( + to_load, workspace, graph, dates, + ) loaded = loader.load_adjusted_array( to_load, mask_dates, assets, mask, ) @@ -326,6 +336,11 @@ class SimplePipelineEngine(object): for loaded_term, adj_array in zip_longest(to_load, loaded): workspace[loaded_term] = adj_array else: + # Asset labels are always the same, but date labels vary by how + # many extra rows are needed. + mask, mask_dates = self._mask_and_dates_for_term( + term, workspace, graph, dates + ) workspace[term] = term._compute( self._inputs_for_term(term, workspace, graph), mask_dates, diff --git a/zipline/pipeline/expression.py b/zipline/pipeline/expression.py index 8ec5200e..b4d100b0 100644 --- a/zipline/pipeline/expression.py +++ b/zipline/pipeline/expression.py @@ -12,7 +12,7 @@ from numpy import ( find_common_type, ) -from zipline.pipeline.term import Term, NotSpecified +from zipline.pipeline.term import Term, NotSpecified, CompositeTerm _VARIABLE_NAME_RE = re.compile("^(x_)([0-9]+)$") @@ -154,7 +154,7 @@ def is_comparison(op): return op in COMPARISONS -class NumericalExpression(Term): +class NumericalExpression(CompositeTerm): """ Term binding to a numexpr expression. diff --git a/zipline/pipeline/factors/factor.py b/zipline/pipeline/factors/factor.py index bc05fcb2..5855777d 100644 --- a/zipline/pipeline/factors/factor.py +++ b/zipline/pipeline/factors/factor.py @@ -21,7 +21,7 @@ from zipline.pipeline.term import ( NotSpecified, RequiredWindowLengthMixin, SingleInputMixin, - Term, + CompositeTerm, ) from zipline.pipeline.expression import ( BadBinaryOperator, @@ -184,7 +184,7 @@ def function_application(func): return mathfunc -class Factor(Term): +class Factor(CompositeTerm): """ Pipeline API expression producing numerically-valued outputs. """ diff --git a/zipline/pipeline/filters/filter.py b/zipline/pipeline/filters/filter.py index 43575750..9869e910 100644 --- a/zipline/pipeline/filters/filter.py +++ b/zipline/pipeline/filters/filter.py @@ -15,7 +15,7 @@ from zipline.errors import ( ) from zipline.pipeline.term import ( SingleInputMixin, - Term, + CompositeTerm, ) from zipline.pipeline.expression import ( BadBinaryOperator, @@ -83,7 +83,7 @@ def binary_operator(op): return binary_operator -class Filter(Term): +class Filter(CompositeTerm): """ Pipeline API expression producing boolean-valued outputs. """ diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py index 428b5424..3a91a37c 100644 --- a/zipline/pipeline/graph.py +++ b/zipline/pipeline/graph.py @@ -104,11 +104,12 @@ class TermGraph(DiGraph): """ out = {} for term in self: - extra_input_rows = term.extra_input_rows - for input_ in term.inputs: - out[term, input_] = self.extra_rows[input_] - extra_input_rows - mask = term.mask - if term.mask is not None: + if not term.atomic: + extra_input_rows = term.extra_input_rows + for input_ in term.inputs: + out[term, input_] = (self.extra_rows[input_] + - extra_input_rows) + mask = term.mask out[term, mask] = self.extra_rows[mask] - extra_input_rows return out @@ -168,6 +169,10 @@ class TermGraph(DiGraph): """ return iter(self._ordered) + @lazyval + def atomic_terms(self): + return tuple(term for term in self if term.atomic) + def _add_to_graph(self, term, parents, extra_rows): """ Add `term` and all its inputs to the graph. @@ -187,21 +192,23 @@ class TermGraph(DiGraph): # Make sure we're going to compute at least `extra_rows` of `term`. self._ensure_extra_rows(term, extra_rows) - # Number of extra rows we need to compute for this term's dependencies. - dependency_extra_rows = extra_rows + term.extra_input_rows + if not term.atomic: + # Number of extra rows we need to compute for this term's + # dependencies. + dependency_extra_rows = extra_rows + term.extra_input_rows - # Recursively add dependencies. - for dependency in term.inputs: - self._add_to_graph( - dependency, - parents, - extra_rows=dependency_extra_rows, - ) - self.add_edge(dependency, term) + # Recursively add dependencies. + for dependency in term.inputs: + self._add_to_graph( + dependency, + parents, + extra_rows=dependency_extra_rows, + ) + self.add_edge(dependency, term) - # Add term's mask, which is really just a specially-enumerated input. - mask = term.mask - if mask is not None: + # Add term's mask, which is really just a specially-enumerated + # input. + mask = term.mask self._add_to_graph(mask, parents, extra_rows=dependency_extra_rows) self.add_edge(mask, term) diff --git a/zipline/pipeline/term.py b/zipline/pipeline/term.py index 56855253..1dc3974a 100644 --- a/zipline/pipeline/term.py +++ b/zipline/pipeline/term.py @@ -43,18 +43,12 @@ class Term(object): Base class for terms in a Pipeline API compute graph. """ # These are NotSpecified because a subclass is required to provide them. - inputs = NotSpecified - window_length = NotSpecified dtype = NotSpecified - mask = NotSpecified domain = NotSpecified _term_cache = WeakValueDictionary() def __new__(cls, - inputs=NotSpecified, - mask=NotSpecified, - window_length=NotSpecified, domain=NotSpecified, dtype=NotSpecified, *args, @@ -72,23 +66,6 @@ class Term(object): # Class-level attributes can be used to provide defaults for Term # subclasses. - if inputs is NotSpecified: - inputs = cls.inputs - # Having inputs = NotSpecified is an error, but we handle it later - # in self._validate rather than here. - if inputs is not NotSpecified: - # Allow users to specify lists as class-level defaults, but - # normalize to a tuple so that inputs is hashable. - inputs = tuple(inputs) - - if mask is NotSpecified: - mask = cls.mask - if mask is NotSpecified: - mask = AssetExists() - - if window_length is NotSpecified: - window_length = cls.window_length - if domain is NotSpecified: domain = cls.domain @@ -96,9 +73,6 @@ class Term(object): dtype = cls.dtype identity = cls.static_identity( - inputs=inputs, - mask=mask, - window_length=window_length, domain=domain, dtype=dtype, *args, **kwargs @@ -109,9 +83,6 @@ class Term(object): except KeyError: new_instance = cls._term_cache[identity] = \ super(Term, cls).__new__(cls)._init( - inputs=inputs, - mask=mask, - window_length=window_length, domain=domain, dtype=dtype, *args, **kwargs @@ -134,10 +105,7 @@ class Term(object): """ pass - def _init(self, inputs, mask, window_length, domain, dtype): - self.inputs = inputs - self.mask = mask - self.window_length = window_length + def _init(self, domain, dtype): self.domain = domain self.dtype = dtype @@ -145,7 +113,7 @@ class Term(object): return self @classmethod - def static_identity(cls, inputs, mask, window_length, domain, dtype): + def static_identity(cls, domain, dtype): """ Return the identity of the Term that would be constructed from the given arguments. @@ -157,80 +125,25 @@ class Term(object): This is a classmethod so that it can be called from Term.__new__ to determine whether to produce a new instance. """ - return (cls, inputs, mask, window_length, domain, dtype) + return (cls, domain, dtype) def _validate(self): """ Assert that this term is well-formed. This should be called exactly once, at the end of Term._init(). """ - if self.inputs is NotSpecified: - raise TermInputsNotSpecified(termname=type(self).__name__) - if self.window_length is NotSpecified: - raise WindowLengthNotSpecified(termname=type(self).__name__) if self.dtype is NotSpecified: raise DTypeNotSpecified(termname=type(self).__name__) - if self.mask is NotSpecified and not self.atomic: - # This isn't user error, this is a bug in our code. - raise AssertionError("{term} has no mask".format(term=self)) - if self.window_length: - for child in self.inputs: - if not child.atomic: - raise InputTermNotAtomic(parent=self, child=child) - - @lazyval + @property def atomic(self): """ Whether or not this term has dependencies. If term.atomic is truthy, it should have dataset and dtype attributes. """ - return len(self.inputs) == 0 - - @lazyval - def windowed(self): - """ - Whether or not this term represents a trailing window computation. - - If term.windowed is truthy, its compute_from_windows method will be - called with instances of AdjustedArray as inputs. - - If term.windowed is falsey, its compute_from_baseline will be called - with instances of np.ndarray as inputs. - """ - return ( - self.window_length is not NotSpecified - and self.window_length > 0 - ) - - @lazyval - def extra_input_rows(self): - """ - The number of extra rows needed for each of our inputs to compute this - term. - """ - return max(0, self.window_length - 1) - - def _compute(self, inputs, dates, assets, mask): - """ - Subclasses should implement this to perform actual computation. - - This is `_compute` rather than just `compute` because `compute` is - reserved for user-supplied functions in CustomFactor. - """ raise NotImplementedError() - def __repr__(self): - return ( - "{type}({inputs}, window_length={window_length})" - ).format( - type=type(self).__name__, - inputs=self.inputs, - window_length=self.window_length, - mask=self.mask, - ) - # TODO: Move mixins to a separate file? class SingleInputMixin(object): @@ -307,7 +220,128 @@ class CustomTermMixin(object): return out -class AssetExists(Term): +class AtomicTerm(Term): + + @property + def atomic(self): + return True + + @property + def dataset(self): + raise NotImplementedError() + + +class CompositeTerm(Term): + inputs = NotSpecified + window_length = NotSpecified + mask = NotSpecified + + def __new__(cls, inputs=NotSpecified, window_length=NotSpecified, + mask=NotSpecified, *args, **kwargs): + + if inputs is NotSpecified: + inputs = cls.inputs + # Having inputs = NotSpecified is an error, but we handle it later + # in self._validate rather than here. + if inputs is not NotSpecified: + # Allow users to specify lists as class-level defaults, but + # normalize to a tuple so that inputs is hashable. + inputs = tuple(inputs) + + if mask is NotSpecified: + mask = cls.mask + if mask is NotSpecified: + mask = AssetExists() + + if window_length is NotSpecified: + window_length = cls.window_length + + return super(CompositeTerm, cls).__new__(cls, inputs=inputs, mask=mask, + window_length=window_length, + *args, **kwargs) + + def _init(self, inputs, window_length, mask, *args, **kwargs): + self.inputs = inputs + self.window_length = window_length + self.mask = mask + return super(CompositeTerm, self)._init(*args, **kwargs) + + @classmethod + def static_identity(cls, inputs, window_length, mask, *args, **kwargs): + return ( + super(CompositeTerm, cls).static_identity(*args, **kwargs), + inputs, + window_length, + mask, + ) + + def _validate(self): + """ + Assert that this term is well-formed. This should be called exactly + once, at the end of Term._init(). + """ + if self.inputs is NotSpecified: + raise TermInputsNotSpecified(termname=type(self).__name__) + if self.window_length is NotSpecified: + raise WindowLengthNotSpecified(termname=type(self).__name__) + if self.mask is NotSpecified: + # This isn't user error, this is a bug in our code. + raise AssertionError("{term} has no mask".format(term=self)) + + if self.window_length: + for child in self.inputs: + if not child.atomic: + raise InputTermNotAtomic(parent=self, child=child) + + return super(CompositeTerm, self)._validate() + + @property + def atomic(self): + return False + + def _compute(self, inputs, dates, assets, mask): + """ + Subclasses should implement this to perform actual computation. + This is `_compute` rather than just `compute` because `compute` is + reserved for user-supplied functions in CustomFactor. + """ + raise NotImplementedError() + + @lazyval + def windowed(self): + """ + Whether or not this term represents a trailing window computation. + + If term.windowed is truthy, its compute_from_windows method will be + called with instances of AdjustedArray as inputs. + + If term.windowed is falsey, its compute_from_baseline will be called + with instances of np.ndarray as inputs. + """ + return ( + self.window_length is not NotSpecified + and self.window_length > 0 + ) + + @lazyval + def extra_input_rows(self): + """ + The number of extra rows needed for each of our inputs to compute this + term. + """ + return max(0, self.window_length - 1) + + def __repr__(self): + return ( + "{type}({inputs}, window_length={window_length})" + ).format( + type=type(self).__name__, + inputs=self.inputs, + window_length=self.window_length, + ) + + +class AssetExists(AtomicTerm): """ Pseudo-filter describing whether or not an asset existed on a given day. This is the default mask for all terms that haven't been passed a mask @@ -321,10 +355,8 @@ class AssetExists(Term): -------- zipline.assets.AssetFinder.lifetimes """ - inputs = () dtype = bool_ - window_length = 0 - mask = None + dataset = None def _compute(self, *args, **kwargs): # TODO: Consider moving the bulk of the logic from @@ -332,3 +364,6 @@ class AssetExists(Term): raise NotImplementedError( "Direct computation of AssetExists is not supported!" ) + + def __repr__(self): + return "AssetExists()" diff --git a/zipline/pipeline/visualize.py b/zipline/pipeline/visualize.py index 03297a27..44ffe530 100644 --- a/zipline/pipeline/visualize.py +++ b/zipline/pipeline/visualize.py @@ -92,7 +92,7 @@ def _render(g, out, format_, include_asset_exists=False): graph_attrs = {'rankdir': 'TB', 'splines': 'ortho'} cluster_attrs = {'style': 'filled', 'color': 'lightgoldenrod1'} - in_nodes = list(node for node in g if node.atomic) + in_nodes = g.atomic_terms out_nodes = list(g.outputs.values()) f = BytesIO()