将系列列表连接到uid

时间:2015-06-16 12:37:01

标签: python python-2.7 pandas pytest

我有一个Pandas数据框,其中有几列组成一个唯一的标识符。我想编写一个通用的测试用例,它允许我将这些列连接成一个列(uid)并测试该列的唯一性。我将以下代码作为非通用测试用例:

import pandas as pd
import pytest
df = pd.DataFrame(columns=['one', 'two', 'three'])
df.one = 'abc', 'def', 'ghi'
df.two = 'jkl', 'mno', 'pqr'
df.three = 'stu', 'vwx', 'yzz'
# Test one
df['uid'] = df.one + df.two
assert len(df.index) == len(df.drop_duplicates(['uid']).index)
# Test two
df['uid'] = df.one + df.three
assert len(df.index) == len(df.drop_duplicates(['uid']).index)

由于我将重复使用它,我需要一个解决方案,允许选择任意列连接在一起,例如在第8行和第11行。

1 个答案:

答案 0 :(得分:1)

假设您要选择要添加的列twothree

col_to_add = ['two', 'three']

使用sum(axis=1)连接这些列:

df['uid'] = df[col_to_add].sum(axis=1)