使用 pyspark 搜索布尔矩阵

Question

我有一个 M x N 的布尔矩阵，其中 M = 6000，N = 1000

1 | 0 1 0 0 0 1 ----> 1000
2 | 1 0 1 0 1 0 ----> 1000
3 | 0 0 1 1 0 0 ----> 1000
  V
6000

现在对于每一列，我想找到第一个值为 1 的地方。对于上面的示例，在前 5 列中，我想要 2 1 2 3 2 1。

现在我的密码是

    sig_matrix = list()
    num_columns = df.columns
    for col_name in num_columns:
        print('Processing column {}'.format(col_name))
        sig_index = df.filter(df[col_name] == 1).\
                    select('perm').limit(1).collect()[0]['perm']
        sig_matrix.append(sig_index)

现在上面的代码真的很慢，我需要 5~7 分钟来解析 1000 列，有没有比我现在做的更快的方法来做到这一点？如果速度更快，我也愿意使用 pandas 数据框而不是 pyspark 数据框。

Answer 1

这是一个对我来说运行 <1s 的 numpy 版本，因此对于这种数据量应该更可取：

arr=np.random.choice([0,1], size=(6000,1000))
[np.argwhere(arr[:,i]==1.)[0][0] for i in range(1000)]

可能会有更高效的 numpy 解决方案。

Answer 2

我最终使用 numpy 解决了我的问题。我是这样做的。

import numpy as np

sig_matrix = list()
    columns = list(df)
    for col_name in columns:
        sig_index = np.argmax(df[col_name]) + 1
        sig_matrix.append(sig_index)

由于我列中的值是 0 和 1，因此 argmax 将 return 第一次出现值 1。