Расширение объекта Dask Dataframe - PullRequest
Новая
       7

Расширение объекта Dask Dataframe

0 голосов
/ 02 июля 2019

Я хочу расширить объект Dataframe dask.dataframe и добавить новые атрибуты / методы, не переписывая большинство существующих функций Dask.

Например, если у меня есть класс: 

from dask.dataframe import DataFrame
class MyExtendedDataframe(DataFrame):
    def __init__(self, df):
        super(DataFrame, self).__init__(df.dask, df._name, df._meta, df.divisions)
        self.myAttribute = "Dog"
    def myMethod(self):
        return self.dropna()

Я могу создать экземпляр этого объекта (ниже), который имеет все атрибуты объекта Dask DataFrame. 

df = dd.read_csv('mydata.csv')
mydf = MyExtendedDataframe(df)
print(type(mydf))

Это печатает: "MyExtendedDataframe"

Затем я могу вызвать функции Pandas-link для объекта MyExtendedDataframe. 

mydf = mydf['field1', 'field2']]
print(type(mydf))

Приведенные выше отпечатки: dask.dataframe.core.DataFrame и, очевидно, атрибут MyAttribute и метод MyMethod не являются членами результирующего экземпляра.

Другими словами, я хочу, чтобы мой новый класс имел все атрибуты класса dask.dataframe.core.DataFrame, но когда я вызываю методы, принадлежащие dask.dataframe.core.DataFrame, я хочу, чтобы он возвращал экземпляр MyExtendedDataframe с данной операцией завершено. Например: 

mydf = mydf[['field1', 'field2']]
print(type(mydf))

должен печатать «MyExtendedDataframe» и 

print(mydf.myAttribute)

должен вывести «Dog» после операции индексации.

Я знаю, что мог бы создать метод-обертку для всех стандартных методов Dask, который бы реконструировал и возвращал экземпляр моего класса. Но это похоже на хлопоты и потребует от меня обновления моего модуля, если Dask изменится. Есть ли более простой способ добиться такого поведения?

Другое возможное решение состоит в том, чтобы сделать мой dask.dataframe.core.DataFrame объект атрибутом моего пользовательского класса, но это выглядит неуклюже, а также для доступа и установки кадра данных. Мне бы постоянно приходилось использовать что-то вроде: 

mydf.dataframe = mydf.dataframe[['field1', 'field2']]
...