from __future__ import division

import pytest
import json
import datetime
import pandas as pd
import great_expectations as ge
from great_expectations.profile import ColumnsExistProfiler

from .test_utils import assertDeepAlmostEqual
from six import PY2


def test_expect_column_values_to_be_dateutil_parseable():

    D = ge.dataset.PandasDataset({
        'c1': ['03/06/09', '23 April 1973', 'January 9, 2016'],
        'c2': ['9/8/2012', 'covfefe', 25],
        'c3': ['Jared', 'June 1, 2013', 'July 18, 1976'],
        'c4': ['1', '2', '49000004632'],
        'already_datetime': [datetime.datetime(2015, 1, 1), datetime.datetime(2016, 1, 1), datetime.datetime(2017, 1, 1)],
    })
    D.set_default_expectation_argument("result_format", "COMPLETE")

    T = [
        {
            'in': {'column': 'c1'},
            'out': {'success': True, 'unexpected_list': [], 'unexpected_index_list': []}},
        {
            'in': {"column": 'c2', "catch_exceptions": True},
            # 'out':{'success':False, 'unexpected_list':['covfefe', 25], 'unexpected_index_list': [1, 2]}},
            'error': {'traceback_substring': 'TypeError: Values passed to expect_column_values_to_be_dateutil_parseable must be of type string'},
        },
        {
            'in': {"column": 'c3'},
            'out': {'success': False, 'unexpected_list': ['Jared'], 'unexpected_index_list': [0]}},
        {
            'in': {'column': 'c3', 'mostly': .5},
            'out': {'success': True, 'unexpected_list': ['Jared'], 'unexpected_index_list': [0]}
        },
        {
            'in': {'column': 'c4'},
            'out': {'success': False, 'unexpected_list': ['49000004632'], 'unexpected_index_list': [2]}
        },
        {
            'in': {'column': 'already_datetime', 'catch_exceptions': True},
            'error': {'traceback_substring': 'TypeError: Values passed to expect_column_values_to_be_dateutil_parseable must be of type string'},
        }
    ]

    for t in T:
        out = D.expect_column_values_to_be_dateutil_parseable(**t['in'])
        if 'out' in t:
            assert t['out']['success'] == out['success']
            assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
            assert t['out']['unexpected_list'] == out['result']['unexpected_list']
        elif 'error' in t:
            assert out['exception_info']['raised_exception'] == True
            assert t['error']['traceback_substring'] in out['exception_info']['exception_traceback']


def test_expect_column_values_to_be_json_parseable():
    d1 = json.dumps({'i': [1, 2, 3], 'j': 35, 'k': {
                    'x': 'five', 'y': 5, 'z': '101'}})
    d2 = json.dumps({'i': 1, 'j': 2, 'k': [3, 4, 5]})
    d3 = json.dumps({'i': 'a', 'j': 'b', 'k': 'c'})
    d4 = json.dumps({'i': [4, 5], 'j': [6, 7], 'k': [
                    8, 9], 'l': {4: 'x', 5: 'y', 6: 'z'}})
    D = ge.dataset.PandasDataset({
        'json_col': [d1, d2, d3, d4],
        'not_json': [4, 5, 6, 7],
        'py_dict': [{'a': 1, 'out': 1}, {'b': 2, 'out': 4}, {'c': 3, 'out': 9}, {'d': 4, 'out': 16}],
        'most': [d1, d2, d3, 'd4']
    })
    D.set_default_expectation_argument("result_format", "COMPLETE")

    T = [
        {
            'in': {'column': 'json_col'},
            'out': {'success': True, 'unexpected_index_list': [], 'unexpected_list':[]}},
        {
            'in': {'column': 'not_json'},
            'out': {'success': False, 'unexpected_index_list': [0, 1, 2, 3], 'unexpected_list':[4, 5, 6, 7]}},
        {
            'in': {'column': 'py_dict'},
            'out': {'success': False, 'unexpected_index_list': [0, 1, 2, 3], 'unexpected_list':[{'a': 1, 'out': 1}, {'b': 2, 'out': 4}, {'c': 3, 'out': 9}, {'d': 4, 'out': 16}]}},
        {
            'in': {'column': 'most'},
            'out': {'success': False, 'unexpected_index_list': [3], 'unexpected_list':['d4']}},
        {
            'in': {'column': 'most', 'mostly': .75},
            'out': {'success': True, 'unexpected_index_list': [3], 'unexpected_list':['d4']}}
    ]

    for t in T:
        out = D.expect_column_values_to_be_json_parseable(**t['in'])
        assert t['out']['success'] == out['success']
        assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
        assert t['out']['unexpected_list'] == out['result']['unexpected_list']


