PySpark 中将数据帧“展平”为每组 1 行的最有效方法 [重复]

Question

我有一种情况，我有一个非常大的数据框，其中包含每个时间戳 5 行的飞机位置数据（请参阅下表以获取 1 个时间戳的示例......其中 3 行只有点，但想象它们有 plane_nums 和地点）

+---------------------------+-----------+-----------+-----------+-----------+
|         utc_time          | plane_num |    lat    |    lon    |  height   |
+---------------------------+-----------+-----------+-----------+-----------+
| 2021-06-02T05:01:40+00:00 |          1| 51.759014 | -1.256688 | 47.337597 |
| 2021-06-02T05:01:40+00:00 |          2| 41.758014 |  1.346678 | 41.632522 |
| ...                       |       ... |       ... |       ... |       ... |
| ...                       |       ... |       ... |       ... |       ... |
| ...                       |       ... |       ... |       ... |       ... |
+---------------------------+-----------+-----------+-----------+-----------+

对于每个时间戳，我想将此数据框折叠到 1 行，其中对一组 2 个平面之间的各种距离进行各种计算，并将它们的所有纬度/经度放在一行中。

我知道我可以在 pyspark 中通过将数据帧过滤到每个平面来轻松完成此操作（每个时间戳总会有一个平面 1、2、3、4、5）。请参阅下面的代码，该代码为 1 个时间戳创建数据帧，但想象一下有数百万个时间戳。

    dfAirplaneData.filter = spark.createDataFrame([
(2021-06-02T05:01:40+00:00, 1, 51.759014, -1.256688, 47.337597),
 (2021-06-02T05:01:40+00:00, 2, 41.758014, 1.346678, 41.632522),
 (2021-06-02T05:01:40+00:00, 3, 41.758014, 1.346678, 11.632522),
 (2021-06-02T05:01:40+00:00, 4, 41.758014, 1.346678, 21.632522),
 (2021-06-02T05:01:40+00:00, 5, 11.758014, 1.346678, 41.632522)
], ("utc_time", "plane_num", "lat", "lon", "height"))

dfPlane1 = dfAirplaneData.filter(F.col('plane_num')==1))
dfPlane2 = dfAirplaneData.filter(F.col('plane_num')==2))
dfPlane3 = dfAirplaneData.filter(F.col('plane_num')==3))

然后将数据框多次连接回自身以获得一行，但这感觉效率低下。在 pyspark 中是否有更好的方法来执行此操作（可能使用带有时间戳的 .groupBy）？

我知道除了加入之外，我还可以使用： w = Window.partitionBy(utc_time) 但我每组只需要 1 行，所以我最终会为每一行进行计算/复制工作，然后过滤到 1 个汇总行。另外，在进行比较时我需要知道哪个 plane_num ，所以不确定使用Window.partitionBy 是否可以轻松实现@

我已经简要了解了Pandas UDF grouped map，如果这是最好的方法，我相信我可以使用分组映射来遍历熊猫数据框。

对解决此问题的最有效、最简单的方法感兴趣。如果加入同一个数据框是最好的，那么我可以使用这种方法。

Answer 1

在 Spark 中，您需要 groupBy DataFrame API 调用。在这种情况下，dfAirplaneData.groupBy('utc_time').agg(collect_list(struct(col('plane_num'),col('lat'),col('lon'),col('height')))) 之类的东西将为您提供每个时间戳的结构数组，然后您可以对其执行计算。我理解你的意思吗？

Answer 2

这是一种在任何计算之前将所有数据放在一行中的方法。您可以使用一个函数来进行计算，而不是使用 lambda 函数来重塑数据。

data='''utc_time, plane_num, lat, lon, height
2021-06-02T05:01:40+00:00, 1, 51.759014, -1.256688, 47.337597
2021-06-02T05:01:40+00:00, 2, 41.758014, 1.346678, 41.632522
2021-06-02T05:01:40+00:00, 3, 41.758014, 1.346678, 11.632522
2021-06-02T05:01:40+00:00, 4, 41.758014, 1.346678, 21.632522
2021-06-02T05:01:40+00:00, 5, 11.758014, 1.346678, 41.632522
2021-07-02T05:01:40+00:00, 1, 51.759014, -1.256688, 47.337597
2021-07-02T05:01:40+00:00, 2, 41.758014, 1.346678, 41.632522
2021-07-02T05:01:40+00:00, 3, 41.758014, 1.346678, 11.632522
2021-07-02T05:01:40+00:00, 4, 41.758014, 1.346678, 21.632522
2021-07-02T05:01:40+00:00, 5, 11.758014, 1.346678, 41.632522
'''
df = pd.read_csv(io.StringIO(data), sep=',', engine='python')
# df.set_index('utc_time', inplace=True)
dft = df.groupby('utc_time').apply(lambda x: pd.DataFrame(x.iloc[:, 1:5].values.reshape(1,-1))).reset_index().drop('level_1', axis=1)
dft.columns = ['utc_time'] + ["plane_num", "lat", "lon", "height"]*5

dft

                    utc_time  plane_num        lat       lon     height  plane_num        lat       lon     height  plane_num  ...       lon     height  plane_num        lat       lon     height  plane_num        lat       lon     height
0  2021-06-02T05:01:40+00:00        1.0  51.759014 -1.256688  47.337597        2.0  41.758014  1.346678  41.632522        3.0  ...  1.346678  11.632522        4.0  41.758014  1.346678  21.632522        5.0  11.758014  1.346678  41.632522
1  2021-07-02T05:01:40+00:00        1.0  51.759014 -1.256688  47.337597        2.0  41.758014  1.346678  41.632522        3.0  ...  1.346678  11.632522        4.0  41.758014  1.346678  21.632522        5.0  11.758014  1.346678  41.632522

[2 rows x 21 columns]