使用 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 将返回值 1 的第一个匹配项。