Modifiche in tempo reale di Bigtable al modello di ricerca vettoriale

Questo modello crea una pipeline di streaming per trasmettere i record delle modifiche dei dati di Bigtable e scriverli in Vertex AI Vector Search utilizzando Dataflow Runner V2.

Requisiti della pipeline

• L'istanza di origine Bigtable deve esistere.
• La tabella di origine Bigtable deve esistere e deve avere gli modifiche in tempo reale abilitati.
• Deve esistere il profilo dell'applicazione Bigtable.
• Il percorso dell'indice Vector Search deve esistere.

Parametri del modello

Parametri obbligatori

• embeddingColumn: il nome completo della colonna in cui sono memorizzati gli incorporamenti. Nel formato cf:col.
• embeddingByteSize: le dimensioni in byte di ogni voce dell'array di incorporamenti. Utilizza 4 per Float e 8 per Double. Il valore predefinito è 4.
• vectorSearchIndex: l'indice Vector Search in cui verranno trasmessi in streaming le modifiche, nel formato "projects/{projectID}/locations/{region}/indexes/{indexID}" (senza spazi iniziali o finali). Ad esempio, projects/123/locations/us-east1/indexes/456.
• bigtableChangeStreamAppProfile: l'ID profilo dell'applicazione Bigtable. Il profilo dell'applicazione deve utilizzare il routing a un singolo cluster e consentire le transazioni su riga singola.
• bigtableReadInstanceId: l'ID dell'istanza Bigtable di origine.
• bigtableReadTableId: l'ID della tabella Bigtable di origine.

Parametri facoltativi

• bigtableMetadataTableTableId: ID tabella utilizzato per creare la tabella dei metadati.
• crowdingTagColumn: il nome completo della colonna in cui è memorizzato il tag di affollamento. Nel formato cf:col.
• allowRestrictsMappings: i nomi di colonna completi separati da virgole delle colonne da utilizzare come restrizioni allow, con il relativo alias. Nel formato cf:col->alias.
• denyRestrictsMappings: i nomi di colonna completi separati da virgole delle colonne che devono essere utilizzate come restrizioni deny, con il relativo alias. Nel formato cf:col->alias.
• intNumericRestrictsMappings: i nomi di colonna completi separati da virgole delle colonne da utilizzare come numeric_restricts intero, con il relativo alias. Nel formato cf:col->alias.
• floatNumericRestrictsMappings: i nomi di colonna completi separati da virgole delle colonne da utilizzare come float (4 byte) numeric_restricts, con il relativo alias. Nel formato cf:col->alias.
• doubleNumericRestrictsMappings: i nomi di colonna completi separati da virgole delle colonne da utilizzare come numeric_restricts double (8 byte), con il relativo alias. Nel formato cf:col->alias.
• upsertMaxBatchSize: il numero massimo di upsert da memorizzare nel buffer prima di eseguire l'upsert del batch nell'indice Vector Search. I batch verranno inviati quando sono pronti upsertBatchSize record o quando è trascorso il tempo di attesa upsertBatchDelay per qualsiasi record. Ad esempio, 10. Il valore predefinito è 10.
• upsertMaxBufferDuration: il ritardo massimo prima che un batch di upsert venga inviato a Vector Search.I batch vengono inviati quando sono pronti upsertBatchSize record o quando è trascorso il tempo di attesa upsertBatchDelay per qualsiasi record. I formati consentiti sono: Ns (per i secondi, ad esempio 5s), Nm (per i minuti, ad esempio 12m), Nh (per le ore, ad esempio 2h). Ad esempio, 10s. Il valore predefinito è 10 secondi.
• deleteMaxBatchSize: il numero massimo di eliminazioni da memorizzare nel buffer prima di eliminare il batch dall'indice Vector Search. I batch vengono inviati quando sono pronti deleteBatchSize record o quando è trascorso il tempo di attesa deleteBatchDelay per qualsiasi record. Ad esempio, 10. Il valore predefinito è 10.
• deleteMaxBufferDuration: il ritardo massimo prima che un batch di eliminazioni venga inviato a Vector Search.I batch vengono inviati quando sono pronti deleteBatchSize record o quando è trascorso il tempo di attesa deleteBatchDelay per un record. I formati consentiti sono: Ns (per i secondi, ad esempio 5s), Nm (per i minuti, ad esempio 12m), Nh (per le ore, ad esempio 2h). Ad esempio, 10s. Il valore predefinito è 10 secondi.
• dlqDirectory: il percorso in cui archiviare i record non elaborati con il motivo per cui non sono stati elaborati. Il valore predefinito è una directory nella posizione temporanea del job Dataflow. Il valore predefinito è sufficiente nella maggior parte delle condizioni.
• bigtableChangeStreamMetadataInstanceId: l'ID istanza dei metadati delle modifiche in tempo reale di Bigtable. Il valore predefinito è vuoto.
• bigtableChangeStreamMetadataTableTableId: l'ID della tabella dei metadati del connettore Bigtable modifiche in tempo reale. Se non specificata, durante l'esecuzione della pipeline viene creata automaticamente una tabella di metadati del connettore di modifiche in tempo reale Bigtable. Il valore predefinito è vuoto.
• bigtableChangeStreamCharset: il nome del set di caratteri dei modifiche in tempo reale Bigtable. Il valore predefinito è UTF-8.
• bigtableChangeStreamStartTimestamp: il timestamp iniziale (https://tools.ietf.org/html/rfc3339), inclusivo, da utilizzare per la lettura dei modifiche in tempo reale. Ad esempio: 2022-05-05T07:59:59Z. Il valore predefinito è il timestamp dell'ora di inizio della pipeline.
• bigtableChangeStreamIgnoreColumnFamilies: un elenco separato da virgole di modifiche ai nomi famiglia di colonne da ignorare. Il valore predefinito è vuoto.
• bigtableChangeStreamIgnoreColumns: un elenco separato da virgole di modifiche ai nomi delle colonne da ignorare. Esempio: "cf1:col1,cf2:col2". Il valore predefinito è vuoto.
• bigtableChangeStreamName: un nome univoco per la pipeline client. Consente di riprendere l'elaborazione dal punto in cui si è interrotta una pipeline in esecuzione in precedenza. Il valore predefinito è un nome generato automaticamente. Consulta i log dei job Dataflow per il valore utilizzato.
• bigtableChangeStreamResume: se impostato su true, una nuova pipeline riprende l'elaborazione dal punto in cui si è interrotta una pipeline in esecuzione in precedenza con lo stesso valore bigtableChangeStreamName. Se la pipeline con il valore bigtableChangeStreamName specificato non è mai stata eseguita, non viene avviata una nuova pipeline. Se impostato su false, viene avviata una nuova pipeline. Se una pipeline con lo stesso valore di bigtableChangeStreamName è già stata eseguita per l'origine specificata, non viene avviata una nuova pipeline. Il valore predefinito è false.
• bigtableReadChangeStreamTimeoutMs: il timeout per le richieste Bigtable ReadChangeStream in millisecondi.
• bigtableReadProjectId: l'ID progetto Bigtable. Il valore predefinito è il progetto per il job Dataflow.

Esegui il modello

Console

1. Vai alla pagina Crea job da modello di Dataflow.
Vai a Crea job da modello
2. Nel campo Nome job, inserisci un nome univoco per il job.
3. (Facoltativo) Per Endpoint a livello di regione, seleziona un valore dal menu a discesa. La regione predefinita è us-central1.

   Per un elenco delle regioni in cui puoi eseguire un job Dataflow, consulta Località di Dataflow.

4. Dal menu a discesa Modello Dataflow, seleziona the Bigtable Change Streams to Vector Search template.
5. Nei campi dei parametri forniti, inserisci i valori dei parametri.
6. Fai clic su Esegui job.

Interfaccia a riga di comando gcloud

Nota:per utilizzare Google Cloud CLI per eseguire modelli flex, devi avere la versione Google Cloud CLI 284.0.0 o successive.

Nella shell o nel terminale, esegui il modello:

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

Sostituisci quanto segue:

• JOB_NAME: un nome univoco del job a tua scelta
• VERSION: la versione del modello che vuoi utilizzare

  Puoi utilizzare i seguenti valori:

  • latest per utilizzare l'ultima versione del modello, disponibile nella cartella principale senza data nel bucket: gs://dataflow-templates-REGION_NAME/latest/
  • il nome della versione, ad esempio 2023-09-12-00_RC00, per utilizzare una versione specifica del modello, che si trova nidificata nella rispettiva cartella principale con data nel bucket: gs://dataflow-templates-REGION_NAME/
  Attenzione:l'ultima versione dei modelli potrebbe essere aggiornata con modifiche che causano interruzioni. Gli ambienti di produzione devono utilizzare i modelli conservati nella cartella principale con data più recente per evitare che queste modifiche che causano interruzioni influiscano sui workflow di produzione.
• REGION_NAME: la regione in cui vuoi eseguire il deployment del job Dataflow, ad esempio us-central1
• EMBEDDING_COLUMN: la colonna Embedding
• EMBEDDING_BYTE_SIZE: la dimensione in byte dell'array di incorporamenti. Può essere 4 o 8.
• VECTOR_SEARCH_INDEX: il percorso dell'indice Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: l'ID profilo di applicazione Bigtable
• BIGTABLE_READ_INSTANCE_ID: l'ID dell'istanza Bigtable di origine
• BIGTABLE_READ_TABLE_ID: l'ID della tabella Bigtable di origine

API

Per eseguire il modello utilizzando l'API REST, invia una richiesta POST HTTP. Per ulteriori informazioni sull'API e sui relativi ambiti di autorizzazione, consulta 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" }
  }
}

