test_dataframe.py 文件源码

def test_group_indices():
    df = pd.DataFrame({'x': [1, 5, 2, 2, 4, 0, 4],
                       'y': [1, 2, 3, 4, 5, 6, 5]})

    results = df >> group_by('x') >> group_indices()
    assert all(results == [1, 4, 2, 2, 3, 0, 3])

    results = df >> group_indices('y % 2')
    assert all(results == [1, 0, 1, 0, 1, 0, 1])

    results = df >> group_indices()
    assert all(results == [1, 1, 1, 1, 1, 1, 1])

    # Branches
    with pytest.warns(UserWarning):
        df >> group_by('x') >> group_indices('y')