寻求AWS Glue脚本优化方案：高效获取多库表行数

任务目标 #

编写AWS Glue的PySpark脚本，统计多源数据库（MySQL、Oracle、PostgreSQL等）中约15000张表的行数，数据源存储在子网内，必须使用DynamicFrame Reader API以复用Glue连接的子网配置。

当前困境 #

• 脚本执行耗时超2小时，效率极低
• 循环创建大量数据库连接，导致源数据库负载过高甚至挂起
• 现有脚本通过遍历S3存储的表清单文件，逐个表执行count查询，实现方式粗放

现有示例代码 #

#Reading file containing the list of database, schemas and tables

s3_input_path = "s3://path/to/csv"
# Create a DynamicFrame from the CSV file in S3

dyf = glueContext.create_dynamic_frame.from_options(
connection_type="s3",
connection_options={"paths": [s3_input_path]},
format="csv",
format_options={"withHeader": True}
)

# Convert the DynamicFrame to a Spark DataFrame
df = dyf.toDF()

#Creating an empty list to collect data from each schema
results = []

schema = StructType([StructField("db_name", StringType()),
StructField("schema_name", StringType()),
StructField("table_name", StringType()),
StructField("row_count", DecimalType()),
StructField("execution_time", StringType())])

for schema_name in df.select('schema_name').distinct().collect():
    schema_name = schema_name['schema_name']

for row in df.filter(df['schema_name'] == schema_name).collect():
    table_name = row['table_name']
    db_name = row['db_name']

    query = f"""Select '{db_name}' as db_name, '{schema_name}' as schema_name, '{table_name}' as table_name, cast(count(*) AS int) as row_count, current_timestamp as execution_time from {schema_name}.{table_name}"""

    try:
     oracle_dyf = glueContext.create_dynamic_frame.from_options(
     connection_type = "oracle",
     connection_options = {
     "useConnectionProperties": "true",
    "dbtable": f"{schema_name}.{table_name}",
    "sampleQuery": query,
    "connectionName": "connection",
    },
transformation_ctx = "src_oracle")

    # Convert to Spark DataFrame
    oracle_df = oracle_dyf.toDF()

    # Collect results to list
    results.extend(oracle_df.collect())

except Exception as e:
    print(f"Error in {schema_name}.{table_name} : {e}")
   # Continue to the next table

# After processing all tables in the current schema, write results to a file
schema_df = spark.createDataFrame(results, schema)
schema_df = schema_df.repartition(1) # Ensure single output file
file_path = "s3://path/to/target"

if schema_df.count() > 0:
    schema_df.write.mode('append').csv(file_path)

# Clear results for the next schema
    results.clear()

核心优化方向 #

1. 复用数据库连接，减少连接创建开销 #

• 按数据库分组处理，同一数据库的所有表复用同一个Glue连接配置，避免每次查询都新建连接
• 配置Glue连接池参数（如Oracle的maxPoolSize），控制并发连接数，避免压垮源库

2. 批量查询，减少单表查询次数 #

• 针对同一schema下的多张表，用UNION ALL合并多个表的count语句，一次查询获取多张表的结果，大幅减少查询请求数：

SELECT 'DB1' as db_name, 'SCHEMA1' as schema_name, 'TABLE1' as table_name, COUNT(*) as row_count, CURRENT_TIMESTAMP as execution_time FROM SCHEMA1.TABLE1
UNION ALL
SELECT 'DB1' as db_name, 'SCHEMA1' as schema_name, 'TABLE2' as table_name, COUNT(*) as row_count, CURRENT_TIMESTAMP as execution_time FROM SCHEMA1.TABLE2

• 控制单批次查询的表数量（建议20-50张，根据数据库性能调整），避免单条SQL过长

3. 利用Spark并行化处理，替代单线程循环 #

• 放弃Python遍历collect()结果的单线程模式，改用Spark分布式并行处理：
  • 将表清单DataFrame按数据库、schema分组，使用mapPartitions或foreachPartition对每个分区内的表批量处理
  • 分区数量控制在源数据库能承受的并发连接数范围内（如10-20个分区）

4. 替换全表count为元数据查询（大幅提速） #

• 利用数据库内置元数据统计信息，避免全表扫描：
  • Oracle：查询ALL_TABLES/DBA_TABLES的NUM_ROWS字段（需定期收集统计信息保证准确性）
  • MySQL：查询INFORMATION_SCHEMA.TABLES的TABLE_ROWS字段
  • PostgreSQL：查询pg_stat_user_tables的n_live_tup字段

5. 优化结果写入逻辑 #

• 去掉repartition(1)，避免强制合并单文件导致的性能瓶颈，允许Spark生成多文件，后续通过Glue Crawler/Athena统一查询
• 改用write.parquet()替代CSV格式，Parquet列式存储的写入、读取效率更高，且支持schema保留

6. 完善错误处理 #

• 添加重试逻辑，针对临时连接失败自动重试
• 将失败的表信息写入单独S3文件，方便后续排查

from pyspark.sql.functions import col, collect_list

# 按数据库、连接信息分组
db_groups = df.groupBy("db_name", "connection_name", "connection_type")

def process_db_group(group_data):
    db_name, conn_name, conn_type, tables = group_data
    # 按schema分组批量处理表
    schema_tables = tables.groupBy("schema_name").agg(collect_list("table_name").alias("table_list"))
    
    results = []
    batch_size = 50  # 每批次处理的表数量
    for schema_row in schema_tables.collect():
        schema_name = schema_row["schema_name"]
        table_list = schema_row["table_list"]
        
        # 分批次构造查询
        for i in range(0, len(table_list), batch_size):
            batch_tables = table_list[i:i+batch_size]
            query_parts = []
            for table in batch_tables:
                # 构造单表count语句
                query_parts.append(f"""
                    SELECT '{db_name}' as db_name, '{schema_name}' as schema_name, '{table}' as table_name, 
                    COUNT(*) as row_count, CURRENT_TIMESTAMP as execution_time 
                    FROM {schema_name}.{table}
                """)
            batch_query = " UNION ALL ".join(query_parts)
            
            try:
                # 复用连接查询批次表
                dyf = glueContext.create_dynamic_frame.from_options(
                    connection_type=conn_type,
                    connection_options={
                        "useConnectionProperties": "true",
                        "sampleQuery": batch_query,
                        "connectionName": conn_name,
                        "connectionPool": "true",
                        "maxPoolSize": 5  # 控制连接池大小
                    },
                    transformation_ctx=f"src_{db_name}_{schema_name}"
                )
                results.extend(dyf.toDF().collect())
            except Exception as e:
                print(f"Batch error in {db_name}.{schema_name}: {e}")
                # 记录失败表
                with open(f"/tmp/failed_tables_{db_name}.txt", "a") as f:
                    f.write("\n".join(batch_tables) + "\n")
    return results

# 并行处理所有数据库分组
final_results = db_groups.flatMap(process_db_group).collect()

# 写入结果到S3（用Parquet格式）
result_schema = StructType([
    StructField("db_name", StringType()),
    StructField("schema_name", StringType()),
    StructField("table_name", StringType()),
    StructField("row_count", IntegerType()),
    StructField("execution_time", TimestampType())
])
final_df = spark.createDataFrame(final_results, schema=result_schema)
final_df.write.mode("append").parquet("s3://path/to/target/row_counts")

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