poyee-data-warehouse/dw_base/spark/spark_sql.py

# -*- coding:utf-8 -*-

import inspect
import re
from importlib import import_module
from typing import List, Optional, Union, Dict, Any, Tuple

from pyspark.sql import Row, SparkSession, DataFrame

from dw_base import *
from dw_base.utils import common_utils
from dw_base.utils.datetime_utils import get_today
from dw_base.utils.file_utils import get_abs_path
from dw_base.utils.log_utils import pretty_print
from dw_base.utils.sql_utils import get_sql_list_from_file, check_parameter_substituted

HDFS_EXPORT_DATA_PATH = '/hdfs-mnt/export-data'


def _load_spark_conf_file(path: str) -> Dict[str, str]:
    """读 Spark 原生 conf（每行 `key value`，# 注释，空白分隔）。文件缺失或行非法时返回空 dict，不抛错。"""
    if not os.path.isfile(path):
        return {}
    config = {}
    with open(path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            parts = line.split(None, 1)
            if len(parts) != 2:
                continue
            config[parts[0]] = parts[1]
    return config


class SparkSQL(object):
    """封装 Spark 会话与 SQL 执行；参数三级覆盖见 conf/spark-defaults.conf + spark-tuning.conf 与 kb/00 §4.2。"""
    REGISTERED_UDF_FILES = []
    ADDED_RESOURCE_FILES = []
    REGISTERED_UDF = {}
    IGNORED_UDF = ['attr', 'udf']

    def __init__(self,
                 session_name: str = 'spark',
                 master: str = 'yarn',
                 spark_yarn_queue: Optional[str] = None,
                 spark_driver_memory: Optional[str] = None,
                 spark_executor_memory: Optional[str] = None,
                 spark_executor_memory_overhead: Optional[str] = None,
                 spark_driver_cores: Optional[int] = None,
                 spark_executor_cores: Optional[int] = None,
                 spark_executor_instances: Optional[int] = None,
                 spark_driver_max_result_size: Optional[str] = None,
                 spark_shuffle_partitions: Optional[int] = None,
                 spark_default_parallelism: Optional[int] = None,
                 extra_spark_config: Dict[str, Any] = None,
                 udf_files: List[str] = None,
                 resource_files: List[str] = None):
        """

        Args:
            session_name:
            master:
            spark_yarn_queue:
            spark_driver_memory:
            spark_executor_memory:
            spark_executor_memory_overhead:
            spark_driver_cores:
            spark_executor_cores:
            spark_executor_instances:
            spark_driver_max_result_size:
            spark_shuffle_partitions:
            spark_default_parallelism:
            extra_spark_config: 额外的Spark配置，优先级高于SQL文件中定义的Spark配置
            udf_files: udf文件列表，必须传相对路径，否则不识别
            resource_files: 资源文件列表，必须传相对路径，否则不识别
        """
        self._spark_session = None
        self._session_name = session_name
        self._master = master
        self._spark_driver_memory = spark_driver_memory
        self._spark_executor_memory = spark_executor_memory
        self._spark_executor_memory_overhead = spark_executor_memory_overhead
        self._spark_driver_cores = spark_driver_cores
        self._spark_executor_cores = spark_executor_cores
        self._spark_executor_instances = spark_executor_instances
        self._spark_yarn_queue = spark_yarn_queue
        self._spark_driver_max_result_size = spark_driver_max_result_size
        self._spark_shuffle_partitions = spark_shuffle_partitions
        self._spark_default_parallelism = spark_default_parallelism
        self._extra_spark_config = extra_spark_config
        self._final_spark_config = {}
        self.limit = 20
        if not udf_files:
            udf_files = []
        self._udf_files = udf_files
        self._py_files = []
        if not resource_files:
            resource_files = []
        self._resource_files = resource_files
        self.global_start_time = None
        self.start_time = None

    def __add_spark_config(self, sql_file: str, spark_config_def: str):
        if not spark_config_def.__contains__('='):
            pretty_print(f'{NORM_YEL}无效的 Spark 配置 {NORM_GRN}{spark_config_def}')
            return
        spark_config, config_value = [e.strip() for e in spark_config_def.split('=')]
        if self._final_spark_config.__contains__(spark_config):
            pretty_print(
                f'{NORM_YEL}SQL 文件 {NORM_GRN}{sql_file}{NORM_YEL} 中定义的 Spark 配置 {NORM_GRN}{spark_config} '
                f'{NORM_YEL}重复提供，原值 {NORM_GRN}{self._final_spark_config[spark_config]} '
                f'{NORM_YEL}将被覆盖为新值 {NORM_GRN}{config_value}')
        else:
            self._final_spark_config[spark_config] = config_value
        self._final_spark_config[spark_config] = config_value

    def __init_spark_session(self):
        if self._spark_session:
            return
        pretty_print(f'{NORM_MGT}基于用户 {NORM_GRN}{USER}{NORM_MGT} 创建 SparkSession')
        builder = SparkSession.builder \
            .appName(self._session_name) \
            .master(self._master)
        # L1：conf/spark-defaults.conf（底层）+ conf/spark-tuning.conf（调优，相同 key 覆盖 defaults）
        l1_defaults = {}
        l1_defaults.update(_load_spark_conf_file(f'{PROJECT_ROOT_PATH}/conf/spark-defaults.conf'))
        l1_defaults.update(_load_spark_conf_file(f'{PROJECT_ROOT_PATH}/conf/spark-tuning.conf'))
        pretty_print(f'{NORM_MGT}L1 加载 {NORM_GRN}{len(l1_defaults)}{NORM_MGT} 条 conf 默认')
        for key, value in l1_defaults.items():
            pretty_print(f'{NORM_MGT}L1 应用 conf 默认 {NORM_GRN}{key} => {str(value)}')
            builder.config(key, value)
        # L2：SQL 内 SET（query() 预扫描，session 启动前写入 builder）
        for key, value in self._final_spark_config.items():
            pretty_print(f'{NORM_MGT}L2 应用 SQL SET {NORM_GRN}{key} => {str(value)}')
            builder.config(key, value)
        # L3：构造函数显式传参 + extra_spark_config
        l3_overrides: Dict[str, Any] = {}
        for conf_key, attr_val in (
                ('spark.yarn.queue', self._spark_yarn_queue),
                ('spark.driver.memory', self._spark_driver_memory),
                ('spark.executor.memory', self._spark_executor_memory),
                ('spark.executor.memoryOverhead', self._spark_executor_memory_overhead),
                ('spark.driver.cores', self._spark_driver_cores),
                ('spark.executor.cores', self._spark_executor_cores),
                ('spark.executor.instances', self._spark_executor_instances),
                ('spark.driver.maxResultSize', self._spark_driver_max_result_size),
                ('spark.sql.shuffle.partitions', self._spark_shuffle_partitions),
                ('spark.default.parallelism', self._spark_default_parallelism),
        ):
            if attr_val is not None:
                l3_overrides[conf_key] = attr_val
        if self._extra_spark_config:
            l3_overrides.update(self._extra_spark_config)
        for key, value in l3_overrides.items():
            pretty_print(f'{NORM_MGT}L3 应用构造参数/extra {NORM_GRN}{key} => {str(value)}')
            builder.config(key, value)
        pretty_print(f'{NORM_MGT}创建 SparkSession')
        self._spark_session = builder.enableHiveSupport().getOrCreate()
        self._spark_session.sparkContext._jsc.hadoopConfiguration().set('mapred.max.split.size', '33554432')
        self._spark_session.sparkContext._jsc.hadoopConfiguration().set(
            'mapreduce.fileoutputcommitter.marksuccessfuljobs', 'false')
        first_py_file = 'dw_base.zip'
        if IS_RUN_IN_RELEASE_DIR and IS_RUN_BY_RELEASE_USER:
            command = f'cd {PROJECT_ROOT_PATH} && if [ ! -f {first_py_file} ];then zip -qr {first_py_file} dw_base; fi'
        else:
            command = f'cd {PROJECT_ROOT_PATH} && rm -f {first_py_file} && zip -qr {first_py_file} dw_base'
        pretty_print(f'{NORM_MGT}执行 Shell 命令 {NORM_GRN}{command}')
        os.system(command)
        self._spark_session.sparkContext.addPyFile(get_abs_path(first_py_file))
        self.register_udf_files(self._udf_files)
        for py_file in self._py_files:
            self.add_py_file(py_file)
        for resource_file in self._resource_files:
            self.add_resource_file(resource_file)

    @staticmethod
    def add_parameters(sql_or_file: str, parameters: Dict, key: str, value: Any):
        if parameters.__contains__(key):
            pretty_print(f'{NORM_YEL}SQL 文件 {NORM_GRN}{sql_or_file} '
                         f'{NORM_YEL}中定义的 SQL 参数 {NORM_GRN}{key} '
                         f'{NORM_YEL}发生变更，将使用新值 {NORM_GRN}{Any} '
                         f'{NORM_YEL}替代原值 {NORM_GRN}{parameters[key]}')
        parameters[key] = value

    def add_py_file(self, py_file: str) -> bool:
        abs_py_file = get_abs_path(py_file)
        if self.REGISTERED_UDF_FILES.__contains__(abs_py_file):
            return False
        if not self._spark_session:
            self._py_files.append(py_file)
            return True
        pretty_print(f'{NORM_MGT}添加 Python 文件 {NORM_GRN}{py_file}')
        self._spark_session.sparkContext.addPyFile(abs_py_file)
        self.REGISTERED_UDF_FILES.append(abs_py_file)
        return True

    def add_resource_file(self, resource_file: str) -> bool:
        abs_resource_file = get_abs_path(resource_file)
        if self.ADDED_RESOURCE_FILES.__contains__(abs_resource_file):
            return False
        if not self._spark_session:
            self._resource_files.append(resource_file)
            return True
        pretty_print(f'{NORM_MGT}添加资源文件 {NORM_GRN}{resource_file}')
        self._spark_session.sparkContext.addFile(abs_resource_file)
        self.ADDED_RESOURCE_FILES.append(abs_resource_file)
        return True

    def execute(self,
                sql_or_file: str,
                check_parameter: bool = False,
                silent: bool = False,
                fill_null=None,
                **kwargs):
        """
        仅执行，不返回任何值
        Args:
            sql_or_file: 需要执行的 sql 语句或者包含 sql 语句的文件路径（相对或绝对都可以）
            check_parameter: 运行前是否进行参数检查
            silent: 执行时是否打印日志，True则不打印
            fill_null: 填充null值
            **kwargs: sql 语句或 sql 文件所用到的参数
        Returns:
        """
        # 运行SQL语句或SQL文件
        start_time = time.time()
        data_frame, is_select = self.query(sql_or_file, check_parameter, silent, fill_null, **kwargs)
        if is_select:
            self.show_data_frame(data_frame, show_number=self.limit, truncate=USER == RELEASE_USER)
        end_time = time.time()
        cost = round(float(end_time - self.start_time), 2)
        pretty_print(f'{NORM_MGT}SQL 语句执行完毕，耗时 {NORM_GRN}{str(cost)}{NORM_MGT} 秒')
        if sql_or_file.endswith('.sql'):
            total_cost = round(float(end_time - self.global_start_time), 2)
            pretty_print(f'{NORM_MGT}SQL 文件 {NORM_GRN}{sql_or_file}'
                         f'{NORM_MGT} 执行完毕，共耗时 {NORM_GRN}{str(total_cost)}{NORM_MGT} 秒')
            self.global_start_time = None

    def export_data(self,
                    data_set_name: str,
                    sql_or_file: str,
                    show_number: int = 100,
                    truncate: Union[bool, int] = 40,
                    delimiter: str = ',',
                    partition: int = 1,
                    **kwargs) -> str:
        """
        导出数据，默认存储于HDFS目录（/hdfs-mnt/export-data），映射本地目录（/opt/hdfs-mnt/export-data）
        Args:
            data_set_name: 导出的数据集名称标识（作为目录）
            sql_or_file: 需要执行的 sql 语句或者包含 sql 语句的文件全路径
            show_number: 显示的行数
            truncate: If set to ``True``, truncate strings longer than 20 chars by default.
                      If set to a number greater than one, truncates long strings to length ``truncate``
                      and align cells right.
            delimiter: 分隔符
            partition: 导出文件个数
            **kwargs: sql 语句或 sql 文件所用到的参数
        Returns:
        """
        start_time = time.time()
        data_frame, _ = self.query(sql_or_file, **kwargs)
        data_frame.persist()
        self.show_data_frame(data_frame, data_set_name, show_number, truncate)
        today = get_today()
        # 注意: 此目录不是linux服务器的目录, 而是hdfs文件系统的目录
        hdfs_directory = f'{data_set_name}.{str(int(time.time()))}'
        hdfs_export_path = os.path.join(HDFS_EXPORT_DATA_PATH, USER, today, hdfs_directory)
        # linux服务器导出文件存放的目录
        # local_export_path = os.path.join(LOCAL_EXPORT_DATA_PATH, USER, today, data_set_name)
        pretty_print(f'{NORM_MGT}准备导出文件到HDFS目录: {NORM_GRN}{hdfs_export_path}')
        if not delimiter:
            delimiter = ','
        # 先将数据保存至hdfs文件系统, 加 option("escape", "\"") 不会导致csv文件“字段文本内容中的英文逗号来分割字段（列膨胀错位）”
        data_frame \
            .repartition(partition) \
            .write \
            .mode('overwrite') \
            .format('com.databricks.spark.csv') \
            .option('header', 'true') \
            .option("escape", "\"") \
            .option('delimiter', delimiter) \
            .save(hdfs_export_path)
        cost = round(float(time.time() - start_time), 2)
        pretty_print(f'{NORM_MGT}数据导出完毕，耗时 {NORM_GRN}{str(cost)}{NORM_MGT} 秒')
        return hdfs_export_path

    def get_columns(self, hive_table_name: str):
        desc_df, _ = self.query(f'desc {hive_table_name}', silent=True)
        desc_rows = desc_df.collect()
        hive_columns = {}
        partition_col_start = False
        for row in desc_rows:  # type: Row
            if row[0] == '# Partition Information':
                partition_col_start = True
                continue
            if row[0] == '# col_name':
                continue
            if partition_col_start:
                del hive_columns[list(hive_columns.keys())[len(hive_columns.keys()) - 1]]
            else:
                hive_columns[row[0]] = (row[1], row[2])
        return hive_columns

    def get_partition_columns(self, hive_table_name: str):
        desc_df, _ = self.query(f'desc {hive_table_name}', silent=True)
        desc_rows = desc_df.collect()
        partition_columns = {}
        partitioned = False
        for row in desc_rows:  # type: Row
            if row[0] != '# Partition Information':
                continue
            if row[0].startswith('#'):
                continue
            partitioned = True
            partition_columns[row[0]] = (row[1], row[2])
        return partitioned, partition_columns

    def get_spark_session(self):
        if not self._spark_session:
            self.__init_spark_session()
        self.start_time = time.time()
        if not self.global_start_time:
            self.global_start_time = time.time()
        return self._spark_session

    def jdbc_read(self, jdbc_url: str, table_name: str) -> DataFrame:
        return self.get_spark_session().read.jdbc(jdbc_url, table_name)

    def list_tables(self, hive_database_name: str = 'default', include_regex: str = None, exclude_regex: str = None):
        if not hive_database_name:
            hive_database_name = 'default'
        tables_df, _ = self.query(f'show tables in {hive_database_name}')
        tables = []
        for row in tables_df.collect():
            if hive_database_name == 'default':
                table_name = row.asDict()['tableName']
            else:
                table_name = f"{hive_database_name}.{row.asDict()['tableName']}"
            if exclude_regex and re.match(exclude_regex, table_name):
                continue
            if include_regex:
                if re.match(include_regex, table_name):
                    tables.append(table_name)
                    continue
            else:
                tables.append(table_name)
        return tables

    def query(self,
              sql_or_file: str,
              check_parameter: bool = False,
              silent: bool = False,
              fill_null=None,
              **kwargs) -> Tuple[DataFrame, bool]:
        """
        执行查询，返回一个DataFrame和该DataFrame是否是SELECT
        Args:
            sql_or_file: 需要执行的 sql 语句或者包含 sql 语句的文件全路径
            check_parameter: 运行前是否进行参数检查
            silent: 执行时是否打印日志，True则不打印
            fill_null: 填充null值
            **kwargs: sql 语句或 sql 文件所用到的参数
        Returns: Tuple[DataFrame, bool]
        """
        if sql_or_file.endswith('.sql'):
            if not silent:
                pretty_print(f'{NORM_MGT}准备执行 SQL 文件 {NORM_GRN}{sql_or_file}')
            sql_list = get_sql_list_from_file(sql_or_file, trim_comment=True)
        else:
            sql_list = [sql_or_file]
        cleaned_sql_list = []
        # 外部传入的参数
        external_parameters = {}
        if kwargs:
            for key, value in kwargs.items():
                if not silent:
                    pretty_print(f'{NORM_MGT}收到外部 SQL 参数 {NORM_GRN}{key} => {value}')
                external_parameters[key] = value
        # SQL文件内部定义的参数
        internal_parameters = {}
        for sql in sql_list:
            upper_sql = sql.upper()
            if upper_sql.startswith('ADD FILES'):
                file_names = sql[len('ADD FILES'):].strip().split(' ')
                for file_name in file_names:
                    file_name = file_name.strip()
                    if not file_name:
                        continue
                    if file_name.endswith('.py'):
                        self.register_udf_file(file_name)
                    else:
                        self.add_resource_file(file_name)
            elif upper_sql.startswith('ADD FILE'):
                file_name = sql[len('ADD FILE'):].strip()
                if file_name.endswith('.py'):
                    self.register_udf_file(file_name)
                else:
                    self.add_resource_file(file_name)
            elif upper_sql.startswith('ADD SPARK_UDF'):
                self.register_udf_file(sql[len('ADD SPARK_UDF'):].strip())
            elif upper_sql.startswith('SET SPARK_VAR:'):
                var_key, var_value = sql[len('SET SPARK_VAR:'):].strip().split('=')
                var_key = var_key.strip()
                var_value = var_value.strip()
                self.add_parameters(sql_or_file, internal_parameters, var_key, var_value)
            elif upper_sql.startswith('SET LIMIT='):
                if not self._spark_session:
                    self.limit = int(sql.split('=')[1].strip())
            elif upper_sql.startswith('SET '):
                set_expr = sql[len('SET '):].strip()
                set_expr_splits = set_expr.split('=')
                set_key = set_expr_splits[0].strip()
                set_value = set_expr_splits[1].strip()
                if re.match(r'.+\..+(\..+)*', set_key):
                    if not self._spark_session:
                        self.__add_spark_config(sql_or_file, sql[len('SET '):].strip())
                else:
                    self.add_parameters(sql_or_file, internal_parameters, set_key, set_value)
            else:
                # 优先使用外部参数
                for key, value in external_parameters.items():
                    sql = sql.replace('${%s}' % key, value)
                # SQL文件里的参数必须在使用前定义（同时也意味着相同的参数可以多次定义不同值，供多个SQL使用）
                for key, value in internal_parameters.items():
                    if external_parameters.__contains__(key):
                        # 如果外部参数中包括当前参数，则使用外部参数（实际上已在上一步替换，此处仅为提示）
                        pretty_print(f'{NORM_YEL}SQL 文件 {NORM_GRN}{sql_or_file} '
                                     f'{NORM_YEL}中定义的 SQL 参数 {NORM_GRN}{key} '
                                     f'{NORM_YEL}已由外部参数抢占，将使用新值 {NORM_GRN}{external_parameters[key]} '
                                     f'{NORM_YEL}替代原值 {NORM_GRN}{internal_parameters[key]}')
                        continue
                    sql = sql.replace('${%s}' % key, value)
                cleaned_sql_list.append(sql)
        if len(cleaned_sql_list) == 0:
            return self.get_spark_session().sql("SELECT ''"), False
        for index in range(len(cleaned_sql_list) - 1):
            sql = cleaned_sql_list[index]
            self.execute(sql)
        last_or_only_sql = cleaned_sql_list[-1]
        check_parameter_substituted(last_or_only_sql, check_parameter)
        self.get_spark_session()
        if not silent:
            if last_or_only_sql.__contains__('\n'):
                pretty_print(f'{NORM_MGT}开始执行 SQL 语句: \n{NORM_GRN}{last_or_only_sql}')
            else:
                pretty_print(f'{NORM_MGT}开始执行 SQL 语句: {NORM_GRN}{last_or_only_sql}')
        data_frame = self._spark_session.sql(last_or_only_sql)
        if fill_null:
            data_frame.fillna(fill_null)
        return data_frame, len(data_frame.schema.fields) > 0

    def query_scalar(self,
                     sql_or_file: str,
                     check_parameter: bool = False,
                     silent: bool = False,
                     fill_null=None,
                     **kwargs):
        data, is_select = self.query(sql_or_file, check_parameter, silent, fill_null, **kwargs)
        if is_select is not True:
            raise Exception('Get scalar from non-query statement')
        res = data.collect()
        res1 = res[0]
        return res[0][0]

    def register_udf_files(self, udf_file_list: List[str]):
        if not udf_file_list:
            return
        for udf_file in udf_file_list:
            self.register_udf_file(udf_file)

    def register_udf_file(self, udf_file: str):
        if not udf_file:
            return
        if not self.add_py_file(udf_file):
            pretty_print(f'{NORM_YEL}Python文件 {NORM_GRN}{udf_file}{NORM_YEL} 已被添加过（重复添加）')
            return
        if not self._spark_session:
            self._udf_files.append(udf_file)
            return
        pretty_print(f'{NORM_MGT}注册文件 {NORM_GRN}{udf_file}{NORM_MGT} 中的 Python UDF')
        udf_module_name = udf_file.replace('/', '.').replace('.py', '')
        module = import_module(udf_module_name)
        for (name, attr) in inspect.getmembers(module):
            if self.IGNORED_UDF.__contains__(name):
                continue
            if self.REGISTERED_UDF.__contains__(name):
                pretty_print(f'{NORM_YEL}名为 {NORM_GRN}{name}{NORM_YEL} 的 Python UDF '
                             f'{NORM_YEL}已在文件 {NORM_GRN}{self.REGISTERED_UDF[name]} '
                             f'{NORM_YEL}被注册（重名UDF以首次出现的为准）')
                continue
            if inspect.isfunction(attr):
                # todo: 注册每个udf的返回值类型
                pretty_print(f'{NORM_MGT}注册 Python UDF {NORM_GRN}{name}')
                self._spark_session.udf.register(name, attr)
                self.REGISTERED_UDF[name] = udf_file

    @staticmethod
    def show_data_frame(data_frame: DataFrame,
                        data_frame_name: str = None,
                        show_number: int = 20,
                        truncate: Union[bool, int] = False):
        """
        打印``DataFrame``的Schema和指定行数的数据
        Args:
            data_frame: DataFrame
            data_frame_name: DataFrame的名称
            show_number: 打印的行数
            truncate: If set to ``True``, truncate strings longer than 20 chars by default.
                      If set to a number greater than one, truncates long strings to length ``truncate``
                      and align cells right.
        Returns:
        """
        if data_frame_name:
            pretty_print(f'{NORM_MGT}%s %s DataFrame Information: %s ' % ('=' * 30, data_frame_name, '=' * 30))
        else:
            pretty_print(f'{NORM_MGT}%s DataFrame Information: %s ' % ('=' * 30, '=' * 30))
        pretty_print(f'{NORM_MGT}Schema :')
        data_frame.printSchema()
        pretty_print(f'{NORM_MGT}Data sample :')
        data_frame.show(show_number, truncate=truncate)
        pretty_print(f'{NORM_MGT}=' * 70)


if __name__ == '__main__':
    with SparkSQL('spark-sql') as spark_sql:
        # spark_sql.execute('workspace/x.sql')
        # spark_sql._spark_session.conf.set('app.name', 'tmp.sa33')
        # spark_sql.execute('workspace/y.sql')
        # spark_sql.query('show databases')[0].show(100, truncate=False)
        # spark_sql.query('show functions')[0].show(1000, truncate=False)
        spark_sql.get_columns('crl_ads.ads_crl_dim_source_stat')