在 pyspark 数据帧数据中查找阈值

Question

总的来说，我对数据帧和 pyspark 完全陌生。

在 python 中，我想做的只是微不足道的 - 但是我似乎找不到使用 pyspark 不需要很长时间的方法。

我有一个大约 4000 行的 pyspark 数据框，其架构如下：

root
 |-- waveformData: struct (nullable = true)
 |    |-- elements: array (nullable = true)
 |    |    |-- element: double (containsNull = true)
 |    |-- dimensions: array (nullable = true)
 |    |    |-- element: integer (containsNull = true)

每个数组大约有 20000 个双精度数。

搜索此数组以查找最大值和阈值（最大值的 50% 的第一个实例）只需要很少的时间 - 但只有在数据采用“正常”格式时（numpy 数组）。

我使用的是基本款：

wav_df = temp_data.select("waveformData").toPandas()
wav = wav_df.to_numpy()[0][0].get("elements")

然后搜索最大值/阈值

但是 'toPandas' 步骤需要很长时间（比如单行需要 30 秒）

为什么？

我一直在尝试使用 .collect 等对 pyspark 数据帧进行操作以避免这种转换，但我尝试的一切都需要很长时间。

如果 pyspark 是针对大数据的，那我肯定做错了，这不可能是处理这么多数据的正常时间。

我错过了什么？

Answer 1

我创建了一些随机测试数据

from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.sql import functions as F

import numpy as np
np.random.seed(123)

spark = SparkSession.builder.getOrCreate()


n = 20000
for i in range(100):
    df = spark.createDataFrame([
        Row(waveFormData=Row(elements=[float(v) for v in np.random.randn(n)], dimensions=[n])) for i in range(40)
    ])
    df.write.parquet('waveFormData.parquet', mode='append')

当我加载数据并选择在 2 秒内运行的数组的最大值时：

df = spark.read.parquet('waveFormData.parquet')

df.select(F.array_max('waveFormData.elements')).toPandas()