# -*- coding:utf-8 -*- import inspect import re from importlib import import_module from typing import List, 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' class SparkSQL(object): """ 封装执行 Spark 相关操作的类, 相关参数说明: In Spark 2.0+ version use spark session variable to set number of executors dynamically (from within program) spark.conf.set("spark.executor.instances', 4) spark.conf.set("spark.executor.cores', 4) In above case maximum 16 tasks will be executed at any given time. other option is dynamic allocation of executors as below - spark.conf.set("spark.dynamicAllocation.enabled', "true') spark.conf.set("spark.executor.cores', 4) spark.conf.set("spark.dynamicAllocation.minExecutors',"1') spark.conf.set("spark.dynamicAllocation.maxExecutors',"5') spark.yarn.executor.memoryOverhead：default is executorMemory * 0.07, with minimum of 384 spark.yarn.driver.memoryOverhead：default is driverMemory * 0.07, with minimum of 384 spark.yarn.am.memoryOverhead：default is AM memory * 0.07, with minimum of 384 """ 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: str = 'spark', spark_driver_memory: str = '2g', spark_executor_memory: str = '6g', spark_executor_memory_overhead: str = '512', spark_driver_cores: int = 2, spark_executor_cores: int = 2, spark_executor_instances: int = 15, spark_driver_max_result_size='4g', spark_shuffle_partitions=200, spark_default_parallelism=200, 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') # for element in os.environ: # pretty_print(f'{NORM_MGT}Environment {NORM_GRN}{element} => {os.environ[element]}') builder = SparkSession.builder \ .appName(self._session_name) \ .master(self._master) \ .config('hive.exec.orc.default.block.size', 134217728) \ .config('spark.debug.maxToStringFields', 5000) \ .config('spark.default.parallelism', self._spark_default_parallelism) \ .config('spark.driver.cores', self._spark_driver_cores) \ .config('spark.driver.maxResultSize', self._spark_driver_max_result_size) \ .config('spark.driver.memory', self._spark_driver_memory) \ .config('spark.dynamicAllocation.enabled', False) \ .config('spark.files.ignoreCorruptFiles', True) \ .config('spark.executor.cores', self._spark_executor_cores) \ .config('spark.executor.instances', self._spark_executor_instances) \ .config('spark.executor.memory', self._spark_executor_memory) \ .config('spark.executor.memoryOverhead', self._spark_executor_memory_overhead) \ .config('spark.sql.adaptive.enabled', 'true') \ .config('spark.sql.broadcastTimeout', -1) \ .config('spark.sql.codegen.wholeStage', 'false') \ .config('spark.sql.execution.arrow.enabled', True) \ .config('spark.sql.execution.arrow.fallback.enabled', True) \ .config('spark.sql.files.ignoreCorruptFiles', True) \ .config('spark.sql.shuffle.partitions', self._spark_shuffle_partitions) \ .config('spark.sql.statistics.fallBackToHdfs', True) \ .config('spark.yarn.queue', self._spark_yarn_queue) \ .config('spark.port.maxRetries', 999) if self._extra_spark_config: for spark_config, config_value in self._extra_spark_config.items(): if self._final_spark_config.__contains__(spark_config): pretty_print(f'{NORM_YEL}构造函数传入的 Spark 配置 {NORM_GRN}{spark_config} => {config_value} ' f'{NORM_YEL}覆盖了在 SQL 文件中定义的配置 ' f'{NORM_GRN}{spark_config} => {self._final_spark_config[spark_config]}') self._final_spark_config[spark_config] = config_value if self._final_spark_config: for key, value in self._final_spark_config.items(): pretty_print(f'{NORM_MGT}添加自定义 Spark 配置 {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')