在PySpark數據框中添加列需要花費時間

Question

我目前正在嘗試將PySpark和Cassandra集成在一起，並且在優化代碼以使其更快執行方面遇到困難。

from pyspark import SparkContext, SparkConf
from pyspark.sql import SQLContext, SparkSession
from pyspark.sql.functions import sum as _sum

def connect_cassandra():
    spark = SparkSession.builder \
      .appName('SparkCassandraApp') \
      .config('spark.cassandra.connection.host', 'localhost') \
      .config('spark.cassandra.connection.port', '9042') \
      .config('spark.cassandra.output.consistency.level','ONE') \
      .master('local[*]') \
      .getOrCreate()

    sqlContext = SQLContext(spark)
    return sqlContext

#--------THIS FUNCTION IS MY CONCERN ACTUALLY------------
def check_ip(ip, df):
    rows= df.filter("src_ip = '"+ip+"' or dst_ip = '"+ip+"'") \
            .agg(_sum('total').alias('data')) \
            .collect()

    print(rows[0][0])
#-----------------------------------------------------------

def load_df(sqlContext):

    df = sqlContext \
      .read \
      .format('org.apache.spark.sql.cassandra') \
      .options(table='acrs_app_data_usage', keyspace='acrs') \
      .load()

    return df

if __name__ == '__main__':
    lists = ['10.8.25.6', '10.8.24.10', '10.8.24.11', '10.8.20.1', '10.8.25.15', '10.8.25.10']
    sqlContext = connect_cassandra()
    df = load_df(sqlContext)
    for ip in lists:
        check_ip(ip, df)

這里的功能check_ip()接受一個ip和一個預加載的數據幀，該數據幀具有3列（ src_ip, dst_ip and total ）和大約250K行，作為參數，然后遍歷總列，將它們添加並返回匯總的數據按提供的IP分組。

但是，當我執行腳本時，每個IP至少要花一秒鍾的時間才能返回總和。 我有超過32K個IP必須相同。 而且時間很長。

任何幫助，將不勝感激。 提前致謝。

Answer 1

簡短的回答：不要使用循環。

可能的解決方案：

• 將lists轉換為數據框。
• 與數據lists_df兩次內部連接lists_df ，第一次在ip == src_ip ，第二次在ip == dst_ip
• 與unionAll串聯
• 最后使用groupBy("ip").agg(_sum("total"))

這使用聯接。 因此，也許有一個更好的解決方案。