def test_expectation_decorator_summary_mode():

    df = ge.dataset.PandasDataset({
        'x': [1, 2, 3, 4, 5, 6, 7, 7, None, None],
    })
    df.set_default_expectation_argument("result_format", "COMPLETE")

    # print '&'*80
    # print json.dumps(df.expect_column_values_to_be_between('x', min_value=1, max_value=5, result_format="SUMMARY"), indent=2)

    exp_output = {
        "success": False,
        "result": {
            "element_count": 10,
            "missing_count": 2,
            "missing_percent": 20.0,
            "unexpected_count": 3,
            "partial_unexpected_counts": [
                {"value": 7.0,
                    "count": 2},
                {"value": 6.0,
                    "count": 1}
            ],
            "unexpected_percent": 30.0,
            "unexpected_percent_nonmissing": 37.5,
            "partial_unexpected_list": [6.0, 7.0, 7.0],
            "partial_unexpected_index_list": [5, 6, 7],
        }
    }
    assert df.expect_column_values_to_be_between('x', min_value=1, max_value=5, result_format="SUMMARY")\
        == exp_output

    exp_output = {
        'success': True,
        'result': {
            'observed_value': 4.375,
            'element_count': 10,
            'missing_count': 2,
            'missing_percent': 20.0
        },
    }

    assert df.expect_column_mean_to_be_between("x", 3, 7, result_format="SUMMARY")\
        == exp_output


def test_positional_arguments():

    df = ge.dataset.PandasDataset({
        'x': [1, 3, 5, 7, 9],
        'y': [2, 4, 6, 8, 10],
        'z': [None, 'a', 'b', 'c', 'abc']
    })
    df.set_default_expectation_argument('result_format', 'COMPLETE')

    exp_output = {'success': True, 'result': {'observed_value': 5, 'element_count': 5,
                                              'missing_count': 0,
                                              'missing_percent': 0.0}}

    assert df.expect_column_mean_to_be_between('x', 4, 6) == exp_output

    out = df.expect_column_values_to_be_between('y', 1, 6)
    t = {'out': {'success': False, 'unexpected_list': [
        8, 10], 'unexpected_index_list': [3, 4]}}
    if 'out' in t:
        assert t['out']['success'] == out['success']
        if 'unexpected_index_list' in t['out']:
            assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
        if 'unexpected_list' in t['out']:
            assert t['out']['unexpected_list'] == out['result']['unexpected_list']

    out = df.expect_column_values_to_be_between('y', 1, 6, mostly=.5)
    t = {'out': {'success': True, 'unexpected_list': [
        8, 10], 'unexpected_index_list': [3, 4]}}
    if 'out' in t:
        assert t['out']['success'] == out['success']
        if 'unexpected_index_list' in t['out']:
            assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
        if 'unexpected_list' in t['out']:
            assert t['out']['unexpected_list'] == out['result']['unexpected_list']

    out = df.expect_column_values_to_be_in_set('z', ['a', 'b', 'c'])
    t = {'out': {'success': False, 'unexpected_list': [
        'abc'], 'unexpected_index_list': [4]}}
    if 'out' in t:
        assert t['out']['success'] == out['success']
        if 'unexpected_index_list' in t['out']:
            assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
        if 'unexpected_list' in t['out']:
            assert t['out']['unexpected_list'] == out['result']['unexpected_list']

    out = df.expect_column_values_to_be_in_set('z', ['a', 'b', 'c'], mostly=.5)
    t = {'out': {'success': True, 'unexpected_list': [
        'abc'], 'unexpected_index_list': [4]}}
    if 'out' in t:
        assert t['out']['success'] == out['success']
        if 'unexpected_index_list' in t['out']:
            assert t['out']['unexpected_index_list'] == out['result']['unexpected_index_list']
        if 'unexpected_list' in t['out']:
            assert t['out']['unexpected_list'] == out['result']['unexpected_list']


