如何在TFX的Transform组件中编写聚合函数实现按时间聚类分组计数

1. TFX Transform 预处理函数实现 #

首先确保你的TFRecord中的datetime特征可被解析为时间格式，预处理函数实现逻辑如下：

import tensorflow as tf
import tensorflow_transform as tft

def preprocessing_fn(inputs):
    outputs = inputs.copy()
    
    # 步骤1：将datetime特征转换为秒级Unix时间戳
    # 若原始特征已是时间戳格式可跳过字符串转datetime的步骤
    pick_datetime = tf.strings.to_datetime(inputs['tpep_pickup_datetime'])
    timestamp_sec = tf.cast(pick_datetime, tf.int64) // 10**9
    
    # 步骤2：生成30分钟间隔的时间桶，等价于pandas的pd.Grouper(freq='30Min')
    TIME_INTERVAL = 30 * 60 # 30分钟对应秒数
    time_bin = tf.floor(timestamp_sec / TIME_INTERVAL) * TIME_INTERVAL
    
    # 步骤3：构造分组组合键：聚类编号 + 时间桶
    cluster_id = tf.cast(inputs['CLUSTER_kmeans40'], tf.string)
    time_bin_str = tf.cast(time_bin, tf.string)
    group_key = tf.strings.join([cluster_id, time_bin_str], separator='_')
    
    # 步骤4：按组合键统计总记录数，等价于pandas的groupby.sum(counts)
    # 统计结果会作为Transform资产存储，也可映射回原始记录
    group_count = tft.count_per_key(group_key, vocab_filename='cluster_time_30min_count')
    
    # 可选：若需要将聚合计数附加到每条原始记录上，新增以下逻辑
    outputs['cluster_30min_demand'] = group_count.lookup(group_key)
    
    return outputs

注意事项 #

• 若仅需要聚合后的统计结果，直接读取Transform输出资产中的cluster_time_30min_count词汇表即可，里面存储了所有分组对应的计数
• 该实现天然支持TFX分布式流水线，不需要全量加载数据到内存，适配超大规模TFRecord数据集

2. TensorFlow 相关聚合API说明 #

TensorFlow没有直接匹配pandas groupby+时间频率分组的一站式内置函数，但可以通过基础API组合实现需求，常用的聚合相关API包括：

• tf.math.unsorted_segment_sum：可按自定义分段ID对张量做聚合求和，适合灵活的自定义分组逻辑
• tf.data.Dataset.group_by_window：在数据加载迭代阶段按指定key分组聚合，适合流式处理超大数据集
• TensorFlow Transform 扩展API：tft.count_per_key、tft.sum_per_key是专门为TFX流水线设计的全量统计API，会利用Transform的全量数据分析阶段完成聚合，是TFX场景下的最优选择

内容的提问来源于stack exchange，提问作者ken koshy varghese