基于Glue Job Spark与tsfresh的时序特征提取性能优化问询

问题背景 #

处理5.85亿行时序数据集，按workorder（WO，工单）分区，每个分区约500万行。目标是处理50个符合以下条件的分区：

• 失真状态均衡（equilibred_workorders）
• 标签数量在120-122之间（limit_tags）
• 使用tsfresh的spark_feature_extraction_on_chunk方法提取特征

单WO测试时计算耗时超1.5小时触发超时，环境为G2X（8 vCPU、32GB内存）+16个worker，已将分区数设为128。疑问：是代码问题？Spark分区优化遗漏？还是tsfresh基于Pandas的转换开销导致？

附代码：

import traceback
import time
import os
import psutil
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, countDistinct, first, count, xxhash64, spark_partition_id
from pyspark.sql.types import *
from awsglue.utils import getResolvedOptions
from tsfresh.convenience.bindings import spark_feature_extraction_on_chunk

def log_timer(start, step_name):
    duration = time.time() - start
    print(f"[TIMER] {step_name} took {duration:.2f}s")
    return time.time()

def log_memory(spark, step):
    sc = spark.sparkContext
    print(f"[MEMORY LOG] {step}")
    print(f"  Executors: {sc._jsc.sc().getExecutorMemoryStatus().keySet()}")

def list_workorders(spark, min_tags, max_tags):
    base_path = "s3://***MASKED_BUCKET***/ml/Join_tags_workorder_data_cleaned/"
    df_all = spark.read.parquet(base_path)
    df_stats = df_all.groupBy("workorder").agg(
        countDistinct("wincctagname").alias("nb_tags"),
        first("isdistorted").alias("isdistorted"),
        count("*").alias("nb_rows")
    )
    workorder_rows = df_stats.filter((col("nb_tags") >= min_tags) & (col("nb_tags") <= max_tags)) \
                             .select("workorder", "isdistorted", "nb_tags", "nb_rows") \
                             .distinct().collect()
    return workorder_rows

def select_workorders(wo_true_list, wo_false_list, max_workorders, equilibred_workorders):
    if max_workorders == 0:
        return [wo for wo in wo_true_list + wo_false_list]
    if equilibred_workorders:
        half = max_workorders // 2
        selected_true = wo_true_list[:half]
        selected_false = wo_false_list[:half]
        selected = selected_true + selected_false
    else:
        combined = wo_true_list + wo_false_list
        selected = combined[:max_workorders]
    return selected

args = getResolvedOptions(sys.argv, ['max_workorders', 'equilibred_workorders', 'limit_tags', 'workorder'])
max_workorders = int(args['max_workorders'])
equilibred_workorders = args['equilibred_workorders'].lower() == 'true'
limit_tags_str = args['limit_tags']
target_workorder = args['workorder']
limit_tags_str = limit_tags_str.strip().strip('"').strip("'")

if "-" in limit_tags_str:
    limit_tags_min, limit_tags_max = map(int, limit_tags_str.split("-"))
else:
    limit_tags_min = int(limit_tags_str)
    limit_tags_max = 126

print(f"Parameters:\n max_workorders={max_workorders}\n equilibred_workorders={equilibred_workorders}\n limit_tags=[{limit_tags_min}, {limit_tags_max}]")

spark = SparkSession.builder.appName("FeatureExtractionByWorkorder").getOrCreate()

workorder_rows = list_workorders(spark, limit_tags_min, limit_tags_max)
wo_true = [row.workorder for row in workorder_rows if row.isdistorted == True]
wo_false = [row.workorder for row in workorder_rows if row.isdistorted == False]
workorders_to_process = select_workorders(wo_true, wo_false, max_workorders, equilibred_workorders)

print(f"Selected workorders ({len(workorders_to_process)}):")
for row in workorder_rows:
    if row.workorder in workorders_to_process:
        print(f"Workorder: {row.workorder}, IsDistorted: {row.isdistorted}, Number of distinct tags: {row.nb_tags}, Number of rows: {row.nb_rows}")

print(f"Workorders to test: {workorders_to_process}")

base_path = "s3://***MASKED_BUCKET***/ml/Join_tags_workorder_data_cleaned/"
global_start = time.time()

