熊貓按組過濾最小

Question

我有一個具有以下格式的數據框：

d = {'id1': ['a', 'a', 'b', 'b',], 'id2': ['a', 'b', 'b', 'c'], 'score': ['1', '2', '3', '4']}
    df = pd.DataFrame(data=d)


print(df)
     id1    id2    score
0     a      a       1        
1     a      b       2             
3     b      b       3        
4     b      c       4

數據框有超過 10 億行，它表示 id1 和 id2 列中對象之間的成對距離分數。 我不需要所有對象對組合，對於 id1 中的每個對象（大約有 40k 個唯一 ID）我只想保留前 100 個最近（最小）距離分數

我正在運行的代碼如下：

df = df.groupby(['id1'])['score'].nsmallest(100)

這段代碼的問題是我每次嘗試運行它時都會遇到內存錯誤

MemoryError: Unable to allocate 8.53 GiB for an array with shape (1144468900,) and data type float64

我假設這是因為在后台熊貓現在正在為 group by 的結果創建一個新的數據框，但現有的數據框仍然保存在內存中。

我只取每個 id 的前 100 個的原因是為了減小數據框的大小，但我似乎在執行該過程時實際上占用了更多空間。

有沒有辦法可以過濾掉這些數據但不占用更多內存？

所需的輸出將是這樣的（假設前 1 名而不是前 100 名）

     id1    id2    score
0     a      a       1        
1     b      b       3            

關於原始 df 的一些附加信息：

df.count()
permid_1    1144468900
permid_2    1144468900
distance    1144468900
dtype: int64

df.dtypes
permid_1      int64
permid_2      int64
distance    float64

df.shape
dtype: object
(1144468900, 3)

id1 & id2 unique value counts: 33,830

Answer 1

由於缺少您的數據，我無法測試此代碼，但也許可以嘗試以下操作：

indicies = []
for the_id in df['id1'].unique():
    scores = df['score'][df['id1'] == the_id]
    min_subindicies = np.argsort(scores.values)[:100]  # numpy is raw index only
    min_indicies = scores.iloc[min_subindicies].index  # convert to pandas indicies
    indicies.extend(min_indicies)

df = df.loc[indicies]

描述性地，在每個唯一 ID ( the_id ) 中，提取匹配的分數。 然后找到最小的 100 個原始索引。選擇這些索引，然后從原始索引映射到 Pandas 索引。 將 Pandas 索引保存到您的列表中。 然后在最后，pandas 索引上的子集。

iloc確實需要一個列表輸入。 some_series.iloc應該與some_series.values正確對齊，這應該允許它工作。 像這樣間接存儲索引應該可以顯着提高內存效率。

df['score'][df['id1'] == the_id]應該比df.loc[df['id1'] == the_id, 'score'] 。 它不是獲取整個數據框並對其進行屏蔽，而是僅獲取數據框的 score 列並將其屏蔽以匹配 ID。 如果您想立即釋放更多內存，您可能希望在每個循環結束時del scores 。

Answer 2

您可以嘗試以下操作：

df.sort_values(["id1", "scores"], inplace=True)
df["dummy_key"] = df["id1"].shift(100).ne(df["id1"])

df = df.loc[df["dummy_key"]]

1. 您按升序排序（最小的在頂部），先分組，然后按分數。

2. 您添加列以指示當前id1是否與后面的 100 行不同（如果不是 - 您的行按順序是 101+）。

3. 您按 2 中的列過濾。

Answer 3

正如 Aryerez 在評論中概述的那樣，您可以執行以下操作：

closest = pd.concat([df.loc[df['id1'] == id1].sort_values(by = 'score').head(100) for 
    id1 in set(df['id1'])])

你也可以這樣做

def get_hundredth(id1):
    sub_df = df.loc[df['id1'] == id1].sort_values(by = 'score')
    return sub_df.iloc[100]['score']

hundredth_dict = {id1: get_hundredth(id1) for id1 in set(df['id1'])}

def check_distance(row):
    return row['score'] <= hundredth_dict[row['id1']]

closest = df.loc[df.apply(check_distance, axis = 1)

另一種策略是查看過濾掉超過閾值的距離如何影響數據幀。 也就是說，取

   low_scores = df.loc[df['score']<threshold]

對於某些合理的閾值，這是否會顯着減小數據幀的大小？ 您需要一個閾值，使數據框足夠小以使用，但為每個id1留下最低的 100 分。

您可能還想研究根據距離度量可以進行哪些優化。 可能有專門針對余弦相似度的算法。

Answer 4

對於具有33,830唯一值計數的給定形狀(1144468900, 3) ， id1和id2列是分類列的良好候選者，將它們轉換為分類數據類型，這將減少大約1144468900/33,830 = 33,830倍的內存需求兩列，然后執行您想要的任何聚合。

df[['id1', 'id2']] = df[['id1', 'id2']].astype('category')
out = df.groupby(['id1'])['score'].nsmallest(100)