如何在Polars或Pyarrow中获取字符串列

Question

我有一个带有字符串键列的pandas dataframe / polars dataframe / pyarrow表。您可以假设字符串是随机的。我想根据此密钥列将该数据框架分为n个较小的数据范围。

使用整数列，我只能使用 df1 = df [df.key％n == 1] ， df2 = df [df.key％n == 2] 等等。

我最好猜测您将如何使用字符串列执行此操作是应用哈希函数（例如，将字符串的ASCII值求和）转换为整数列，然后使用模量。

请让我知道在Pandas，Polars或Pyarrow中可以做到的最有效的方法，理想情况下，API中的纯柱状操作。对于我的用例来说，做DF。可能太慢了。

Answer 1

我对@cbilots的答案有很小的补充。 Polars具有 Hash 表达式，因此计算分区ID将是微不足道的。

如果将其与 partition_by 相结合，则可以以较高的速度创建以下速度分区的分区：

df = pl.DataFrame({
    "strings": ["A", "A", "B", "A"],
    "payload": [1, 2, 3, 4]
})


N = 2
(df.with_columns(
     (pl.col("strings").hash() % N).alias("partition_id")
).partition_by("partition_id"))

[shape: (3, 3)
 ┌─────────┬─────────┬──────────────┐
 │ strings ┆ payload ┆ partition_id │
 │ ---     ┆ ---     ┆ ---          │
 │ str     ┆ i64     ┆ u64          │
 ╞═════════╪═════════╪══════════════╡
 │ A       ┆ 1       ┆ 0            │
 │ A       ┆ 2       ┆ 0            │
 │ A       ┆ 4       ┆ 0            │
 └─────────┴─────────┴──────────────┘,
 shape: (1, 3)
 ┌─────────┬─────────┬──────────────┐
 │ strings ┆ payload ┆ partition_id │
 │ ---     ┆ ---     ┆ ---          │
 │ str     ┆ i64     ┆ u64          │
 ╞═════════╪═════════╪══════════════╡
 │ B       ┆ 3       ┆ 1            │
 └─────────┴─────────┴──────────────┘]

分组和分区的实现将并行完成。

Answer 2

我会尝试使用 查看其在数据集和计算平台上的性能。 （请注意，在计算中，我有效地仅选择键字段并在此上运行 hash_rows ）

N = 50
df = df.with_columns(
    pl.lit(df.select('key').hash_rows() % N).alias('hash')
)

我只是在数据集中运行了此数据集，其中有近4900万条记录一个32核系统，并在几秒钟内完成。 （我的数据集中的“键”字段是人的姓氏。）

我还应该注意，有一个 partition_by 可能在分区中有帮助的方法。