如何处理polars中用户定义函数的多行结果？

Question

我想在polars中把文本行解析成多个列和行，并使用用户定义的函数。

import polars as pl
df = pl.DataFrame({'file': ['aaa.txt','bbb.txt'], 'text': ['my little pony, your big pony','apple+banana, cake+coke']})

def myfunc(p_str: str) -> list:
    res = []
    for line in p_str.split(','):
        x = line.strip().split(' ')
        res.append({f'word{e+1}': w for e, w in enumerate(x)})
    return res

如果我只是运行一个测试，结果是好的，会创建一个字典的列表：

myfunc(df['text'][0])

[{'word1': 'my', 'word2': 'little', 'word3': 'pony'},
 {'word1': 'your', 'word2': 'big', 'word3': 'pony'}]

甚至创建一个数据框也很简单：

pl.DataFrame(myfunc(df['text'][0]))

但是尝试使用map_elements()时就失败了：

(df.with_columns(pl.struct(['text']).map_elements(lambda x: myfunc(x['text'])).alias('aaa')
                 )
 )

线程 '' 在 crates/polars-core/src/chunked_array/builder/list/anonymous.rs:161:69 处崩溃： 在一个 Err 值上调用了 Result::unwrap()：InvalidOperation(ErrString("无法连接不同数据类型的数组。")) --- PyO3 在从 Python 获取 PanicException 后恢复了一个崩溃。 ---

我想要的结果是这样的：

file     word1         word2   word3
aaa.txt  my            little  pony
aaa.txt  your          big     pony
bbb.txt  apple+banana
bbb.txt  cake+coke

有什么想法吗？

Answer 1

你可以完全不使用用户自定义函数（UDF），可以按照以下步骤操作。

1. 把文本分割成一行一行的列表
2. 把这些行的列表展开成单独的行
3. 去掉每行的多余空格，然后把每行分割成一个个单词
4. 把这些单词的列表转换成一个结构体，并把字段命名为符合格式 "word_IDX" 的样子

(
    df
    .with_columns(
        pl.col("text").str.split(",")
    )
    .explode("text")
    .with_columns(
        pl.col("text").str.strip_chars().str.split(" ")
        .list.to_struct()
        .name.map_fields(lambda s: s.replace("field_", "word_"))
    )
    .unnest("text")
)

shape: (5, 4)
┌─────────┬──────────────┬────────┬────────┐
│ file    ┆ word_0       ┆ word_1 ┆ word_2 │
│ ---     ┆ ---          ┆ ---    ┆ ---    │
│ str     ┆ str          ┆ str    ┆ str    │
╞═════════╪══════════════╪════════╪════════╡
│ aaa.txt ┆ my           ┆ little ┆ pony   │
│ aaa.txt ┆ your         ┆ big    ┆ pony   │
│ aaa.txt ┆ your         ┆ big    ┆ pony   │
│ bbb.txt ┆ apple+banana ┆ null   ┆ null   │
│ bbb.txt ┆ cake+coke    ┆ null   ┆ null   │
└─────────┴──────────────┴────────┴────────┘

Answer 2

这个错误很可能是因为用户自定义函数（UDF）返回的类型比较复杂。Polars会快速推断内部字典的数据类型。但是，后面被评估的字典可能字段数量不同，这就导致了错误的发生。

一个更简单的例子可以是下面这个。

import random

import polars as pl

def ufunc(x: int):
    return [
        {f"word_{i}": "elephant" for i in range(random.randint(1, 4))}
        for _ in range(random.randint(1, 4))
    ]

pl.DataFrame({"id": [1, 2]}).with_columns(pl.col("id").map_elements(ufunc))

最开始，我以为正确设置pl.Expr.map_elements的returned_dtype参数就能解决这个问题。然而，我没有找到一种方法来指定字段数量不固定的结构类型。

不过，当把返回值放进另一个列表里时，Polars似乎能够正确推断返回的数据类型。接着，我们可以用.list.first获取第一个（也是唯一的）列表元素，然后使用pl.DataFrame.explode将行展开成多行，最后再把内部的结构解开。

(
    pl.DataFrame({"id": [1, 2]})
    .with_columns(
        pl.col("id").map_elements(lambda x: [ufunc(x)]).list.first()
    )
    .explode("id")
    .unnest("id")
)

shape: (7, 4)
┌──────────┬──────────┬──────────┬──────────┐
│ word_0   ┆ word_1   ┆ word_2   ┆ word_3   │
│ ---      ┆ ---      ┆ ---      ┆ ---      │
│ str      ┆ str      ┆ str      ┆ str      │
╞══════════╪══════════╪══════════╪══════════╡
│ elephant ┆ elephant ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ elephant ┆ elephant │
│ elephant ┆ null     ┆ null     ┆ null     │
│ elephant ┆ null     ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ null     ┆ null     │
│ elephant ┆ elephant ┆ elephant ┆ null     │
│ elephant ┆ null     ┆ null     ┆ null     │
└──────────┴──────────┴──────────┴──────────┘