如何实现支持重复项选择、严格比较及UTF-8的高级Asof Join？

可以实现满足以下需求的类Asof Join操作：

• 存在重复匹配项时，可指定选择first或last匹配元素
• 支持严格小于（lt）、严格大于（gt）、小于等于（leq）、大于等于（geq）四种连接逻辑
• 完全兼容UTF-8编码

基础数据准备 #

import polars as pl

df1 = pl.DataFrame({
    'by_1': ['X', 'X', 'Y', 'Y'] * 8,
    'by_2': ['X', 'Y', 'X', 'Y'] * 8,
    'on_1': ['A'] * 16 + ['C'] * 16,
    'on_2': (['A'] * 8 + ['C'] * 8) * 2,
    '__index__': list(range(32))
})

df2 = pl.DataFrame([
    { 'by_1': 'Y', 'by_2': 'Y', 'on_1': 'B', 'on_2': 'A' },
    { 'by_1': 'Y', 'by_2': 'Y', 'on_1': 'C', 'on_2': 'A' },
    { 'by_1': 'Y', 'by_2': 'Z', 'on_1': 'A', 'on_2': 'A' },
])

df1数据结构 #

┌──────┬──────┬──────┬──────┬───────────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 ┆ __index__ │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---       │
│ str  ┆ str  ┆ str  ┆ str  ┆ i64       │
╞══════╪══════╪══════╪══════╪═══════════╡
│ X    ┆ X    ┆ A    ┆ A    ┆ 0         │
│ X    ┆ Y    ┆ A    ┆ A    ┆ 1         │
│ Y    ┆ X    ┆ A    ┆ A    ┆ 2         │
│ Y    ┆ Y    ┆ A    ┆ A    ┆ 3         │
│ X    ┆ X    ┆ A    ┆ A    ┆ 4         │
│ X    ┆ Y    ┆ A    ┆ A    ┆ 5         │
│ Y    ┆ X    ┆ A    ┆ A    ┆ 6         │
│ Y    ┆ Y    ┆ A    ┆ A    ┆ 7         │
│ X    ┆ X    ┆ A    ┆ C    ┆ 8         │
│ X    ┆ Y    ┆ A    ┆ C    ┆ 9         │
│ Y    ┆ X    ┆ A    ┆ C    ┆ 10        │
│ Y    ┆ Y    ┆ A    ┆ C    ┆ 11        │
│ X    ┆ X    ┆ A    ┆ C    ┆ 12        │
│ X    ┆ Y    ┆ A    ┆ C    ┆ 13        │
│ Y    ┆ X    ┆ A    ┆ C    ┆ 14        │
│ Y    ┆ Y    ┆ A    ┆ C    ┆ 15        │
│ X    ┆ X    ┆ C    ┆ A    ┆ 16        │
│ X    ┆ Y    ┆ C    ┆ A    ┆ 17        │
│ Y    ┆ X    ┆ C    ┆ A    ┆ 18        │
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 19        │
│ X    ┆ X    ┆ C    ┆ A    ┆ 20        │
│ X    ┆ Y    ┆ C    ┆ A    ┆ 21        │
│ Y    ┆ X    ┆ C    ┆ A    ┆ 22        │
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 23        │
│ X    ┆ X    ┆ C    ┆ C    ┆ 24        │
│ X    ┆ Y    ┆ C    ┆ C    ┆ 25        │
│ Y    ┆ X    ┆ C    ┆ C    ┆ 26        │
│ Y    ┆ Y    ┆ C    ┆ C    ┆ 27        │
│ X    ┆ X    ┆ C    ┆ C    ┆ 28        │
│ X    ┆ Y    ┆ C    ┆ C    ┆ 29        │
│ Y    ┆ X    ┆ C    ┆ C    ┆ 30        │
│ Y    ┆ Y    ┆ C    ┆ C    ┆ 31        │
└──────┴──────┴──────┴──────┴───────────┘

df2数据结构 #

┌──────┬──────┬──────┬──────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 │
│ ---  ┆ ---  ┆ ---  ┆ ---  │
│ str  ┆ str  ┆ str  ┆ str  │
╞══════╪══════╪══════╪══════╡
│ Y    ┆ Y    ┆ B    ┆ A    │
│ Y    ┆ Y    ┆ C    ┆ A    │
│ Y    ┆ Z    ┆ A    ┆ A    │
└──────┴──────┴──────┴──────┘

