Repository: unionai-oss/pandera

Issue:
add support for grouping column of interest by another column.
if the user needs to validate some column `x` conditioned on the values of column `y`, the user can specify multiple column names argument

```
DataFrameSchema({
    ("x", "y"): Columns([Int, String], Check(lambda x, y: x[y == "foo"] > 10))
})
```


Discussion:
What's your view on handling the `nullable`, `strict`, `allow_duplicates` when supplying a tuple of columns for checking?

Do you see something like this?:

```
DataFrameSchema({
    "a": Column(Float, nullable=True),
    ("b", "c", "d"): Columns(
        (
            (Int, nullable=True, strict=True),
            (Float, nullable=True, allow_duplicates=True),
            (String, nullable=True, strict=False),
        ),
        Check(lambda b, c, d: b*c == d > 10),
    "e": Column(Bool, nullable=True),
})
```
this is an interesting problem... I'm wondering if it would make more sense to extend the `Check` class so specifying a tuple of column names that the column `depends_on` that modifies the signature of the check function, like so:

```
DataFrameSchema({
    "a": Column(Float, nullable=True, checks=[
        Check(lambda a: a > 0),
        Check(lambda a, b, c: a[(b == "foo") & (c == "bar")] > 10, depends_on=("b", "c"))
    ]),
    "b": Column(String, nullable=True),
    "c": Column(String, nullable=True),
})
```

This way the user only needs to specify column properties once, and modify the function signature of a `Check`.

I'm concerned that my original proposal will make schemas more convoluted to read.

One problem with this above proposal that I haven't thought about too much is whether the order of execution of `Check`s matter. Before each check was atomic in the sense that it only depends on one column... but now with `depends_on`, I wonder if under the hood we'd need to run checks on "b" and "c" before executing the check on `a[(b == "foo") & (c == "bar")] > 10`.

Can probably check at schema initialization whether there are circular dependencies, and error out immediately if there are any, e.g.

```
DataFrameSchema({
    "a": Column(Float, nullable=True, checks=[
        Check(lambda a, b: a[(b == "foo")] > 10, depends_on=("b"))
    ]),
    "b": Column(String, nullable=True, checks=[
        Check(lambda b, a: b[(a > 5)] == "bar", depends_on=("a"))
    ]),
})

# SchemaError: circular dependency in columns "a" and "b"
```
Specifying `depends_on` before those columns are defined seems a bit odd on first read (it makes me think of the classic IDE error "variable referenced before assignment").

Maybe the use of `b` in the `check` of `a` could be handled implicitly, i.e. the fact that `b` exists is enough for building a `depends_on` dependency graph behind the scenes?

That still suffers from the readability issue. 


One option could be having an entirely separate `checks` list separate from the `columns` for readability, something like this:
```
DataFrameSchema(
    columns={
        "a": Column(Float, nullable=True),
        "b": Column(String, nullable=True),
        "c": Column(String, nullable=True)
    },
    checks=[
        Check(lambda a: a > 0),
        Check(lambda a, b, c: a[(b == "foo") & (c == "bar")] > 10)
    ]
)
```
the separate `checks` idea is challenging because I think checks should be bound to the data provided by the `columns` keys. There really isn't a straightforward way to let a `Check` know what the function signature is besides associating a `Check` with column names in an explicit way.

Really, the intent behind checks that involve multiple columns is to enable things like hypothesis testing https://github.com/cosmicBboy/pandera/issues/17, which I think is the strongest use case for a solution where additional column dependencies are specified at the `Check` level, one possible proposal being the `depends_on` kwarg.

Basically the backend implementation of the one_sided_t_test described in #17 would be something like:

```
from scipy import stats

def two_sample_one_sided_ttest(group1, group2, relationship="gt", alpha=0.05, **kwargs):
    ttest = stats.ttest_ind(group1, group2, **kwargs)
    is_significant = ttest.pvalue / 2 < alpha
    # `relationship` refers to the relationship of group1 w.r.t. group2
    if relationship == "gt":
        return ttest.statistic > 0 and is_significant
    elif relationship == "lt":
        return ttest.statistic < 0 and is_significant
    else:
        raise ValueError("relationship %s not recognized" % relationship)

DataFrameSchema({
    "weight": Column(
        Float,
        checks=[
            Check(
                lambda w, s: two_sample_one_sided_ttest(w[s == 'men'], w[s == 'women']),
                depends_on=("sex")
            )
        ]
    ),
    "sex": Column(String, Check(lambda s: s.isin(["men", "women"]).all()))
})
```

