Repository: great-expectations/great_expectations

Issue:
When calling groupby on subclassed PandasDataset, groups revert to PandasDataset
It's possible I'm doing something wrong when subclassing, but: 

I'm subclassing PandasDataset to create custom expectations. When calling the groupby method on these custom datasets, the groups revert to the base PandasDataset type (and lose the custom expectation methods). 

```
import great_expectations as ge
import pandas as pd
from great_expectations.dataset import Dataset, PandasDataset

class SubclassPandasDataset(PandasDataset, pd.DataFrame):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

df = pd.DataFrame([{'a': 1}, {'a': 2}, {'a':1}])
sdf = ge._convert_to_dataset_class(df, SubclassPandasDataset)
print(type(sdf))

for ix, g in sdf.groupby('a'):
    print(type(g))
```
output:

<class '__main__.SubclassPandasDataset'>
<class 'great_expectations.dataset.pandas_dataset.PandasDataset'>
<class 'great_expectations.dataset.pandas_dataset.PandasDataset'>

Discussion:

Patch:
diff --git a/great_expectations/dataset/pandas_dataset.py b/great_expectations/dataset/pandas_dataset.py
index 5b3b5447e..c53c23a11 100644
--- a/great_expectations/dataset/pandas_dataset.py
+++ b/great_expectations/dataset/pandas_dataset.py
@@ -250,7 +250,7 @@ class PandasDataset(MetaPandasDataset, pd.DataFrame):
 
     @property
     def _constructor(self):
-        return PandasDataset
+        return self.__class__
 
 # Do we need to define _constructor_sliced and/or _constructor_expanddim? See http://pandas.pydata.org/pandas-docs/stable/internals.html#subclassing-pandas-data-structures
 


Test patch:
diff --git a/tests/test_pandas_dataset.py b/tests/test_pandas_dataset.py
index 250c1e6e5..1bd3cc4c4 100644
--- a/tests/test_pandas_dataset.py
+++ b/tests/test_pandas_dataset.py
@@ -1032,5 +1032,30 @@ class TestPandasDataset(unittest.TestCase):
         self.assertIsInstance(sub1, ge.dataset.PandasDataset)
         self.assertEqual(sub1.find_expectations(), exp1)
 
+    def test_subclass_pandas_subset_retains_subclass(self):
+        """A subclass of PandasDataset should still be that subclass after a Pandas subsetting operation"""
+        class CustomPandasDataset(ge.dataset.PandasDataset):
+
+            @ge.dataset.MetaPandasDataset.column_map_expectation
+            def expect_column_values_to_be_odd(self, column):
+                return column.map(lambda x: x % 2 )
+
+            @ge.dataset.MetaPandasDataset.column_map_expectation
+            def expectation_that_crashes_on_sixes(self, column):
+                return column.map(lambda x: (x-6)/0 != "duck")
+
+        df = CustomPandasDataset({
+            'all_odd': [1, 3, 5, 5, 5, 7, 9, 9, 9, 11],
+            'mostly_odd': [1, 3, 5, 7, 9, 2, 4, 1, 3, 5],
+            'all_even': [2, 4, 4, 6, 6, 6, 8, 8, 8, 8],
+            'odd_missing': [1, 3, 5, None, None, None, None, 1, 3, None],
+            'mixed_missing': [1, 3, 5, None, None, 2, 4, 1, 3, None],
+            'all_missing': [None, None, None, None, None, None, None, None, None, None]
+        })
+
+        df2 = df.sample(frac=0.5)
+        self.assertTrue(type(df2) == type(df))
+
+
 if __name__ == "__main__":
     unittest.main()