如何指定列数据类型

Question

我有以下代码：

import polars as pl
from typing import NamedTuple


class Event(NamedTuple):
    name: str
    description: str


def event_table(num) -> list[Event]:
    events = []
    for i in range(num):
        events.append(Event("name", "description"))
    return events


data = {"events": [1, 2]}
df = pl.DataFrame(data).select(events=pl.col("events").map_elements(event_table))

"""
shape: (2, 1)
┌───────────────────────────────────┐
│ events                            │
│ ---                               │
│ list[struct[2]]                   │
╞═══════════════════════════════════╡
│ [{"name","description"}]          │
│ [{"name","description"}, {"name"… │
└───────────────────────────────────┘
"""

但是如果第一个列表为空，我会得到一个

list[list[str]]

而不是我需要的

list[struct[2]]

：

data = {"events": [0, 1, 2]}
df = pl.DataFrame(data).select(events=pl.col("events").map_elements(event_table))
print(df)

"""
shape: (3, 1)
┌───────────────────────────────────┐
│ events                            │
│ ---                               │
│ list[list[str]]                   │
╞═══════════════════════════════════╡
│ []                                │
│ [["name", "description"]]         │
│ [["name", "description"], ["name… │
└───────────────────────────────────┘
"""

我尝试使用

return_dtype

函数的

map_elements

，例如：

data = {"events": [0, 1, 2]}
df = pl.DataFrame(data).select(
    events=pl.col("events").map_elements(
        event_table,
        return_dtype=pl.List(pl.Struct({"name": pl.String, "description": pl.String})),
    )
)

但这失败了：

Traceback (most recent call last):
  File "script.py", line 18, in <module>
    df = pl.DataFrame(data).select(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/polars/dataframe/frame.py", line 8193, in select
    return self.lazy().select(*exprs, **named_exprs).collect(_eager=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".venv/lib/python3.11/site-packages/polars/lazyframe/frame.py", line 1943, in collect
    return wrap_df(ldf.collect())
                   ^^^^^^^^^^^^^
polars.exceptions.SchemaError: expected output type 'List(Struct([Field { name: "name", dtype: String }, Field { name: "description", dtype: String }]))', got 'List(List(String))'; set `return_dtype` to the proper datatype

我怎样才能让它发挥作用？如果第一个列表为空，我需要此列的类型为

list[struct[2]]

事件。

Answer 1

正如评论中所述，这似乎是极地的一个错误。但是，我设法通过在

pl.Field

dtype 的定义中使用

pl.Struct

来运行它。

data = {"events": [0, 1, 2]}
pl.DataFrame(data).select(
    events=pl.col("events").map_elements(
        event_table,
        return_dtype=pl.List(pl.Struct([pl.Field("name", pl.String), pl.Field("description", pl.String)])),
    )
)

shape: (3, 1)
┌───────────────────────────────────┐
│ events                            │
│ ---                               │
│ list[struct[2]]                   │
╞═══════════════════════════════════╡
│ []                                │
│ [{"my_name","my_desc"}]           │
│ [{"my_name","my_desc"}, {"my_nam… │
└───────────────────────────────────┘

有趣的是，在 Jupyter 中运行以下单元格之后不再引发错误。

data = {"events": [0, 1, 2]}
df = pl.DataFrame(data).select(
    events=pl.col("events").map_elements(
        event_table,
        return_dtype=pl.List(pl.Struct({"name": pl.String, "description": pl.String})),
    )
)