如何將函數應用於 dask 數據幀並返回多個值？

Question

在 Pandas 中，我使用下面的典型模式將矢量化函數應用於 df 並返回多個值。 只有當所述函數從單個任務產生多個獨立輸出時，這才是真正必要的。 看我過於瑣碎的例子：

import pandas as pd
df = pd.DataFrame({'val1': [1, 2, 3, 4, 5],
                   'val2': [1, 2, 3, 4, 5]})

def myfunc(in1, in2):
    out1 = in1 + in2
    out2 = in1 * in2
    return (out1, out2)

df['out1'], df['out2'] = zip(*df.apply(lambda x: myfunc(x['val1'], x['val2']), axis=1))

目前，我編寫了一個單獨的函數來對 Pandas df 進行分塊，並使用多處理來提高效率，但我想使用 dask 來完成此任務。 繼續這個例子，這是我在使用 dask 時如何運行向量化函數以返回單個值：

import dask.dataframe as dd
ddf = dd.from_pandas(df, npartitions=2)

def simple_func(in1, in2):
    out1 = in1 + in2
    return out1

df['out3'] = ddf.map_partitions(lambda x: simple_func(x['val1'], x['val2']), meta=(None, 'i8')).compute()

現在我想使用 dask 並返回兩個值，如熊貓示例中所示。 我試圖向元添加一個列表並返回一個元組，但只是得到錯誤。 這在 dask 中可能嗎？

Answer 1

我認為這里的問題源於您組合結果的方式不是很好。 理想情況下，您將df.apply與result_expand參數一起使用，然后使用df.merge 。 將此代碼從 Pandas 移植到 Dask 是微不足道的。 對於熊貓，這將是：

熊貓

import pandas as pd

def return_two_things(x, y):
    return (
        x + y,
        x * y,
    )

def pandas_wrapper(row):
    return return_two_things(row['val1'], row['val2'])

df = pd.DataFrame({
    'val1': range(1, 6),
    'val2': range(1, 6),
})

res = df.apply(pandas_wrapper, axis=1, result_type='expand')
res.columns = ['out1', 'out2']
full = df.merge(res, left_index=True, right_index=True)
print(full)

哪些輸出：

   val1  val2  out1  out2
0     1     1     2     1
1     2     2     4     4
2     3     3     6     9
3     4     4     8    16
4     5     5    10    25

達斯克

對於 Dask，將函數應用於數據並整理結果實際上是相同的：

import dask.dataframe as dd

ddf = dd.from_pandas(df, npartitions=2)
# here 0 and 1 refer to the default column names of the resulting dataframe
res = ddf.apply(pandas_wrapper, axis=1, result_type='expand', meta={0: int, 1: int})
# which are renamed out1, and out2 here
res.columns = ['out1', 'out2']
# this merge is considered "embarrassingly parallel", as a worker does not need to contact 
# any other workers when it is merging the results (that it created) with the input data it used.
full = ddf.merge(res, left_index=True, right_index=True)

print(full.compute())

輸出：

   val1  val2  out1  out2
0     1     1     2     1
1     2     2     4     4
2     3     3     6     9
3     4     4     8    16
4     5     5    10    25

Answer 2

聚會遲到了。 也許在提出問題時這是不可能的。

我不喜歡結束分配模式。 據我所知，dask 不允許像 Pandas 那樣進行新的列分配。

您需要將元值設置為您要返回的基本類型。 你可以很簡單地從我的測試中返回一個字典、元組、集合或列表。 無論如何，元似乎並不關心類型是否與返回對象的類型相匹配。

import pandas
import dask.dataframe

def myfunc(in1, in2):
    out1 = in1 + in2
    out2 = in1 * in2
    return (out1, out2)

df = pandas.DataFrame({'val1': [1, 2, 3, 4, 5],
                   'val2': [1, 2, 3, 4, 5]})
ddf = dask.dataframe.from_pandas(df, npartitions=2)

df['out1'], df['out2'] = zip(*df.apply(lambda x: myfunc(x['val1'], x['val2']), axis=1))


output = ddf.map_partitions(lambda part: part.apply(lambda x: myfunc(x['val1'], x['val2']), axis=1), meta=tuple).compute()

out1, out2 = zip(*output)

ddf = ddf.assign(out1 = pandas.Series(out1))
ddf = ddf.assign(out2 = pandas.Series(out2))

print('\nPandas\n',df)
print('\nDask\n',ddf.compute())
print('\nEqual\n',ddf.eq(df).compute().all())

輸出：

Pandas
    val1  val2  out1  out2
0     1     1     2     1
1     2     2     4     4
2     3     3     6     9
3     4     4     8    16
4     5     5    10    25

Dask
    val1  val2  out1  out2
0     1     1     2     1
1     2     2     4     4
2     3     3     6     9
3     4     4     8    16
4     5     5    10    25

Equal
val1    True
val2    True
out1    True
out2    True
dtype: bool

注意到 map_partition 的 lambda 返回是較大數據幀的一個分區（在這種情況下，基於您的 npartitions 值）會有所幫助。 然后你會像對待任何其他數據框一樣對待你的 .apply()。