> Specifying depends_on before those columns are defined seems a bit odd on first read (it makes me think of the classic IDE error "variable referenced before assignment").

I think this makes sense for imperative program models, but the goal with pandera is that the schema definition is (somewhat) declarative, such that the properties of a dataframe are asserted up-front, and the details of how validation happens is abstracted away for the user... I say somewhat because the validators are still fairly low-level, e.g. asserting the values of a column are positive via `lambda s: s > 0`.

In this way we're not really defining columns in an imperative way, we're declaring that a `Check` associated with the column `weight` depends on this other column `sex`. What are the properties of the `sex` column? Well, the user would be expected to define these properties with `{"sex": Column(...)}`

Not sure if this thought process makes sense... we can continue mulling this over, but two of the things I'd like to shoot for are:
- keeping the API as clear and concise as possible
- discouraging an API that leads to over-specification of column properties
I like your thinking - column wise checks follow the logic of "this column depends on these others", and fits the primary use case. The syntax gets clunky if a user writes too many overlapping checks, and that can easily be handled outside of Pandera if required.

I looked at implementing it and one observation: if the user made a mistake in specifying the `lambda` function or `depends_on`, it would be difficult to diagnose/debug:, e.g.:
```
DataFrameSchema({
    "weight": Column(
        Float,
        checks=[
            Check(
                lambda w, s, h, a: w+s+a+h,
                depends_on=("sex","age","height")
            )
        ]
    )
})
```
In the above function, the positional order of `depends_on` sets `h=age`, `height=a` which is probably not what the user intended. Not sure if or how to make this better.

One possible implementation leaves open the possibility of changing syntax in future relatively flexibly: passing a dataframe through `SeriesSchemaBase` and into `Check`.

So the call of `Check` goes from:
```
    def __call__(self, parent_schema, series, index):
```
and becomes this:
```
    def __call__(self, parent_schema, dataframe, index):
```

To make this work the whole dataframe could be passed into `SeriesSchemaBase`, and the existing code be adapted to only reference the named column, i.e. using `df[self._name]` currently in `column`.

Then modifying the way the checks are called in  `SeriesSchemaBase` from:
```
        for i, check in enumerate(self._checks):
            check_results.append(check(self, series, i))
```
into _something like_:
```
        for i, check in enumerate(self._checks):
            check_results.append(check(self, df[this_column,depend_on], i))
```

Effectively by passing a Dataframe into `Check`, either the subsetted dataframe, or the full dataframe - instead of multiple series leaves open the possibility of extending DataFrameSchema to also have checks, or changing the syntax..
> In the above function, the positional order of depends_on sets h=age, height=a which is probably not what the user intended. Not sure if or how to make this better.

good point, I do like the direction of making the API for specifying checks clearer, and I'd like to maintain the separation of the primary column (the one that is specified in the `columns` key) and the `depends_on` columns.

For example:

```
DataFrameSchema({
    "weight": Column(
        Float,
        checks=[
            Check(
                lambda w, d: w + d["sex"] + d["age"] + d["height"],
                depends_on=("sex","age","height")
            )
        ]
    )
})
```

basically specifying `depends_on` modifies the expected `Check` function signature to
```
(series, dataframe) -> bool|pd.Series[bool]
```

where `dataframe` contains the columns specified in `depends_on`.

I like this because `weight` is still a first-class argument to the check function and `sex`, `age`, `height` are contained within the second argument.

It's then the user's responsibility to use the contents in the dependencies dataframe to make an assertion about the primary column.

On a side-note, more complex `check` functions (with many dependencies and complex logic) should probably be defined before hand with the `def myfunc(x, df)` pattern.

