将Databricks SQL代码转换为带类与函数的PySpark代码并示例

一、需创建的类、函数与对象 #

核心类 #

• HospitalEfficiencyAnalyzer：封装医院效率分析的全流程逻辑，作为核心业务类，包含初始化和数据转换方法。

类内核心函数 #

• __init__：初始化SparkSession，配置源表所属schema（支持灵活适配不同环境），完成源表数据的预加载。
• transform：对外暴露的主方法，串联所有数据处理步骤，返回最终分析结果DataFrame。
• _load_source_data（私有方法）：统一加载指定schema下的源数据表，封装数据读取逻辑，避免重复代码。
• _prepare_part1_data（私有方法）：实现原SQL中part1的关联与筛选逻辑，得到基础医院数据及对应可比医院列表。
• _calculate_comparator_metrics（私有方法）：计算可比医院的市场份额、成本节省额等衍生指标，处理窗口函数与关联逻辑。

对象 #

• 创建HospitalEfficiencyAnalyzer类的实例，传入SparkSession即可调用transform方法执行完整分析流程。

二、PySpark示例代码 #

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

class HospitalEfficiencyAnalyzer:
    def __init__(self, spark: SparkSession, schema_name: str = "myvdi"):
        """
        初始化医院效率分析器
        :param spark: SparkSession实例
        :param schema_name: 源数据表所在的schema名称，默认myvdi
        """
        self.spark = spark
        self.schema_name = schema_name
        # 预加载所有源表
        self.hospital_flag = self._load_source_data("hospital_list_1_flag")
        self.hospital_efficiency = self._load_source_data("hospital_Efficiency")
        self.hospital_distance = self._load_source_data("hospital_distance_Dev")

    def _load_source_data(self, table_name: str):
        """私有方法：加载指定schema下的源数据表"""
        return self.spark.table(f"{self.schema_name}.{table_name}")

    def _prepare_part1_data(self):
        """私有方法：生成原SQL中part1的数据集"""
        part1_df = self.hospital_flag.filter(F.col("Keep_flag") == 0) \
            .join(self.hospital_efficiency, on="HOSPITAL_PRACTICE_NO", how="left") \
            .join(
                self.hospital_distance,
                (self.hospital_distance["HospitalName"] == self.hospital_efficiency["HOSPITAL_PRACTICE_NO"]) &
                (self.hospital_distance["HospitalName_comb"] != self.hospital_efficiency["HOSPITAL_PRACTICE_NO"]) &
                (F.trim(self.hospital_distance["StateName_comb"]) == self.hospital_efficiency["PROVINCE_DESCR"]) &
                (self.hospital_distance["Distance"] <= 15),
                how="left"
            ) \
            .select(
                self.hospital_efficiency["*"],
                self.hospital_distance["HospitalName_comb"].alias("COMPARATOR_HOSPITAL")
            )
        return part1_df

    def _calculate_comparator_metrics(self, part1_df):
        """私有方法：计算可比医院的衍生指标"""
        # 关联可比医院的效率数据及有效标记
        joined_df = part1_df.join(
            self.hospital_flag.filter(F.col("Keep_flag") == 1),
            on=part1_df["COMPARATOR_HOSPITAL"] == self.hospital_flag["HOSPITAL_PRACTICE_NO"],
            how="inner"
        ) \
        .join(
            self.hospital_efficiency.alias("d1"),
            on=F.col("COMPARATOR_HOSPITAL") == F.col("d1.HOSPITAL_PRACTICE_NO"),
            how="left"
        )

        # 定义窗口：按主医院分组计算聚合值
        window_spec = Window.partitionBy("HOSPITAL_PRACTICE_NO")

        # 计算所有衍生指标
        final_df = joined_df.withColumn(
            "COMPARATOR_EVENTS", F.col("d1.EVENTS")
        ).withColumn(
            "COMAPRATOR_CAPABILITY", F.col("d1.CAPABILITY")
        ).withColumn(
            "COMPARATOR_EFFICIENCY", F.col("d1.HOSPITAL_EFFICIENCY")
        ).withColumn(
            "COMPARATOR_CPE_CMA", F.col("d1.CPE_CMA_HOSPITAL")
        ).withColumn(
            "COMARRATOR_MARKET_SHARE",
            F.col("d1.EVENTS") / F.sum(F.col("d1.EVENTS")).over(window_spec).cast("float")
        ).withColumn(
            "SAVINGS",
            (F.col("d1.EVENTS") / F.sum(F.col("d1.EVENTS")).over(window_spec).cast("float")) *
            F.col("HOSPITAL_COST_WEIGHT") *
            (F.col("CPE_CMA_HOSPITAL") - F.col("d1.CPE_CMA_HOSPITAL"))
        ).select(
            part1_df["*"],
            "COMPARATOR_EVENTS", "COMAPRATOR_CAPABILITY", "COMPARATOR_EFFICIENCY",
            "COMPARATOR_CPE_CMA", "COMARRATOR_MARKET_SHARE", "SAVINGS"
        )

        return final_df

    def transform(self):
        """对外主方法：执行全流程分析，返回最终结果"""
        part1_df = self._prepare_part1_data()
        final_df = self._calculate_comparator_metrics(part1_df)
        return final_df

# 使用示例
if __name__ == "__main__":
    spark = SparkSession.builder.appName("HospitalEfficiencyAnalysis").getOrCreate()
    analyzer = HospitalEfficiencyAnalyzer(spark)
    result_df = analyzer.transform()
    # 可根据需求展示或保存结果
    result_df.show()
    spark.stop()

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