From f7e9281b14a9f63f5067e5f155dccdda27a1992a Mon Sep 17 00:00:00 2001
From: Scott Sanderson <ssanderson@quantopian.com>
Date: Tue, 10 May 2016 12:40:36 -0400
Subject: [PATCH] BUG: Fix groupby with string columns.

The previous algorithm assumed that the group labels were integers. It
produced nonsense with LabelArrays (though sadly didn't crash because
numpy promotes None and void to object).
---
 tests/pipeline/test_factor.py      | 13 +++++++++++++
 zipline/lib/labelarray.py          |  4 ++++
 zipline/pipeline/factors/factor.py | 26 +++++++++++++++++++-------
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/tests/pipeline/test_factor.py b/tests/pipeline/test_factor.py
index c5873c1f..4eee9f2e 100644
--- a/tests/pipeline/test_factor.py
+++ b/tests/pipeline/test_factor.py
@@ -23,6 +23,7 @@ from numpy import (
 from numpy.random import randn, seed
 
 from zipline.errors import UnknownRankMethod
+from zipline.lib.labelarray import LabelArray
 from zipline.lib.rank import masked_rankdata_2d
 from zipline.lib.normalize import naive_grouped_rowwise_apply as grouped_apply
 from zipline.pipeline import Classifier, Factor, Filter, TermGraph
@@ -38,6 +39,7 @@ from zipline.testing import (
 )
 from zipline.utils.functional import dzip_exact
 from zipline.utils.numpy_utils import (
+    categorical_dtype,
     datetime64ns_dtype,
     float64_dtype,
     int64_dtype,
@@ -442,6 +444,7 @@ class FactorTestCase(BasePipelineTestCase):
         f = self.f
         m = Mask()
         c = C()
+        str_c = C(dtype=categorical_dtype)
 
         factor_data = array(
             [[1.0, 2.0, 3.0, 4.0],
@@ -463,12 +466,18 @@ class FactorTestCase(BasePipelineTestCase):
              [1, 1, 2, 2]],
             dtype=int64_dtype,
         )
+        string_classifier_data = LabelArray(
+            classifier_data.astype(str).astype(object),
+            missing_value=None,
+        )
 
         terms = {
             'vanilla': f.demean(),
             'masked': f.demean(mask=m),
             'grouped': f.demean(groupby=c),
+            'grouped_str': f.demean(groupby=str_c),
             'grouped_masked': f.demean(mask=m, groupby=c),
+            'grouped_masked_str': f.demean(mask=m, groupby=str_c),
         }
         expected = {
             'vanilla': array(
@@ -496,6 +505,9 @@ class FactorTestCase(BasePipelineTestCase):
                  [-0.500,  0.500,  0.000,    nan]]
             )
         }
+        # Changing the classifier dtype shouldn't affect anything.
+        expected['grouped_str'] = expected['grouped']
+        expected['grouped_masked_str'] = expected['grouped_masked']
 
         graph = TermGraph(terms)
         results = self.run_graph(
@@ -503,6 +515,7 @@ class FactorTestCase(BasePipelineTestCase):
             initial_workspace={
                 f: factor_data,
                 c: classifier_data,
+                str_c: string_classifier_data,
                 m: filter_data,
             },
             mask=self.build_mask(self.ones_mask(shape=factor_data.shape)),
diff --git a/zipline/lib/labelarray.py b/zipline/lib/labelarray.py
index b707f0cd..a3ef5027 100644
--- a/zipline/lib/labelarray.py
+++ b/zipline/lib/labelarray.py
@@ -211,6 +211,10 @@ class LabelArray(ndarray):
         # This is a property because it should be immutable.
         return self._missing_value
 
+    @property
+    def missing_value_code(self):
+        return self.reverse_categories[self.missing_value]
+
     def has_label(self, value):
         return value in self.reverse_categories
 
diff --git a/zipline/pipeline/factors/factor.py b/zipline/pipeline/factors/factor.py
index a76002a9..6a0b0b08 100644
--- a/zipline/pipeline/factors/factor.py
+++ b/zipline/pipeline/factors/factor.py
@@ -47,6 +47,7 @@ from zipline.utils.input_validation import expect_types
 from zipline.utils.math_utils import nanmean, nanstd
 from zipline.utils.numpy_utils import (
     bool_dtype,
+    categorical_dtype,
     coerce_to_dtype,
     datetime64ns_dtype,
     float64_dtype,
@@ -939,15 +940,26 @@ class GroupedRowTransform(Factor):
 
     def _compute(self, arrays, dates, assets, mask):
         data = arrays[0]
-        null_group_value = self.inputs[1].missing_value
-        group_labels = where(
-            mask,
-            arrays[1],
-            null_group_value,
-        )
+        groupby_expr = self.inputs[1]
+        if groupby_expr.dtype == int64_dtype:
+            group_labels = arrays[1]
+            null_label = self.inputs[1].missing_value
+        elif groupby_expr.dtype == categorical_dtype:
+            # Coerce our LabelArray into an isomorphic array of ints.  This is
+            # necessary because np.where doesn't know about LabelArrays or the
+            # void dtype.
+            group_labels = arrays[1].as_int_array()
+            null_label = arrays[1].missing_value_code
+        else:
+            raise TypeError(
+                "Unexpected groupby dtype: %s." % groupby_expr.dtype
+            )
+
+        # Make a copy with the null code written to masked locations.
+        group_labels = where(mask, group_labels, null_label)
 
         return where(
-            group_labels != null_group_value,
+            group_labels != null_label,
             naive_grouped_rowwise_apply(
                 data=data,
                 group_labels=group_labels,