June 6, 2021

Капризы условных конструкций в Pandas и NumPy

Рассмотрим, как избежать ключевую проблему, которая может возникнуть при использовании излюбленной условной конструкции - where пакета NumPy со значениями из DataFrame.

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

import pandas as pd
import numpy as np
df = pd.DataFrame({'data':np.arange(1,11),
                   'date_time':pd.date_range(start='2019-01', periods=10, freq='M')})

Попытка заменить в столбце date_time значения даты приведет к следующему результату:

np.where(df['date_time']>pd.Timestamp('2019-03-31'), df['date_time'], pd.Timestamp('2019-03-31'))

При этом типы данных отдельных элементов в столбце и у нового значения совпадают, а тип данных всего столбца - datetime64 модуля NumPy:

Для решения проблемы можно преобразовать значение новой даты в datetime64:

np.where(df['date_time']>pd.Timestamp('2019-03-31'), df['date_time'], np.datetime64('2019-03-31'))

Либо можно использовать аналогичные where методы библиотеки Pandas. В частности, существуют два метода датафрейма, позволяющие добиться цели - where и mask:

df['date_time'].where(df['date_time']>pd.Timestamp('2019-03-31'), pd.Timestamp('2019-03-31'))

where применяет заменяемое значение из второго аргумента, когда условие ложно, а mask - когда истинно.

df['date_time'].mask(df['date_time']<=pd.Timestamp('2019-03-31'), pd.Timestamp('2019-03-31'))