Resolving patch:
diff --git a/pandera/__init__.py b/pandera/__init__.py
--- a/pandera/__init__.py
+++ b/pandera/__init__.py
@@ -1,6 +1,7 @@
 from .pandera import DataFrameSchema, Column, Index, PandasDtype, \
-    SeriesSchema, SchemaError, Check, check_input, check_output, \
-    Bool, DateTime, Category, Float, Int, Object, String, Timedelta
+    SeriesSchema, SchemaError, SchemaInitError, Check, check_input, \
+    check_output, Bool, DateTime, Category, Float, Int, Object, String, \
+    Timedelta
 
 
 __version__ = "0.1.2"
diff --git a/pandera/pandera.py b/pandera/pandera.py
--- a/pandera/pandera.py
+++ b/pandera/pandera.py
@@ -8,6 +8,11 @@
 
 from collections import OrderedDict
 from enum import Enum
+from functools import partial
+
+
+class SchemaInitError(Exception):
+    pass
 
 
 class SchemaError(Exception):
@@ -39,7 +44,8 @@ class PandasDtype(Enum):
 
 class Check(object):
 
-    def __init__(self, fn, element_wise=True, error=None, n_failure_cases=10):
+    def __init__(self, fn, element_wise=False, error=None, n_failure_cases=10,
+                 groupby=None, groups=None):
         """Check object applies function element-wise or series-wise
 
         Parameters
@@ -59,29 +65,66 @@ def __init__(self, fn, element_wise=True, error=None, n_failure_cases=10):
         n_failure_cases : int|None
             report the top n failure cases. If None, then report all failure
             cases.
+        groupby : str|list[str]|callable|None
+            Only applies to Column Checks. If a string or list of strings is
+            provided, then these columns are used to group the Column Series by
+            `groupby`. If a callable is passed, the expected signature is
+            DataFrame -> DataFrameGroupby. The function has access to the
+            entire dataframe, but the Column.name is selected from
+            this DataFrameGroupby object so that a SeriesGroupBy object is
+            passed into `fn`.
+
+            Specifying this argument changes the `fn` signature to:
+
+            dict[str|tuple[str], Series] -> bool|pd.Series[bool]
+
+            Where specific groups can be obtained from the input dict.
+
+        groups : str|list[str]|None
+            The dict input to the `fn` callable will be constrained to the
+            groups specified by `groups`.
         """
+        if element_wise and groupby is not None:
+            raise SchemaInitError("Cannot use groupby when element_wise=True.")
         self.fn = fn
         self.element_wise = element_wise
         self.error = error
         self.n_failure_cases = n_failure_cases
 
+        if groupby is None and groups is not None:
+            raise ValueError(
+                "`groupby` argument needs to be provided when `groups` "
+                "argument is defined")
+
+        if isinstance(groupby, str):
+            groupby = [groupby]
+        self.groupby = groupby
+        if isinstance(groups, str):
+            groups = [groups]
+        self.groups = groups
+
+    @staticmethod
+    def _check_groupby(column_name, df):
+        return df.groupby(column_name)
+
     @property
     def error_message(self):
         if self.error:
             return "%s: %s" % (self.fn.__name__, self.error)
         return "%s" % self.fn.__name__
 
