Apache spark 矢量化以长格式来源的高维数据

我有一个非常大的长格式数据帧（几十亿行），我想对其进行透视和矢量化（使用VectorAssembler），目的是使用类似于TF-IDF的东西来降低维度。一旦数据透视，数据帧将有约1.3亿列

源长格式架构如下所示：

root
 |-- entity_id: long (nullable = false)
 |-- attribute_id: long (nullable = false)
 |-- event_count: integer (nullable = true)

按照以下步骤旋转失败，耗尽执行器和驱动程序内存。我可以增加这两个，但我的感觉是，旋转然后使用矢量汇编程序不是解决这个问题的正确方法

wide_frame = (
    long_frame.groupBy("entity_id")
    .pivot("attribute_id")
    .agg(F.first("event_count"))
)

为了实现每个实体的属性向量的最终目标，我是否应该尝试其他火花模式