map_batches在极坐标中的表达有什么意义？

Question

我正在从官方文档中阅读有关map_batches的内容 这里：

尽管示例描述表明以下示例代码应该为map_elements和map_batches返回不同的结果，但我得到了相同的结果

out = df.group_by("keys", maintain_order=True).agg(
    pl.col("values").map_batches(lambda s: s.shift()).alias("shift_map_batches"),
    pl.col("values").shift().alias("shift_expression"),
)
print(out)
output:
shape: (2, 3)
┌──────┬───────────────────┬──────────────────┐
│ keys ┆ shift_map_batches ┆ shift_expression │
│ ---  ┆ ---               ┆ ---              │
│ str  ┆ list[i64]         ┆ list[i64]        │
╞══════╪═══════════════════╪══════════════════╡
│ a    ┆ [null, 10]        ┆ [null, 10]       │
│ b    ┆ [null]            ┆ [null]           │
└──────┴───────────────────┴──────────────────┘

您可以看到代表两种方法结果的两列是相同的

因此，尽管文档说两个结果应该不同，但似乎没有任何区别。其次，在非分组场景中，我们将整个系列传递给map_batches，并按元素传递给map_elements，因此，在这里我们似乎也可以通过函数传递系列并完成工作。

仅以文档为例：

counter = 0


def add_counter(val: int) -> int:
    global counter
    counter += 1
    return counter + val


out1 = df.select(
    pl.col("values").map_batches(add_counter).alias("solution_map_elements")
)

counter = 0
out2 = df.select(
    add_counter(pl.col("values"))
)
print(out1, out2)

所以，我没有看到 map_batches 的用例。我想了解我们可以在哪里使用map_batches？

Answer 1

文档需要更新。

map_batches

过去的行为有所不同，该示例需要修改。

• https://github.com/pola-rs/polars/issues/14699#issuecomment-1969240018

• https://github.com/pola-rs/polars/issues/14521

---

用例 #1 - 性能

文档中的另一个提及：

• “映射函数的一个合理用例是使用第三方库转换表达式表示的值。”

如果我们使用

map_elements

作为基本情况，我们在每个

“Row”

上调用 Python 的 .swapcase()。

df = pl.select(text = pl.repeat(pl.lit("áéíóúüñ" * 1000), n=100_000))

df.with_columns(out = pl.col("text").map_elements(lambda text: 
    text.swapcase(),
    return_dtype = pl.String
))

Elapsed time: 11.93447 seconds

例如，

PyArrow 有一个

utf8_swapcase()

计算函数。

使用

map_batches

我们可以传递整个“列”并且utf8_swapcase被调用一次。

import pyarrow.compute

df.with_columns(out = pl.col("text").map_batches(lambda text: 
    pl.Series(pyarrow.compute.utf8_swapcase(text.to_arrow()))
))

Elapsed time: 4.05791 seconds

---

用例 #2 - 访问表达式内的 DataFrame API

最初引入

Expr.replace()

时，是用 Python 实现并使用

.map_batches

一个过于简单的例子：

def my_custom_function(col, mapping, default=None):
    df = pl.LazyFrame({
        "old": mapping.keys(),
        "new": mapping.values()
    })
   
    select_expr = pl.col("new")
    if default is not None:
        select_expr = select_expr.fill_null(default)
        
    return (
        col.to_frame("old")
           .lazy()
           .join(df, on="old", how="left")
           .select(select_expr)
           .collect()
           .to_series()
    )

我们在函数内调用

Series.to_frame()

来获取 DataFrame。

这使我们能够访问 DataFrame 方法，例如

.join()

用于映射旧的 -> 值。

lf = pl.LazyFrame({
   "a": [1, 2, 3, 4],
   "b": [5, 6, 7, 8]
})

lf.with_columns(
   c = pl.col("a").map_batches(lambda col: 
      my_custom_function(
         col, 
         mapping = {2: "TWO", 4: "FOUR"}, 
         default = "FOUR O FOUR"
      )
   )
).collect()

shape: (4, 3)
┌─────┬─────┬─────────────┐
│ a   ┆ b   ┆ c           │
│ --- ┆ --- ┆ ---         │
│ i64 ┆ i64 ┆ str         │
╞═════╪═════╪═════════════╡
│ 1   ┆ 5   ┆ FOUR O FOUR │
│ 2   ┆ 6   ┆ TWO         │
│ 3   ┆ 7   ┆ FOUR O FOUR │
│ 4   ┆ 8   ┆ FOUR        │
└─────┴─────┴─────────────┘

“最终用户”可能不会使用太多。