Мутирующие ячейки в большом DataFrame Polars (Python) с использованием iter_rows вызывают ошибку сегментации

Содержание

Вопрос или проблема
Ответ или решение
Решение
Код
Объяснение
Заключение

Вопрос или проблема

Мутирующие ячейки в большом DataFrame Polars (Python) с использованием iter_rows вызывают ошибку сегментации

У меня есть большой DataFrame, который выглядит следующим образом:

df_large = pl.DataFrame({'x':['h1','h2','h2','h3'],
                         'y':[1,2,3,4], 
                         'ind1':['0/0','1/0','1/1','0/1'], 
                         'ind2':['0/1','0/2','1/1','0/0'] }).lazy()
df_large.collect()

|   x   |   y   | ind_1 | ind_2   | 
|_______|_______|_______|_________|
| "h1"  |   1   | "0/0" |  '0/1'  |       
| "h2"  |   2   | "1/0" |  '0/2'  | 
| "h2"  |   3   | "1/1" |  '1/1'  | 
| "h3"  |   4   | "0/1" |  '0/0'  |

df_large содержит координаты [x (строка), y (целое число)] и строковые значения для многих индивидуумов [ind_1, ind_2,…]. Он очень большой, поэтому мне нужно читать CSV файл как ленивый DataFrame. Кроме того, у меня есть маленький DataFrame, который выглядит так:

df_rep = pl.DataFrame({'x':['h1','h2','h2'],
                       'y':[1,2,2], 
                       'ind':['ind1','ind1','ind2']})
df_rep

|   x   |   y   |  indvs  | 
|_______|_______|_________|
| "h1"  |   1   | "ind_1" |     
| "h2"  |   2   | "ind_1" |
| "h2"  |   2   | "ind_2" |

Мне нужно изменить значения для столбцов ind_k в df_large, когда они появляются в df_rep.

Я написал следующий код для этого:

for row in df_rep.iter_rows():
    df_large = df_large.with_columns(
                    pl.when(pl.col('x') == row[0],
                       pl.col('y')    == row[1])
                       .then(pl.col(row[2]).str.replace_all('(.)/(.)','./.'))
                       .otherwise(pl.col(row[2]))
                      .alias(row[2])
                  )
df_large.collect()

|   x   |   y   | ind_1 | ind_2   | 
|_______|_______|_______|_________|
| "h1"  |   1   | "./." |  '0/1'  |       
| "h2"  |   2   | "./." |  './.'  | 
| "h2"  |   3   | "1/1" |  '1/1'  | 
| "h3"  |   4   | "0/1" |  '0/0'  |

Этот метод, хоть и медленный, работает для подмножества большого набора данных. Однако Polars вызывает ошибку сегментации, когда применяется к полному набору данных. Я надеялся, что вы сможете дать мне совет, как решить эту проблему. Альтернативный метод достижения моей цели без использования iter_rows() был бы идеален!

Я новичок в Polars и буду очень признателен за любые комментарии. Я довольно давно застрял на этой проблеме 🙁

Ответ или решение

Для решения проблемы, связанной с замещением значений в больших датафреймах Polars без использования iter_rows(), который может вызвать сегментационные ошибки, мы можем воспользоваться возможностями языка, предоставляемыми Polars, для более эффективной обработки данных векторизированным способом.

Решение

Вместо итерации по строкам с использованием iter_rows(), мы можем применить методы join и conditional update. Идея заключается в том, чтобы объединить два датафрейма, а затем изменять значения в целевом датафрейме на основе условий, указанных в исходном датафрейме.

Вот шаги, чтобы это выполнить:

Сначала подготовим оба датафрейма.
Затем сделаем join датафрейма df_large с df_rep, а затем применим необходимые изменения.

Код

import polars as pl

# Создание большого датафрейма
df_large = pl.DataFrame({
    'x': ['h1', 'h2', 'h2', 'h3'],
    'y': [1, 2, 3, 4],
    'ind1': ['0/0', '1/0', '1/1', '0/1'],
    'ind2': ['0/1', '0/2', '1/1', '0/0']
}).lazy()

# Создание меньшего датафрейма
df_rep = pl.DataFrame({
    'x': ['h1', 'h2', 'h2'],
    'y': [1, 2, 2],
    'ind': ['ind1', 'ind1', 'ind2']
})

# Преобразуем df_large в eagerness (жадный) режим
df_large = df_large.collect()

# Объединяем датафреймы по столбцам 'x' и 'y', а затем применяем метод замены
df_result = df_large.join(df_rep, on=['x', 'y'], how='left') \
    .with_columns(
        pl.when(pl.col('ind').is_not_null())  # Проверяем на совпадение
        .then(pl.col('ind').str.replace_all('(.)/(.)', './.'))  # Заменяем
        .otherwise(pl.col('ind1')).alias('ind1'),  # Если нет, оставляем исходные значения
        pl.when(pl.col('ind').is_not_null())
        .then(pl.col('ind').str.replace_all('(.)/(.)', './.'))
        .otherwise(pl.col('ind2')).alias('ind2')
    ).select(['x', 'y', 'ind1', 'ind2'])

# Итоговый результат
print(df_result.collect())

Объяснение

Объединение (Join): Мы объединили два датафрейма по столбцам x и y, результаты объединения содержат информацию о том, какой ind относится к каждой записи в df_large.
Замена на основе условия: С помощью pl.when(...).then(...).otherwise(...) мы производим проверку. Для каждого индикатора (ind1, ind2) мы заменяем значение в зависимости от того, присутствует ли оно в df_rep.
Избор столбцов: Наконец, мы выбираем необходимые столбцы, чтобы вернуть конечный результат.

Заключение

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

Если у вас возникли дополнительные вопросы или потребуется дальнейшая помощь, пожалуйста, не стесняйтесь обращаться!