dask 数据框中的 df.groupby(...).apply(...) 函数

Question

我使用 Python dask 处理大型 csv 面板数据集 (15+GB)，我需要执行 groupby(...).apply(...) 函数来删除每天每只股票的最后观察值。我的数据集看起来像

 stock     date     time   spread  time_diff 
  VOD      01-01    9:05    0.01     0:07     
  VOD      01-01    9:12    0.03     0:52     
  VOD      01-01   10:04    0.02     0:11
  VOD      01-01   10:15    0.01     0:10     
  VOD      01-01   10:25    0.03     0:39  
  VOD      01-01   11:04    0.02    22:00 
  VOD      01-02    9:04    0.02     0:05
  ...       ...     ...     ....     ...
  BAT      01-01    13:05   0.04    10:02
  BAT      01-02    9:07    0.05     0:03
  BAT      01-02    9:10    0.06     0:04
  ...       ...     ...     ....     ...

如果数据框在 pandas 中，那么这可以通过

df_new=df_have.groupby(['stock','date'], as_index=False).apply(lambda x: x.iloc[:-1])

此代码适用于 pandas df。但是，我无法在 dask 数据框中执行此代码。我做了以下尝试。

ddf_new=ddf_have.groupby(['stock','date']).apply(lambda x: x.iloc[:-1]).compute()

或

ddf_new=ddf_have.groupby(['stock','date']).apply(lambda x: x.iloc[:-1], meta=('stock' : 'f8')).compute()

或

ddf_new=ddf_have.groupby(['stock','date']).apply(lambda x: x.iloc[:-1], meta=meta).compute()

不幸的是，它们都不起作用。谁能帮我为 dask dataframe 获取正确的代码？谢了

Answer 1

我认为对于您的具体情况，问题是您分配的meta。这应该可以。

import pandas as pd
import numpy as np
import dask.dataframe as dd

dates = pd.date_range(start='2019-01-01',
                      end='2019-12-31',
                      freq='5T')

out = []
for stock in list("abcdefgh"):
    df = pd.DataFrame({"stock":[stock]*len(dates),
                       "date":dates,
                       "spread":np.random.randn(len(dates))})
    df["time_diff"] = df["date"].diff().shift(-1)
    df["time"] = df["date"].dt.time.astype(str)
    df["date"] = df["date"].dt.date.astype(str)
    out.append(df)
df = pd.concat(out, ignore_index=True)

del out

ddf = dd.from_pandas(df, npartitions=4)

out = ddf.groupby(['stock','date']).apply(lambda x: x[:-1],
                                          meta={"stock":"str",
                                                "date":"str",
                                                "spread":"f8",
                                                "time_diff":"str",
                                                "time":"str"})
out = out.compute().reset_index(drop=True)

如果您可以按股票日期很好地对文件进行分区并保存在to_parquet 中，那么您可以获得更好的性能，因为您可以使用map_partitions 而不是apply。