如何在PySpark类中通过self变量读取CSV文件生成DataFrame

你需要通过self访问类内初始化的SparkSession实例和类定义的路径属性即可完成逻辑，完整可运行的实现如下：

from pyspark.sql import SparkSession
import pyspark.sql.functions as F

class ChargePointsETLJob:
    # 类属性，可通过self直接访问
    input_path = 'data/input/electric-chargepoints-2017.csv'
    output_path = 'data/output/chargepoints-2017-analysis'

    def __init__(self):
        # 初始化的SparkSession存在self实例变量中，所有实例方法都可通过self访问
        self.spark_session = (SparkSession.builder
                                          .master("local[*]")
                                          .appName("ElectricChargePointsETLJob")
                                          .getOrCreate())
    def extract(self):
        # 读取CSV文件，直接用self调用spark_session和input_path
        return self.spark_session.read.option("header", True)\
                                      .option("inferSchema", True)\
                                      .csv(self.input_path)

    def transform(self, df):
        # 示例转换逻辑：统计每个充电点的总充电时长、平均充电时长、充电次数
        # 你可以替换为自己需要的业务逻辑
        return df.groupBy("chargepoint_id")\
                 .agg(
                     F.sum("charge_duration").alias("total_duration_hour"),
                     F.avg("charge_duration").alias("avg_duration_hour"),
                     F.count("*").alias("charge_times")
                 )

    def load(self, df):
        # 写入结果到输出路径，用self访问output_path
        df.write.option("header", True)\
                .mode("overwrite")\
                .parquet(self.output_path)

    def run(self):
        self.load(self.transform(self.extract()))

# 执行任务
if __name__ == "__main__":
    etl_job = ChargePointsETLJob()
    etl_job.run()

关键逻辑说明 #

• 实例初始化时创建的SparkSession存在self.spark_session中，所有实例方法都可以直接通过self调用，不需要重复创建SparkSession
• 类中定义的输入输出路径是类属性，同样可以通过self.input_path、self.output_path直接访问
• 你可以根据自己的业务需求修改transform方法的逻辑，也可以调整extract、load的读写参数，比如指定Schema、调整输出格式为CSV等

内容的提问来源于stack exchange，提问作者Developer Rajinikanth