Mark rows of one dataframe based on values from another dataframe

I have following problem. Let's say I have two dataframes

df1 = pl.DataFrame({'a': range(10)})
df2 = pl.DataFrame({'b': [[1, 3], [5,6], [8, 9]], 'tags': ['aa', 'bb', 'cc']})
print(df1)
print(df2)

shape: (10, 1)
┌─────┐
│ a   │
│ --- │
│ i64 │
╞═════╡
│ 0   │
│ 1   │
│ 2   │
│ 3   │
│ 4   │
│ 5   │
│ 6   │
│ 7   │
│ 8   │
│ 9   │
└─────┘
shape: (3, 2)
┌───────────┬──────┐
│ b         ┆ tags │
│ ---       ┆ ---  │
│ list[i64] ┆ str  │
╞═══════════╪══════╡
│ [1, 3]    ┆ aa   │
│ [5, 6]    ┆ bb   │
│ [8, 9]    ┆ cc   │
└───────────┴──────┘

I need to mark/tag rows in dataframe df1 based on values of dataframe df2, so I can get following dataframe

print(pl.DataFrame({'a': range(10), 'tag': ['NA', 'aa', 'aa', 'aa', 'NA', 'bb', 'bb', 'NA', 'cc', 'cc']}))

shape: (10, 2)
┌─────┬─────┐
│ a   ┆ tag │
│ --- ┆ --- │
│ i64 ┆ str │
╞═════╪═════╡
│ 0   ┆ NA  │
│ 1   ┆ aa  │
│ 2   ┆ aa  │
│ 3   ┆ aa  │
│ 4   ┆ NA  │
│ 5   ┆ bb  │
│ 6   ┆ bb  │
│ 7   ┆ NA  │
│ 8   ┆ cc  │
│ 9   ┆ cc  │
└─────┴─────┘

So list in column b of df2 indicates start and end values for column a of df1 that needs to be tagged with what's in column tags.

Thanks

You could .explode the .arange and use a left join.

df1.join(
   df2.with_columns(
      pl.arange(pl.col("b").arr.first(), pl.col("b").arr.last() + 1)
   ).explode("b"),
   left_on="a",
   right_on="b",
   how="left"
)

shape: (10, 2)
┌─────┬──────┐
│ a   ┆ tags │
│ --- ┆ ---  │
│ i64 ┆ str  │
╞═════╪══════╡
│ 0   ┆ null │
│ 1   ┆ aa   │
│ 2   ┆ aa   │
│ 3   ┆ aa   │
│ 4   ┆ null │
│ 5   ┆ bb   │
│ 6   ┆ bb   │
│ 7   ┆ null │
│ 8   ┆ cc   │
│ 9   ┆ cc   │
└─────┴──────┘

If the ranges don't overlap, another option is to reshape df2:

df2.with_columns(
   pl.col("b").arr.to_struct()
     .struct.rename_fields(["start", "end"])
).unnest("b")

shape: (3, 3)
┌───────┬─────┬──────┐
│ start ┆ end ┆ tags │
│ ---   ┆ --- ┆ ---  │
│ i64   ┆ i64 ┆ str  │
╞═══════╪═════╪══════╡
│ 1     ┆ 3   ┆ aa   │
│ 5     ┆ 6   ┆ bb   │
│ 8     ┆ 9   ┆ cc   │
└───────┴─────┴──────┘

And use .join_asof

df1.join_asof(
   df2.with_columns(
      pl.col("b").arr.to_struct().struct.rename_fields(["start", "end"])
   ).unnest("b"),
   left_on="a",
   right_on="end",
   strategy="forward"
).with_columns(
   pl.when(pl.col("a").is_between("start", "end"))
     .then(pl.col("tags"))
)

shape: (10, 4)
┌─────┬───────┬─────┬──────┐
│ a   ┆ start ┆ end ┆ tags │
│ --- ┆ ---   ┆ --- ┆ ---  │
│ i64 ┆ i64   ┆ i64 ┆ str  │
╞═════╪═══════╪═════╪══════╡
│ 0   ┆ 1     ┆ 3   ┆ null │
│ 1   ┆ 1     ┆ 3   ┆ aa   │
│ 2   ┆ 1     ┆ 3   ┆ aa   │
│ 3   ┆ 1     ┆ 3   ┆ aa   │
│ 4   ┆ 5     ┆ 6   ┆ null │
│ 5   ┆ 5     ┆ 6   ┆ bb   │
│ 6   ┆ 5     ┆ 6   ┆ bb   │
│ 7   ┆ 8     ┆ 9   ┆ null │
│ 8   ┆ 8     ┆ 9   ┆ cc   │
│ 9   ┆ 8     ┆ 9   ┆ cc   │
└─────┴───────┴─────┴──────┘