在PySpark中一次比较两行

时间:2018-05-08 21:00:40

标签: apache-spark join pyspark

我是Spark的新手,正在寻求最佳实践方面的帮助。我有一个大的DataFrame,需要一次两行输入一个比较它们的函数。

actual_data是一个DataFrame,其中包含id列和多个值列。

rows_to_compare是一个包含两列的DataFrame:left_idright_id

对于rows_to_compare中的每一对,我想将actual_data中的两个相应行提供给函数。

我的实际数据非常大(~30GB)并且有很多列,所以我把它简化为这个更简单的例子:

import pandas as pd
from pyspark.sql import SQLContext
from pyspark.sql.functions import col
import builtins

sqlContext = SQLContext(sc)
# Build DataFrame of Actual Data
data = { 
    'id': [1,2,3,4,5],
    'value': [11,12,13,14,15]}
actual_data_df = sqlContext.createDataFrame(
    pd.DataFrame(data, columns=data.keys()))

# Build DataFrame of Rows To Compare
rows_to_compare = {
    'left_id': [1,2,3,4,5],
    'right_id': [1,1,1,1,1]}
rows_to_compare_df = 
    sqlContext.createDataFrame(
        pd.DataFrame(rows_to_compare, columns=rows_to_compare.keys()))

result = (
    rows_to_compare_df
        .join(
            actual_data_df.alias('a'),
            col('left_id') == col('a.id'))
        .join(
            actual_data_df.alias('b'),
            col('right_id') == col('b.id'))
        .withColumn(
            'total',
            builtins.sum(
                [col('a.value'), 
                 col('b.value')]))
        .select('a.id', 'b.id', 'total')
        .collect())

返回所需的输出:

[Row(id=2, id=1, total=23), Row(id=5, id=1, total=26), Row(id=4, id=1, total=25), Row(id=1, id=1, total=22), Row(id=3, id=1, total=24)]

当我运行它时,即使是这个玩具问题,它看起来也很慢。这是解决这个问题的最佳方法吗?我能想到的最清晰的替代方法是让我的DataFrame的每一行都包含我想要比较的两行的值。我担心这种方法,因为它会涉及大量的数据重复。

非常感谢任何帮助,谢谢。

0 个答案:

没有答案