如何基于列值从DataFrame中选择行?

39 浏览
0 Comments

如何基于列值从DataFrame中选择行?

如何在 Pandas 中根据某列的值选择 DataFrame 中的行?

在 SQL 中,我会使用:

SELECT *
FROM table
WHERE column_name = some_value

admin 更改状态以发布 2023年5月22日
0
0 Comments

有几种方法可以从Pandas数据框中选择行:

  1. 布尔索引 (df[df['col'] == value] )
  2. 位置索引 (df.iloc[...])
  3. 标签索引 (df.xs(...))
  4. df.query(...) API

下面我将展示每种方法的示例,并给出何时使用特定技术的建议。假设我们的条件是列 'A' == 'foo'

(关于性能的说明:对于每种基本类型,我们可以使用Pandas API来简化操作,或者可以进入NumPy等外部库来加速操作。)


设置

首先,我们需要确定一个条件作为选择行的标准。我们将从OP的情况 column_name == some_value 开始,包括一些其他常见用例。

借鉴@unutbu的方法:

import pandas as pd, numpy as np
df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(),
                   'B': 'one one two three two two one three'.split(),
                   'C': np.arange(8), 'D': np.arange(8) * 2})


1. 布尔索引

...布尔索引需要找到每一行 'A' 列等于 'foo' 的真值,然后使用这些真值来识别要保留哪些行。通常,我们将这个系列,即真值数组,命名为 mask。我们在这里也这样做。

mask = df['A'] == 'foo'

然后我们可以使用这个掩码来切片或索引数据框。

df[mask]
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14

这是完成此任务最简单的方法之一,如果性能或直观性不是问题,那么这应该是您选择的方法。但是,如果性能是一个问题,那么您可能需要考虑使用替代方法来创建 mask


2. 位置索引

位置索引(df.iloc[...])有其使用场景,但这不是其中之一。为了确定要切片的位置,我们首先需要执行上述相同的布尔分析。这使我们执行了一个额外的步骤来完成相同的任务。

mask = df['A'] == 'foo'
pos = np.flatnonzero(mask)
df.iloc[pos]
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14

3. 标签索引

标签索引非常方便,但在这种情况下,我们又做了更多的工作,却没有得到任何好处。

df.set_index('A', append=True, drop=False).xs('foo', level=1)
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14

4. df.query() API

pd.DataFrame.query 是一种非常优雅/直观的执行此任务的方法,但速度通常较慢。 然而,如果您注意下面的计时,对于大数据,查询非常高效。 比标准方法更高效,并且与我的最佳建议相当。

df.query('A == "foo"')
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14


我更喜欢使用 Boolean mask

通过修改如何创建我们的 Boolean mask,可以实现实际改进。

mask 替代方法 1
使用底层的 NumPy 数组,避免创建另一个 pd.Series 的开销

mask = df['A'].values == 'foo'

我将在最后展示更完整的时间测试,但是先来看看使用示例数据框获得的性能提升。首先,我们比较创建 mask 的不同之处。

%timeit mask = df['A'].values == 'foo'
%timeit mask = df['A'] == 'foo'
5.84 µs ± 195 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
166 µs ± 4.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

使用 NumPy 数组计算 mask 大约快了 30 倍。这在一定程度上是由于 NumPy 的计算速度通常更快。另一方面,由于无需构建索引和相应的 pd.Series 对象,也会导致更快的速度。

接下来,我们将比较使用两种 mask 进行切片的时间。

mask = df['A'].values == 'foo'
%timeit df[mask]
mask = df['A'] == 'foo'
%timeit df[mask]
219 µs ± 12.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
239 µs ± 7.03 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

性能提升并不明显。我们将在更为全面的测试中看到这是否成立。


mask 替代方案 2
我们也可以重新构建数据框。但要注意重新构建数据框时必须处理好 dtypes

我们将不再使用df[mask]

pd.DataFrame(df.values[mask], df.index[mask], df.columns).astype(df.dtypes)

如果数据帧是混合类型的,比如我们的例子,那么当我们获得df.values的结果数组的dtypeobject,因此新数据帧的所有列也都是dtypeobject。因此需要astype(df.dtypes),这将消除任何潜在的性能增益。

%timeit df[m]
%timeit pd.DataFrame(df.values[mask], df.index[mask], df.columns).astype(df.dtypes)
216 µs ± 10.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.43 ms ± 39.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