-    def vectorized_error_message(self, parent_schema, index, failure_cases):
+    def vectorized_error_message(
+            self, parent_schema, check_index, failure_cases):
         return (
                 "%s failed element-wise validator %d:\n"
                 "%s\nfailure cases:\n%s" %
-                (parent_schema, index,
+                (parent_schema, check_index,
                  self.error_message,
                  self._format_failure_cases(failure_cases)))
 
-    def generic_error_message(self, parent_schema, index):
+    def generic_error_message(self, parent_schema, check_index):
         return "%s failed series validator %d: %s" % \
-               (parent_schema, index, self.error_message)
+               (parent_schema, check_index, self.error_message)
 
     def _format_failure_cases(self, failure_cases):
         failure_cases = (
@@ -98,43 +141,79 @@ def _format_failure_cases(self, failure_cases):
         else:
             return failure_cases.head(self.n_failure_cases)
 
-    def __call__(self, parent_schema, series, index):
+    def prepare_input(self, series, dataframe):
+        # used by Column.__call__ method to prepare series/SeriesGroupBy input
+        if dataframe is None or self.groupby is None:
+            return series
+        elif isinstance(self.groupby, list):
+            groupby_obj = (
+                pd.concat([series, dataframe[self.groupby]], axis=1)
+                .groupby(self.groupby)[series.name]
+            )
+        elif callable(self.groupby):
+            groupby_obj = self.groupby(
+                pd.concat([series, dataframe], axis=1))[series.name]
+        else:
+            raise TypeError("Type %s not recognized for `groupby` argument.")
+
+        if self.groups is None:
+            return {g: s for g, s in groupby_obj}
+
+        group_keys = set(g for g, _ in groupby_obj)
+        invalid_groups = [g for g in self.groups if g not in group_keys]
+        if invalid_groups:
+            raise KeyError(
+                "groups %s provided in `groups` argument not a valid group "
+                "key. Valid group keys: %s" % (invalid_groups, group_keys))
+        fn_input = {}
+        for group, series in groupby_obj:
+            if group in self.groups:
+                fn_input[group] = series
+        return fn_input
+
+    def _vectorized_series_check(self, parent_schema, check_index, check_obj):
+        val_result = self.fn(check_obj)
+        if isinstance(val_result, pd.Series):
+            if not val_result.dtype == PandasDtype.Bool.value:
+                raise TypeError(
+                    "validator %d: %s must return bool or Series of type "
+                    "bool, found %s" %
+                    (check_index, self.fn.__name__, val_result.dtype))
+            if val_result.all():
+                return True
+            elif isinstance(check_obj, dict) or \
+                    check_obj.shape[0] != val_result.shape[0] or \
+                    (check_obj.index != val_result.index).all():
+                raise SchemaError(
+                    self.generic_error_message(parent_schema, check_index))
+            else:
+                raise SchemaError(self.vectorized_error_message(
+                    parent_schema, check_index, check_obj[~val_result]))
+        else:
+            if val_result:
+                return True
+            raise SchemaError(
+                self.generic_error_message(parent_schema, check_index))
+
+    def __call__(self, parent_schema, check_index, check_obj):
+        _vcheck = partial(
+            self._vectorized_series_check, parent_schema, check_index)
         if self.element_wise:
-            val_result = series.map(self.fn)
+            val_result = check_obj.map(self.fn)
             if val_result.all():
                 return True
             raise SchemaError(self.vectorized_error_message(
-                parent_schema, index, series[~val_result]))
-        else:
-            # series-wise validator can return either a boolean or a
-            # pd.Series of booleans.
-            val_result = self.fn(series)
-            if isinstance(val_result, pd.Series):
-                if not val_result.dtype == PandasDtype.Bool.value:
-                    raise TypeError(
-                        "validator %d: %s must return bool or Series of type "
-                        "bool, found %s" %
-                        (index, self.fn.__name__, val_result.dtype))
-                if val_result.all():
-                    return True
-                try:
-                    raise SchemaError(self.vectorized_error_message(
-                        parent_schema, index, series[~val_result]))
-                except pd.core.indexing.IndexingError:
-                    raise SchemaError(
-                        self.generic_error_message(parent_schema, index))
-            else:
-                if val_result:
-                    return True
-                else:
-                    raise SchemaError(
-                        self.generic_error_message(parent_schema, index))
+                parent_schema, check_index, check_obj[~val_result]))
+        elif isinstance(check_obj, (pd.Series, dict)):
+            return _vcheck(check_obj)
 
 
 class DataFrameSchema(object):
     """A light-weight pandas DataFrame validator."""
 
