使用 dask DataFrame 有效地加载和操作 csv

Question

我正在尝试使用dask.dataframe操作https://www.kaggle.com/raymondsunartio/6000-nasdaq-stocks-historical-daily-prices中的 csv 文件。 原来的 dataframe 有列'date'、'ticker'、'open'、'close'等...

我的目标是创建一个新的数据框，其中索引“日期”和列作为每个唯一代码的收盘价。

以下代码可以解决问题，但速度很慢， N = 6使用了将近一分钟。 我怀疑 dask 尝试在 for 循环中多次读取 CSV 文件，但我不知道如何让 go 更快。 我最初的猜测是在某处使用df.groupby('ticker')会有所帮助，但我对 pandas 不够熟悉。

import dask.dataframe as dd
from functools import reduce

def load_and_fix_csv(path: str, N: int, tickers: list = None) -> dd.DataFrame:
    raw = dd.read_csv(path, parse_dates=["date"])
    if tickers is None:
        tickers = raw.ticker.unique().compute()[:N] # Get unique tickers
    dfs = []
    for tick in tickers:
        tmp = raw[raw.ticker == tick][["date", "close"]] # Temporary dataframe from specific ticker with columns date, close
        dfs.append(tmp)
    df = reduce(lambda x, y: dd.merge(x, y, how="outer", on="date"), dfs) # Merge all dataframes on date
    df = df.set_index("date").compute()
    return df

各种帮助表示赞赏。 谢谢你。

Answer 1

我很确定你是对的，Dask 很可能会在每个循环中“回到井中”。 这是因为 Dask 构建了一个操作图并尝试将计算推迟到强制或必要时。 我喜欢做的一件事是使用Client.persist图形的读取操作：

from distributed import Client

client = Client()


def persist_load_and_fix_csv(path: str, N: int, tickers: list = None) -> dd.DataFrame:
    raw = dd.read_csv(path, parse_dates=["date"])

    # This "cuts the graph" prior operations (just the `read_csv` here)
    raw = client.persist(raw)
    if tickers is None:
        tickers = raw.ticker.unique().compute()[:N] # Get unique tickers
    dfs = []
    for tick in tickers:
        tmp = raw[raw.ticker == tick][["date", "close"]] # Temporary dataframe from specific ticker with columns date, close
        dfs.append(tmp)
    df = reduce(lambda x, y: dd.merge(x, y, how="outer", on="date"), dfs) # Merge all dataframes on date
    df = df.set_index("date").compute()
    return df

在 Kaggle session 中，我使用persist_load_and_fix_csv(csv_path, N=3)测试了这两个函数，并设法将时间缩短了一半。 通过仅保留最终使用的列，您还将获得更好的性能。

（注意：我发现，至少对我和我的代码而言，如果我开始看到.compute()出现在我应该退后一步并重新评估代码路径的函数中；我将其视为代码气味）