Bigtable change streams to Vector Search 模板

此模板会创建一个流处理流水线，用于流式传输 Bigtable 数据更改记录，并使用 Dataflow Runner V2 将其写入 Vertex AI Vector Search。

流水线要求

• Bigtable 源实例必须存在。
• Bigtable 源表必须已存在，且 必须启用变更数据流。
• Bigtable 应用配置文件必须存在。
• 向量搜索索引路径必须存在。

模板参数

参数	说明
embeddingColumn	在其中存储嵌入的完全限定列名。格式为 cf:col。
embeddingByteSize	嵌入数组中每个条目的字节大小。对于 float，请使用 4；对于 double，请使用 8。默认值为 4。
vectorSearchIndex	将流式传输更改的 Vector Search 索引，格式为“projects/{projectID}/locations/{region}/indexes/{indexID}”（不含前后空格）。例如：projects/123/locations/us-east1/indexes/456。
bigtableChangeStreamAppProfile	用于区分 Bigtable 中工作负载的应用配置文件。
bigtableReadInstanceId	表所属的 Bigtable 实例的 ID。
bigtableReadTableId	要读取的 Bigtable 表。
bigtableMetadataTableTableId	可选：要创建的元数据表的 ID。如果未设置，Bigtable 会生成一个 ID。
crowdingTagColumn	可选：存储拥挤标记的完全限定列名称，格式为 cf:col。
allowRestrictsMappings	可选：要用作 allow 限制的列的以逗号分隔完全限定列名及其别名。每个列名称都必须采用 cf:col->alias 格式。
denyRestrictsMappings	可选：要用作 deny 限制的列的以逗号分隔的完全限定列名称及其别名。每个列名称都必须采用 cf:col->alias 格式。
intNumericRestrictsMappings	可选：要用作整数 numeric_restricts 的列的以逗号分隔完全限定列名及其别名。每个列名称都必须采用 cf:col->alias 格式。
floatNumericRestrictsMappings	可选：要用作浮点数（4 字节）numeric_restricts 的列的以逗号分隔完全限定列名及其别名。每个列名称都必须采用 cf:col->alias 格式
doubleNumericRestrictsMappings	可选：要用作双精度（8 字节）numeric_restricts 的列的英文逗号分隔的完全限定列名称及其别名。每个列名称都必须采用 cf:col->alias 格式
upsertMaxBatchSize	可选：在将批次更新/插入向量搜索索引之前要缓冲的更新/插入数量上限。当有 upsertBatchSize 条记录准备就绪时，系统会发送批量数据。 示例：10。
upsertMaxBufferDuration	可选：将一批更新/插入操作发送到向量搜索之前的最大延迟。当有 upsertBatchSize 条记录准备就绪时，系统会发送批量数据。允许的格式如下：Ns（表示秒，例如：5s）、Nm（表示分钟，例如 12m）、Nh（表示小时，例如 2h）。默认值：10s。
deleteMaxBatchSize	可选：从向量搜索索引中删除批次之前要缓冲的最大删除次数。当有 deleteBatchSize 条记录准备就绪时，系统会发送批量数据。例如：10。
deleteMaxBufferDuration	可选：在向 Vector Search 发送批量删除操作之前的最大延迟时间。当有 deleteBatchSize 条记录准备就绪时，系统会发送批量数据。允许的格式如下：Ns（表示秒，例如：5s）、Nm（表示分钟，例如 12m）、Nh（表示小时，例如 2h）。默认值：10s。
dlqDirectory	可选：用于存储未处理记录以及无法处理原因的路径。默认值为 Dataflow 作业的临时位置下的目录。默认值适用于大多数情况。
bigtableChangeStreamMetadataInstanceId	可选：要用于变更数据流连接器元数据表的 Bigtable 实例。默认值为空。
bigtableChangeStreamMetadataTableTableId	可选：要使用的 Bigtable 变更数据流连接器元数据表 ID。如果未提供，系统会在流水线流期间自动创建 Bigtable 变更数据流连接器元数据表。默认值为空。
bigtableChangeStreamCharset	可选：读取值和列限定符时的 Bigtable 变更数据流字符集名称。默认值为 UTF-8。
bigtableChangeStreamStartTimestamp	可选：要用于读取变更数据流的起始 DateTime（含边界值）(https://tools.ietf.org/html/rfc3339)。例如，2022-05-05T07:59:59Z。默认为流水线启动时的时间戳。
bigtableChangeStreamIgnoreColumnFamilies	可选：系统不会捕获对列族名称所做更改的逗号分隔列表。默认值为空。
bigtableChangeStreamIgnoreColumns	可选：要忽略的列名称变更的英文逗号分隔列表。默认值为空。
bigtableChangeStreamName	可选：客户端流水线的唯一名称。此参数允许您从先前运行的流水线停止的位置继续处理。默认为自动生成的名称。如需了解所用的值，请参阅 Dataflow 作业日志。
bigtableChangeStreamResume	可选：设置为 true 时，新流水线将从具有相同名称的先前运行的流水线停止时的点开始处理。如果具有该名称的流水线过去从未运行过，则新流水线将无法启动。 使用 bigtableChangeStreamName 参数 指定流水线行。 如果设置为 false，系统会启动新的流水线。如果以给定来源为来源，过去已运行过与 bigtableChangeStreamName 同名的流水线，则新流水线无法启动。 默认值为 false。
bigtableReadProjectId	可选：要从中读取 Bigtable 数据的项目。此参数的默认项目是 Dataflow 流水线在其中运行的项目。

运行模板

控制台

1. 转到 Dataflow 基于模板创建作业页面。
转到“基于模板创建作业”
2. 在作业名称字段中，输入唯一的作业名称。
3. 可选：对于区域性端点，从下拉菜单中选择一个值。默认区域为 us-central1。

   如需查看可以在其中运行 Dataflow 作业的区域列表，请参阅 Dataflow 位置。

4. 从 Dataflow 模板下拉菜单中，选择 the Bigtable Change Streams to Vector Search template。
5. 在提供的参数字段中，输入您的参数值。
6. 点击运行作业。

gcloud CLI

注意：如需使用 Google Cloud CLI 运行 flex模板，您必须具有 Google Cloud CLI 284.0.0 或更高版本。

在 shell 或终端中，运行模板：

gcloud dataflow flex-template run JOB_NAME \
    --template-file-gcs-location=gs://dataflow-templates-REGION_NAME/VERSION/flex/Bigtable_Change_Streams_to_Vector_Search \
    --project=PROJECT_ID \
    --region=REGION_NAME \
    --parameters \
       embeddingColumn=EMBEDDING_COLUMN,\
       embeddingByteSize=EMBEDDING_BYTE_SIZE,\
       vectorSearchIndex=VECTOR_SEARCH_INDEX,\
       bigtableChangeStreamAppProfile=BIGTABLE_CHANGE_STREAM_APP_PROFILE,\
       bigtableReadInstanceId=BIGTABLE_READ_INSTANCE_ID,\
       bigtableReadTableId=BIGTABLE_READ_TABLE_ID,\

替换以下内容：

• JOB_NAME：您选择的唯一性作业名称
• VERSION：您要使用的模板的版本

  您可使用以下值：

  • latest，以使用模板的最新版本，该模板在存储桶的未标示日期的父文件夹 (gs://dataflow-templates-REGION_NAME/latest/) 中可用
  • 版本名称（如 2023-09-12-00_RC00），以使用模板的特定版本，该版本嵌套在存储桶的相应日期父文件夹 (gs://dataflow-templates-REGION_NAME/) 中
  注意：最新版模板可能会随着重大更改而更新。为了防止这些重大更改影响您的生产工作流程，生产环境应使用有最近标示日期的父文件夹中保存的模板。
• REGION_NAME：要在其中部署 Dataflow 作业的区域，例如 us-central1
• EMBEDDING_COLUMN：嵌入列
• EMBEDDING_BYTE_SIZE：嵌入数组的字节大小。可以是 4 或 8。
• VECTOR_SEARCH_INDEX：向量搜索索引路径
• BIGTABLE_CHANGE_STREAM_APP_PROFILE：Bigtable 应用配置文件 ID
• BIGTABLE_READ_INSTANCE_ID：源 Bigtable 实例 ID
• BIGTABLE_READ_TABLE_ID：源 Bigtable 表 ID

API

如需使用 REST API 来运行模板，请发送 HTTP POST 请求。如需详细了解 API 及其授权范围，请参阅 projects.templates.launch。

POST https://dataflow.googleapis.com/v1b3/projects/PROJECT_ID/locations/LOCATION/flexTemplates:launch
{
   "launchParameter": {
     "jobName": "JOB_NAME",
     "parameters": {
       "embeddingColumn": "EMBEDDING_COLUMN",
       "embeddingByteSize": "EMBEDDING_BYTE_SIZE",
       "vectorSearchIndex": "VECTOR_SEARCH_INDEX",
       "bigtableChangeStreamAppProfile": "BIGTABLE_CHANGE_STREAM_APP_PROFILE",
       "bigtableReadInstanceId": "BIGTABLE_READ_INSTANCE_ID",
       "bigtableReadTableId": "BIGTABLE_READ_TABLE_ID",
     },
     "containerSpecGcsPath": "gs://dataflow-templates-LOCATION/VERSION/flex/Bigtable_Change_Streams_to_Vector_Search",
     "environment": { "maxWorkers": "10" }
  }
}

替换以下内容：

• PROJECT_ID：您要在其中运行 Dataflow 作业的 Google Cloud 项目的 ID
• JOB_NAME：您选择的唯一性作业名称
• VERSION：您要使用的模板的版本

  您可使用以下值：

  • latest，以使用模板的最新版本，该模板在存储桶的未标示日期的父文件夹 (gs://dataflow-templates-REGION_NAME/latest/) 中可用
  • 版本名称（如 2023-09-12-00_RC00），以使用模板的特定版本，该版本嵌套在存储桶的相应日期父文件夹 (gs://dataflow-templates-REGION_NAME/) 中
  注意：最新版模板可能会随着重大更改而更新。为了防止这些重大更改影响您的生产工作流程，生产环境应使用有最近标示日期的父文件夹中保存的模板。
• LOCATION：要在其中部署 Dataflow 作业的区域，例如 us-central1
• EMBEDDING_COLUMN：嵌入列
• EMBEDDING_BYTE_SIZE：嵌入数组的字节大小。可以是 4 或 8。
• VECTOR_SEARCH_INDEX：向量搜索索引路径
• BIGTABLE_CHANGE_STREAM_APP_PROFILE：Bigtable 应用配置文件 ID
• BIGTABLE_READ_INSTANCE_ID：源 Bigtable 实例 ID
• BIGTABLE_READ_TABLE_ID：源 Bigtable 表 ID

模板源代码

Java

/*
 * Copyright (C) 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations under
 * the License.
 */
package com.google.cloud.teleport.v2.templates.bigtablechangestreamstovectorsearch;

import com.google.cloud.Timestamp;
import com.google.cloud.aiplatform.v1.IndexDatapoint;
import com.google.cloud.teleport.metadata.Template;
import com.google.cloud.teleport.metadata.TemplateCategory;
import com.google.cloud.teleport.v2.bigtable.options.BigtableCommonOptions.ReadChangeStreamOptions;
import com.google.cloud.teleport.v2.bigtable.options.BigtableCommonOptions.ReadOptions;
import com.google.cloud.teleport.v2.cdc.dlq.DeadLetterQueueManager;
import com.google.cloud.teleport.v2.options.BigtableChangeStreamsToVectorSearchOptions;
import com.google.cloud.teleport.v2.transforms.DLQWriteTransform;
import com.google.cloud.teleport.v2.utils.DurationUtils;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.GroupIntoBatches;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Values;
import org.apache.beam.sdk.transforms.WithKeys;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.commons.lang3.StringUtils;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

@Template(
    name = "Bigtable_Change_Streams_to_Vector_Search",
    category = TemplateCategory.STREAMING,
    displayName = "Bigtable Change Streams to Vector Search",
    description =
        "Streaming pipeline. Streams Bigtable data change records and writes them into Vertex AI Vector Search using Dataflow Runner V2.",
    optionsClass = BigtableChangeStreamsToVectorSearchOptions.class,
    optionsOrder = {
      BigtableChangeStreamsToVectorSearchOptions.class,
      ReadChangeStreamOptions.class,
      ReadOptions.class
    },
    skipOptions = {
      "bigtableReadAppProfile",
      "bigtableAdditionalRetryCodes",
      "bigtableRpcAttemptTimeoutMs",
      "bigtableRpcTimeoutMs"
    },
    documentation =
        "https://cloud.google.com/dataflow/docs/guides/templates/provided/bigtable-change-streams-to-vector-search",
    flexContainerName = "bigtable-changestreams-to-vector-search",
    contactInformation = "https://cloud.google.com/support",
    streaming = true)
public final class BigtableChangeStreamsToVectorSearch {
  private static final Logger LOG =
      LoggerFactory.getLogger(BigtableChangeStreamsToVectorSearch.class);

  private static final String USE_RUNNER_V2_EXPERIMENT = "use_runner_v2";

  /**
   * Main entry point for executing the pipeline.
   *
   * @param args The command-line arguments to the pipeline.
   */
  public static void main(String[] args) throws Exception {
    LOG.info("Starting replication from Cloud Bigtable Change Streams to Vector Search");

    BigtableChangeStreamsToVectorSearchOptions options =
        PipelineOptionsFactory.fromArgs(args)
            .withValidation()
            .as(BigtableChangeStreamsToVectorSearchOptions.class);

    run(options);
  }

  public static PipelineResult run(BigtableChangeStreamsToVectorSearchOptions options)
      throws IOException {
    options.setStreaming(true);
    options.setEnableStreamingEngine(true);

    List<String> experiments = options.getExperiments();
    if (experiments == null) {
      experiments = new ArrayList<>();
    }
    boolean hasUseRunnerV2 = false;
    for (String experiment : experiments) {
      if (experiment.equalsIgnoreCase(USE_RUNNER_V2_EXPERIMENT)) {
        hasUseRunnerV2 = true;
        break;
      }
    }
    if (!hasUseRunnerV2) {
      experiments.add(USE_RUNNER_V2_EXPERIMENT);
    }
    options.setExperiments(experiments);

    Instant startTimestamp =
        options.getBigtableChangeStreamStartTimestamp().isEmpty()
            ? Instant.now()
            : toInstant(Timestamp.parseTimestamp(options.getBigtableChangeStreamStartTimestamp()));

    String bigtableProjectId = getBigtableProjectId(options);

    LOG.info("  - startTimestamp {}", startTimestamp);
    LOG.info("  - bigtableReadInstanceId {}", options.getBigtableReadInstanceId());
    LOG.info("  - bigtableReadTableId {}", options.getBigtableReadTableId());
    LOG.info("  - bigtableChangeStreamAppProfile {}", options.getBigtableChangeStreamAppProfile());
    LOG.info("  - embeddingColumn {}", options.getEmbeddingColumn());
    LOG.info("  - crowdingTagColumn {}", options.getCrowdingTagColumn());
    LOG.info("  - project {}", options.getProject());
    LOG.info("  - indexName {}", options.getVectorSearchIndex());

    String indexName = options.getVectorSearchIndex();

    String vertexRegion = Utils.extractRegionFromIndexName(indexName);
    String vertexEndpoint = vertexRegion + "-aiplatform.googleapis.com:443";

    final Pipeline pipeline = Pipeline.create(options);

    DeadLetterQueueManager dlqManager = buildDlqManager(options);

    BigtableIO.ReadChangeStream readChangeStream =
        BigtableIO.readChangeStream()
            .withChangeStreamName(options.getBigtableChangeStreamName())
            .withExistingPipelineOptions(
                options.getBigtableChangeStreamResume()
                    ? BigtableIO.ExistingPipelineOptions.RESUME_OR_FAIL
                    : BigtableIO.ExistingPipelineOptions.FAIL_IF_EXISTS)
            .withProjectId(bigtableProjectId)
            .withAppProfileId(options.getBigtableChangeStreamAppProfile())
            .withInstanceId(options.getBigtableReadInstanceId())
            .withTableId(options.getBigtableReadTableId())
            .withMetadataTableInstanceId(options.getBigtableChangeStreamMetadataInstanceId())
            .withMetadataTableTableId(options.getBigtableMetadataTableTableId())
            .withStartTime(startTimestamp);

    PCollectionTuple results =
        pipeline
            .apply("Read from Cloud Bigtable Change Streams", readChangeStream)
            .apply("Create Values", Values.create())
            .apply(
                "Converting to Vector Search Datapoints",
                ParDo.of(
                        new ChangeStreamMutationToDatapointOperationFn(
                            options.getEmbeddingColumn(),
                            options.getEmbeddingByteSize(),
                            options.getCrowdingTagColumn(),
                            Utils.parseColumnMapping(options.getAllowRestrictsMappings()),
                            Utils.parseColumnMapping(options.getDenyRestrictsMappings()),
                            Utils.parseColumnMapping(options.getIntNumericRestrictsMappings()),
                            Utils.parseColumnMapping(options.getFloatNumericRestrictsMappings()),
                            Utils.parseColumnMapping(options.getDoubleNumericRestrictsMappings())))
                    .withOutputTags(
                        ChangeStreamMutationToDatapointOperationFn.UPSERT_DATAPOINT_TAG,
                        TupleTagList.of(
                            ChangeStreamMutationToDatapointOperationFn.REMOVE_DATAPOINT_TAG)));
    results
        .get(ChangeStreamMutationToDatapointOperationFn.UPSERT_DATAPOINT_TAG)
        .apply("Add placeholer keys", WithKeys.of("placeholder"))
        .apply(
            "Batch Contents",
            GroupIntoBatches.<String, IndexDatapoint>ofSize(
                    bufferSizeOption(options.getUpsertMaxBatchSize()))
                .withMaxBufferingDuration(
                    bufferDurationOption(options.getUpsertMaxBufferDuration())))
        .apply("Map to Values", Values.create())
        .apply(
            "Upsert Datapoints to VectorSearch",
            ParDo.of(new UpsertDatapointsFn(vertexEndpoint, indexName)))
        .apply(
            "Write errors to DLQ",
            DLQWriteTransform.WriteDLQ.newBuilder()
                .withDlqDirectory(dlqManager.getSevereDlqDirectory() + "YYYY/MM/dd/HH/mm/")
                .withTmpDirectory(dlqManager.getSevereDlqDirectory() + "tmp/")
                .setIncludePaneInfo(true)
                .build());

    results
        .get(ChangeStreamMutationToDatapointOperationFn.REMOVE_DATAPOINT_TAG)
        .apply("Add placeholder keys", WithKeys.of("placeholer"))
        .apply(
            "Batch Contents",
            GroupIntoBatches.<String, String>ofSize(
                    bufferSizeOption(options.getDeleteMaxBatchSize()))
                .withMaxBufferingDuration(
                    bufferDurationOption(options.getDeleteMaxBufferDuration())))
        .apply("Map to Values", Values.create())
        .apply(
            "Remove Datapoints From VectorSearch",
            ParDo.of(new RemoveDatapointsFn(vertexEndpoint, indexName)))
        .apply(
            "Write errors to DLQ",
            DLQWriteTransform.WriteDLQ.newBuilder()
                .withDlqDirectory(dlqManager.getSevereDlqDirectory() + "YYYY/MM/dd/HH/mm/")
                .withTmpDirectory(dlqManager.getSevereDlqDirectory() + "tmp/")
                .setIncludePaneInfo(true)
                .build());

    return pipeline.run();
  }

  private static String getBigtableProjectId(BigtableChangeStreamsToVectorSearchOptions options) {
    return StringUtils.isEmpty(options.getBigtableReadProjectId())
        ? options.getProject()
        : options.getBigtableReadProjectId();
  }

  private static Instant toInstant(Timestamp timestamp) {
    if (timestamp == null) {
      return null;
    } else {
      return Instant.ofEpochMilli(timestamp.getSeconds() * 1000 + timestamp.getNanos() / 1000000);
    }
  }

  private static int bufferSizeOption(int size) {
    if (size < 1) {
      size = 1;
    }

    return size;
  }

  private static Duration bufferDurationOption(String duration) {
    if (duration.isEmpty()) {
      return Duration.standardSeconds(1);
    }

    return DurationUtils.parseDuration(duration);
  }

  private static DeadLetterQueueManager buildDlqManager(
      BigtableChangeStreamsToVectorSearchOptions options) {
    String dlqDirectory = options.getDlqDirectory();
    if (dlqDirectory.isEmpty()) {
      LOG.info("Falling back to temp dir for DLQ");

      String tempLocation = options.as(DataflowPipelineOptions.class).getTempLocation();

      LOG.info("Have temp location {}", tempLocation);
      if (tempLocation == null || tempLocation.isEmpty()) {
        tempLocation = "/";
      } else if (!tempLocation.endsWith("/")) {
        tempLocation += "/";
      }

      dlqDirectory = tempLocation + "dlq";
    }

    LOG.info("Writing dead letter queue to: {}", dlqDirectory);

    return DeadLetterQueueManager.create(dlqDirectory, 1);
  }
}