将列表列转换为 Python Polars 中的虚拟列？

Question

我有一个非常大的数据框，其中有一列是代表类别成员资格的数字列表。

这是一个虚拟版本

import pandas as pd 
import numpy as np 

segments = [str(i) for i in range(1_000)]

# My real data is ~500m rows
nums = np.random.choice(segments, (100_000,10))

df = pd.DataFrame({'segments': [','.join(n) for n in nums]})

用户ID	细分
0	885,106,49,138,295,254,26,460,0,844
1	908,709,454,966,151,922,666,886,65,708
2	664,713,272,241,301,498,630,834,702,289
3	60,880,906,471,437,383,878,369,556,876
4	817,183,365,171,23,484,934,476,273,230
...	...

请注意，有一个已知的段列表（示例中为 0-999）

我想将其转换为虚拟列，指示每个段的成员资格。

我找到了几种方法：

在

pandas

：

df_one_hot_encoded = (df['segments']
    .str.split(',')
    .explode()
    .reset_index()
    .assign(__one__=1)
    .pivot_table(index='index', columns='segments', values='__one__', fill_value=0)
)

（100k 行样本需要 8 秒）

还有

polars

df_ans = (df2
          .with_columns(
              pl.int_range(pl.len()).alias('row_index'), 
              pl.col('segments').str.split(','), 
              pl.lit(1).alias('__one__')
          )
          .explode('segments')
          .pivot(on='segments', index='row_index', values='__one__', aggregate_function='first')
          .fill_null(0)
     )
df_one_hot_encoded = df_ans.to_pandas()

（需要 1.5 秒，包括与 pandas 之间的转换，不包括 0.9 秒）

但是，我听说 .pivot 效率不高，并且它不能很好地处理惰性框架。

我尝试了

polars

中的其他解决方案，但它们慢得多：

_ = df2.lazy().with_columns(**{segment: pl.col('segments').str.contains(segment) for segment in segments}).collect()

（2秒）

(df2
 .with_columns(
     pl.int_range(pl.len()).alias('row_index'), 
     pl.col('segments').str.split(',')
 )
 .explode('segments')
 .to_dummies(columns=['segments'])
 .group_by('row_index')
 .sum()
)

（4秒）

有人知道比 .9s 枢轴更好的解决方案吗？

Answer 1

这种方法最终比枢轴慢，但它有一个不同的技巧，所以我将包括它。

df2=pl.from_pandas(df)
df2_ans=(df2.with_row_count('userId').with_column(pl.col('segments').str.split(',')).explode('segments') \
    .with_columns([pl.when(pl.col('segments')==pl.lit(str(i))).then(pl.lit(1,pl.Int32)).otherwise(pl.lit(0,pl.Int32)).alias(str(i)) for i in range(1000)]) \
    .groupby('userId')).agg(pl.exclude('segments').sum())
df_one_hot_encoded = df2_ans.to_pandas()   

其他一些观察结果。 我不确定您是否检查了

str.contains

方法的输出，但我认为这行不通，因为例如，在查看字符串时，15 包含在 154 中。

另一件事，我想这只是一种偏好，是

with_row_count

语法与

pl.arrange

。 我不认为两者的性能更好（至少不是很明显），但您不必引用 df 名称来获取它的

len

，这很好。

我尝试了其他一些更糟糕的事情，包括不进行爆炸而只进行

is_in

，但速度较慢。 我尝试使用布尔值而不是 1 和 0，然后使用

any

进行聚合，但速度较慢。