pyspark火花RDD到3D矩阵

Question

我正在使用pyspark来执行矩阵运算。我一直在尝试将火花数据帧转换为3D矩阵。 我的数据框看起来像这样（x_input）：

Prod  | Location  |    Day |    Sales  
---------------------------------
0       0           0          65.4
1       0           0          96.1
2       0           0          98.1
3       0           0          9.5
0       0           1          16.1
1       0           1          59.4
2       0           1          17
3       0           1          92.8
0       1           0          52.3
1       1           0          41.1
2       1           0          14.6
3       1           0          92.2
0       1           1          27.7
1       1           1          69
2       1           1          50.7
3       1           1          53.4
0       2           0          54
1       2           0          50.4
2       2           0          54.6
3       2           0          1.3
0       2           1          2.7
1       2           1          60.6
2       2           1          77.7
3       2           1          21.7

Prod，Location和Day列是矩阵的索引，我想把Sales值放在一起，输出应该是这样的：

[
[
  [65.4, 96.1, 98.1, 9.5],
  [16.1, 59.4, 17, 92.8]
],
[
  [52.3, 41.1, 14.6, 92.2],
  [27.7, 69, 50.7, 53.4]
],
[
  [54, 50.4, 54.6, 1.3],
  [2.7, 60.6, 77.7, 21.7]
]
]

我尝试关注this post并提出了这个解决方案：

x_rdd = x_input.rdd

rows = (x_rdd
       .zipWithIndex()
       .groupBy(lambda (x,i): i/4)
       .mapValues(lambda vals: [x.Sales for (x,i) in sorted(vals, key=lambda (x, i): i)]))

f = rows.map(lambda x: x[1])

这样做之后，f现在是一个2D矩阵，如何将其转换为如上所述的3D矩阵？
由于3D矩阵具有较大的尺寸，因此无法实现for循环方法。