try:
    t0 = time.time()
    df_all = spark.read.parquet(base_path).select("timestamp", "wincctagname", "valuefloat", "isdistorted", "workorder")
    t0 = log_timer(t0, "Read parquet")

    df_filtered = df_all.filter(col("workorder").isin(workorders_to_process))
    count_filtered = df_filtered.count()
    df_filtered = df_filtered.repartition(128)
    log_memory(spark, "After reading and caching filtered data")
    t0 = log_timer(t0, f"Filter workorders → {count_filtered} rows")

    df_filtered = df_filtered.withColumn("id", xxhash64("workorder"))
    df_id_workorder = df_filtered.select("workorder", "id", "isdistorted").distinct()
    count_id_workorder = df_id_workorder.count()
    t0 = log_timer(t0, f"Build df_id_workorder → {count_id_workorder} rows")

    df_tsfresh = df_filtered.select(
        "id",
        col("wincctagname").alias("kind"),
        col("valuefloat").alias("value"),
        col("timestamp").alias("sort")
    )
    count_tsfresh = df_tsfresh.count()
    t0 = log_timer(t0, f"Prepare df_tsfresh → {count_tsfresh} rows")

    log_memory(spark, "Before feature extraction")

    fc_parameters = {
        "mean": None,
        "kurtosis": None,
        "standard_deviation": None,
        "skewness": None,
        "autocorrelation": [{"lag": 1}, {"lag": 5}],
        "time_reversal_asymmetry_statistic": [{"lag": 1}],
        "longest_strike_below_mean": None,
        "absolute_sum_of_changes": None,
        "cid_ce": [{"normalize": True}],
        "mean_abs_change": None,
        "mean_change": None,
        "sample_entropy": None,
        "linear_trend": [{"attr": "slope"}, {"attr": "rvalue"}],
        "fft_coefficient": [{"coeff": 3, "attr": "abs"}],
        "range_count": [{"min": -1e9, "max": 1e9}],
        "quantile": [{"q": 0.25}, {"q": 0.75}]
    }

    df_grouped = df_tsfresh.groupBy("id", "kind")
    df_features = spark_feature_extraction_on_chunk(
        df_grouped,
        column_id="id",
        column_kind="kind",
        column_value="value",
        column_sort="sort",
        default_fc_parameters=fc_parameters
    )

    df_with_pid = df_features.withColumn("partition_id", spark_partition_id())
    counts = df_with_pid.groupBy("partition_id").count().orderBy("count", ascending=False)
    counts.show(100)

    count_features = df_features.count()
    t0 = log_timer(t0, f"Feature extraction → {count_features} rows")

    log_memory(spark, "After feature extraction")

    spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")
    df_features_with_wo = df_features.join(df_id_workorder, on="id", how="left")
    t0 = log_timer(t0, "Join features with workorder")

except Exception as e:
    print(f"An error occurred: {e}")
    traceback.print_exc()

finally:
    total_duration = time.time() - global_start
    print(f"[INFO] Execution completed in {total_duration:.2f}s")

核心问题根源 #

1. tsfresh的Spark绑定本质开销 #

spark_feature_extraction_on_chunk会将每个Spark分区的数据转换为Pandas DataFrame处理，涉及JVM到Python的序列化（Py4J），500万行/WO的量级下，序列化和Pandas计算的开销会被放大。同时每个WO对应120+个标签，按id+kind分组后，Python GIL会限制单进程的并行效率。

2. Spark分区策略不合理 #

直接repartition(128)未利用数据的天然WO分区特性，可能导致单个WO数据被拆分到多分区，或多个WO挤入同一分区，引发数据倾斜或资源浪费。虽然分区数匹配总核心数，但数据分布不均仍会导致部分任务超时。

3. 代码层面冗余操作 #

• 读取全量数据后再过滤WO，未利用Parquet的分区路径优势，引发全表扫描
• 多次调用count()触发不必要的Action，增加重复计算开销
• 未缓存df_filtered，导致后续操作重复读取处理数据

针对性优化方案 #

1. 利用数据分区减少IO #

如果数据按workorder分区存储（路径形如s3://.../workorder=XXX/），直接读取指定WO的路径，避免全表扫描：

# 替换原df_all读取逻辑
selected_paths = [f"{base_path}/workorder={wo}" for wo in workorders_to_process]
df_all = spark.read.parquet(*selected_paths).select("timestamp", "wincctagname", "valuefloat", "isdistorted", "workorder")

2. 优化Spark分区与资源配置 #

• 按workorder分区，确保单个WO数据集中在少数分区：

# 按WO分区，分区数设为WO数量的1-2倍（比如50个WO设为80-100）
df_filtered = df_filtered.repartition(100, "workorder")

• 调整Spark资源参数，适配G2X实例：

spark = SparkSession.builder.appName("FeatureExtractionByWorkorder") \
    .config("spark.executor.memory", "16g") \
    .config("spark.driver.memory", "16g") \
    .config("spark.executor.cores", 4) \
    .config("spark.python.worker.memory", "8g") \
    .getOrCreate()

3. 减少冗余计算与缓存优化 #

• 缓存过滤后的DataFrame，避免重复计算：

df_filtered = df_all.filter(col("workorder").isin(workorders_to_process)).cache()
count_filtered = df_filtered.count()  # 触发缓存

• 移除非必需的count()调用，仅在需要验证数据量时保留

4. 优化tsfresh计算参数 #

• 移除计算密集型非必需特征（如sample_entropy），先测试性能变化
• 调整n_jobs参数，平衡Python进程并行度与GIL竞争：

df_features = spark_feature_extraction_on_chunk(
    df_grouped,
    column_id="id",
    column_kind="kind",
    column_value="value",
    column_sort="sort",
    default_fc_parameters=fc_parameters,
    n_jobs=2  # 根据Executor核心数调整
)

5. 排查与解决数据倾斜 #

查看分区计数结果，如果某分区行数远高于其他，对大WO单独拆分分区：

# 假设600万行是大WO阈值
large_wos = [wo for wo in workorders_to_process if get_wo_row_count(spark, wo) > 6000000]
# 大WO拆分多分区，小WO按WO分区
df_large = df_filtered.filter(col("workorder").isin(large_wos)).repartition(len(large_wos)*2, "workorder")
df_small = df_filtered.filter(~col("workorder").isin(large_wos)).repartition(len(workorders_to_process)-len(large_wos), "workorder")
df_filtered = df_large.union(df_small)

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