Самый эффективный способ условно выбрать мой большой df
У меня большой DF (~ 35 миллионов строк), и я пытаюсь создать новый df, случайным образом выбирая две строки из каждого уникального идентификатора кластера (~ 1,8 миллиона уникальных идентификаторов кластера) - одна строка должна иметь метку 0 и одна строка. должна иметь метку 1 (иногда метка только одна, поэтому сначала нужно проверить, присутствуют ли обе метки в кластере). Для справки, мой набор данных имеет 3 основных столбца: «вложения», «cluster_ID», «метка».
Я обнаружил, что этот процесс занимает гораздо больше времени, чем я ожидал (более 15 часов), и хочу знать, есть ли способ оптимизировать мой код.
Пожалуйста, обрати внимание: В моем исходном df есть только две метки: 1 и 0. В итоге у меня должно получиться более 1,8 миллиона строк. Я хотел бы выбрать 2 строки для каждого кластера (одна метка 1, а другая — метка 0). В некоторых кластерах есть смешанный набор строк с обеими метками, но в некоторых есть только строки с меткой 1. Если это так: мне нужны 2 строки из этого кластера (по одной для каждой метки) - если верно последнее : Мне нужна только одна строка из этого кластера. Следовательно: в итоге у меня должно получиться от 1,8 до 3,6 миллиона строк.
import pandas as pd
import random
# Create an empty list to store selected rows
selected_rows = []
# Iterate over unique cluster IDs
for cluster_id in result_df_30['cluster_ID'].unique():
# Filter the DataFrame for the current cluster ID
df_cluster = result_df_30[result_df_30['cluster_ID'] == cluster_id]
# Filter rows with label 0 and 1
df_label_0 = df_cluster[df_cluster['label'] == 0]
df_label_1 = df_cluster[df_cluster['label'] == 1]
# Sample rows if they exist
if not df_label_0.empty:
sample_label_0 = df_label_0.sample(n=1, random_state=42)
selected_rows.append(sample_label_0)
if not df_label_1.empty:
sample_label_1 = df_label_1.sample(n=1, random_state=42)
selected_rows.append(sample_label_1)
# Concatenate the selected rows into a single DataFrame
selected_rows_df = pd.concat(selected_rows)
selected_rows_df
Это крайне неэффективно: вы сканируете весь фрейм данных почти 2 миллиона раз только для того, чтобы получить группы, по сути, это квадратичное время. Вместо этого используйте groupby
@juanpa.arrivillaga Не могли бы вы написать ответ, чтобы я мог увидеть ваше решение?
Ответы 1
Обновленные требования: выберите по 1 строке из каждой комбинации меток кластера.
sample = (
df.sample(frac=1)
.groupby(["cluster_ID", "label"], as_index=False)
.first()
)
Исходные требования: выборка только из кластеров с метками 0 и 1.
Для этого требуется, чтобы ваш фрейм данных был последовательно проиндексирован (0, 1,..., 35m). С df = df.reset_index(drop=True) это очень просто.
# Some test data
n = 35_000_000
df = pd.DataFrame(
{
"embeddings": np.random.rand(n),
"cluster_ID": np.random.randint(0, 1_800_000, n),
"label": np.random.randint(0, 2, n),
}
)
tmp = (
# since you only care about rows with label 0 or 1
df[df["label"].isin([0, 1])]
# shuffle the rows. This is where the randomness comes from
.sample(frac=1)
# reset the index
.reset_index()
# for each Cluster ID and label, get the number of the first row
.pivot_table(index = "cluster_ID", columns = "label", values = "index", aggfunc = "first")
)
# Filter for clusters that have both 0 and 1, then get their row numbers
idx = tmp[tmp.notna().all(axis=1)].to_numpy("int").flatten()
# and your random sample
sample = df.iloc[idx]
В моем исходном df есть только две метки: 1 и 0. Я изменил код соответствующим образом, но после запуска вашего кода у меня должно получиться более 1,8 миллиона строк, но в итоге у меня получается именно такое количество. Мне нужно взять по 2 (или 1, если необходимо) образца для каждого кластера. В некоторых кластерах есть смешанный набор строк с обеими метками, в некоторых есть только строки с меткой 1. Если первый случай: мне нужны 2 строки из этого кластера, если второй: мне нужна только одна строка из этот кластер. Следовательно: в итоге у меня должно получиться от 1,8 до 3,6 миллиона строк.
Мой sample.shape возвращает примерно то, что вы описали. Если быть точным, в нем 3 599 554 строки. Что sample.shape возвращает вам?
Используя мой реальный df, я получил 31 386 строк.
Я думаю, проблема в этой строке: # Отфильтруйте кластеры, которые имеют как 0, так и 1, затем получите номера их строк idx = tmp[tmp.notna().all(axis=1)].to_numpy("int").flatten () Я не хочу фильтровать только кластеры, имеющие как 0, так и 1. Я хочу, чтобы каждый кластер был представлен максимум двумя метками, если это возможно (если нет, то 1 метка подойдет). Как мне изменить эту последнюю строку, чтобы она была точной?
Из вашего первоначального вопроса у меня сложилось впечатление, что вам нужны только кластеры, имеющие метки 0 и 1. Вопрос обновлен, хотя вы отметили его как отвеченный.
Чем вы занимаетесь гораздо больше времени, чем предполагалось?