主页 / coding / 问题 / 79161450
Accepted
apostofes
Asked: 2024-11-06 13:47:43 +0800 CST

在 polars 中 join_where 与 starts_with

  • 772

我有两个数据框，

df = pl.DataFrame({'url': ['https//abc.com', 'https//abcd.com', 'https//abcd.com/aaa', 'https//abc.com/abcd']})


conditions_df = pl.DataFrame({'url': ['https//abc.com', 'https//abcd.com', 'https//abcd.com/aaa', 'https//abc.com/aaa'], 'category': [['a'], ['b'], ['c'], ['d']]})

现在我想要一个 df，用于根据第二个 df 中以 url 开头的第一个匹配项为第一个 df 分配类别，即输出应该是，

网址 类别
https//abc.com ['一个']
https//abcd.com ['b']
https//abcd.com/aaa ['b'] - 这个以 https//abcd.com 开头，这是第一个匹配
https//abc.com/abcd ['a'] - 这个以 https//abc.com 开头，这是第一个匹配

目前有效的代码是这样的，

def add_category_column(df: pl.DataFrame, conditions_df) -> pl.DataFrame:
    
    # Initialize the category column with empty lists
    df = df.with_columns(pl.Series("category", [[] for _ in range(len(df))], dtype=pl.List(pl.String)))
    
    # Apply the conditions to populate the category column
    for row in conditions_df.iter_rows():
        url_start, category = row
        df = df.with_columns(
            pl.when(
                (pl.col("url").str.starts_with(url_start)) & (pl.col("category").list.len() == 0)
            )
            .then(pl.lit(category))
            .otherwise(pl.col("category"))
            .alias("category")
        )
    
    return df

但是有没有办法在不使用 for 循环的情况下实现相同的效果，我们可以在这里使用 join_where 吗，但在我的尝试中 join_where 对 starts_with 不起作用

python-polars

3 个回答

  1. 不幸的是，目前看起来并非如此。我刚刚在 Polars 问题跟踪器上提出了这个问题并请求它。

    这是@roman 答案的一个细微变化，它在连接之前准备一个行索引。

    (
    df.join(conditions_df.with_row_index(), how="cross")
    .filter(pl.col("url").str.starts_with(pl.col("url_right")))
    # pick only the first match based on lowest row number
    # needs to happen after the starts with filter
    .filter(pl.col("index") == pl.col("index").min().over("url"))
    .select("url", "category")
)

    或者，如果您有更大的数据或性能问题，则通过 DuckDB 的解决方案将进行连接starts_with（Polars SQL 似乎还不支持它）

    sql = """
select df.url, c.category
from df
inner join (select *, row_number() over() as index from conditions_df) as c
on starts_with(df.url, c.url)
-- pick only the first match based on lowest row number
-- this happens after the join is already done
qualify index = min(index) over(partition by df.url)
"""
duckdb.query(sql).pl()
    • 2
  2. Best Answer

    我希望pl.DataFrame.join_where()能工作，但显然它还不允许pl.Expr.str.starts_with()条件——我得到了only 1 binary comparison allowed as join condition错误。

    因此你可以pl.DataFrame.join()改用pl.DataFrame.filter()：

    (
    df
    .join(conditions_df, how="cross")
    .filter(pl.col("url").str.starts_with(pl.col("url_right")))
    .sort("url")
    .group_by("url", maintain_order=True)
    .agg(pl.col.category.first())
)

    shape: (4, 2)
┌─────────────────────┬───────────┐
│ url                 ┆ category  │
│ ---                 ┆ ---       │
│ str                 ┆ list[str] │
╞═════════════════════╪═══════════╡
│ https//abc.com      ┆ ["a"]     │
│ https//abc.com/abcd ┆ ["a"]     │
│ https//abcd.com     ┆ ["b"]     │
│ https//abcd.com/aaa ┆ ["b"]     │
└─────────────────────┴───────────┘

    您还可以将DuckDB 与 Polars 集成并使用lateral join：

    import duckdb

duckdb.sql("""
    select
        d.url,
        c.category
    from df as d,
        lateral (
            select c.category
            from conditions_df as c
            where
                starts_with(d.url, c.url)
            limit 1
        ) as c
""")

    ┌─────────────────────┬───────────┐
