Vertex AI批量预测时如何将结果映射至含record_id的原记录

方案一：直接保留record_id提交预测（推荐） #

若自定义模型的输入签名仅指定了所需特征字段，Vertex AI批量预测会自动忽略输入表中额外的record_id字段，且该字段会被保留在预测结果中，无需提前剔除。修改代码如下：

from google.cloud import bigquery
from google.cloud import aiplatform

aiplatform.init(project=project_id)
client = bigquery.Client(project=project_id)

# 直接使用包含record_id的原表提交预测
clf = aiplatform.Model(model_id='custom_model')
clf.batch_predict(job_display_name='custom model batch prediction',
                 bigquery_source='bq://your_original_table',
                 instances_format='bigquery',
                 bigquery_destination_prefix='bq://prediction_result_table',
                 predictions_format='bigquery',
                 machine_type='n1-standard-4',
                 max_replica_count=1
                 )

预测完成后，prediction_result_table将包含原表的record_id字段，可直接关联原记录。

方案二：通过临时映射表关联（适用于必须剔除record_id的场景） #

如果模型限制必须剔除record_id才能提交预测，可通过创建临时映射表的方式恢复关联：

步骤1：生成特征表与映射表 #

from google.cloud import bigquery
from google.cloud import aiplatform

aiplatform.init(project=project_id)
client = bigquery.Client(project=project_id)

# 创建带临时唯一ID的特征表，以及record_id与临时ID的映射表
query = '''
CREATE OR REPLACE TABLE `table_with_temp_id` AS
SELECT *, GENERATE_UUID() AS temp_id FROM `your_original_table`;

CREATE OR REPLACE TABLE `record_id_mapping` AS
SELECT record_id, temp_id FROM `table_with_temp_id`;

# 生成用于预测的无record_id表，保留temp_id用于后续关联
CREATE OR REPLACE TABLE `table_for_prediction` AS
SELECT * EXCEPT(record_id) FROM `table_with_temp_id`;
'''
client.query(query).result()

# 提交批量预测任务
clf = aiplatform.Model(model_id='custom_model')
clf.batch_predict(job_display_name='custom model batch prediction',
                 bigquery_source='bq://table_for_prediction',
                 instances_format='bigquery',
                 bigquery_destination_prefix='bq://prediction_result_table',
                 predictions_format='bigquery',
                 machine_type='n1-standard-4',
                 max_replica_count=1
                 )

步骤2：关联恢复record_id #

预测完成后，执行BigQuery SQL将预测结果与原record_id关联：

CREATE OR REPLACE TABLE `final_prediction_with_record_id` AS
SELECT p.*, m.record_id
FROM `prediction_result_table` p
INNER JOIN `record_id_mapping` m ON p.temp_id = m.temp_id;

方案三：利用行号关联（简单但有局限性） #

若原表数据在生成特征表期间无变更，且行顺序严格一致，可通过行号快速关联：

步骤1：生成带行号的特征表与映射表 #

query = '''
CREATE OR REPLACE TABLE `table_with_row_num` AS
SELECT *, ROW_NUMBER() OVER() AS row_num FROM `your_original_table`;

CREATE OR REPLACE TABLE `record_id_mapping` AS
SELECT record_id, row_num FROM `table_with_row_num`;

CREATE OR REPLACE TABLE `table_for_prediction` AS
SELECT * EXCEPT(record_id) FROM `table_with_row_num`;
'''
client.query(query).result()

步骤2：关联恢复record_id #

CREATE OR REPLACE TABLE `final_prediction_with_record_id` AS
SELECT p.*, m.record_id
FROM `prediction_result_table` p
INNER JOIN `record_id_mapping` m ON p.row_num = m.row_num;

注意：此方法仅适用于原表数据固定、行顺序无变化的场景，否则会出现关联错误。

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