主页 / coding / 问题 / 79148243
Accepted
gernophil
Asked: 2024-11-01 21:53:47 +0800 CST

在 Polars df 中转置会导致“TypeError：尚未实现：嵌套对象类型”

  • 772

我有这些数据：

┌────────────┬─────────────────────────────────────┐
│ col1       ┆ col2                                │
│ ---        ┆ ---                                 │
│ list[str]  ┆ list[list[str]]                     │
╞════════════╪═════════════════════════════════════╡
│ ["a"]      ┆ [["a"]]                             │
│ ["b", "c"] ┆ [["b", "c"], ["b", "c"], ["b", "c"] │
│ ["d"]      ┆ [["d"]]                             │
└────────────┴─────────────────────────────────────┘

我希望第 2 行中同一列表中包含所有 b 和所有 c，但如您所见，第 2 行中未维护 b 到 b 和 c 到 c 的关联。对于 pandas，我使用了：

import pandas as pd

pddf = pd.DataFrame({"col1": [["a"], ["b", "c"], ["d"]],
                     "col2": [[["a"]], [["b", "c"], ["b", "c"], ["b", "c"]], [["d"]]]})

pddf["col2"] = pddf["col2"].apply(lambda listed: pd.DataFrame(listed).transpose().values.tolist())

print(pddf)

#      col1                    col2
# 0     [a]                   [[a]]
# 1  [b, c]  [[b, b, b], [c, c, c]]
# 2     [d]                   [[d]]

这是期望的结果。我试图用 替换 polars 来实现同样的效果pddf.transpose().values.tolist()，pldf.transpose().to_numpy().tolist()但总是得到 和TypeError: not yet implemented: Nested object types。有什么解决方法吗？以下是完整的 polars 代码：

import polars as pl

pldf = pl.DataFrame({"col1": [["a"], ["b", "c"], ["d"]],
                     "col2": [[["a"]], [["b", "c"], ["b", "c"], ["b", "c"]], [["d"]]]})

pldf = pldf.with_columns(pl.col("col2").map_elements(lambda listed: pl.DataFrame(listed).transpose().to_numpy().tolist()))

print(pldf)

# TypeError: not yet implemented: Nested object types
# Hint: Try setting `strict=False` to allow passing data with mixed types.

我需要在哪里应用上述内容strict=False？

在更简单的 df 上pddf.transpose().values.tolist()并且pldf.transpose().to_numpy().tolist()是相同的：

import pandas as pd
import polars as pl

pd.DataFrame(
    {"col1": ["a", "b", "c"],
     "col2": ["d", "e", "f"]}
     ).transpose().values.tolist() == pl.DataFrame(
                                          {"col1": ["a", "b", "c"],
                                           "col2": ["d", "e", "f"]}
                                           ).transpose().to_numpy().tolist()

# True

请尽可能接近代码，即使使用.apply()或并不理想.map_elements()，但这是一个更大的项目，我不想破坏任何其他东西:)。

(编辑：我稍微简化了代码，因为第二个 lambda 对于这个问题来说并不是真正必要的。)

python

3 个回答

  1. 这是一个完全符合 polars 表达式 API 的解决方案，它不会委托给（缓慢的）本机 python。

    这个想法是小心地将嵌套列表列展开pl.DataFrame.explode为一个长字符串列（使用），并在此过程中小心地添加索引列。

    这使我们能够跟踪每个元素的行/列表/内部列表索引，并将元素聚合回预期结果。

    (
    df
    .with_row_index(name="row_idx")
    .explode("col2")
    .with_columns(
        pl.int_range(pl.len()).over("row_idx").alias("outer_list_idx")
    )
    .explode("col2")
    .with_columns(
        pl.int_range(pl.len()).over("row_idx", "outer_list_idx").alias("inner_list_idx")
    )
    .group_by("row_idx", "inner_list_idx", maintain_order=True)
    .agg(
        pl.col("col1").first(),
        "col2"
    )
    .group_by("row_idx")
    .agg(
        pl.col("col1").first(),
        "col2"
    )
    .drop("row_idx")
)

    shape: (3, 2)