def test_result_format_argument_in_decorators():
    df = ge.dataset.PandasDataset({
        'x': [1, 3, 5, 7, 9],
        'y': [2, 4, 6, 8, 10],
        'z': [None, 'a', 'b', 'c', 'abc']
    })
    df.set_default_expectation_argument('result_format', 'COMPLETE')

    # Test explicit Nones in result_format

    exp_output = {'success': True, 'result': {'observed_value': 5, 'element_count': 5,
                                              'missing_count': 0,
                                              'missing_percent': 0.0
                                              }}
    assert df.expect_column_mean_to_be_between('x', 4, 6, result_format=None)\
        == exp_output

    exp_output = {'result': {'element_count': 5,
                             'missing_count': 0,
                             'missing_percent': 0.0,
                             'partial_unexpected_counts': [{'count': 1, 'value': 8},
                                                           {'count': 1, 'value': 10}],
                             'partial_unexpected_index_list': [3, 4],
                             'partial_unexpected_list': [8, 10],
                             'unexpected_count': 2,
                             'unexpected_index_list': [3, 4],
                             'unexpected_list': [8, 10],
                             'unexpected_percent': 40.0,
                             'unexpected_percent_nonmissing': 40.0},
                  'success': False}

    assert df.expect_column_values_to_be_between('y', 1, 6, result_format=None)\
        == exp_output

    # Test unknown output format
    with pytest.raises(ValueError):
        df.expect_column_values_to_be_between('y', 1, 6, result_format="QUACK")

    with pytest.raises(ValueError):
        df.expect_column_mean_to_be_between('x', 4, 6, result_format="QUACK")


def test_from_pandas():
    pd_df = pd.DataFrame({
        'x': [1, 3, 5, 7, 9],
        'y': [2, 4, 6, 8, 10],
        'z': [None, 'a', 'b', 'c', 'abc']
    })

    ge_df = ge.from_pandas(pd_df)
    assert isinstance(ge_df, ge.data_asset.DataAsset)
    assert list(ge_df.columns) == ['x', 'y', 'z']
    assert list(ge_df['x']) == list(pd_df['x'])
    assert list(ge_df['y']) == list(pd_df['y'])
    assert list(ge_df['z']) == list(pd_df['z'])

    # make an empty subclass to test dataset_class argument
    class CustomPandasDataset(ge.dataset.PandasDataset):
        pass

    ge_df_custom = ge.from_pandas(pd_df, dataset_class=CustomPandasDataset)

    assert not isinstance(ge_df, CustomPandasDataset)
    assert isinstance(ge_df_custom, CustomPandasDataset)
    assert list(ge_df_custom.columns)==['x', 'y', 'z']
    assert list(ge_df_custom['x'])==list(pd_df['x'])
    assert list(ge_df_custom['y'])==list(pd_df['y'])
    assert list(ge_df_custom['z'])==list(pd_df['z'])


def test_ge_pandas_concatenating_no_autoinspect():
    df1 = ge.dataset.PandasDataset({
        'A': ['A0', 'A1', 'A2'],
        'B': ['B0', 'B1', 'B2']
    })

    df1.expect_column_to_exist('A')
    df1.expect_column_to_exist('B')
    df1.expect_column_values_to_match_regex('A', '^A[0-2]$')
    df1.expect_column_values_to_match_regex('B', '^B[0-2]$')

    df2 = ge.dataset.PandasDataset({
        'A': ['A3', 'A4', 'A5'],
        'B': ['B3', 'B4', 'B5']
    })

    df2.expect_column_to_exist('A')
    df2.expect_column_to_exist('B')
    df2.expect_column_values_to_match_regex('A', '^A[3-5]$')
    df2.expect_column_values_to_match_regex('B', '^B[3-5]$')

    df = pd.concat([df1, df2])

    exp_c = []

    # The concatenated data frame will:
    #
    #   1. Be a ge.dataset.PandaDataSet
    #   2. Have no expectations (since no default expectations are created), even expectations that were common
    #      to the concatenated dataframes and still make sense (since no autoinspection happens).

    assert isinstance(df, ge.dataset.PandasDataset)
    assert df.find_expectations() == exp_c


def test_ge_pandas_joining():
    df1 = ge.dataset.PandasDataset({
        'A': ['A0', 'A1', 'A2'],
        'B': ['B0', 'B1', 'B2']},
        index=['K0', 'K1', 'K2'])

    df1.expect_column_values_to_match_regex('A', '^A[0-2]$')
    df1.expect_column_values_to_match_regex('B', '^B[0-2]$')

    df2 = ge.dataset.PandasDataset({
        'C': ['C0', 'C2', 'C3'],
        'D': ['C0', 'D2', 'D3']},
        index=['K0', 'K2', 'K3'])

    df2.expect_column_values_to_match_regex('C', '^C[0-2]$')
    df2.expect_column_values_to_match_regex('D', '^D[0-2]$')

    df = df1.join(df2)

    exp_j = [
        # No autoinspection is default 20180920
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'A'}},
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'B'}},
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'C'}},
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'D'}}
    ]

    # The joined data frame will:
    #
    #   1. Be a ge.dataset.PandaDataSet
    #   2. Have no expectations (no autoinspection)

    assert isinstance(df, ge.dataset.PandasDataset)
    assert df.find_expectations() == exp_j