但是,如果数据帧不是混合类型,那么这是一种非常有用的方法。

给定

np.random.seed([3,1415])
d1 = pd.DataFrame(np.random.randint(10, size=(10, 5)), columns=list('ABCDE'))
d1
   A  B  C  D  E
0  0  2  7  3  8
1  7  0  6  8  6
2  0  2  0  4  9
3  7  3  2  4  3
4  3  6  7  7  4
5  5  3  7  5  9
6  8  7  6  4  7
7  6  2  6  6  5
8  2  8  7  5  8
9  4  7  6  1  5


%%timeit
mask = d1['A'].values == 7
d1[mask]
179 µs ± 8.73 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

%%timeit
mask = d1['A'].values == 7
pd.DataFrame(d1.values[mask], d1.index[mask], d1.columns)
87 µs ± 5.12 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

我们将时间减半。


mask alternative 3

@unutbu也向我们展示了如何使用pd.Series.isin来考虑df['A']的每个元素是否在一组值中。如果我们的值集合是一个值的集合,即'foo',那么这将得到相同的结果。但是如果需要,它也可以推广到包括更大的值集合。结果表明,即使这是一种更普遍的解决方案,它仍然非常快速。唯一的损失是对于不熟悉这个概念的人来说,它不够直观。

mask = df['A'].isin(['foo'])
df[mask]
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14

然而,与之前一样,我们可以利用NumPy来提高性能,而几乎不损失任何东西。我们将使用np.in1d

mask = np.in1d(df['A'].values, ['foo'])
df[mask]
     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14


计时

我还将包括其他帖子中提到的概念以供参考。

下面是代码

这个表中的每一列代表在不同长度的数据框上测试每个函数所需的相对时间。每一列显示相对于最快的函数所花费的时间,最快的函数被赋予基本指数1.0

res.div(res.min())
                         10        30        100       300       1000      3000      10000     30000
mask_standard         2.156872  1.850663  2.034149  2.166312  2.164541  3.090372  2.981326  3.131151
mask_standard_loc     1.879035  1.782366  1.988823  2.338112  2.361391  3.036131  2.998112  2.990103
mask_with_values      1.010166  1.000000  1.005113  1.026363  1.028698  1.293741  1.007824  1.016919
mask_with_values_loc  1.196843  1.300228  1.000000  1.000000  1.038989  1.219233  1.037020  1.000000
query                 4.997304  4.765554  5.934096  4.500559  2.997924  2.397013  1.680447  1.398190
xs_label              4.124597  4.272363  5.596152  4.295331  4.676591  5.710680  6.032809  8.950255
mask_with_isin        1.674055  1.679935  1.847972  1.724183  1.345111  1.405231  1.253554  1.264760
mask_with_in1d        1.000000  1.083807  1.220493  1.101929  1.000000  1.000000  1.000000  1.144175

你会注意到,最快的时间似乎在mask_with_valuesmask_with_in1d之间共享。

res.T.plot(loglog=True)

Enter image description here

把下面的内容翻译成中文如果有html标签和引号或双引号内的内容请不要翻译:

Functions

def mask_standard(df):
    mask = df['A'] == 'foo'
    return df[mask]
def mask_standard_loc(df):
    mask = df['A'] == 'foo'
    return df.loc[mask]
def mask_with_values(df):
    mask = df['A'].values == 'foo'
    return df[mask]
def mask_with_values_loc(df):
    mask = df['A'].values == 'foo'
    return df.loc[mask]
def query(df):
    return df.query('A == "foo"')
def xs_label(df):
    return df.set_index('A', append=True, drop=False).xs('foo', level=-1)
def mask_with_isin(df):
    mask = df['A'].isin(['foo'])
    return df[mask]
def mask_with_in1d(df):
    mask = np.in1d(df['A'].values, ['foo'])
    return df[mask]


Testing

res = pd.DataFrame(
    index=[
        'mask_standard', 'mask_standard_loc', 'mask_with_values', 'mask_with_values_loc',
        'query', 'xs_label', 'mask_with_isin', 'mask_with_in1d'
    ],
    columns=[10, 30, 100, 300, 1000, 3000, 10000, 30000],
    dtype=float
)
for j in res.columns:
    d = pd.concat([df] * j, ignore_index=True)
    for i in res.index:a
        stmt = '{}(d)'.format(i)
        setp = 'from __main__ import d, {}'.format(i)
        res.at[i, j] = timeit(stmt, setp, number=50)


