使用一个数据帧中的值从另一个数据帧中选择子字符串

Question

给定一个包含字符串的数据帧，如何使用第二个数据帧从第一个数据帧中获取切片作为新列？

df_strs = pl.DataFrame({
    "Col_A": ["AAABBCCCC", "DDDEEFFFF"],
    "Col_B": ["AAB", "DDE"]
})
print(df_strs)

df_offsets = pl.DataFrame({
    "Reference_Col": ["Col_A", "Col_A", "Col_A", "Col_B", "Col_B"],
    "Offset": [0, 3, 5, 0, 2],
    "Length": [3, 2, 4, 2, 1],
    "Alias": ["Col_A_1", "Col_A_2", "Col_A_3", "Col_B_1", "Col_B_2"]
})
print(df_offsets)

输出：

df_strs
shape: (2, 2)
┌───────────┬───────┐
│ Col_A     ┆ Col_B │
│ ---       ┆ ---   │
│ str       ┆ str   │
╞═══════════╪═══════╡
│ AAABBCCCC ┆ AAB   │
│ DDDEEFFFF ┆ DDE   │
└───────────┴───────┘

df_offsets
shape: (5, 4)
┌───────────────┬────────┬────────┬─────────┐
│ Reference_Col ┆ Offset ┆ Length ┆ Alias   │
│ ---           ┆ ---    ┆ ---    ┆ ---     │
│ str           ┆ i64    ┆ i64    ┆ str     │
╞═══════════════╪════════╪════════╪═════════╡
│ Col_A         ┆ 0      ┆ 3      ┆ Col_A_1 │
│ Col_A         ┆ 3      ┆ 2      ┆ Col_A_2 │
│ Col_A         ┆ 5      ┆ 4      ┆ Col_A_3 │
│ Col_B         ┆ 0      ┆ 2      ┆ Col_B_1 │
│ Col_B         ┆ 2      ┆ 1      ┆ Col_B_2 │
└───────────────┴────────┴────────┴─────────┘

因此生成的数据框有五列，其中每列都是

reference[offset:offset+length]

，名称由

alias

?

定义

在上面的例子中，结果将是：

shape: (5, 4)
┌─────────┬─────────┬─────────┬─────────┬─────────┐
│ Col_A_1 ┆ Col_A_2 ┆ Col_A_3 ┆ Col_B_1 ┆ Col_B_2 │
│ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
│ str     ┆ str     ┆ str     ┆ str     ┆ str     │
╞═════════╪═════════╪═════════╪═════════╪═════════╡
│ AAA     ┆ BB      ┆ CCCC    ┆ AA      ┆ B       │
│ DDD     ┆ EE      ┆ FFFF    ┆ DD      ┆ E       │
└─────────┴─────────┴─────────┴─────────┴─────────┘

我可以使用循环来做到这一点，但我觉得必须有一种更惯用的方法来使用极坐标。我目前的解决方案是：

df_offsets = df_offsets.transpose() # this converts our ints to strs
for s in df_offsets:
    df_strs = df_strs.with_columns(
        pl.col(s[0]).str.slice(int(s[1]), int(s[2])).alias(s[3])
    )

输出：

>>>print(df_strs)
shape: (2, 7)
┌───────────┬───────┬─────────┬─────────┬─────────┬─────────┬─────────┐
│ Col_A     ┆ Col_B ┆ Col_A_0 ┆ Col_A_1 ┆ Col_A_2 ┆ Col_B_0 ┆ Col_B_1 │
│ ---       ┆ ---   ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
│ str       ┆ str   ┆ str     ┆ str     ┆ str     ┆ str     ┆ str     │
╞═══════════╪═══════╪═════════╪═════════╪═════════╪═════════╪═════════╡
│ AAABBCCCC ┆ AAB   ┆ AAA     ┆ BB      ┆ CCCC    ┆ AA      ┆ B       │
│ DDDEEFFFF ┆ DDE   ┆ DDD     ┆ EE      ┆ FFFF    ┆ DD      ┆ E       │
└───────────┴───────┴─────────┴─────────┴─────────┴─────────┴─────────┘

有更优化的方法吗？

Answer 1

我可以想到两个潜在的优化。

主要：并行运行表达式

您的代码正在运行您的表达式，一次一个。 一个简单的优化是首先将表达式收集到列表中，然后使用

with_columns

并行运行所有表达式。

此外，让我们在惰性模式下运行它，这将应用进一步的优化。

请注意，我们不希望列表理解在惰性模式下完成——只有

df_strs

。 （LazyFrame 没有

transpose

方法。）

(
    df_strs
    .lazy()
    .with_columns(
        pl.col(s[0]).str.slice(int(s[1]), int(s[2])).alias(s[3])
        for s in df_offsets.transpose()
    )
    .collect()
)

shape: (2, 7)
┌───────────┬───────┬─────────┬─────────┬─────────┬─────────┬─────────┐
│ Col_A     ┆ Col_B ┆ Col_A_1 ┆ Col_A_2 ┆ Col_A_3 ┆ Col_B_1 ┆ Col_B_2 │
│ ---       ┆ ---   ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
│ str       ┆ str   ┆ str     ┆ str     ┆ str     ┆ str     ┆ str     │
╞═══════════╪═══════╪═════════╪═════════╪═════════╪═════════╪═════════╡
│ AAABBCCCC ┆ AAB   ┆ AAA     ┆ BB      ┆ CCCC    ┆ AA      ┆ B       │
│ DDDEEFFFF ┆ DDE   ┆ DDD     ┆ EE      ┆ FFFF    ┆ DD      ┆ E       │
└───────────┴───────┴─────────┴─────────┴─────────┴─────────┴─────────┘

仅此一项就可以提供巨大的加速。

次要：生成表达式列表。

您可以尝试更快速地生成表达式列表的方法。 以下内容可能会也可能不会加快表达式的创建速度（而不是转置

df_offsets

）。 你当然可以尝试一下。

(
    df_strs
    .lazy()
    .with_columns(
        eval(expr) for expr in
        df_offsets.select(
            pl.format(
                r"pl.col('{}').str.slice({}, {}).alias('{}')",
                *df_offsets.columns
            )
        )
        .to_series()
    )
    .collect()
)

权衡是我们在 Polars 内部生成表达式（作为字符串）而不转置 DataFrame ...但代价是使用 Python 的

eval

表达式将字符串转换为 Polars 表达式。