Sostituisci quanto segue:

• PROJECT_ID: l'ID progetto Google Cloud in cui vuoi eseguire il job Dataflow
• JOB_NAME: un nome univoco del job a tua scelta
• VERSION: la versione del modello che vuoi utilizzare

  Puoi utilizzare i seguenti valori:

  • latest per utilizzare l'ultima versione del modello, disponibile nella cartella principale senza data nel bucket: gs://dataflow-templates-REGION_NAME/latest/
  • il nome della versione, ad esempio 2023-09-12-00_RC00, per utilizzare una versione specifica del modello, che si trova nidificata nella rispettiva cartella principale con data nel bucket: gs://dataflow-templates-REGION_NAME/
  Attenzione:l'ultima versione dei modelli potrebbe essere aggiornata con modifiche che causano interruzioni. Gli ambienti di produzione devono utilizzare i modelli conservati nella cartella principale con data più recente per evitare che queste modifiche che causano interruzioni influiscano sui workflow di produzione.
• LOCATION: la regione in cui vuoi eseguire il deployment del job Dataflow, ad esempio us-central1
• EMBEDDING_COLUMN: la colonna Embedding
• EMBEDDING_BYTE_SIZE: la dimensione in byte dell'array di incorporamenti. Può essere 4 o 8.
• VECTOR_SEARCH_INDEX: il percorso dell'indice Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: l'ID profilo di applicazione Bigtable
• BIGTABLE_READ_INSTANCE_ID: l'ID dell'istanza Bigtable di origine
• BIGTABLE_READ_TABLE_ID: l'ID della tabella Bigtable di origine

Codice sorgente del modello

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