Streams de alterações do Bigtable para o modelo de pesquisa vetorial

Este modelo cria um pipeline de streaming para transmitir registos de alterações de dados do Bigtable e escrevê-los na pesquisa vetorial da Vertex AI através do Dataflow Runner V2.

Requisitos do pipeline

• A instância de origem do Bigtable tem de existir.
• A tabela de origem do Bigtable tem de existir e a tabela tem de ter streams de alterações ativadas.
• O perfil da aplicação Bigtable tem de existir.
• O caminho do índice do Vector Search tem de existir.

Parâmetros de modelos

Parâmetros obrigatórios

• embeddingColumn: o nome da coluna totalmente qualificado onde os encaixes estão armazenados. No formato cf:col.
• embeddingByteSize: o tamanho em bytes de cada entrada na matriz de incorporações. Use 4 para Float e 8 para Double. A predefinição é: 4.
• vectorSearchIndex: o índice do Vector Search onde as alterações vão ser transmitidas em fluxo, no formato "projects/{projectID}/locations/{region}/indexes/{indexID}" (sem espaços no início nem no fim). Por exemplo, projects/123/locations/us-east1/indexes/456.
• bigtableChangeStreamAppProfile: o ID do perfil da aplicação do Bigtable. O perfil da aplicação tem de usar o encaminhamento de cluster único e permitir transações de linha única.
• bigtableReadInstanceId: o ID da instância do Bigtable de origem.
• bigtableReadTableId: o ID da tabela do Bigtable de origem.

Parâmetros opcionais

• bigtableMetadataTableTableId: ID da tabela usado para criar a tabela de metadados.
• crowdingTagColumn: o nome da coluna totalmente qualificado onde a etiqueta de sobreposição está armazenada. No formato cf:col.
• allowRestrictsMappings: os nomes de colunas totalmente qualificados separados por vírgulas das colunas que devem ser usadas como allow restrições, com o respetivo alias. No formato cf:col->alias.
• denyRestrictsMappings: os nomes das colunas totalmente qualificados separados por vírgulas das colunas que devem ser usadas como deny restrições, com o respetivo alias. No formato cf:col->alias.
• intNumericRestrictsMappings: os nomes das colunas totalmente qualificados separados por vírgulas das colunas que devem ser usados como número inteiro numeric_restricts, com o respetivo alias. No formato cf:col->alias.
• floatNumericRestrictsMappings: os nomes das colunas totalmente qualificados separados por vírgulas das colunas que devem ser usados como float (4 bytes) numeric_restricts, com o respetivo alias. No formato cf:col->alias.
• doubleNumericRestrictsMappings: os nomes das colunas totalmente qualificados separados por vírgulas das colunas que devem ser usados como duplos (8 bytes) numeric_restricts, com o respetivo alias. No formato cf:col->alias.
• upsertMaxBatchSize: o número máximo de inserções/atualizações a armazenar em buffer antes de inserir/atualizar o lote no índice do Vector Search. Os lotes são enviados quando existem registos upsertBatchSize prontos ou quando passou o tempo upsertBatchDelay de espera de qualquer registo. Por exemplo, 10. A predefinição é: 10.
• upsertMaxBufferDuration: o atraso máximo antes de um lote de inserções/atualizações ser enviado para a pesquisa vetorial.Os lotes são enviados quando existem registos upsertBatchSize prontos ou quando qualquer registo está à espera há upsertBatchDelay. Os formatos permitidos são: Ns (para segundos, exemplo: 5s), Nm (para minutos, exemplo: 12m), Nh (para horas, exemplo: 2h). Por exemplo, 10s. A predefinição é: 10 s.
• deleteMaxBatchSize: o número máximo de eliminações a colocar em buffer antes de eliminar o lote do índice de pesquisa vetorial. Os lotes são enviados quando existem registos deleteBatchSize prontos ou quando o tempo de espera deleteBatchDelay de qualquer registo tiver passado. Por exemplo, 10. A predefinição é: 10.
• deleteMaxBufferDuration: o atraso máximo antes de um lote de eliminações ser enviado para a pesquisa vetorial.Os lotes são enviados quando existem registos deleteBatchSize prontos ou quando qualquer registo está à espera há deleteBatchDelay. Os formatos permitidos são: Ns (para segundos, exemplo: 5s), Nm (para minutos, exemplo: 12m), Nh (para horas, exemplo: 2h). Por exemplo, 10s. A predefinição é: 10 s.
• dlqDirectory: o caminho para armazenar todos os registos não processados com o motivo pelo qual não foram processados. A predefinição é um diretório na localização temporária da tarefa do Dataflow. O valor predefinido é suficiente na maioria das condições.
• bigtableChangeStreamMetadataInstanceId: o ID da instância de metadados das streams de alterações do Bigtable. A predefinição é vazio.
• bigtableChangeStreamMetadataTableTableId: o ID da tabela de metadados do conetor de streams de alterações do Bigtable. Se não for fornecida, é criada automaticamente uma tabela de metadados do conector de streams de alterações do Bigtable durante a execução do pipeline. A predefinição é vazio.
• bigtableChangeStreamCharset: o nome do conjunto de carateres das streams de alterações do Bigtable. A predefinição é: UTF-8.
• bigtableChangeStreamStartTimestamp: a data/hora de início (https://tools.ietf.org/html/rfc3339), inclusive, a usar para ler streams de alterações. Por exemplo, 2022-05-05T07:59:59Z. O valor predefinido é a data/hora de início do pipeline.
• bigtableChangeStreamIgnoreColumnFamilies: uma lista separada por vírgulas de alterações de nomes de famílias de colunas a ignorar. A predefinição é vazio.
• bigtableChangeStreamIgnoreColumns: uma lista separada por vírgulas de alterações de nomes de colunas a ignorar. Exemplo: "cf1:col1,cf2:col2". A predefinição é vazio.
• bigtableChangeStreamName: um nome exclusivo para o pipeline do cliente. Permite-lhe retomar o processamento a partir do ponto em que um pipeline em execução anterior parou. A predefinição é um nome gerado automaticamente. Consulte os registos da tarefa do Dataflow para ver o valor usado.
• bigtableChangeStreamResume: quando definido como true, um novo pipeline retoma o processamento a partir do ponto em que um pipeline executado anteriormente com o mesmo valor bigtableChangeStreamName parou. Se o pipeline com o valor bigtableChangeStreamName especificado nunca tiver sido executado, não é iniciado um novo pipeline. Quando definido como false, é iniciado um novo pipeline. Se já tiver sido executado um pipeline com o mesmo valor de bigtableChangeStreamName para a origem especificada, não é iniciado um novo pipeline. A predefinição é false.
• bigtableReadChangeStreamTimeoutMs: o limite de tempo para pedidos Bigtable ReadChangeStream em milissegundos.
• bigtableReadProjectId: o ID do projeto do Bigtable. A predefinição é o projeto da tarefa do Dataflow.

Execute o modelo

Consola

1. Aceda à página Dataflow Criar tarefa a partir de um modelo.
Aceda a Criar tarefa a partir de modelo
2. No campo Nome da tarefa, introduza um nome exclusivo para a tarefa.
3. Opcional: para Ponto final regional, selecione um valor no menu pendente. A região predefinida é us-central1.

   Para ver uma lista das regiões onde pode executar uma tarefa do Dataflow, consulte o artigo Localizações do Dataflow.

4. No menu pendente Modelo do fluxo de dados, selecione the Bigtable Change Streams to Vector Search template.
5. Nos campos de parâmetros fornecidos, introduza os valores dos parâmetros.
6. Clique em Executar tarefa.

CLI gcloud

Nota: para usar a CLI do Google Cloud para executar modelos, tem de ter a versão 284.0.0 ou posterior da CLI do Google Cloud. flex

Na shell ou no terminal, execute o modelo:

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,\

Substitua o seguinte:

• JOB_NAME: um nome de tarefa exclusivo à sua escolha
• VERSION: a versão do modelo que quer usar

  Pode usar os seguintes valores:

  • latest para usar a versão mais recente do modelo, que está disponível na pasta principal sem data no contentor: gs://dataflow-templates-REGION_NAME/latest/
  • o nome da versão, como 2023-09-12-00_RC00, para usar uma versão específica do modelo, que pode ser encontrada aninhada na pasta principal com a data correspondente no contentor: gs://dataflow-templates-REGION_NAME/
  Atenção: a versão mais recente dos modelos pode ser atualizada com alterações destrutivas. Os seus ambientes de produção devem usar modelos mantidos na pasta principal datada mais recente para evitar que estas alterações significativas afetem os seus fluxos de trabalho de produção.
• REGION_NAME: a região onde quer implementar a tarefa do Dataflow, por exemplo, us-central1
• EMBEDDING_COLUMN: a coluna Incorporação
• EMBEDDING_BYTE_SIZE: o tamanho em bytes da matriz de incorporações. Pode ser 4 ou 8.
• VECTOR_SEARCH_INDEX: o caminho do índice do Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: o ID do perfil da aplicação do Bigtable
• BIGTABLE_READ_INSTANCE_ID: o ID da instância do Bigtable de origem
• BIGTABLE_READ_TABLE_ID: o ID da tabela do Bigtable de origem

API

Para executar o modelo através da API REST, envie um pedido HTTP POST. Para mais informações sobre a API e os respetivos âmbitos de autorização, consulte 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" }
  }
}

