Polars中如何无需循环实现基于starts_with前缀的首匹配分类赋值

现有两个Polars DataFrame：

import polars as pl

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的每条url，匹配conditions_df中**按行顺序第一个满足前缀匹配（starts_with）**的url，将对应的category赋值给该url。预期输出：

目前已有基于循环的实现：

def add_category_column(df: pl.DataFrame, conditions_df) -> pl.DataFrame:
    # 初始化category列为空列表
    df = df.with_columns(pl.Series("category", [[] for _ in range(len(df))], dtype=pl.List(pl.String)))
    
    # 应用条件填充category列
    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

提问：是否可以无需循环实现该需求？能否使用join_where？尝试用join_where处理前缀匹配时未能成功。

可以不用循环实现，直接利用Polars的向量化操作完成。join_where本身可以实现条件匹配，但它会返回所有符合条件的结果，无法直接返回「第一个匹配项」，需要结合排序或窗口函数进一步处理。

方法1：交叉连接+排序取首个匹配 #

核心思路是给条件表添加顺序标记，交叉连接后过滤出前缀匹配的行，再按条件表的顺序排序，最后分组取每个url的第一个匹配结果：

import polars as pl

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"]]
})

# 给条件表添加匹配顺序列，标记行的先后顺序
conditions_with_order = conditions_df.with_columns(
    pl.int_range(0, len(conditions_df)).alias("match_order")
)

# 交叉连接后过滤前缀匹配，按匹配顺序排序，分组取首个结果
result = df.join(conditions_with_order, how="cross") \
          .filter(pl.col("url").str.starts_with(pl.col("url_right"))) \
          .sort("match_order") \
          .group_by("url") \
          .first() \
          .select("url", "category")

# 左连接回原表，确保所有原始url都被保留
final_result = df.join(result, on="url", how="left")

print(final_result)

输出结果：

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

方法2：交叉连接+窗口函数取最小匹配顺序 #

另一种方式是用窗口函数找到每个url对应的最小匹配顺序，再过滤出对应行：

result = df.join(conditions_with_order, how="cross") \
          .filter(pl.col("url").str.starts_with(pl.col("url_right"))) \
          .with_columns(pl.min("match_order").over("url").alias("min_order")) \
          .filter(pl.col("match_order") == pl.col("min_order")) \
          .select("url", "category")

final_result = df.join(result, on="url", how="left")

关于join_where的说明 #

join_where可以用来实现前缀匹配的连接，但它会返回所有符合条件的匹配对。例如：

temp = df.join(conditions_df, how="left", on=pl.col("url").str.starts_with(pl.col("url_right")))

这段代码会给https//abcd.com/aaa返回两个匹配结果（https//abcd.com和https//abcd.com/aaa），无法直接得到我们需要的「第一个匹配项」。因此需要额外步骤（如排序、窗口函数）来筛选出优先级最高的匹配。

内容的提问来源于stack exchange，提问作者apostofes