эффективно запрашивать большое количество логических столбцов в pandas.DataFrame

Question

В моем фрейме данных каждая строка представляет набор кодов с горячим кодированием, поэтому в фрейме данных содержится большое количество логических столбцов.

Я хочу выбрать все строки, которые содержат подмножество кодов, то есть все строки, в которых значения для данного набора столбцов равны 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 с).Для больших наборов происходит сбой с ошибкой памяти:

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 с.Преимущество состоит в том, что он может запрашивать большие наборы без ошибок памяти.

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

Работает, прост и лаконичен и занимает: Время стены: 30 с.Полагаю, трудно быть ниже этого времени только с пандами.

Answer 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],:]

Answer 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]

Код, подобный этому, должен дать вам нужные вам строки.

Answer 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)