现在我有一个数据框架:df = pd.DataFrame({"a":[1,2,3,4,5],"b":[2,3,4,5,6],"c":[3,4,5,6,7]})
功能:
def fun(a,b,shift_len):
return a+b*shift_len,b-shift_len
我可以通过以下方式得到结果:
df[["d","e"]] = df.apply(lambda row:fun(row["a"],row["b"],3),axis=1,result_type="expand")
我想知道如何使用极值来得到相同的结果?
答案取决于你是否可以使用polar表达式重写你的函数。
使用极性表达式
要获得polar的最佳性能，请尝试使用表达式对计算进行编码。表达式产生了性能最好的、令人尴尬的并行解。
例如，函数可以表示为:
shift_len = 3
df.with_columns(
[
(pl.col("a") + (pl.col("b") * shift_len)).alias("d"),
(pl.col("b") - shift_len).alias("e"),
]
)
shape: (5, 5)
┌─────┬─────┬─────┬─────┬─────┐
│ a ┆ b ┆ c ┆ d ┆ e │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╡
│ 1 ┆ 2 ┆ 3 ┆ 7 ┆ -1 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2 ┆ 3 ┆ 4 ┆ 11 ┆ 0 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 3 ┆ 4 ┆ 5 ┆ 15 ┆ 1 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 4 ┆ 5 ┆ 6 ┆ 19 ┆ 2 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 5 ┆ 6 ┆ 7 ┆ 23 ┆ 3 │
└─────┴─────┴─────┴─────┴─────┘
polar将并行运行两个表达式，产生非常快的结果。
使用apply
让我们假设你不能将你的函数编码为极性表达式(例如，你需要使用一个外部库)。由于您的函数接受多个参数并返回多个值，因此我们将分步骤进行。
传递多个值
我们可以通过"图章耦合"将多个值传递给apply中的fun函数;使用polars.struct将多个列合并为单个序列。在lambda函数中，这些值作为Pythondict传递，列的名称作为键。因此，例如，我们将下面lambda中a列中的值访问为cols["a"]。
df.with_column(
pl.struct(["a", "b"])
.apply(lambda cols: fun(cols["a"], cols["b"], 3))
.alias("result")
)
shape: (5, 4)
┌─────┬─────┬─────┬─────────┐
│ a ┆ b ┆ c ┆ result │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ object │
╞═════╪═════╪═════╪═════════╡
│ 1 ┆ 2 ┆ 3 ┆ (7, -1) │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 2 ┆ 3 ┆ 4 ┆ (11, 0) │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 4 ┆ 5 ┆ (15, 1) │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 4 ┆ 5 ┆ 6 ┆ (19, 2) │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌┤
│ 5 ┆ 6 ┆ 7 ┆ (23, 3) │
└─────┴─────┴─────┴─────────┘
result列包含fun函数返回的元组。但是，请注意result列的类型:object。object类型的列在polar中不是很有用，并且功能有限。
处理多个返回值
接下来我们将fun返回的tuple函数转换成更有用的东西:字典的键值对,其中键是所需的列名(d和e例子)。
为此,我们将使用Python的zip函数和一个元组所需的名字。
当我们运行这段代码时，将得到一个类型为struct的列。
df.with_column(
pl.struct(["a", "b"])
.apply(lambda cols: dict(zip(("d", "e"), fun(cols["a"], cols["b"], 3))))
.alias("result")
)
df
shape: (5, 4)
┌─────┬─────┬─────┬───────────┐
│ a ┆ b ┆ c ┆ result │
│ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ struct[2] │
╞═════╪═════╪═════╪═══════════╡
│ 1 ┆ 2 ┆ 3 ┆ {7,-1} │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 2 ┆ 3 ┆ 4 ┆ {11,0} │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 3 ┆ 4 ┆ 5 ┆ {15,1} │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 4 ┆ 5 ┆ 6 ┆ {19,2} │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┤
│ 5 ┆ 6 ┆ 7 ┆ {23,3} │
└─────┴─────┴─────┴───────────┘
名称d和e没有出现在result列的输出中，但是它们存在。
使用unnest
在最后一步中，我们将使用unnest函数将结构体分成两个新列。
df.with_column(
pl.struct(["a", "b"])
.apply(lambda cols: dict(zip(("d", "e"), fun(cols["a"], cols["b"], 3))))
.alias("result")
).unnest("result")
shape: (5, 5)
┌─────┬─────┬─────┬─────┬─────┐
│ a ┆ b ┆ c ┆ d ┆ e │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╪═════╪═════╡
│ 1 ┆ 2 ┆ 3 ┆ 7 ┆ -1 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 2 ┆ 3 ┆ 4 ┆ 11 ┆ 0 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 3 ┆ 4 ┆ 5 ┆ 15 ┆ 1 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 4 ┆ 5 ┆ 6 ┆ 19 ┆ 2 │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌┤
│ 5 ┆ 6 ┆ 7 ┆ 23 ┆ 3 │
└─────┴─────┴─────┴─────┴─────┘
一个谨慎使用apply与外部库和/或自定义Python字节码使您的代码服从Python GIL。结果是非常缓慢的单线程性能-无论如何编码。因此，我强烈建议避免使用apply和自定义Python函数，如果可以的话，尝试只使用polar表达式来编写算法。
用args传递参数
import pandas as pd
df1 = pd.DataFrame({"a":[1,2,3,4,5],"b":[2,3,4,5,6],"c":[3,4,5,6,7]})
def t(df, row1, row2, shift_len):
return df[row1] + df[row2] * shift_len, df[row2] - shift_len
df1[["d", "e"]] = df1.apply(t, args=("a", "b", 3), axis=1, result_type="expand")
print(df1)
OUTPUT:
a b c d e
0 1 2 3 7 -1
1 2 3 4 11 0
2 3 4 5 15 1
3 4 5 6 19 2
4 5 6 7 23 3