用一对特征序列计算Spark上的Chi Square

Posted 2021-05-01

实际上，有几天我坚持这个挑战，我想知道你的想法。我在pyspark中有一个大数据框，有大约150个功能。每行属于一个人，每个单元格的值是0或1.我有兴趣知道这些功能之间是否存在任何相关性。因为这些功能是分类的，所以我选择了Chi Square Independence测试。

对于每一对，我使用交叉表创建了列联表，然后将corsstab输出转换为密集矩阵并计算每对的p值。对于小规模，它可以工作，但对于大数据量，代码会长时间运行，在我看来效率不高。问题是因为我需要所有功能的组合，它会以指数方式增加计算大小。

您建议采用什么方法以更好的性能完成这项工作？

def calculate_chi(df, col1, col2):
   test = df.crosstab(col2, col1) 
   vals = test.rdd.map(lambda x:[x[1],x[2]]).flatMap(lambda y: y).collect()
   Matrix = Matrices.dense(2, 2, vals)
   return Statistics.chiSqTest(Matrix).pValue

if __name__=="__main__":
   spark = SparkSession.builder.getOrCreate()
   df=spark.read.option("header",True).csv("s3a://*.csv")
   df = df.select(selected_cols)
   data = []
   cols =  combinations(df.columns,2)
   for c in cols:
      data.append((str(c), calculate_chi(df, c[0],c[1])))
   df_results = spark.createDataFrame(data, ["feature_pair","pvalue"])
   df_results.coalesce(1).write.format('csv').options(header='true').save("s3a://chiSqTest")

答案

IMO，问题在于calculate_chi函数的第3行。您使用了未分布的矩阵表示，因此驱动程序节点在内存和CPU方面都是算法的瓶颈。我建议你考虑使用here的分布式矩阵表示法。由于可以独立地对每列进行χ2测试，我认为您应该能够更好地利用群集的功能。