Dask 将列存储数据高效索引为 dask 数组

Question

我正在尝试找到一种有效的方法来索引数据集，该数据集按使用 Dask (Python) 加载的 zarr 数组中的列存储。我必须对其执行一些操作，其中根据另一个数组对所有列进行索引（即采用

column[ind] for column in columns

）。我以为这是一件微不足道的事情，但我还没有找到有效的解决方案。

1. 我尝试的第一件事就是简单地使用 dask 数据框，因为这似乎就是我所需要的。因此，我将二维数组中的列连接起来，并从中创建了一个数据框。但是，我经常访问每列的数据，并且我注意到从数据帧访问单个列仍然会从磁盘加载所有数据（显然二维数组中的所有列都是按块加载的，然后只保留选定的列） 。这对我来说是绝对不行的。我大约有10000列，效率太低了。

2. 我尝试的第二件事是简单地坚持dask数组（列）的字典。这显然对我来说很好，因为我主要必须对数组执行操作。请记住，了解块大小对我来说非常重要，因为我经常对数组执行重塑。因此，在索引数组时（如

   df = {k : v[ind] for k, v in df.items()}

   ），我遇到了以下问题：
   • 如果

     ind

     是一个 numpy 数组，则索引操作只会耗尽我机器的 RAM (64GB)，尽管块的大小约为 20MB。
   • 如果

     ind

     是一个 dask 数组，我需要在每个数组上调用

     compute_chunk_sizes

     来计算结果形状。然而，这是非常低效的，因为 dask 对数据集中的每个数组重复相同的块计算，而它可以简单地一次性完成。

总而言之，我的问题是：你将如何解决这个问题？我愿意考虑其他方式，请记住，柱状访问对我来说很重要。

Answer 1

我通过将所有一维数组连接成一个二维数组，然后对第一维进行切片，最后从该数组中提取列来解决这个问题。一次访问 1 个数组仍然高效，不涉及从磁盘加载所有数据。