Special Timing

观察当整个数据帧都有一个非对象dtype的特殊情况。

代码如下

spec.div(spec.min())
                     10        30        100       300       1000      3000      10000     30000
mask_with_values  1.009030  1.000000  1.194276  1.000000  1.236892  1.095343  1.000000  1.000000
mask_with_in1d    1.104638  1.094524  1.156930  1.072094  1.000000  1.000000  1.040043  1.027100
reconstruct       1.000000  1.142838  1.000000  1.355440  1.650270  2.222181  2.294913  3.406735

结果表明,在几百行以上,重构已经不值得了。

spec.T.plot(loglog=True)

Enter image description here

Functions

np.random.seed([3,1415])
d1 = pd.DataFrame(np.random.randint(10, size=(10, 5)), columns=list('ABCDE'))
def mask_with_values(df):
    mask = df['A'].values == 'foo'
    return df[mask]
def mask_with_in1d(df):
    mask = np.in1d(df['A'].values, ['foo'])
    return df[mask]
def reconstruct(df):
    v = df.values
    mask = np.in1d(df['A'].values, ['foo'])
    return pd.DataFrame(v[mask], df.index[mask], df.columns)
spec = pd.DataFrame(
    index=['mask_with_values', 'mask_with_in1d', 'reconstruct'],
    columns=[10, 30, 100, 300, 1000, 3000, 10000, 30000],
    dtype=float
)

Testing

for j in spec.columns:
    d = pd.concat([df] * j, ignore_index=True)
    for i in spec.index:
        stmt = '{}(d)'.format(i)
        setp = 'from __main__ import d, {}'.format(i)
        spec.at[i, j] = timeit(stmt, setp, number=50)


0
0 Comments

要选择其列值等于标量some_value的行,请使用==

df.loc[df['column_name'] == some_value]

要选择其列值在可迭代对象some_values中的行,请使用isin

df.loc[df['column_name'].isin(some_values)]

使用&将多个条件组合:

df.loc[(df['column_name'] >= A) & (df['column_name'] <= B)]

请注意括号。由于 Python 的运算符优先级规则&的优先级高于<=>=。因此,最后一个示例中的括号是必需的。如果没有括号,则

df['column_name'] >= A & df['column_name'] <= B

将被解析为

df['column_name'] >= (A & df['column_name']) <= B

这将导致“系列的真值是模棱两可的”错误


要选择其列值不等于some_value的行,请使用!=

df.loc[df['column_name'] != some_value]

isin返回一个布尔值系列,因此要选择其值不在some_values中的行,请使用~来取反布尔值系列:

df.loc[~df['column_name'].isin(some_values)]


例如,

import pandas as pd
import numpy as np
df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(),
                   'B': 'one one two three two two one three'.split(),
                   'C': np.arange(8), 'D': np.arange(8) * 2})
print(df)
#      A      B  C   D
# 0  foo    one  0   0
# 1  bar    one  1   2
# 2  foo    two  2   4
# 3  bar  three  3   6
# 4  foo    two  4   8
# 5  bar    two  5  10
# 6  foo    one  6  12
# 7  foo  three  7  14
print(df.loc[df['A'] == 'foo'])

将产生:

     A      B  C   D
0  foo    one  0   0
2  foo    two  2   4
4  foo    two  4   8
6  foo    one  6  12
7  foo  three  7  14


如果你有多个要包括的值,将它们放在列表(或更普遍的可迭代对象)中并使用isin

print(df.loc[df['B'].isin(['one','three'])])

将产生:

     A      B  C   D
0  foo    one  0   0
1  bar    one  1   2
3  bar  three  3   6
6  foo    one  6  12
7  foo  three  7  14


但请注意,如果您想多次执行此操作,则更高效的方法是首先创建索引,然后使用df.loc

df = df.set_index(['B'])
print(df.loc['one'])

将产生:

       A  C   D
B              
one  foo  0   0
one  bar  1   2
one  foo  6  12

或者,要从索引中包括多个值,请使用df.index.isin

df.loc[df.index.isin(['one','two'])]

产生:

       A  C   D
B              
one  foo  0   0
one  bar  1   2
two  foo  2   4
two  foo  4   8
two  bar  5  10
one  foo  6  12

0