│         url         │ category  │
│       varchar       │ varchar[] │
├─────────────────────┼───────────┤
│ https//abc.com      │ [a]       │
│ https//abc.com/abcd │ [a]       │
│ https//abcd.com/aaa │ [b]       │
│ https//abcd.com     │ [b]       │
└─────────────────────┴───────────┘

    但是，您必须小心，因为在标准 SQL 规范中行集合是无序的，因此如果不在order by侧面添加明确的子句，我不会在生产中这样做。

    • 1

  3. 可以concat先进行水平操作（而不是初始操作）来找到匹配项join。

    它需要更多的手动步骤 - 但在处理较大的输入时给了我最快的结果，所以可能会感兴趣。

    (df.with_row_index()
   .join(
       pl.concat(
           [
               df.with_row_index(), 
               conditions_df.rename({"url": "condition_url"}).drop("category")
           ],
           how = "horizontal"
       )
       .with_columns(
           pl.col.url.str.extract_many(pl.col.condition_url, overlapping=True)
             .alias("condition_url")
       )
       .explode("condition_url")
       .filter(pl.col.url.str.starts_with(pl.col.condition_url))
       .group_by("index")
       .agg(pl.col.condition_url.first())
       .join(
           conditions_df.rename({"url": "condition_url"}),
           on = "condition_url"
       ),
       on = "index",
       how = "left"
  )
)

    shape: (5, 4)
┌───────┬─────────────────────┬─────────────────┬───────────┐
│ index ┆ url                 ┆ condition_url   ┆ category  │
│ ---   ┆ ---                 ┆ ---             ┆ ---       │
│ u32   ┆ str                 ┆ str             ┆ list[str] │
╞═══════╪═════════════════════╪═════════════════╪═══════════╡
│ 0     ┆ https//abc.com      ┆ https//abc.com  ┆ ["a"]     │
│ 1     ┆ https//abcd.com     ┆ https//abcd.com ┆ ["b"]     │
│ 2     ┆ https//abcd.com/aaa ┆ https//abcd.com ┆ ["b"]     │
│ 3     ┆ https//abc.com/abcd ┆ https//abc.com  ┆ ["a"]     │
│ 4     ┆ nomatch             ┆ null            ┆ null      │
└───────┴─────────────────────┴─────────────────┴───────────┘

    解释

    我们水平排列.concat()框架并用来.str.extract_many()获取所有子字符串匹配的列表，然后将其explode分成几行。

    possible_matches = (
    pl.concat(
        [
            df.with_row_index(), 
            conditions_df.rename({"url": "condition_url"}).drop("category")
        ],
        how = "horizontal"
    )
    .with_columns(
        pl.col.url.str.extract_many(pl.col.condition_url, overlapping=True)
          .alias("condition_url")
    )
    .explode("condition_url")
)

    shape: (6, 3)
┌───────┬─────────────────────┬─────────────────────┐
│ index ┆ url                 ┆ condition_url       │
│ ---   ┆ ---                 ┆ ---                 │
│ u32   ┆ str                 ┆ str                 │
╞═══════╪═════════════════════╪═════════════════════╡
│ 0     ┆ https//abc.com      ┆ https//abc.com      │
│ 1     ┆ https//abcd.com     ┆ https//abcd.com     │
│ 2     ┆ https//abcd.com/aaa ┆ https//abcd.com     │
│ 2     ┆ https//abcd.com/aaa ┆ https//abcd.com/aaa │
│ 3     ┆ https//abc.com/abcd ┆ https//abc.com      │
│ 4     ┆ nomatch             ┆ null                │
└───────┴─────────────────────┴─────────────────────┘

    然后我们应用starts_with过滤器和first约束。

    matches = (
    possible_matches
     .filter(pl.col.url.str.starts_with(pl.col.condition_url))
     .group_by("index")
     .agg(pl.col.condition_url.first())
)

    shape: (4, 2)
┌───────┬─────────────────┐
│ index ┆ condition_url   │
│ ---   ┆ ---             │
│ u32   ┆ str             │
╞═══════╪═════════════════╡
│ 0     ┆ https//abc.com  │
│ 1     ┆ https//abcd.com │
│ 2     ┆ https//abcd.com │
│ 3     ┆ https//abc.com  │
└───────┴─────────────────┘

    A.join()用于获取对应的类别。

    matches.join(
    conditions_df.rename({"url": "condition_url"}), 
    on = "condition_url"
)

    shape: (4, 3)
┌───────┬─────────────────┬───────────┐
│ index ┆ condition_url   ┆ category  │
│ ---   ┆ ---             ┆ ---       │
│ u32   ┆ str             ┆ list[str] │
╞═══════╪═════════════════╪═══════════╡
│ 0     ┆ https//abc.com  ┆ ["a"]     │
│ 3     ┆ https//abc.com  ┆ ["a"]     │
│ 1     ┆ https//abcd.com ┆ ["b"]     │
│ 2     ┆ https//abcd.com ┆ ["b"]     │
└───────┴─────────────────┴───────────┘

    最后.join()根据匹配的行将其添加回原始框架index。

    • 0

相关问题