Polars 惰性框架 - 将其他惰性框架中的字段添加为不带“collect”的结构

我正在尝试填充一个包含所有结构的新字段other来自另一个基于谓词的惰性帧的字段。

虽然示例是用 python 编写的，但我愿意接受 python 或 rust 中的答案。

companies = pl.DataFrame({
  "id": [1], 
  "name": ["google"], 
  "industry": [1001]
}).lazy()

industries = pl.DataFrame({
  "id": [1001],
  "name": ["tech"],
  "sectors": [[10011, 10012]]
}).lazy()

expected = pl.DataFrame({
  "id": [1],
  "name": ["polars"],
  "industry": [{
    "name": "tech",
    "sectors": [[10011, 10012]]
  }]
})

我可以通过 join -> select -> to_struct -> lit 来天真地做到这一点，但这并不理想，因为我必须在惰性框架上执行收集才能获得我想要的结果。

right_columns = industries.schema.keys()
industry_matches = companies.join(
    industries,
    left_on="industry",
    right_on="id",
    how="inner"
).collect().select(right_columns).drop(["id"]).to_struct("industry")

df = companies.with_column(pl.lit(industry_matches))

df.collect()

理想情况下，我想做一些类似在表达式中加入之类的事情。

companies.select([
  pl.col("*"),
  pl.col("industry").join(industries, on="id").exclude(["id"]),
])

但我对任何不需要的替代方案持开放态度collect

这对你有用吗：

(
    companies
    .join(
        industries
        .select([
            pl.col('id').alias('industry'),
            pl.struct(pl.exclude('id')).alias('industry_struct'),
        ]),
        on="industry",
        how="inner"
    )
    .drop('industry')
    .collect()
)

shape: (1, 3)
┌─────┬────────┬─────────────────────────┐
│ id  ┆ name   ┆ industry_struct         │
│ --- ┆ ---    ┆ ---                     │
│ i64 ┆ str    ┆ struct[2]               │
╞═════╪════════╪═════════════════════════╡
│ 1   ┆ google ┆ {"tech",[10011, 10012]} │
└─────┴────────┴─────────────────────────┘

本质上，我们是在连接本身内创建结构。