def test_ge_pandas_merging():
    df1 = ge.dataset.PandasDataset({
        'id': [1, 2, 3, 4],
        'name': ['a', 'b', 'c', 'd']
    })

    df1.expect_column_values_to_match_regex('name', '^[A-Za-z ]+$')

    df2 = ge.dataset.PandasDataset({
        'id': [1, 2, 3, 4],
        'salary': [57000, 52000, 59000, 65000]
    })

    df2.expect_column_values_to_match_regex('salary', '^[0-9]{4,6}$')

    df = df1.merge(df2, on='id')

    exp_m = [
        # No autoinspection as of 20180920
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'id'}},
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'name'}},
        # {'expectation_type': 'expect_column_to_exist',
        #  'kwargs': {'column': 'salary'}}
    ]

    # The merged data frame will:
    #
    #   1. Be a ge.dataset.PandaDataSet
    #   2. Have no expectations (no autoinspection is now default)

    assert isinstance(df, ge.dataset.PandasDataset)
    assert df.find_expectations() == exp_m


def test_ge_pandas_sampling():
    df = ge.dataset.PandasDataset({
        'A': [1, 2, 3, 4],
        'B': [5, 6, 7, 8],
        'C': ['a', 'b', 'c', 'd'],
        'D': ['e', 'f', 'g', 'h']
    })

    # Put some simple expectations on the data frame
    df.profile(profiler=ColumnsExistProfiler)
    df.expect_column_values_to_be_in_set("A", [1, 2, 3, 4])
    df.expect_column_values_to_be_in_set("B", [5, 6, 7, 8])
    df.expect_column_values_to_be_in_set("C", ['a', 'b', 'c', 'd'])
    df.expect_column_values_to_be_in_set("D", ['e', 'f', 'g', 'h'])

    exp1 = df.find_expectations()

    # The sampled data frame should:
    #
    #   1. Be a ge.dataset.PandaDataSet
    #   2. Inherit ALL the expectations of the parent data frame

    samp1 = df.sample(n=2)
    assert isinstance(samp1, ge.dataset.PandasDataset)
    if PY2:
        assert sorted(samp1.find_expectations()) == sorted(exp1)
    else:
        assert samp1.find_expectations() == exp1

    samp1 = df.sample(frac=0.25, replace=True)
    assert isinstance(samp1, ge.dataset.PandasDataset)
    assert samp1.find_expectations() == exp1

    # Change expectation on column "D", sample, and check expectations.
    # The failing expectation on column "D" is NOT automatically dropped
    # in the sample.
    df.expect_column_values_to_be_in_set("D", ['e', 'f', 'g', 'x'])
    samp1 = df.sample(n=2)
    exp1 = [
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'A'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'B'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'C'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'D'}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'A', 'value_set': [1, 2, 3, 4]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'B', 'value_set': [5, 6, 7, 8]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'C', 'value_set': ['a', 'b', 'c', 'd']}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'D', 'value_set': ['e', 'f', 'g', 'x']}}
    ]
    if PY2:
        assert sorted(samp1.find_expectations()) == sorted(exp1)
    else:
        assert samp1.find_expectations() == exp1


def test_ge_pandas_subsetting():
    df = ge.dataset.PandasDataset({
        'A': [1, 2, 3, 4],
        'B': [5, 6, 7, 8],
        'C': ['a', 'b', 'c', 'd'],
        'D': ['e', 'f', 'g', 'h']
    })

    # Put some simple expectations on the data frame
    df.expect_column_values_to_be_in_set("A", [1, 2, 3, 4])
    df.expect_column_values_to_be_in_set("B", [5, 6, 7, 8])
    df.expect_column_values_to_be_in_set("C", ['a', 'b', 'c', 'd'])
    df.expect_column_values_to_be_in_set("D", ['e', 'f', 'g', 'h'])

    # The subsetted data frame should:
    #
    #   1. Be a ge.dataset.PandaDataSet
    #   2. Inherit ALL the expectations of the parent data frame

    exp1 = df.find_expectations()

    sub1 = df[['A', 'D']]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df[['A']]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df[:3]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df[1:2]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df[:-1]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df[-1:]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df.iloc[:3, 1:4]
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1

    sub1 = df.loc[0:, 'A':'B']
    assert isinstance(sub1, ge.dataset.PandasDataset)
    assert sub1.find_expectations() == exp1


