基於多列條目過濾行

Question

我有一個包含數百萬個條目的數據框，看起來像這樣：

鉻	開始	替代
1	21651521	一種
1	41681521	噸
1	41681521	噸
...	...	...
X	423565	噸

我目前正在嘗試計算同時匹配多個條件的行數，即Chr==1 、 Start==41681521和Alt==T 。 現在我正在使用這種語法，它工作得很好，但似乎不是 Pythonic 並且我認為也很慢。

num_occurrence = sum((df["Chr"] == chrom) & 
                      (df["Start"] == int(position)) & 
                      (df["Alt"] == allele))

有沒有人有比我更合適的方法？ 任何幫助深表感謝！

干杯!

Answer 1

使用DataFrame.all + Series.sum ：

res = (df[["Chr", "Start", "Alt"]] == [chrom, int(position), allele]).all(1).sum()

例如：

import pandas as pd

# toy data
df = pd.DataFrame(data=[[1, 21651521, "A"], [1, 41681521, "T"], [1, 41681521, "T"]], columns=["Chr", "Start", "Alt"])
chrom, position, allele = 1, "21651521", "A"


res = (df[["Chr", "Start", "Alt"]] == [chrom, int(position), allele]).all(1).sum()
print(res)

輸出

1

Answer 2

備選方案 1： pd.DataFrame.query()

您可以使用查詢（另請參閱此處的說明性示例）：

expr = "Chr=={chr} & Start=={pos} & Alt=='{alt}'"
ret = df.query(expr.format(chr=chrom, pos=int(position), alt=allele))

在我的實驗中，這已經導致了相當大的加速。

進一步優化這需要有關所涉及數據類型的附加信息。 您可以嘗試以下幾種方法：

方案二：查詢排序后的數據

如果您有能力在查詢之前對 DataFrame 進行排序，則可以使用pd.Series.searchsorted() 。 這是一種可能的方法：

def query_sorted(df, chrom, position, allele):
    """
    Returns index of the matches.
    """
    assert df["Start"].is_monotonic_increasing
    i_min, i_max = df["Start"].searchsorted([position, position+1])
    df = df.iloc[i_min:i_max]
    return df[(df["Chr"] == chrom) & (df["Alt"] == allele)].index

# Usage: first sort df by column "Start", then query:
df = df.sort_values("Start")
ret_index = query_sorted(df, chrom, position, allele)
print(len(ret_index))

備選方案 3：使用哈希

另一個想法是使用哈希。 同樣，這需要預先進行一些計算，但它會大大加快查詢速度。 下面是一個基於pd.util.hash_pandas_object()的例子：

def query_hash(df, chrom, position, allele):
    """
    Returns a view on df
    """
    assert "hash" in df
    dummy = pd.DataFrame([[chrom, position, allele]])
    query_hash = pd.util.hash_pandas_object(dummy, index=False).squeeze()
    return df[df["hash"] == query_hash].index

# Usage: first compute hashes over the columns of interest, then query
df["hash"] = pd.util.hash_pandas_object(df[["Chr", "Start", "Alt"]], 
                                        index=False)
ret_index = query_hash(df, chrom, position, allele)
print(len(ret_index))

備選方案 4：使用多索引

在通過索引訪問行時，Pandas 也使用散列操作。 因此，無需像前面的替代方案那樣顯式計算哈希值，而是可以在查詢之前簡單地設置 DataFrame 的索引。 （因為將所有列設置為索引會導致一個空的 DataFrame，我首先創建一個虛擬列。對於具有附加列的真實 DataFrame，這可能沒有必要。）

df["dummy"] = None
df = df.set_index(["Chr", "Start", "Alt"])
df = df.sort_index()  # Improves performance
print(len(df.loc[(chrom, position, allele)])
# Interestingly, chaining .loc[] is about twice as fast
print(len(df.loc[chrom].loc[position].loc[allele]))

請注意，使用一個索引值映射到許多記錄的索引並不總是一個好主意。 此外，這種方法比替代方法 3 慢，這表明 Pandas 在這里做了一些額外的工作。

當然還有更多方法可以改善這一點，但替代方法將取決於您的特定需求。

結果

我在運行 Python 3.8、Pandas 1.2.4 和 IPython 7.24.1 的 MacBook Pro（2015 年中）上用 n=10M 樣本進行了測試。 請注意，性能評估取決於問題的大小。 因此，對於不同的問題規模，方法的相對評估會發生變化。

# original (sum(s)):  1642.0 ms ± 19.1 ms
# original (s.sum()):  639.0 ms ± 21.9 ms
# query():             175.0 ms ±  1.1 ms 
# query_sorted():       17.5 ms ± 60.4 µs
# query-hash():         10.6 ms ± 62.5 µs
# multi-index:          71.5 ms ±  0.7 ms 
# multi-index (seq.):   36.5 ms ±  0.6 ms

---

執行

這就是我構建數據並比較不同方法的方式。

import numpy as np 
import pandas as pd

# Create test data
n = int(10*1e6)
df = pd.DataFrame({"Chr": np.random.randint(1,23+1,n),
                   "Start": np.random.randint(100,999, n),
                   "Alt": np.random.choice(list("ACTG"), n)})

# Query point
chrom, position, allele = 1, 142, "A"

# Create test data
n = 10000000
df = pd.DataFrame({"Chr": np.random.randint(1,23+1,n),
                   "Start": np.random.randint(100,999, n),
                   "Alt": np.random.choice(list("ACTG"), n)})

# Query point
chrom, position, allele = 1, 142, "A"

# Measure performance in IPython
print("original (sum(s)):")
%timeit sum((df["Chr"] == chrom) & \
            (df["Start"] == int(position)) & \
            (df["Alt"] == allele))

print("original (s.sum()):")
%timeit ((df["Chr"] == chrom) & \
         (df["Start"] == int(position)) & \
         (df["Alt"] == allele)).sum()

print("query():")
%timeit len(df.query(expr.format(chr=chrom, \
                                 pos=position, \
                                 alt=allele)))

print("query_sorted():")
df_sorted = df.sort_values("Start")
%timeit query_sorted(df_sorted, chrom, position, allele)

print("query-hash():")
df_hash = df.copy()
df_hash["hash"] = pd.util.hash_pandas_object(df_hash[["Chr", "Start", "Alt"]], 
                                             index=False)
%timeit query_hash(df_hash, chrom, position, allele)

print("multi-index:")
df_multi = df.copy()
df_multi["dummy"] = None
df_multi = df_multi.set_index(["Chr", "Start", "Alt"]).sort_index()
%timeit df_multi.loc[(chrom, position, allele)]
print("multi-index (seq.):")
%timeit len(df_multi.loc[chrom].loc[position].loc[allele])