如何让PySpark任务在资源可用时自动并行执行而非串行

嘿，针对你现在用串行for循环处理多个查询导致CPU资源浪费的问题，我有两个适合你场景的解决方案，都能实现自动根据剩余资源并行启动处理任务，不用手动指定进程数。

方案一：利用Spark分布式能力并行处理查询任务 #

这个方案适合你有Spark集群的场景（单机Spark集群也适用），Spark会自动根据集群的executor资源分配并行任务，完全不用手动管理进程数量。

关键步骤说明 #

1. 确保Executor能访问MySQL：把MySQL驱动包放到所有Executor节点的classpath中，或者启动Spark时通过--jars /path/to/mysql-connector-java-5.1.38.jar参数传递驱动，同时配置spark.executor.extraClassPath指向驱动路径。
2. 封装单查询处理逻辑：把单个查询的读取、过滤、保存逻辑打包成函数，这个函数会在Executor节点上独立执行。
3. 转RDD并行执行：将查询列表转为Spark RDD，Spark会自动根据集群资源调度并行任务。

修改后的代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import isnull, when, count, countDistinct
from time import time

# 封装单个查询的完整处理逻辑，将在Executor节点上运行
def process_query(query):
    # 每个任务初始化独立的SparkSession（适配PySpark 2.2.1版本）
    spark = SparkSession \
        .builder \
        .appName('Filtering Columns - Worker') \
        .config('spark.driver.memory', '16g') \
        .config('spark.executor.memory', '16g') \
        .config('spark.driver.extraClassPath', '/path/to/mysql-connector-java-5.1.38.jar') \
        .config('spark.executor.extraClassPath', '/path/to/mysql-connector-java-5.1.38.jar') \
        .getOrCreate()

    jdbc_config = {
        'url': 'jdbc:mysql://my_db_ip_address',
        'properties': {
            'user': 'my_db_user',
            'password': 'my_db_password'
        }
    }

    def drop_constant_columns(df):
        cols_to_drop_map = df.select([
            when(countDistinct(column_name) == 1, True).alias(column_name)
            for column_name in df.columns
        ]).first().asDict()
        cols_to_drop = [
            col for col, should_drop in cols_to_drop_map.iteritems() if should_drop
        ]
        return df.drop(*cols_to_drop)

    print('Querying {}'.format(query['table_name']))
    df = spark.read.jdbc(table=query['query'], **jdbc_config)
    print('Filtering {}'.format(query['table_name']))
    n_cols = len(df.columns)
    start = time()
    df = drop_constant_columns(df)
    elapsed = time() - start
    n_cols_filtered = n_cols - len(df.columns)
    print('Filtered {} of {} columns in {:.2f} secs'.format(n_cols_filtered, n_cols, elapsed))
    print('Persisting {}'.format(query['table_name']))
    df.write.mode('overwrite').parquet('./{}_test.parquet'.format(query['table_name']))
    
    # 关闭当前任务的SparkSession
    spark.stop()
    return f"Processed {query['table_name']} successfully"

if __name__ == '__main__':
    # 初始化Driver端的SparkSession
    spark = SparkSession \
        .builder \
        .appName('Filtering Columns - Driver') \
        .config('spark.driver.memory', '16g') \
        .config('spark.executor.memory', '16g') \
        .config('spark.driver.extraClassPath', '/path/to/mysql-connector-java-5.1.38.jar') \
        .config('spark.executor.extraClassPath', '/path/to/mysql-connector-java-5.1.38.jar') \
        .getOrCreate()

    queries_to_run = [
        {'table_name': 'table1', 'query': '''(some query) as tmp '''},
        {'table_name': 'table2', 'query': '''(some query) as tmp '''},
        {'table_name': 'table3', 'query': '''(some query) as tmp '''},
        # 更多查询...
    ]

    # 将查询列表转为RDD，Spark自动调度并行任务
    queries_rdd = spark.sparkContext.parallelize(queries_to_run)
    results = queries_rdd.map(process_query).collect()

    # 打印处理结果
    for res in results:
        print(res)
    
    spark.stop()

注意事项 #

• 要确保MySQL数据库允许所有Executor节点的IP访问，否则会出现连接拒绝错误。
• 如果是单机Spark集群，这个方案同样能利用多核CPU并行处理任务。

方案二：Python多进程并行处理（适合单机多核场景） #

如果你的环境是单机没有Spark集群，用Python的multiprocessing模块可以自动根据CPU核心数启动并行进程，充分利用单机CPU资源。

关键步骤说明 #

1. 封装单查询处理逻辑：每个进程需要独立初始化SparkSession，所以把单个查询的完整逻辑（包括Session初始化、查询、过滤、保存）打包成函数。
2. 使用进程池调度：multiprocessing.Pool默认会用CPU核心数作为进程数，当一个进程完成任务后，会自动启动下一个任务，充分利用空闲资源。

修改后的代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import isnull, when, count, countDistinct
from time import time
import multiprocessing

# 单个查询的完整处理逻辑，每个进程独立执行
def process_query(query):
    # 每个进程初始化独立的SparkSession
    spark = SparkSession \
        .builder \
        .appName('Filtering Columns - Process') \
        .config('spark.driver.memory', '16g') \
        .config('spark.executor.memory', '16g') \
        .config('spark.driver.extraClassPath', '/path/to/mysql-connector-java-5.1.38.jar') \
        .getOrCreate()

    jdbc_config = {
        'url': 'jdbc:mysql://my_db_ip_address',
        'properties': {
            'user': 'my_db_user',
            'password': 'my_db_password'
        }
    }

    def drop_constant_columns(df):
        cols_to_drop_map = df.select([
            when(countDistinct(column_name) == 1, True).alias(column_name)
            for column_name in df.columns
        ]).first().asDict()
        cols_to_drop = [
            col for col, should_drop in cols_to_drop_map.iteritems() if should_drop
        ]
        return df.drop(*cols_to_drop)

    print('Querying {}'.format(query['table_name']))
    df = spark.read.jdbc(table=query['query'], **jdbc_config)
    print('Filtering {}'.format(query['table_name']))
    n_cols = len(df.columns)
    start = time()
    df = drop_constant_columns(df)
    elapsed = time() - start
    n_cols_filtered = n_cols - len(df.columns)
    print('Filtered {} of {} columns in {:.2f} secs'.format(n_cols_filtered, n_cols, elapsed))
    print('Persisting {}'.format(query['table_name']))
    df.write.mode('overwrite').parquet('./{}_test.parquet'.format(query['table_name']))
    
    # 关闭当前进程的SparkSession
    spark.stop()
    return f"Processed {query['table_name']} successfully"

if __name__ == '__main__':
    queries_to_run = [
        {'table_name': 'table1', 'query': '''(some query) as tmp '''},
        {'table_name': 'table2', 'query': '''(some query) as tmp '''},
        {'table_name': 'table3', 'query': '''(some query) as tmp '''},
        # 更多查询...
    ]

    # 创建进程池，默认使用CPU核心数作为并行数
    pool = multiprocessing.Pool()
    # 并行处理所有查询
    results = pool.map(process_query, queries_to_run)
    # 关闭进程池并等待所有任务完成
    pool.close()
    pool.join()

    # 打印处理结果
    for res in results:
        print(res)

注意事项 #

• 每个进程会占用独立的内存资源，要根据机器实际内存调整spark.driver.memory参数（比如4核机器配4个进程，每个进程分配4G内存会更合理）。
• 这个方案完全依赖单机CPU资源，任务调度由Python自动完成，不用手动干预。

内容的提问来源于stack exchange，提问作者Álvaro Lemos