-    def __init__(self, columns, index=None, transformer=None, coerce=False, strict=False):
+    def __init__(
+            self, columns, index=None, transformer=None, coerce=False,
+            strict=False):
         """Initialize pandas dataframe schema.
 
         Parameters
@@ -151,14 +230,28 @@ def __init__(self, columns, index=None, transformer=None, coerce=False, strict=F
         coerce : bool
             whether or not to coerce all of the columns on validation.
         strict : bool
-            whether or not to accept columns in the dataframe that aren't in the
-            DataFrame Schema.
+            whether or not to accept columns in the dataframe that aren't in
+            the DataFrame Schema.
         """
         self.index = index
         self.columns = columns
         self.transformer = transformer
         self.coerce = coerce
         self.strict = strict
+        self._validate_schema()
+
+    def _validate_schema(self):
+        for column_name, column in self.columns.items():
+            for check in column._checks:
+                if check.groupby is None or callable(check.groupby):
+                    continue
+                nonexistent_dep_columns = [
+                    c for c in check.groupby if c not in self.columns]
+                if nonexistent_dep_columns:
+                    raise SchemaInitError(
+                        "groupby argument %s in Check for Column %s not "
+                        "specified in the DataFrameSchema." %
+                        (nonexistent_dep_columns, column_name))
 
     def validate(self, dataframe):
         # Check if all columns in the dataframe have a corresponding column in
