pandas - 使用大型数据集在 pyspark 中获取相关矩阵

Question

我想计算大型数据集(100 万行)的相关矩阵。这个想法是计算产品销售的相关性。如果两种产品的销售额同比增减相似，则可能存在相关性。

我已经试过这里的帖子了:

• How to get correlation matrix values pyspark
• How to get the correlation matrix of a pyspark data frame?
• How to plot correlation heatmap when using pyspark+databricks
• https://gist.github.com/cameres/bc24ac6711c9e537dd20be47b2a83558

它们或多或少都做同样的事情，但它们会在驱动程序处收集相关矩阵。这是一个问题，因为大型数据集使该集合RAM 密集。我正在寻找一种方法将这个问题分解成多个部分并利用 Spark 的分布式计算。有 170,000 种不同的产品，因此作业运行了 170,000 次，并且有 29B 种组合。

我的想法是逐列计算相关性(交叉应用)，然后将其收集到数据框(或 RDD)中以对其运行过滤器(仅相关性 > 0.8)。但我没有好的想法开始。

数据集基本上是这样的。

d = {'Product': ['A', 'B', 'C','A', 'B', 'C','A', 'B', 'C'],\
     'Year': [2010, 2010, 2010, 2011, 2011, 2011, 2012, 2012, 2012],\
     'Revenue': [100, 200, 300, 110, 190, 320, 120, 220, 350]}
df = pd.DataFrame(data=d)

我将数据转置为列中的年份。

df = df.pivot(index='Product', columns='Year', values='Revenue').fillna(0)

我计算出 pct_change 具有逐年的相对变化。

df_diff = df.pct_change(axis=1).replace([np.inf, -np.inf], np.nan).fillna(0)

Year     2010      2011      2012
Product                          
A         0.0  0.100000  0.090909
B         0.0 -0.050000  0.157895
C         0.0  0.066667  0.093750

我需要相关性...有了 Pandas 很容易

# change structure
df_diff = df_diff.stack().unstack(level=0)
# get correlation
df_diff = df_diff.corr().abs()
# change structure back
df_diff = df_diff.unstack().to_frame(name='value')
df_diff.index = df_diff.index.set_names(['Product_1', 'Product_2'])
df_diff.reset_index(inplace=True)

    Product_1 Product_2     value
0         A         A  1.000000
1         A         B  0.207317
2         A         C  0.933485
3         B         A  0.207317
4         B         B  1.000000
5         B         C  0.544352
6         C         A  0.933485
7         C         B  0.544352
8         C         C  1.000000

Answer 1

我使用了一个 udf 并将其映射到 spark df。使用 numOfPartitions，您可以控制生成并分发到工作节点的任务数。在我的示例中，我使用了 16 个节点，每个节点有 8 个 cpu，并将 df 分成 10000 个分区。

import pandas as pd
import numpy as np

d = {'Product': ['A', 'B', 'C','A', 'B', 'C','A', 'B', 'C'],\
     'Year': [2010, 2010, 2010, 2011, 2011, 2011, 2012, 2012, 2012],\
     'Revenue': [100, 200, 300, 110, 190, 320, 120, 220, 350]}
df = pd.DataFrame(data=d)

df = df.pivot(index='Product', columns='Year', values='Revenue').fillna(0)

df_diff = df.pct_change(axis=1, limit=1).replace([np.inf, -np.inf], np.nan).fillna(0)
df_diff = df_diff.dropna(how='all')

# pivot columns and rows to have year on rows and product on columns
df_diff_piv = df_diff.stack().unstack(level=0).sort_index()

# bring to spark df
df_diff_spark = spark.createDataFrame(df_diff.reset_index())

# correlate on at least x periods
correlation_min_periods = 1 # I used 10 for a 20 periods dataset

# set num of partitions to parallelize on tasks 
numOfPartitions = 200 #200 is default

from pyspark.sql.functions import udf, struct
from pyspark.sql.types import StringType, ArrayType, StructType, StructField, FloatType

schema = StructType(
    [
    StructField("Product_1", StringType()),
    StructField("Product_2", StringType()),
    StructField("corr", StringType()) #cant get it to work on FloatType()
    ]
)

def calculate_correlation(product):
  data = df_diff_piv
  arr = []
  for col in data.columns:
    m1 = product
    m2 = data[col].name
    c = np.absolute(data[product].corr(data[col])) #, min_periods=correlation_min_periods
    arr.append([m1, m2, str(c)]) #cant get it to work on FloatType()
  return arr

#register udf
spark.udf.register("calculate_correlation_udf", calculate_correlation)
calculate_correlation_udf = udf(calculate_correlation, ArrayType(schema))

#apply udf to distinct product
distinct_product = df_diff_spark.select("Product").distinct().repartition(numOfPartitions)
res = distinct_product.select("Product", calculate_correlation_udf("Product").alias("corr_matrix"))

from pyspark.sql.functions import explode

# explode (flatten) array and struct back to dataframe
expl = res.select(explode("corr_matrix").alias("corr_row"))
rowlevel = expl.select("corr_row.Product_1","corr_row.Product_2","corr_row.corr")

# convert string to float
rowlevel = rowlevel.withColumn("corr", rowlevel["corr"].cast(FloatType()))

rowlevel.show()