PERF: Batch load atomic terms by dataset

Added CompositeTerm and now we dispatch more generally on atomic

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):
 

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

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,

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,

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.
 

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.
     """

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.
     """

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)
 

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()"

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()