Substitua o seguinte:

• PROJECT_ID: o ID do projeto onde quer executar a tarefa do Dataflow Google Cloud
• JOB_NAME: um nome de tarefa exclusivo à sua escolha
• VERSION: a versão do modelo que quer usar

  Pode usar os seguintes valores:

  • latest para usar a versão mais recente do modelo, que está disponível na pasta principal sem data no contentor: gs://dataflow-templates-REGION_NAME/latest/
  • o nome da versão, como 2023-09-12-00_RC00, para usar uma versão específica do modelo, que pode ser encontrada aninhada na pasta principal com a data correspondente no contentor: gs://dataflow-templates-REGION_NAME/
  Atenção: a versão mais recente dos modelos pode ser atualizada com alterações destrutivas. Os seus ambientes de produção devem usar modelos mantidos na pasta principal datada mais recente para evitar que estas alterações significativas afetem os seus fluxos de trabalho de produção.
• LOCATION: a região onde quer implementar a tarefa do Dataflow, por exemplo, us-central1
• EMBEDDING_COLUMN: a coluna Embedding
• EMBEDDING_BYTE_SIZE: o tamanho em bytes da matriz de incorporações. Pode ser 4 ou 8.
• VECTOR_SEARCH_INDEX: o caminho do índice do Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: o ID do perfil da aplicação do Bigtable
• BIGTABLE_READ_INSTANCE_ID: o ID da instância do Bigtable de origem
• BIGTABLE_READ_TABLE_ID: o ID da tabela do Bigtable de origem

Código-fonte do modelo

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);
  }
}