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

Question

我想计算一个大型数据集（1M 行）的相关矩阵。 这个想法是计算产品销售的相关性。如果两种产品的销售额同比增长/下降幅度相似，则可能存在相关性。

我已经尝试过这里的帖子：

• 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 的分布式计算。有 170k 个独特的产品，因此作业运行 170k 次，并且有 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

而且我需要相关性... 用熊猫很容易

# 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()

Answer 2

所以以下应该有效（至少对于玩具示例）：我很想听听它是如何缩放的：

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql.window import Window

# define pyspark df
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 = spark.createDataFrame(pd.DataFrame(data=d))

# Define a window over products and calc %-changes over time
win = Window().partitionBy("Product").orderBy("Year")
df = df.withColumn("pct_change", 
      F.col("Revenue")/F.lag(F.col("Revenue")).over(win) - 1
)

# replace nulls with 0
df = df.na.fill(0)

# pivot
df = (df.groupBy("Product")
    .pivot("Year")
    .agg(F.first("pct_change"))
    .orderBy("Product"))

# Get pair-RDD of (product, %-changes) and cross-join
numerical_cols = df.columns[1:]
rdd = df.rdd.map(lambda x: (x['Product'], [x[col] for col in numerical_cols]))
rdd = rdd.cartesian(rdd)

# correlation helper function
def corr(pair):
    (prod1, series1), (prod2, series2) = pair
    corr = pd.Series(series1).corr(pd.Series(series2))
    return (prod1, prod2, float(corr))

# pairwise correlation DF
corr_df = rdd.map(corr).toDF(schema=['Product_1', 'Product_2', 'value'])
corr_df.show(5)