Dask：并行分组，按顺序保存

Question

总结一下：如何对有限数量的组并行执行groupby操作，但将每个组的应用函数的结果写入磁盘？

我的问题：我正在尝试根据许多客户的信息为回归模型创建一个监督结构，这些信息分为几年。我必须从相同的客户端构建不同的模型，使用不同的输入 X 和标签 Y，因此我的想法是创建一个 X 和 Y 数据框，同时保存所有变量，并根据任务对每个变量进行切片。例如，X 可以保存来自薪水、年龄或性别的信息，但模型 1 将仅使用年龄和性别，而模型 2 仅使用薪水。

由于客户不是每年都在场，我只能使用从一个时期到下一个时期都在场的客户。 我没有为每对连续年份选择客户的交集，而是尝试连接整个信息并按客户 ID 执行 groupby 操作（然后按年份顺序过滤，例如使用期间差为 1 的行）。使用 Dask 执行此任务的问题是分布式工作人员的内存不足（即使将每个限制增加到 30Gb 之后）。 请注意，我正在为每个组创建一个新数据帧，因此我不会将计算减少到每个组的单个数字，因此会占用大量内存。

我目前正在做的是执行 groupby 操作，然后遍历 groupby 对象并按顺序写入磁盘：例如：

x_file=open('X.csv', 'w')
for name, group in concatenated_data.groupby('ID'):
   data_x=my_func(group) # In my real code, my_func returns x and y dataframes
   data_x.to_csv(x_file, header=None)
x_file.close()

依次应用my_func 写入数据，该my_func 为每个组选择x 和y。
我想要的是对受控数量的组执行操作（当时说3 个） ，并将每个组的结果写入磁盘（可能使用data_x.to_csv(x_file, single_file=True)）。
当然，我可以对 dask 数据帧执行相同的操作，并使用 get_group() 遍历 groupbpy 对象，但我不相信它会在保持内存检查的同时并行运行。

编辑：示例

# Lets say I have 3 csv files:
data=['./data_2016', './data_2017', './data_2018'] # Each file contains millions of rows (1 per client ID) and like 85 columns
# and certains variables
x_vars=['x1', 'x2', 'x3'] # x variables
y_vars= ['y1', 'y2', 'x1'] # note than some variables can be among x and y (like using today's salary to predict tomorrows salary)
data=[pd.read_csv(x) for x in data]

def func1(df_):
   # do some preprocessing stuff
   return df_
data=map(func1, data) # Some preprocessing and adding some columns (for example adding column for year)

concatenated_data=pd.concat(data, axis=1) # Big file, all clients from 2016-2018

def my_func(df_): # function applied above
   # order by year
   df_['Diff']=df_.year.diff() # calculating the difference among years 
   df['shifted']=df.Diff.shift(-1) # calculate shift of difference
   # For exammple, *client z* may be on 2016 and 2018, thus his year difference is 2. 
   # I can't use *clien z* x_vars to predict y (only a single period ahead regression)
   x=df_.loc[df_['shifted']==1, x_vars] # select only contigous years
   y=df_.loc[df_['Diff']==1, y_vars] # the same, but a year ahead of x
   return (x, y)

# ... Iteration over groupby object

我没有使用 groupby() 来减少，而是将单个大文件扩展为 x 和 y 数据帧，其中 y 在 x 之前保存了一段时间的信息。
如您所见，使用 dask 数据帧 groupby（为简单起见省略）将并行化 my_func 操作，但据我所知，也会等到所有操作节点完成，从而耗尽我的内存。 我想要对某些组执行my_func（理想情况下是内存可以容纳的最多），完成它们，保存到磁盘（没有与并行保存相关的问题），最后继续进行下一批组。

也许我可以使用一些 dask 延迟的对象，但我认为如果手动设置批次不会很好地利用我的内存。

Answer 1

我不确定这是不是你要找的东西

生成数据

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

n = 200
df = pd.DataFrame({"grp":np.random.choice(list("abcd"), n),
                   "x":np.random.randn(n),
                   "y":np.random.randn(n),
                   "z":np.random.randn(n)})

df.to_csv("file.csv", index=False)

# we will need later on
df.to_parquet("file.parquet", index=False)

熊猫解决方案

# we save our files on a given folder
fldr = "output1"
os.makedirs(fldr, exist_ok=True)

# we read the columns we need only
cols2read = ["grp", "x", "y"]

df = pd.read_csv("file.csv")
df = df[cols2read]

def write_file(x, fldr):
    name = x["grp"].iloc[0]
    x.to_csv(f"{fldr}/{name}.csv", index=False)

df.groupby("grp")\
  .apply(lambda x: write_file(x, fldr))

简单的解决方案

这基本上是一样的，但是我们需要将meta添加到我们的apply和compute中

# we save our files on a given folder
fldr = "output2"
os.makedirs(fldr, exist_ok=True)

# we read the columns we need only
cols2read = ["grp", "x", "y"]

df = pd.read_csv("file.csv")
df = df[cols2read]

def write_file(x, fldr):
    name = x["grp"].iloc[0]
    x.to_csv(f"{fldr}/{name}.csv", index=False)

df.groupby("grp")\
  .apply(lambda x: write_file(x, fldr), meta='f8')\
  .compute()

使用镶木地板

在这里我建议您使用镶木地板，因为这样会更有效率

cols2read = ["grp", "x", "y"]
df = dd.read_parquet("file.parquet",
                     columns=cols2read)

df.to_parquet("output3/",
              partition_on="grp")

在output3 中，您可以找到多个名为grp=a 的文件夹，依此类推。它们中的每一个最终都可能包含几个文件。但你可以用pd.read_parquet("output3/grp=a)阅读所有这些