From 3babc380385564e0b65930e69d6b971852bf1a3e Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Tue, 13 Sep 2016 23:28:25 -0400
Subject: [PATCH] PERF: Release unneeded pipeline terms.

Refcount pipeline terms during execution and release terms once they're
no longer needed.

This dramatically reduces memory usage on large pipelines.
---
 zipline/pipeline/engine.py | 10 ++++++++++
 zipline/pipeline/graph.py  | 13 +++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/zipline/pipeline/engine.py b/zipline/pipeline/engine.py
index 2fc54c5a..93d296ed 100644
--- a/zipline/pipeline/engine.py
+++ b/zipline/pipeline/engine.py
@@ -344,6 +344,8 @@ class SimplePipelineEngine(object):
         loader_group_key = juxt(get_loader, getitem(graph.extra_rows))
         loader_groups = groupby(loader_group_key, graph.loadable_terms)
 
+        refcounts = graph.initial_refcounts()
+
         for term in graph.ordered():
             # `term` may have been supplied in `initial_workspace`, and in the
             # future we may pre-compute loadable terms coming from the same
@@ -380,6 +382,14 @@ class SimplePipelineEngine(object):
                 else:
                     assert workspace[term].shape == (mask.shape[0], 1)
 
+                # Decref any term we depended on.
+                for (parent, _) in graph.in_edges(term):
+                    refcounts[parent] -= 1
+                    # No one else depends on this term. Remove it from the
+                    # workspace to conserve memory.
+                    if refcounts[parent] == 0:
+                        del workspace[parent]
+
         out = {}
         graph_extra_rows = graph.extra_rows
         for name, term in iteritems(graph.outputs):
diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py
index 9ff09d45..9078c7ea 100644
--- a/zipline/pipeline/graph.py
+++ b/zipline/pipeline/graph.py
@@ -119,6 +119,19 @@ class TermGraph(DiGraph):
     def _repr_png_(self):
         return self.png.data
 
+    def initial_refcounts(self):
+        """
+        Calculate initial refcounts for execution of this graph.
+
+        Each node starts with a refcount equal to its outdegree, and output
+        nodes get one extra reference to ensure that they're still in the graph
+        at the end of execution.
+        """
+        refcounts = self.out_degree()
+        for t in self.outputs.values():
+            refcounts[t] += 1
+        return refcounts
+
 
 class ExecutionPlan(TermGraph):
     """