从 Pandas 到 Polars 十七：Polars 中的嵌套数据类型：pl.List 数据类型

Polars 使用 Apache Arrow 在内存中存储其数据。Arrow 的一个巨大优势是它支持各种嵌套数据类型（或“dtypes”）。在本文中，我们将更详细地分析 pl.List 数据类型：

1. 我们首先概述 pl.List 数据类型
2. 我们在 pl.List 列的每一行上调用表达式
3. 我们使用神经网络嵌入进行聚合
4. 我们进行简单的文本分析

pl.List 数据类型的概述

pl.List 数据类型允许我们在每一行上存储一个值的数组。关键点是，每个数组内部的值类型必须相同，并且这些类型在所有行上都必须相同。

在这个例子中，我们创建了一个包含整数、浮点数和字符串 pl.List 列的 DataFrame。请注意：

1. 在 floats 列中，我们有一行同时包含浮点数和整数，因此 Polars 会将所有值转换为浮点数类型
2. 数组的长度在列内部可以有所不同

import polars as pl
dfLists = pl.DataFrame({
    'ints':[ [0,1], [4,3,2]],
    'floats':[ [0.0,1], [2,3]],
'strings':[ ["0","1"],["2","3"]]})
dfLists
shape: (2, 3)

关于 pl.List 数据类型需要理解的关键点是，每一行在底层都是一个 pl.Series。这意味着对 pl.List 列的操作将是快速的矢量化操作。

数组内的表达式

在本文后面的用例中，我们将看到如何对整个数组应用表达式。但是，我们也可以在 pl.List 列上逐行应用表达式。

在这个例子中，我们对每个数组内的元素进行排名。

(
    dfLists
    .with_columns(
        pl.col("ints").arr.eval(
            pl.element().rank(method="ordinal")
        )
))
shape: (2, 1)

要在每个数组内部调用 rank 表达式，我们

1. 在 ints 列上调用 arr.eval
2. 在 arr.eval 内部，我们使用 pl.element 来为每一行启动表达式，并
3. 然后在 pl.element 上调用 rank 来对每一行执行 rank 表达式

用例

嵌入分析

当你与神经网络模型的嵌入以及其他元数据一起工作时，pl.List 数据类型是一个很好的选择。

在下面的示例中，我们有一个 doc_id 列来标识每一行来自哪个文档，一个 text 列显示每个文档的文本块，以及一个 embeddings 列，其中包含该文本的嵌入。

df = pl.DataFrame(
    {
    "doc_id":[0,0,1,1,2,2],
    "text":
        [
                "Polars is a dataframe library",
                "Polars is written in Rust",
                "Expressions allow you to transform data",
                "Expressions run in paralell",
                "Apache Arrow supports nested data",
                "There are three nested dtypes"
        ]
    }
   )
   .with_columns(
       pl.Series(
        "embeddings",
        [pl.Series("",np.random.randint(0,5,3)) for _ in range(6)]
        ))
shape: (6, 3)

然后，我们通过对 doc_id 列进行分组并取 embeddings 的平均值，来获得文档平均嵌入向量。

(
    df
    .groupby(
        "doc_id"
        )
    .agg(
        pl.col("embeddings").arr.mean()
        )
)

我们使用 arr.mean 而不是简单的 mean 来执行聚合操作。通过使用 arr.mean，我们利用了 arr 命名空间中针对 pl.List 数据类型的数组表达式。您可以在这里查看完整的表达式集。

单词计数

数组的另一个用例是当我们拆分字符串时。在这个例子中，我们按空白字符拆分 text 列以获得单个单词。这将 text 列转换为包含字符串数组的列。

df2 = (
    df
    .with_columns(
        pl.col("text").str.split(" ")
    )
    .select(
        ["doc_id","text"]
))
shape: (6, 2)

使用这个字符串数组，我们就可以使用 arr.lengths 来计算每一行的单词数量了。

(
    df2
    .with_columns(
        pl.col("text").arr.lengths()
))
shape: (6, 2)

或者，我们可以计算每个单词的出现次数。我们通过调用 explode 将字符串数组转换为单独的行来实现这一点。

(
    df2
    .select(["doc_id","text"])
.explode("text"))
shape: (30, 2)

从这个展开（explode）的格式中，我们可以计算单词的出现次数

(
    df2
    .explode("text")
    ["text"]
.value_counts(sort=True))
shape: (24, 2)

当然，这只是我们使用 pl.List 数据类型的开始。