Pyspark / Databricks。 Kolmogorov - Smirnov 随着时间的推移。 有效率的。 在平行下

Question

你好 StackOverflowers。

我有一个 pyspark dataframe ，它由一个 time_column 和一个包含值的列组成。 例如

+----------+--------------------+
|  snapshot|              values|
+----------+--------------------+
|2005-01-31| 0.19120256617637743|
|2005-01-31|  0.7972692479278891|
|2005-02-28|0.005236883665445502|
|2005-02-28|  0.5474099672222935|
|2005-02-28| 0.13077227571485905|
+----------+--------------------+

我想用前一个快照值对每个快照值进行 KS 测试。

我试图用 for 循环来做到这一点。

import numpy as np
from scipy.stats import ks_2samp
import pyspark.sql.functions as F

def KS_for_one_snapshot(temp_df, snapshots_list, j, var = "values"):
  sample1=temp_df.filter(F.col("snapshot")==snapshots_list[j])
  sample2=temp_df.filter(F.col("snapshot")==snapshots_list[j-1]) # pick the last snapshot as the one to compare with

  if (sample1.count() == 0 or sample2.count() == 0 ):
    ks_value = -1  # previously "0 observations" which gave type error
  else:
    ks_value, p_value = ks_2samp(  np.array(sample1.select(var).collect()).reshape(-1)
                                 , np.array(sample2.select(var).collect()).reshape(-1)
                                 , alternative="two-sided"
                                 , mode="auto")



  return ks_value

results = []
snapshots_list = df.select('snapshot').dropDuplicates().sort('snapshot').rdd.flatMap(lambda x: x).collect()
for j in range(len(snapshots_list) - 1 ):
  results.append(KS_for_one_snapshot(df, snapshots_list, j+1))
results

但现实中的数据是巨大的，所以它需要很长时间。 我正在使用 databricks 和 pyspark，所以我想知道通过避免 for 循环和利用可用的工作程序来运行它的更有效方法是什么。

我试图通过使用 udf 来做到这一点，但徒劳无功。

有任何想法吗？

PS。 您可以使用以下代码生成数据。

from random import randint

    df = (spark.createDataFrame( range(1,1000), T.IntegerType())
          .withColumn('snapshot' ,F.array(F.lit("2005-01-31"), F.lit("2005-02-28"),F.lit("2005-03-30")  ).getItem((F.rand()*3).cast("int")))
          .withColumn('values', F.rand()).drop('value')
    )

更新：

我使用 UDF 尝试了以下操作。

var_used = 'values'
data_input_1 = df.groupBy('snapshot').agg(F.collect_list(var_used).alias('value_list'))
data_input_2 = df.groupBy('snapshot').agg(F.collect_list(var_used).alias("value_list_2"))
windowSpec = Window.orderBy("snapshot")
data_input_2 = data_input_2.withColumn('snapshot_2', F.lag("snapshot", 1).over(Window.orderBy('snapshot'))).filter('snapshot_2 is not NULL')
data_input_final = data_input_final = data_input_1.join(data_input_2, data_input_1.snapshot == data_input_2.snapshot_2)

def KS_one_snapshot_general(sample_in_list_1, sample_in_list_2):
  
  if (len(sample_in_list_1) == 0 or len(sample_in_list_2) == 0 ):
    ks_value = -1  # previously "0 observations" which gave type error
  else:
    print('something')
    ks_value, p_value = ks_2samp(  sample_in_list_1
                                 , sample_in_list_2
                                 , alternative="two-sided"
                                 , mode="auto")
  
  return ks_value

import pyspark.sql.types as T
KS_one_snapshot_general_udf = udf(KS_one_snapshot_general, T.FloatType())
data_input_final.select( KS_one_snapshot_general_udf('value_list', 'value_list_2')).display()

如果数据集（每个快照）很小，则效果很好。 但是如果我增加行数，那么我最终会出错。 PickleException：预期零 arguments 用于构造 ClassDict（对于 numpy.dtype）