def test_ge_pandas_automatic_failure_removal():
    df = ge.dataset.PandasDataset({
        'A': [1, 2, 3, 4],
        'B': [5, 6, 7, 8],
        'C': ['a', 'b', 'c', 'd'],
        'D': ['e', 'f', 'g', 'h']
    })

    # Put some simple expectations on the data frame
    df.profile(ge.profile.ColumnsExistProfiler)
    df.expect_column_values_to_be_in_set("A", [1, 2, 3, 4])
    df.expect_column_values_to_be_in_set("B", [5, 6, 7, 8])
    df.expect_column_values_to_be_in_set("C", ['w', 'x', 'y', 'z'])
    df.expect_column_values_to_be_in_set("D", ['e', 'f', 'g', 'h'])

    # First check that failing expectations are NOT automatically
    # dropped when sampling.
    # For this data frame, the expectation on column "C" above fails.
    exp1 = [
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'A'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'B'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'C'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'D'}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'A', 'value_set': [1, 2, 3, 4]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'B', 'value_set': [5, 6, 7, 8]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'C', 'value_set': ['w', 'x', 'y', 'z']}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'D', 'value_set': ['e', 'f', 'g', 'h']}}
    ]
    samp1 = df.sample(n=2)
    if PY2:
        assert sorted(samp1.find_expectations()) == sorted(exp1)
    else:
        assert samp1.find_expectations() == exp1

    # Now check subsetting to verify that failing expectations are NOT
    # automatically dropped when subsetting.
    sub1 = df[['A', 'D']]
    if PY2:
        assert sorted(samp1.find_expectations()) == sorted(exp1)
    else:
        assert samp1.find_expectations() == exp1

    # Set property/attribute so that failing expectations are
    # automatically removed when sampling or subsetting.
    df.discard_subset_failing_expectations = True

    ###
    # Note: Order matters in this test, and a validationoperator may change order
    ###

    exp_samp = [
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'A'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'B'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'C'}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'D'}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'A', 'value_set': [1, 2, 3, 4]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'B', 'value_set': [5, 6, 7, 8]}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'D', 'value_set': ['e', 'f', 'g', 'h']}}
    ]

    samp2 = df.sample(n=2)
    if PY2:
        assert sorted(samp2.find_expectations()) == sorted(exp_samp)
    else:
        assert samp2.find_expectations() == exp_samp

    # Now check subsetting. In additional to the failure on column "C",
    # the expectations on column "B" now fail since column "B" doesn't
    # exist in the subset.
    sub2 = df[['A', 'D']]
    exp_sub = [
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'A'}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'A', 'value_set': [1, 2, 3, 4]}},
        {'expectation_type': 'expect_column_to_exist',
         'kwargs': {'column': 'D'}},
        {'expectation_type': 'expect_column_values_to_be_in_set',
         'kwargs': {'column': 'D', 'value_set': ['e', 'f', 'g', 'h']}}
    ]
    if PY2:
        assert sorted(samp2.find_expectations()) == sorted(exp_samp)
    else:
        assert samp2.find_expectations() == exp_samp


def test_subclass_pandas_subset_retains_subclass():
    """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)
    assert type(df2) == type(df)

    def test_validate_map_expectation_on_categorical_column(self):
        """Map expectations should work on categorical columns"""

        D = ge.dataset.PandasDataset({
            'cat_column_1': ['cat_one', 'cat_two', 'cat_one', 'cat_two', 'cat_one', 'cat_two', 'cat_one', 'cat_two'],
        })

        D['cat_column_1'] = D['cat_column_1'].astype('category')

        D.set_default_expectation_argument("result_format", "COMPLETE")

        out = D.expect_column_value_lengths_to_equal('cat_column_1', 7)

        self.assertEqual(out['success'], True)


def test_pandas_deepcopy():
    import copy

    df = ge.dataset.PandasDataset({"a": [1, 2, 3]})
    df2 = copy.deepcopy(df)

    df["a"] = [2, 3, 4]

    # Our copied dataframe should not be affected
    assert df2.expect_column_to_exist("a")["success"] == True
    assert list(df["a"]) == [2, 3, 4]
    assert list(df2["a"]) == [1, 2, 3]