┌────────────┬──────────────────────────────────┐
│ col1       ┆ col2                             │
│ ---        ┆ ---                              │
│ list[str]  ┆ list[list[str]]                  │
╞════════════╪══════════════════════════════════╡
│ ["a"]      ┆ [["a"]]                          │
│ ["b", "c"] ┆ [["b", "b", "b"], ["c", "c", "c" │
│ ["d"]      ┆ [["d"]]                          │
└────────────┴──────────────────────────────────┘
    • 1
  2. Best Answer

    以下是我的做法。我会定义一个函数transpose_nested(listed)来 zip 转置每个内部列表。基本上，和 pandas 中一样transpose().values.tolist()。然后，我会 map_elements(transpose_nested)。

    import polars as pl

pldf = pl.DataFrame({
    "col1": [["a"], ["b", "c"], ["d"]],
    "col2": [[["a"]], [["b", "c"], ["b", "c"], ["b", "c"]], [["d"]]]
})

print(pldf)

def transpose_nested(listed):
    return list(map(list, zip(*listed)))

pldf = pldf.with_columns(
    pl.col("col2").map_elements(transpose_nested)
)

print(pldf)

    返回原始数据和转换后的数据：

    shape: (3, 2)
┌────────────┬─────────────────────────────────┐
│ col1       ┆ col2                            │
│ ---        ┆ ---                             │
│ list[str]  ┆ list[list[str]]                 │
╞════════════╪═════════════════════════════════╡
│ ["a"]      ┆ [["a"]]                         │
│ ["b", "c"] ┆ [["b", "c"], ["b", "c"], ["b",… │
│ ["d"]      ┆ [["d"]]                         │
└────────────┴─────────────────────────────────┘
shape: (3, 2)
┌────────────┬─────────────────────────────────┐
│ col1       ┆ col2                            │
│ ---        ┆ ---                             │
│ list[str]  ┆ list[list[str]]                 │
╞════════════╪═════════════════════════════════╡
│ ["a"]      ┆ [["a"]]                         │
│ ["b", "c"] ┆ [["b", "b", "b"], ["c", "c", "… │
│ ["d"]      ┆ [["d"]]                         │
└────────────┴─────────────────────────────────┘
    • 0

  3. 正如您在评论中提到的性能 - 这也可能引起人们的兴趣。

    基于@Herick 的回答- 当您添加索引并展开时，保证数据已排序。

    对于已排序的数据，您可以使用 - 添加“每组的行号”， rle()这样速度会快得多。

    (df.with_row_index("row")
   .select("row", "col2")
   .explode("col2")
   .with_row_index()
   .explode("col2")
   .with_columns(
      pl.int_ranges(pl.col("index").rle().struct.field("len"))
        .flatten()
        .alias("list_idx")
   )
)

    shape: (8, 4)
┌───────┬─────┬──────┬──────────┐
│ index ┆ row ┆ col2 ┆ list_idx │
│ ---   ┆ --- ┆ ---  ┆ ---      │
│ u32   ┆ u32 ┆ str  ┆ i64      │
╞═══════╪═════╪══════╪══════════╡
│ 0     ┆ 0   ┆ a    ┆ 0        │
│ 1     ┆ 1   ┆ b    ┆ 0        │
│ 1     ┆ 1   ┆ c    ┆ 1        │
│ 2     ┆ 1   ┆ b    ┆ 0        │
│ 2     ┆ 1   ┆ c    ┆ 1        │
│ 3     ┆ 1   ┆ b    ┆ 0        │
│ 3     ┆ 1   ┆ c    ┆ 1        │
│ 4     ┆ 2   ┆ d    ┆ 0        │
└───────┴─────┴──────┴──────────┘

    在我的测试中，只选择要爆炸的列表列然后将结果水平连接回原始框架也更快。

    pl.concat(
    [
        df.drop("col2"),
        (df.with_row_index("row")
           .select("row", "col2")
           .explode("col2")
           .with_row_index()
           .explode("col2")
           .with_columns(
              pl.int_ranges(pl.col("index").rle().struct.field("len"))
                .flatten()
                .alias("index")
           )
           .group_by("index", "row", maintain_order=True) # order required for horizontal concat
           .agg("col2")
           .group_by("row", maintain_order=True)          # order required for horizontal concat
           .agg("col2")
           .drop("row")
        )
    ],
    how = "horizontal"
)

    shape: (3, 2)
┌────────────┬────────────────────────────────────┐
│ col1       ┆ col2                               │
│ ---        ┆ ---                                │
│ list[str]  ┆ list[list[str]]                    │
╞════════════╪════════════════════════════════════╡
│ ["a"]      ┆ [["a"]]                            │
│ ["b", "c"] ┆ [["b", "b", "b"], ["c", "c", "c"]] │
│ ["d"]      ┆ [["d"]]                            │
└────────────┴────────────────────────────────────┘

    作为基本比较，如果我对5_000_000行进行抽样，我会得到：

    姓名 时间
    地图元素 61.19秒
    爆炸 + 连接 7.14秒
    • 0

相关问题