@@ -209,6 +302,8 @@ def __init__(self, pandas_dtype, checks=None, nullable=False,
         checks : Check|list[Check]
         nullable : bool
             Whether or not column can contain null values.
+        allow_duplicates : bool
+            Whether or not to allow duplicated values in the series.
         """
         self._pandas_dtype = pandas_dtype
         self._nullable = nullable
@@ -219,12 +314,20 @@ def __init__(self, pandas_dtype, checks=None, nullable=False,
             checks = [checks]
         self._checks = checks
 
-    def __call__(self, series):
+        for check in self._checks:
+            if check.groupby is not None and not isinstance(self, Column):
+                raise SchemaInitError(
+                    "Can only use `groupby` with a pandera.Column, found %s" %
+                    type(self))
+
+    def __call__(self, series, dataframe=None):
         """Validate a series."""
         expected_dtype = _dtype = self._pandas_dtype if \
             isinstance(self._pandas_dtype, str) else self._pandas_dtype.value
         if self._nullable:
             series = series.dropna()
+            if dataframe is not None:
+                dataframe = dataframe.loc[series.index]
             if _dtype in ["int_", "int8", "int16", "int32", "int64", "uint8",
                           "uint16", "uint32", "uint64"]:
                 _dtype = Float.value
@@ -238,8 +341,7 @@ def __call__(self, series):
         else:
             nulls = series.isnull()
             if nulls.sum() > 0:
-                type_val_result = series.dtype == _dtype
-                if not type_val_result:
+                if series.dtype != _dtype:
                     raise SchemaError(
                         "expected series '%s' to have type %s, got %s and "
                         "non-nullable series contains null values: %s" %
@@ -260,16 +362,14 @@ def __call__(self, series):
                     (series.name,
                      series[duplicates].head(N_FAILURE_CASES).to_dict()))
 
-        type_val_result = series.dtype == _dtype
-        if not type_val_result:
+        if series.dtype != _dtype:
             raise SchemaError(
                 "expected series '%s' to have type %s, got %s" %
                 (series.name, expected_dtype, series.dtype))
 
-        check_results = []
-        for i, check in enumerate(self._checks):
-            check_results.append(check(self, series, i))
-        return all([type_val_result] + check_results)
+        return all(
+            check(self, check_index, check.prepare_input(series, dataframe))
+            for check_index, check in enumerate(self._checks))
 
 
 class SeriesSchema(SeriesSchemaBase):
@@ -366,7 +466,8 @@ def __call__(self, df):
         if self._name is None:
             raise RuntimeError(
                 "need to `set_name` of column before calling it.")
-        return super(Column, self).__call__(df[self._name])
+        return super(Column, self).__call__(
+            df[self._name], df.drop(self._name, axis=1))
 
     def __repr__(self):
         if isinstance(self._pandas_dtype, PandasDtype):


Test patch:
diff --git a/tests/test_pandera.py b/tests/test_pandera.py
--- a/tests/test_pandera.py
+++ b/tests/test_pandera.py
@@ -5,13 +5,13 @@
 import pytest
 
 from pandera import Column, DataFrameSchema, Index, PandasDtype, \
-    SeriesSchema, Check, Float, Int, DateTime, String, check_input, \
-    check_output, SchemaError
+    SeriesSchema, Check, Bool, Float, Int, DateTime, String, check_input, \
+    check_output, SchemaError, SchemaInitError
 
 
 def test_column():
     schema = DataFrameSchema({
-        "a": Column(PandasDtype.Int, Check(lambda x: x > 0))
+        "a": Column(PandasDtype.Int, Check(lambda x: x > 0, element_wise=True))
     })
     data = pd.DataFrame({"a": [1, 2, 3]})
     assert isinstance(schema.validate(data), pd.DataFrame)
@@ -19,7 +19,7 @@ def test_column():
 
 def test_series_schema():
     schema = SeriesSchema(
-        PandasDtype.Int, Check(lambda x: 0 <= x <= 100))
+        PandasDtype.Int, Check(lambda x: 0 <= x <= 100, element_wise=True))
     validated_series = schema.validate(pd.Series([0, 30, 50, 100]))
     assert isinstance(validated_series, pd.Series)
 
@@ -35,7 +35,7 @@ def test_series_schema():
     non_duplicate_schema = SeriesSchema(
         PandasDtype.Int, allow_duplicates=False)
     with pytest.raises(SchemaError):
-        non_duplicate_schema.validate(pd.Series([0,1,2,3,4,1]))
+        non_duplicate_schema.validate(pd.Series([0, 1, 2, 3, 4, 1]))
 
 
 def test_vectorized_checks():
@@ -53,8 +53,8 @@ def test_vectorized_checks():
 def test_series_schema_multiple_validators():
     schema = SeriesSchema(
         PandasDtype.Int, [
-            Check(lambda x: 0 <= x <= 50),
-            Check(lambda s: (s == 21).any(), element_wise=False)])
+            Check(lambda x: 0 <= x <= 50, element_wise=True),
+            Check(lambda s: (s == 21).any())])
     validated_series = schema.validate(pd.Series([1, 5, 21, 50]))
     assert isinstance(validated_series, pd.Series)
 
@@ -66,24 +66,24 @@ def test_series_schema_multiple_validators():
 def test_dataframe_schema():
     schema = DataFrameSchema(
         {
-            "a": Column(PandasDtype.Int, Check(lambda x: x > 0)),
-            "b": Column(PandasDtype.Float, Check(lambda x: 0 <= x <= 10)),
+            "a": Column(PandasDtype.Int,
+                        Check(lambda x: x > 0, element_wise=True)),
+            "b": Column(PandasDtype.Float,
+                        Check(lambda x: 0 <= x <= 10, element_wise=True)),
             "c": Column(PandasDtype.String,
-                        Check(lambda x: set(x) == {"x", "y", "z"},
-                              element_wise=False)),
+                        Check(lambda x: set(x) == {"x", "y", "z"})),
             "d": Column(PandasDtype.Bool,
-                        Check(lambda x: x.mean() > 0.5,
-                              element_wise=False)),
+                        Check(lambda x: x.mean() > 0.5)),
             "e": Column(PandasDtype.Category,
-                        Check(lambda x: set(x) == {"c1", "c2", "c3"},
-                              element_wise=False)),
+                        Check(lambda x: set(x) == {"c1", "c2", "c3"})),
             "f": Column(PandasDtype.Object,
-                        Check(lambda x: x.isin([(1,), (2,), (3,)]),
-                              element_wise=False)),
+                        Check(lambda x: x.isin([(1,), (2,), (3,)]))),
             "g": Column(PandasDtype.DateTime,
-                        Check(lambda x: x >= pd.Timestamp("2015-01-01"))),
+                        Check(lambda x: x >= pd.Timestamp("2015-01-01"),
+                              element_wise=True)),
             "i": Column(PandasDtype.Timedelta,
-                        Check(lambda x: x < pd.Timedelta(10, unit="D")))
+                        Check(lambda x: x < pd.Timedelta(10, unit="D"),
+                              element_wise=True))
         })
     df = pd.DataFrame({
         "a": [1, 2, 3],
@@ -108,7 +108,8 @@ def test_dataframe_schema():
     with pytest.raises(SchemaError):
         schema.validate(df.assign(a=[-1, -2, -1]))
 
-    # checks if 'a' is converted to float, while schema says int, will a schema error be thrown
+    # checks if 'a' is converted to float, while schema says int, will a schema
+    # error be thrown
     with pytest.raises(SchemaError):
         schema.validate(df.assign(a=[1.7, 2.3, 3.1]))
 
@@ -128,8 +129,8 @@ def test_index_schema():
         columns={},
         index=Index(
             PandasDtype.Int, [
-                Check(lambda x: 1 <= x <= 11),
-                Check(lambda index: index.mean() > 1, element_wise=False)]
+                Check(lambda x: 1 <= x <= 11, element_wise=True),
+                Check(lambda index: index.mean() > 1)]
         ))
     df = pd.DataFrame(index=range(1, 11), dtype="int64")
     assert isinstance(schema.validate(df), pd.DataFrame)
@@ -142,14 +143,17 @@ def test_check_function_decorators():
     in_schema = DataFrameSchema(
         {
             "a": Column(PandasDtype.Int, [
-                Check(lambda x: x >= 1),
-                Check(lambda s: s.mean() > 0, element_wise=False)]),
+                Check(lambda x: x >= 1, element_wise=True),
+                Check(lambda s: s.mean() > 0)]),
             "b": Column(PandasDtype.String,
-                        Check(lambda x: x in ["x", "y", "z"])),
+                        Check(lambda x: x in ["x", "y", "z"],
+                              element_wise=True)),
             "c": Column(PandasDtype.DateTime,
-                        Check(lambda x: pd.Timestamp("2018-01-01") <= x)),
+                        Check(lambda x: pd.Timestamp("2018-01-01") <= x,
+                              element_wise=True)),
             "d": Column(PandasDtype.Float,
-                        Check(lambda x: np.isnan(x) or x < 3),
+                        Check(lambda x: np.isnan(x) or x < 3,
+                              element_wise=True),
                         nullable=True)
         },
         transformer=lambda df: df.assign(e="foo")
@@ -157,9 +161,9 @@ def test_check_function_decorators():
     out_schema = DataFrameSchema(
         {
             "e": Column(PandasDtype.String,
-                        Check(lambda s: s == "foo", element_wise=False)),
+                        Check(lambda s: s == "foo")),
             "f": Column(PandasDtype.String,
-                        Check(lambda x: x in ["a", "b"]))
+                        Check(lambda x: x in ["a", "b"], element_wise=True))
         })
 
     # case 1: simplest path test - df is first argument and function returns
@@ -351,3 +355,154 @@ def test_required():
 
     with pytest.raises(Exception):
         schema.validate(df_not_ok)
+
+
+def test_check_groupby():
+    schema = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s["foo"] > 10, groupby="col2"),
+            Check(lambda s: s["bar"] < 10, groupby=["col2"]),
+            Check(lambda s: s["foo"] > 10,
+                  groupby=lambda df: df.groupby("col2")),
+            Check(lambda s: s["bar"] < 10,
+                  groupby=lambda df: df.groupby("col2"))
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+    })
+
+    df_pass = pd.DataFrame({
+        "col1": [7, 8, 9, 11, 12, 13],
+        "col2": ["bar", "bar", "bar", "foo", "foo", "foo"],
+    })
+
+    df = schema.validate(df_pass)
+    assert isinstance(df, pd.DataFrame)
+    assert len(df.columns) == 2
+    assert set(df.columns) == {"col1", "col2"}
+
+    # raise SchemaError when Check fails
+    df_fail_on_bar = pd.DataFrame({
+        "col1": [7, 8, 20, 11, 12, 13],
+        "col2": ["bar", "bar", "bar", "foo", "foo", "foo"],
+    })
+    df_fail_on_foo = pd.DataFrame({
+        "col1": [7, 8, 9, 11, 1, 13],
+        "col2": ["bar", "bar", "bar", "foo", "foo", "foo"],
+    })
+    # raise SchemaError when groupby column doesn't exist
+    df_fail_no_column = pd.DataFrame({
+        "col1": [7, 8, 20, 11, 12, 13],
+    })
+
+    for df in [df_fail_on_bar, df_fail_on_foo, df_fail_no_column]:
+        with pytest.raises(SchemaError):
+            schema.validate(df)
+
+
+def test_check_groupby_multiple_columns():
+    schema = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s[("bar", True)].sum() == 16,  # 7 + 9
+                  groupby=["col2", "col3"]),
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+        "col3": Column(Bool),
+    })
+
+    df_pass = pd.DataFrame({
+        "col1": [7, 8, 9, 11, 12, 13],
+        "col2": ["bar", "bar", "bar", "foo", "foo", "foo"],
+        "col3": [True, False, True, False, True, False],
+    })
+
+    df = schema.validate(df_pass)
+    assert isinstance(df, pd.DataFrame)
+    assert len(df.columns) == 3
+    assert set(df.columns) == {"col1", "col2", "col3"}
+
+
+def test_check_groups():
+    schema = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s["foo"] > 10, groupby="col2", groups=["foo"]),
+            Check(lambda s: s["foo"] > 10, groupby="col2", groups="foo"),
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+    })
+
+    df = pd.DataFrame({
+        "col1": [7, 8, 9, 11, 12, 13],
+        "col2": ["bar", "bar", "bar", "foo", "foo", "foo"],
+    })
+
+    validated_df = schema.validate(df)
+    assert isinstance(validated_df, pd.DataFrame)
+    assert len(validated_df.columns) == 2
+    assert set(validated_df.columns) == {"col1", "col2"}
+
+    # raise KeyError when groups does not include a particular group name
+    schema_fail_key_error = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s["bar"] > 10, groupby="col2", groups="foo"),
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+    })
+    with pytest.raises(KeyError, match="^'bar'"):
+        schema_fail_key_error.validate(df)
+
+    # raise KeyError when the group does not exist in the groupby column when
+    # referenced in the Check function
+    schema_fail_nonexistent_key_in_fn = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s["baz"] > 10, groupby="col2", groups=["foo"]),
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+    })
+    with pytest.raises(KeyError, match="^'baz'"):
+        schema_fail_nonexistent_key_in_fn.validate(df)
+
+    # raise KeyError when the group does not exist in the groups argument.
+    schema_fail_nonexistent_key_in_groups = DataFrameSchema({
+        "col1": Column(Int, [
+            Check(lambda s: s["foo"] > 10, groupby="col2", groups=["baz"]),
+        ]),
+        "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+    })
+    with pytest.raises(KeyError):
+        schema_fail_nonexistent_key_in_groups.validate(df)
+
+
+def test_groupby_init_exceptions():
+    def init_schema_element_wise():
+        DataFrameSchema({
+            "col1": Column(Int, [
+                Check(lambda s: s["foo"] > 10,
+                      element_wise=True,
+                      groupby=["col2"]),
+            ]),
+            "col2": Column(String, Check(lambda s: s.isin(["foo", "bar"]))),
+        })
+
+    # can't use groupby in Checks where element_wise == True
+    with pytest.raises(
+            SchemaInitError,
+            match=r"^Cannot use groupby when element_wise=True."):
+        init_schema_element_wise()
+
+    # raise SchemaInitError even when the schema doesn't specify column key for
+    # groupby column
+    def init_schema_no_groupby_column():
+        DataFrameSchema({
+            "col1": Column(Int, [
+                Check(lambda s: s["foo"] > 10, groupby=["col2"]),
+            ]),
+        })
+    with pytest.raises(SchemaInitError):
+        init_schema_no_groupby_column()
+
+    # can't use groupby argument in SeriesSchema or Index objects
+    for SchemaClass in [SeriesSchema, Index]:
+        with pytest.raises(
+                SchemaInitError,
+                match="^Can only use `groupby` with a pandera.Column, found"):
+            SchemaClass(Int, Check(lambda s: s["bar"] == 1, groupby="foo"))