如何根据条件连接两个 Polars 数据框？

Question

给定以下 3 个 Polars 数据框

journeys = pl.DataFrame({'id':[1,2,3,4,5,6],'order_id':[11,12,13,14,14,15],'order_type':['restaurant','restaurant','restaurant','restaurant','grocery','grocery']})

restaurant_orders = pl.DataFrame({'id':[11,12,13,14],'item_count':[4,7,3,5]})

grocery_orders = pl.DataFrame({'id':[14,15],'item_count':[23,21]})

journeys：

restaurant_orders：

grocery_orders：

我想将其引入item_count数据journeys框。

journeys最简单的方法是根据列过滤数据帧order_type，对每个过滤后的数据帧执行连接，最后将它们连接在一起。

journeys是否有一种 Polars 惯用方法可以根据 的值对数据帧执行条件（多态？）连接order_type。

Answer 1

我认为不可能以 的条件加入polars。

实现所需结果的一种方法是将两个数据框连接起来并将它们添加为一列order_type。

restaurant_orders = restaurant_orders.with_columns(order_type = pl.lit("restaurant"))
grocery_orders = grocery_orders.with_columns(order_type = pl.lit("grocery"))
orders = pl.concat([restaurant_orders, grocery_orders])

journeys.join(orders, left_on=['order_id', 'order_type'], right_on=['id', 'order_type'], how='left')

如果您有多个数据帧要像这样连接，您可以使用一个函数根据数据帧的名称来执行此操作：

def concat_df(list_df: list) -> pl.DataFrame:
  list_newdf = []

  for df in list_df:
    name = [x for x in globals() if globals()[x] is df][0]
    name = name[:-7]

    list_newdf.append(df.with_columns(order_type = pl.lit(name)))
    
  return pl.concat(list_newdf)


list_df = [restaurant_orders, grocery_orders]
orders = concat_df(list_df)

journeys.join(orders, left_on=['order_id', 'order_type'], right_on=['id', 'order_type'], how='left')

它将产生相同的结果。请注意，第二个解决方案取决于 DataFrame 名称。