Dask 将 dtype 设置为整数数组

Question

使用 Dask，我尝试创建一个具有整数类型列表的列。例如：

import dask.dataframe as dd
import pandas as pd

# Have an example Dask Dataframe
ddf = dd.from_pandas(pd.DataFrame({
    'id': [1, 2, 3, 4, 5],
    'name': ['Alice', 'Bob', 'Charlie', 'David', 'Emma'],
    'age': [25, 30, 35, 40, 45]
}), npartitions=1)

# now create an array type column
ddf["alist"] = ddf.apply(
    lambda k: [1, 0, 0], axis=1, meta=("alist", "list<item: int64>")
)

这个特殊案例失败了，因为：

类型错误：数据类型“列表”无法理解

最终我想写信给镶木地板：

ddf.to_parquet(
    "example",
    engine="pyarrow",
    compression="snappy",
    overwrite=True,
)

如果我指定的数据类型不正确，则会引发：

ValueError: Failed to convert partition to expected pyarrow schema:
    `ArrowInvalid('Could not convert [1, 2, 3] with type list: tried to convert to int64', 'Conversion failed for column alist with type object')`

Expected partition schema:
    id: int64
    name: large_string
    age: int64
    alist: int64
    __null_dask_index__: int64

Received partition schema:
    id: int64
    name: large_string
    age: int64
    alist: list<item: int64>
      child 0, item: int64
    __null_dask_index__: int64

This error *may* be resolved by passing in schema information for
the mismatched column(s) using the `schema` keyword in `to_parquet`.

Answer 1

也遇到过在 dask 数据框中直接创建和定义列表列并遇到相同问题的问题，因此我选择了另一种方法，其中涉及将列表序列化为 json 字符串以进行存储。

这是我的建议

import dask.dataframe as dd
import pandas as pd
import json

pdf = pd.DataFrame({
    'id': [1, 2, 3, 4, 5],
    'name': ['Alice', 'Bob', 'Charlie', 'David', 'Emma'],
    'age': [25, 30, 35, 40, 45],
    'alist': [[1, 0, 0] for _ in range(5)]  
})

pdf['alist'] = pdf['alist'].apply(json.dumps)

ddf = dd.from_pandas(pdf, npartitions=1)

import pyarrow as pa
schema = pa.schema([
    ('id', pa.int64()),
    ('name', pa.string()),
    ('age', pa.int64()),
    ('alist', pa.string()), 
])

ddf.to_parquet(
    "example",
    engine="pyarrow",
    compression="snappy",
    overwrite=True,
    schema=schema
)