使用带时间戳键的feature_store_lookups时AutoML Forecast失败

环境与场景 #

我们在Databricks Runtime 15.4 ML LTS和16.4 ML LTS上运行AutoML Forecast，使用包含Feature Store时间协变量（如corona_dummy）的时序数据集，通过lookup_key和timestamp_lookup_key配置feature_store_lookups关联特征表。

特征表定义 #

fs.create_table(
  name="...features.example_corona_features",
  primary_keys=["Monat", "Produkt", "Vertriebstyp_Art"],
  df=...,
  timestamp_keys="Monat",
  ...
)

AutoML调用代码 #

feature_store_lookups=[{
  "table_name": "...features.example_corona_features",
  "lookup_key": ["Produkt", "Vertriebstyp_Art"],
  "timestamp_lookup_key": "Monat"
}]

预期结果 #

AutoML能在数据集和特征表之间通过时间戳与键执行时间连接，将corona_dummy协变量纳入训练并顺利完成训练。

实际问题 #

AutoML启动后，在内部applyInPandas()或.toPandas()转换阶段失败，抛出错误：

ValueError: Length mismatch: Expected axis has 8 elements, new values have 11 elements

该错误发生在特征连接和训练集加载完成后，即AutoML内部训练循环执行期间。

观察结果 #

• 移除feature_store_lookups配置后，AutoML可无错误完成运行。
• 仅当时间戳列（Monat）同时满足以下两个条件时才会触发问题：
  • 该列属于特征表的primary_keys；
  • 同时被指定为timestamp_lookup_key传入。

问题 #

这是否是已知问题？在AutoML中使用带timestamp_lookup_key的feature_store_lookups的正确方式是什么？

问题分析与解决方案 #

这是Databricks AutoML Forecast在处理时间戳同时作为主键和时间查找键时的已知兼容问题，核心原因是AutoML内部在特征连接后的数据处理逻辑中，对重复的时间戳列（主数据集和特征表的Monat列）处理不当，导致列数匹配错误。

正确的使用方式需要调整特征表定义和AutoML调用配置：

1. 调整特征表主键：
   将时间戳列（Monat）从primary_keys中移除，仅保留非时间维度的键作为主键，同时保留timestamp_keys配置。修改后的特征表定义示例：

   fs.create_table(
     name="...features.example_corona_features",
     primary_keys=["Produkt", "Vertriebstyp_Art"],
     df=...,
     timestamp_keys="Monat",
     ...
   )
   

   这样特征表会通过timestamp_keys维护时间维度，同时主键仅用于非时间维度的关联。

2. 保持AutoML调用配置不变：
   原有的feature_store_lookups配置无需修改，仍然指定lookup_key为非时间主键，timestamp_lookup_key为Monat即可。

3. 验证数据唯一性：
   调整主键后，需确保在相同的Produkt、Vertriebstyp_Art和Monat组合下，特征表中的数据是唯一的，避免时间连接时出现重复行导致后续训练错误。

如果无法修改现有特征表结构，可临时通过以下方式规避：

• 在主数据集中重命名时间戳列，例如将主数据集的Monat改为ds，然后在feature_store_lookups中将timestamp_lookup_key指定为ds，同时确保特征表的timestamp_keys仍为Monat，AutoML会自动处理时间映射。

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