Bigtable 變更串流到 Vector Search 範本

這個範本會建立串流管道，使用 Dataflow Runner V2 串流 Bigtable 資料變更記錄，並將這些記錄寫入 Vertex AI Vector Search。

管道相關規定

• Bigtable 來源執行個體必須存在。
• Bigtable 來源資料表必須存在，且該資料表必須啟用變更串流。
• Bigtable 應用程式設定檔必須存在。
• Vector Search 索引路徑必須存在。

範本參數

必要參數

• embeddingColumn：儲存嵌入內容的完整欄名。格式為 cf:col。
• embeddingByteSize：embeddings 陣列中每個項目的位元組大小。使用 4 表示 Float，8 表示 Double。預設值為 4。
• vectorSearchIndex：要串流變更的 Vector Search 索引，格式為「projects/{projectID}/locations/{region}/indexes/{indexID}」(開頭或結尾不得有空格)，例如 projects/123/locations/us-east1/indexes/456。
• bigtableChangeStreamAppProfile：Bigtable 應用程式設定檔 ID。應用程式設定檔必須使用單叢集轉送，並允許單列交易。
• bigtableReadInstanceId：來源 Bigtable 執行個體 ID。
• bigtableReadTableId：來源 Bigtable 資料表 ID。

選用參數

• bigtableMetadataTableTableId：用於建立中繼資料表的資料表 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：將批次 upsert 至 Vector Search 索引前，要緩衝處理的 upsert 數量上限。當有 upsertBatchSize 筆記錄準備就緒，或任何記錄等待 upsertBatchDelay 時間已過時，系統就會傳送批次。例如，10。預設值為 10。
• upsertMaxBufferDuration：將一批 upsert 傳送至 Vector Search 前的最大延遲時間。當有 upsertBatchSize 筆記錄準備就緒，或任何記錄等待 upsertBatchDelay 時間已過，系統就會傳送批次。允許的格式為 Ns (以秒為單位例如 5s)、Nm (以分鐘為單位例如 12m)、Nh (以小時為單位 2h)。例如，10s。預設值為 10 秒。
• deleteMaxBatchSize：從 Vector Search 索引刪除批次前，要緩衝處理的刪除作業數量上限。當有 deleteBatchSize 筆記錄準備就緒，或任何記錄等待 deleteBatchDelay 時間已過時，系統就會傳送批次。例如，10。預設值為 10。
• deleteMaxBufferDuration：將一批刪除作業傳送至 Vector Search 前的最大延遲時間。當有 deleteBatchSize 筆記錄準備就緒，或任何記錄等待 deleteBatchDelay 時間已過時，系統就會傳送批次作業。允許的格式為 Ns (以秒為單位例如 5s)、Nm (以分鐘為單位例如 12m)、Nh (以小時為單位 2h)。例如，10s。預設值為 10 秒。
• dlqDirectory：儲存任何未處理記錄的路徑，並附上處理失敗的原因。預設值是 Dataflow 工作暫時位置下的目錄。在大多數情況下，預設值就已足夠。
• bigtableChangeStreamMetadataInstanceId：Bigtable 變更串流中繼資料執行個體 ID。預設為空白。
• bigtableChangeStreamMetadataTableTableId：Bigtable 變更串流連接器中繼資料表的 ID。如果未提供，管道執行期間會自動建立 Bigtable 變更串流連接器中繼資料表。預設為空白。
• bigtableChangeStreamCharset：Bigtable 變更串流字元集名稱。預設值為 UTF-8。
• bigtableChangeStreamStartTimestamp：用於讀取變更串流的起始時間戳記 (https://tools.ietf.org/html/rfc3339)，包含在內。例如 2022-05-05T07:59:59Z。預設為管道開始時間的時間戳記。
• bigtableChangeStreamIgnoreColumnFamilies：以逗號分隔的清單，列出要忽略的資料欄系列名稱變更。預設為空白。
• bigtableChangeStreamIgnoreColumns：以半形逗號分隔的清單，列出要忽略的資料欄名稱變更。例如：「cf1:col1,cf2:col2」。預設為空白。
• bigtableChangeStreamName：用戶端管道的專屬名稱。可從先前執行管道停止的點繼續處理。預設為自動產生的名稱。如要瞭解使用的值，請參閱 Dataflow 工作記錄。
• bigtableChangeStreamResume：設為 true 時，新管道會從先前執行管道 (具有相同 bigtableChangeStreamName 值) 停止處理的點繼續處理。如果具有指定 bigtableChangeStreamName 值的管道從未執行，系統就不會啟動新管道。設為 false 時，系統會啟動新的管道。如果具有相同 bigtableChangeStreamName 值的管道已針對指定來源執行，系統就不會啟動新管道。預設值為 false。
• bigtableReadChangeStreamTimeoutMs：Bigtable ReadChangeStream 要求的逾時時間 (以毫秒為單位)。
• bigtableReadProjectId：Bigtable 專案 ID。預設值為 Dataflow 工作的專案。

執行範本

控制台

1. 前往 Dataflow 的「Create job from template」(透過範本建立工作) 頁面。
前往「依據範本建立工作」
2. 在「工作名稱」欄位中，輸入不重複的工作名稱。
3. 選用：如要使用區域端點，請從下拉式選單中選取值。預設區域為 us-central1。

   如需可執行 Dataflow 工作的地區清單，請參閱「Dataflow 位置」。

4. 從「Dataflow template」(Dataflow 範本) 下拉式選單中選取 the Bigtable Change Streams to Vector Search template。
5. 在提供的參數欄位中輸入參數值。
6. 按一下「Run Job」(執行工作)。

gcloud CLI

注意：如要使用 Google Cloud CLI 執行 flex 範本，您必須安裝 Google Cloud CLI 284.0.0 版或更新版本。 138.0.0

在殼層或終端機中執行範本：

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 資料欄
• EMBEDDING_BYTE_SIZE：嵌入陣列的位元組大小。可以是 4 或 8。
• VECTOR_SEARCH_INDEX：Vector Search 索引路徑
• 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 工作的專案 ID Google Cloud
• 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 資料欄
• EMBEDDING_BYTE_SIZE：嵌入陣列的位元組大小。可以是 4 或 8。
• VECTOR_SEARCH_INDEX：Vector Search 索引路徑
• 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);
  }
}