场景1：严格小于（lt）连接 #

选择首个匹配项 #

df2.join_asof_lt(
    df1,
    by=['by_1', 'by_2'],
    on=['on_1', 'on_2'],
    lt_select_eq = 'first',
)

输出结果：

┌──────┬──────┬──────┬──────┬───────────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 ┆ __index__ │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---       │
│ str  ┆ str  ┆ str  ┆ str  ┆ i64       │
╞══════╪══════╪══════╪══════╪═══════════╡
│ Y    ┆ Y    ┆ B    ┆ A    ┆ 11        │ # 严格小于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，取首个匹配项11
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 11        │ # 严格小于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，取首个匹配项11
│ Y    ┆ Z    ┆ A    ┆ A    ┆ null      │ # 无匹配分组，返回null
└──────┴──────┴──────┴──────┴───────────┘

选择最后一个匹配项 #

df2.join_asof_lt(
    df1,
    by=['by_1', 'by_2'],
    on=['on_1', 'on_2'],
    lt_select_eq = 'last',
)

输出结果：

┌──────┬──────┬──────┬──────┬───────────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 ┆ __index__ │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---       │
│ str  ┆ str  ┆ str  ┆ str  ┆ i64       │
╞══════╪══════╪══════╪══════╪═══════════╡
│ Y    ┆ Y    ┆ B    ┆ A    ┆ 15        │ # 严格小于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，取最后一个匹配项15
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 15        │ # 严格小于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，取最后一个匹配项15
│ Y    ┆ Z    ┆ A    ┆ A    ┆ null      │ # 无匹配分组，返回null
└──────┴──────┴──────┴──────┴───────────┘

场景2：小于等于（leq）连接 #

严格小于取首个，等于取最后一个 #

df2.join_asof_leq(
    df1,
    by=['by_1', 'by_2'],
    on=['on_1', 'on_2'],
    lt_select_eq = 'first',
    eq_select_eq = 'last',
)

输出结果：

┌──────┬──────┬──────┬──────┬───────────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 ┆ __index__ │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---       │
│ str  ┆ str  ┆ str  ┆ str  ┆ i64       │
╞══════╪══════╪══════╪══════╪═══════════╡
│ Y    ┆ Y    ┆ B    ┆ A    ┆ 11        │ # 小于等于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，严格小于场景取首个匹配项11
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 23        │ # 小于等于的匹配组为('Y', 'Y'), ('C', 'A')，对应索引19和23，等于场景取最后一个匹配项23
│ Y    ┆ Z    ┆ A    ┆ A    ┆ null      │ # 无匹配分组，返回null
└──────┴──────┴──────┴──────┴───────────┘

严格小于取最后一个，等于取首个 #

df2.join_asof_leq(
    df1,
    by=['by_1', 'by_2'],
    on=['on_1', 'on_2'],
    lt_select_eq = 'last',
    eq_select_eq = 'first',
)

输出结果：

┌──────┬──────┬──────┬──────┬───────────┐
│ by_1 ┆ by_2 ┆ on_1 ┆ on_2 ┆ __index__ │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---       │
│ str  ┆ str  ┆ str  ┆ str  ┆ i64       │
╞══════╪══════╪══════╪══════╪═══════════╡
│ Y    ┆ Y    ┆ B    ┆ A    ┆ 15        │ # 小于等于的匹配组为('Y', 'Y'), ('A', 'C')，对应索引11和15，严格小于场景取最后一个匹配项15
│ Y    ┆ Y    ┆ C    ┆ A    ┆ 19        │ # 小于等于的匹配组为('Y', 'Y'), ('C', 'A')，对应索引19和23，等于场景取首个匹配项19
│ Y    ┆ Z    ┆ A    ┆ A    ┆ null      │ # 无匹配分组，返回null
└──────┴──────┴──────┴──────┴───────────┘

上述lt/leq的逻辑可直接扩展至gt/geq场景，实现严格大于、大于等于的自定义规则连接。

内容的提问来源于stack exchange，提问作者T.H Rice