如何在Dask DataFrame中读取并存储向量（List[float]）

我需要在Dask DataFrame中创建一个名为vector的列，数据源是一个规模为500k * 1536的大型numpy数组。用Pandas实现的代码如下：

import pandas as pd
import numpy as np

vectors = np.array([
    np.array([1, 2, 3]), 
    np.array([4, 5, 6]), 
    np.array([7, 8, 9])
])

df = pd.DataFrame({
    "vector": vectors.tolist()
})

df

生成的DataFrame结构符合预期，但仅加载就占用了34GB内存，结构如下：

我尝试了两种方案，但都存在问题：

方案一 #

import dask.dataframe as dd
import dask.array as da
import numpy as np

vectors = np.array([
    np.array([1, 2, 3]), 
    np.array([4, 5, 6]), 
    np.array([7, 8, 9])
])

vectors = da.from_array(vectors)
df = dd.from_dask_array(vectors)

df

该方案生成的DataFrame会将向量的每个元素拆分为单独一列，不符合需求。

方案二 #

import dask.dataframe as dd
import dask.array as da
import numpy as np

# vectors = np.load(dataset_path / "vectors.npy")

vectors = np.array([
    np.array([1, 2, 3]), 
    np.array([4, 5, 6]), 
    np.array([7, 8, 9])
])

df = dd.from_dask_array(da.from_array(vectors))
columns_to_drop = df.columns.tolist()
df["vector"] = df.apply(lambda row: tuple(row), axis=1, meta=(None, 'object'))
df = df.drop(columns=columns_to_drop)

df

该方案能得到正确结果，但实现繁琐且效率极低。

高效解决方案 #

可以借助Dask Array的map_blocks方法，直接将数组的每一行转换为列表/元组，再构造Dask DataFrame，避免逐行apply的开销：

import dask.dataframe as dd
import dask.array as da
import numpy as np

# 加载大型numpy数组，这里用示例数据代替
vectors = np.array([
    np.array([1, 2, 3]), 
    np.array([4, 5, 6]), 
    np.array([7, 8, 9])
])

# 转换为Dask Array，根据内存情况调整chunks大小
da_vectors = da.from_array(vectors, chunks=(10000, 1536))

# 定义每个块的处理函数，将每行转为列表
def process_chunk(chunk):
    return pd.DataFrame({"vector": chunk.tolist()})

# 用map_blocks生成Dask DataFrame
ddf = da_vectors.map_blocks(process_chunk, meta={"vector": object})

# 查看结果
ddf.compute()

这种方式利用了Dask的分块并行处理能力，既简洁又高效，同时能控制内存占用。

内容的提问来源于stack exchange，提问作者Mike Chaliy