有效地查询pandas.DataFrame

时间:2018-12-11 14:02:51

标签: python sql pandas dataframe

在我的数据框中,每一行代表一组代码,经过一次热编码,因此数据框中有大量的布尔列。

enter image description here

我要选择包含代码子集的所有行,即所有给定列集的值为True的行。

示例集可能是:

code_selection = {"H045027", "S100031", "G121001", "S456005", "M743110"} 

我的第一次尝试依靠DataFrame.query并从给定的集合构建查询字符串:

def filter_codeset_1(codesets_onehot, code_selection):
    """Return only code sets that contain all of the codes in the code selection"""
    query_string = " & ".join(code_selection)
    return codesets_onehot.query(query_string)

这适用于小型设备,但需要相当长的时间(耗时:31.8 s)。对于大型集,它会因内存错误而崩溃:

MemoryError                               Traceback (most recent call last)
<ipython-input-86-8fb45d40b678> in <module>
----> 1 filtered = filter_codeset(codesets_onehot, code_selection)

<ipython-input-71-ca3fccfa21ba> in filter_codeset(codesets_onehot, code_selection)
      2     """Return only code sets that contain all of the codes in the code selection"""
      3     query_string = " & ".join(code_selection)
----> 4     return codesets_onehot.query(query_string)

~/anaconda3/lib/python3.7/site-packages/pandas/core/frame.py in query(self, expr, inplace, **kwargs)
   2845         kwargs['level'] = kwargs.pop('level', 0) + 1
   2846         kwargs['target'] = None
-> 2847         res = self.eval(expr, **kwargs)
   2848 
   2849         try:

~/anaconda3/lib/python3.7/site-packages/pandas/core/frame.py in eval(self, expr, inplace, **kwargs)
   2960             kwargs['target'] = self
   2961         kwargs['resolvers'] = kwargs.get('resolvers', ()) + tuple(resolvers)
-> 2962         return _eval(expr, inplace=inplace, **kwargs)
   2963 
   2964     def select_dtypes(self, include=None, exclude=None):

~/anaconda3/lib/python3.7/site-packages/pandas/core/computation/eval.py in eval(expr, parser, engine, truediv, local_dict, global_dict, resolvers, level, target, inplace)
    294         eng = _engines[engine]
    295         eng_inst = eng(parsed_expr)
--> 296         ret = eng_inst.evaluate()
    297 
    298         if parsed_expr.assigner is None:

~/anaconda3/lib/python3.7/site-packages/pandas/core/computation/engines.py in evaluate(self)
     74 
     75         # make sure no names in resolvers and locals/globals clash
---> 76         res = self._evaluate()
     77         return _reconstruct_object(self.result_type, res, self.aligned_axes,
     78                                    self.expr.terms.return_type)

~/anaconda3/lib/python3.7/site-packages/pandas/core/computation/engines.py in _evaluate(self)
    121             truediv = scope['truediv']
    122             _check_ne_builtin_clash(self.expr)
--> 123             return ne.evaluate(s, local_dict=scope, truediv=truediv)
    124         except KeyError as e:
    125             # python 3 compat kludge

~/anaconda3/lib/python3.7/site-packages/numexpr/necompiler.py in evaluate(ex, local_dict, global_dict, out, order, casting, **kwargs)
    814     expr_key = (ex, tuple(sorted(context.items())))
    815     if expr_key not in _names_cache:
--> 816         _names_cache[expr_key] = getExprNames(ex, context)
    817     names, ex_uses_vml = _names_cache[expr_key]
    818     arguments = getArguments(names, local_dict, global_dict)

~/anaconda3/lib/python3.7/site-packages/numexpr/necompiler.py in getExprNames(text, context)
    705 
    706 def getExprNames(text, context):
--> 707     ex = stringToExpression(text, {}, context)
    708     ast = expressionToAST(ex)
    709     input_order = getInputOrder(ast, None)

~/anaconda3/lib/python3.7/site-packages/numexpr/necompiler.py in stringToExpression(s, types, context)
    282         else:
    283             flags = 0
--> 284         c = compile(s, '<expr>', 'eval', flags)
    285         # make VariableNode's for the names
    286         names = {}

MemoryError: 

对于更具扩展性的实现(在不超过几秒钟的时间内用成百上千的代码查询数十万行),我有哪些选择?应该有可能非常有效地执行此操作,因为基本上需要为每一行选择一组固定的布尔值并将其与and连接。

以下是替代实现,包括答案中建议的实现:

def filter_codeset_2(codesets_onehot, code_selection):
    column_mask = codesets_onehot.columns.isin(code_selection)
    return codesets_onehot[codesets_onehot.apply(lambda row: row[column_mask].all(), axis=1)]

似乎可以工作,但需要更长的时间:挂墙时间:1分22秒

def filter_codesets_3(codesets_onehot, code_selection):
    codesets_onehot = codesets_onehot.reset_index(drop=True)
    return codesets_onehot.loc[[set(codesets_onehot.columns[i]) == code_selection for i in codesets_onehot.values],:]

需要更长的时间才能得出空的结果:挂墙时间:1分钟5秒

def filter_codesets_4(codesets_onehot, code_selection):
    columns_of_interest = list(code_selection)
    len_coi = len(columns_of_interest)
    return codesets_onehot.loc[codesets_onehot[columns_of_interest].sum(axis=1) == len_coi]

这有效并且与第一个版本一样快:挂墙时间:28.7 s。优点是它可以查询更大的集而不会出现内存错误。

def filter_codesets_5(codesets_onehot, code_selection):
    return codesets_onehot[codesets_onehot[list(code_selection)].all(1)]

工作原理简单明了,需要:挂墙时间:30 s。我想很难仅靠熊猫来达到这个运行时间。

3 个答案:

答案 0 :(得分:3)

再次考虑这一点,看起来很简单,只需选择感兴趣的列并调用DataFrame.all

df_filtered = df[df[list(code_selection)].all(1)]

通过调用np.ndarray.all而不是DataFrame.all可以更快。

df_filtered = df[df[list(code_selection)].values.all(1)]

使用numba,我们可以走得更快:

from numba import njit, prange

@njit(parallel=True)
def get_mask(v, pos):
    mask = [True] * v.shape[0]
    for i in prange(v.shape[0]):
        for j in pos:
            mask[i] &= v[i, j]

    return np.array(mask)

性能

np.random.seed(0)
df = pd.DataFrame(np.random.choice(2, (100000, 1000), p=[0.1, 0.9]))
code_selection = set(np.random.choice(df.columns, 20))

%timeit df[df[list(code_selection)].all(1)]
%timeit df[df[list(code_selection)].values.all(1)]
%timeit df[get_mask(df.values, df.columns.get_indexer(code_selection))]

61.2 ms ± 2.02 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
52.6 ms ± 435 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
36.1 ms ± 460 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

答案 1 :(得分:2)

我会做这样的事情-

data = [
    [True, False, True],
    [False, True, False],
    [True, True, True],
    [True, True, False],
    [False, True, True]
]

df = pd.DataFrame(data, columns=['a', 'b', 'c'])

columns_of_interest = ['b', 'c']
len_coi = len(columns_of_interest)

df.loc[df[columns_of_interest].sum(axis=1) == len_coi]

这样的代码应该为您提供所需的行。

答案 2 :(得分:1)

这是一种方法:

df.loc[[set(df.columns[i]) == code_selection for i in df.values],:]

如果索引不起作用,请尝试将其删除:

df = df.reset_index(drop=True)
df.loc[[set(df.columns[i]) == code_selection for i in df.values],:]