Pyspark数据帧转换为pandas会丢失数据吗?

时间:2018-05-01 17:13:47

标签: pandas pyspark

我有一个相当复杂的过程,即创建一个pyspark数据帧,将其转换为pandas数据帧,并将结果输出到一个平面文件。我不确定错误的引入点,所以我将描述整个过程。

开始我有一个pyspark数据框,其中包含id组的成对相似性。它看起来像这样:

  +------+-------+-------------------+
  |  ID_A|   ID_B|  EuclideanDistance|
  +------+-------+-------------------+
  |     1|      1|                0.0|
  |     1|      2|0.13103884200454394|
  |     1|      3| 0.2176246463836219|
  |     1|      4|  0.280568636550471|
 ...

我喜欢用ID_A对它进行分组,按EuclideanDistance对每个组进行排序,并且只抓取每组的前N对。所以首先我这样做:

from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col, row_number

window = Window.partitionBy(df['ID_A']).orderBy(df_sim['EuclideanDistance'])
result = (df.withColumn('row_num', row_number().over(window)))

我确保ID_A = 1仍然在“结果”数据帧中。然后我这样做将每组限制为只有20行:

result1 = result.where(result.row_num<20)
result1.toPandas().to_csv("mytest.csv")

和ID_A = 1不在结果.csv文件中(虽然它仍然存在于result1中)。这个转换链中的某个地方是否存在可能导致数据丢失的问题?

2 个答案:

答案 0 :(得分:0)

您在解决方案的窗口中引用了2个数据框。不确定这是否会导致您的错误,但值得清理。无论如何,您不需要在window definition中引用特定的数据框。无论如何,请尝试

window = Window.partitionBy('ID_A').orderBy('EuclideanDistance')

答案 1 :(得分:0)

正如David所说,你在窗口函数中引用了第二个数据帧“df_sim”。

我测试了以下内容,它可以在我的机器上工作(着名的最后一句话):

from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col, row_number
import pandas as pd

#simulate some data
df = pd.DataFrame({'ID_A': pd.np.arange(100)%5, 
    'ID_B': pd.np.repeat(pd.np.arange(20),5), 
    'EuclideanDistance': pd.np.random.rand(100)*5}
    )
#artificially set distance between point and self to 0
df['EuclideanDistance'][df['ID_A'] == df['ID_B']] = 0
df = spark.createDataFrame(df)
#end simulation
window = Window.partitionBy(df['ID_A']).orderBy(df['EuclideanDistance'])
output = df.select('*', row_number().over(window).alias('rank')).filter(col('rank') <= 10)
output.show(50)

模拟代码只是为了使它成为一个独立的例子。您当然可以使用实际的数据帧,并在测试时忽略模拟。希望有用!