Modifiche in tempo reale di Bigtable per il modello Vector Search

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

Requisiti della pipeline

• Deve esistere l'istanza di origine Bigtable.
• La tabella di origine Bigtable deve esistere e la tabella deve avere le modifiche in tempo reale abilitate.
• Il profilo dell'applicazione Bigtable deve esistere.
• Il percorso dell'indice di Vector Search deve esistere.

Parametri del modello

Parametro	Descrizione
embeddingColumn	Il nome completo della colonna in cui sono archiviati gli incorporamenti. Nel formato cf:col.
embeddingByteSize	La dimensione in byte di ogni voce nell'array degli incorporamenti. Usa 4 per float e 8 per il doppio. Il valore predefinito è 4.
vectorSearchIndex	L'indice di Vector Search in cui verranno trasmesse 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	Il profilo dell'applicazione utilizzato per distinguere i carichi di lavoro in Bigtable.
bigtableReadInstanceId	L'ID dell'istanza Bigtable che contiene la tabella.
bigtableReadTableId	La tabella Bigtable da cui leggere.
bigtableMetadataTableTableId	Facoltativo: ID per la tabella di metadati creata. Se il criterio non viene configurato, Bigtable genera un ID.
crowdingTagColumn	(Facoltativo) Il nome completo della colonna in cui è archiviato il tag di crowding, nel formato cf:col.
allowRestrictsMappings	Facoltativo: i nomi delle colonne completi e separati da virgole delle colonne da utilizzare come limitazioni di allow, oltre ai relativi alias. Ogni nome colonna deve essere nel formato cf:col->alias.
denyRestrictsMappings	Facoltativo: i nomi delle colonne completi e separati da virgole delle colonne da utilizzare come limitazioni di deny, oltre ai relativi alias. Ogni nome colonna deve essere nel formato cf:col->alias.
intNumericRestrictsMappings	Facoltativo: i nomi delle colonne completi e separati da virgole delle colonne da utilizzare come numero intero numeric_restricts, più i relativi alias. Ogni nome colonna deve essere nel formato cf:col->alias.
floatNumericRestrictsMappings	Facoltativo: i nomi delle colonne completi e separati da virgole delle colonne da utilizzare come numeri in virgola mobile (4 byte) numeric_restricts, più i relativi alias. Ogni nome colonna deve essere nel formato cf:col->alias
doubleNumericRestrictsMappings	Facoltativo: i nomi delle colonne completi e separati da virgole delle colonne da utilizzare come doppie (8 byte) numeric_restricts, più i relativi alias. Ogni nome colonna deve essere nel formato cf:col->alias
upsertMaxBatchSize	(Facoltativo) Il numero massimo di upsert di cui eseguire il buffer prima di eseguire l'upsert del batch all'indice di Vector Search. I batch vengono inviati quando sono pronti upsertBatchSize record. Esempio: 10.
upsertMaxBufferDuration	(Facoltativo) Il ritardo massimo prima che un batch di upsert venga inviato a Vector Search. I batch vengono inviati quando sono pronti upsertBatchSize record. I formati consentiti sono: Ns per secondi (ad esempio: 5 s), Nm per minuti (ad esempio: 12 min) e Nh per le ore (esempio: 2 h). Valore predefinito: 10s.
deleteMaxBatchSize	(Facoltativo) Il numero massimo di eliminazioni nel buffer prima di eliminare il batch dall'indice di Vector Search. I batch vengono inviati quando sono pronti deleteBatchSize record. Ad esempio: 10.
deleteMaxBufferDuration	(Facoltativo) Il ritardo massimo prima che un batch di eliminazioni venga inviato a Vector Search. I batch vengono inviati quando sono pronti deleteBatchSize record. I formati consentiti sono: Ns per i secondi (ad esempio: 5 sec), Nm per i minuti (ad esempio: 12 m) e Nh per le ore (ad esempio: 2 h). Valore predefinito: 10s.
dlqDirectory	(Facoltativo) Il percorso in cui archiviare tutti i record non elaborati con il motivo della mancata elaborazione. Il valore predefinito è una directory all'interno della posizione temporanea del job Dataflow. Il valore predefinito è appropriato per la maggior parte degli scenari.
bigtableChangeStreamMetadataInstanceId	(Facoltativo) L'istanza Bigtable da utilizzare per la tabella dei metadati del connettore delle modifiche in tempo reale. Il campo predefinito è vuoto.
bigtableChangeStreamMetadataTableTableId	(Facoltativo) L'ID tabella dei metadati del connettore delle modifiche in tempo reale di Bigtable da utilizzare. Se non viene fornita, viene creata automaticamente una tabella di metadati del connettore delle modifiche in tempo reale di Bigtable durante il flusso della pipeline. Il campo predefinito è vuoto.
bigtableChangeStreamCharset	(Facoltativo) Bigtable modifica il nome del set di caratteri durante la lettura dei valori e dei qualificatori di colonna. Il valore predefinito è UTF-8.
bigtableChangeStreamStartTimestamp	Facoltativo: il valore DateTime di inizio, incluso, da utilizzare per la lettura delle modifiche in tempo reale (https://tools.ietf.org/html/rfc3339). Ad esempio, 2022-05-05T07:59:59Z. Il valore predefinito è il timestamp all'avvio della pipeline.
bigtableChangeStreamIgnoreColumnFamilies	(Facoltativo) Un elenco separato da virgole di nomi di famiglia di colonne che non verranno acquisiti. Il campo predefinito è vuoto.
bigtableChangeStreamIgnoreColumns	(Facoltativo) Un elenco separato da virgole di nomi di colonne che non verranno acquisiti. Il campo predefinito è vuoto.
bigtableChangeStreamName	(Facoltativo) Un nome univoco per la pipeline client. Questo parametro consente di riprendere l'elaborazione dal punto in cui è stata arrestata una pipeline in esecuzione in precedenza. Il nome predefinito è generato automaticamente. Controlla i log del job Dataflow per il valore utilizzato.
bigtableChangeStreamResume	(Facoltativo) Se viene impostato su true, l'elaborazione di una nuova pipeline riprende dal punto in cui è stata arrestata una pipeline precedentemente in esecuzione con lo stesso nome. Se una pipeline con questo nome non è mai stata eseguita in passato, non viene avviata la nuova pipeline. Utilizza il parametro bigtableChangeStreamName per specificare la linea della pipeline. Se il criterio viene impostato su false, viene avviata una nuova pipeline. Se una pipeline con lo stesso nome di bigtableChangeStreamName è già stata eseguita in passato per l'origine specificata, la nuova pipeline non viene avviata. Il valore predefinito è false.
bigtableReadProjectId	(Facoltativo) Progetto da cui leggere i dati Bigtable. Il valore predefinito per questo parametro è il progetto in cui è in esecuzione la pipeline 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 di regioni in cui è possibile 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. Inserisci i valori parametro negli appositi campi.
6. Fai clic su Esegui job.

Interfaccia a riga di comando gcloud

Nota: per utilizzare Google Cloud CLI per eseguire flex modelli, devi disporre di Google Cloud CLI 284.0.0 o successiva.

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 job univoco a tua scelta
• VERSION: la versione del modello che vuoi utilizzare

  Puoi utilizzare i seguenti valori:

  • latest per utilizzare la versione più recente del modello, disponibile nella cartella padre non con data del 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 è possibile trovare nidificata nella rispettiva cartella principale con data nel bucket: gs://dataflow-templates-REGION_NAME/
  Attenzione:la versione più recente dei modelli potrebbe essere aggiornata con modifiche che provocano errori. I tuoi ambienti di produzione devono utilizzare i modelli conservati nella cartella principale con data più recente per evitare che queste modifiche che provocano un errore influiscano sui flussi di lavoro di produzione.
• REGION_NAME: la regione in cui vuoi eseguire il deployment del job Dataflow, ad esempio us-central1
• EMBEDDING_COLUMN: la colonna Incorporamento
• EMBEDDING_BYTE_SIZE: la dimensione in byte dell'array degli incorporamenti. Può essere 4 o 8.
• VECTOR_SEARCH_INDEX: il percorso dell'indice di Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: ID profilo applicazione Bigtable
• BIGTABLE_READ_INSTANCE_ID: l'ID dell'istanza Bigtable di origine
• BIGTABLE_READ_TABLE_ID: l'ID tabella Bigtable di origine

API

Per eseguire il modello utilizzando l'API REST, invia una richiesta POST HTTP. Per maggiori 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 del progetto Google Cloud in cui vuoi eseguire il job Dataflow
• JOB_NAME: un nome job univoco a tua scelta
• VERSION: la versione del modello che vuoi utilizzare

  Puoi utilizzare i seguenti valori:

  • latest per utilizzare la versione più recente del modello, disponibile nella cartella padre non con data del 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 è possibile trovare nidificata nella rispettiva cartella principale con data nel bucket: gs://dataflow-templates-REGION_NAME/
  Attenzione:la versione più recente dei modelli potrebbe essere aggiornata con modifiche che provocano errori. I tuoi ambienti di produzione devono utilizzare i modelli conservati nella cartella principale con data più recente per evitare che queste modifiche che provocano un errore influiscano sui flussi di lavoro di produzione.
• LOCATION: la regione in cui vuoi eseguire il deployment del job Dataflow, ad esempio us-central1
• EMBEDDING_COLUMN: la colonna Incorporamento
• EMBEDDING_BYTE_SIZE: la dimensione in byte dell'array degli incorporamenti. Può essere 4 o 8.
• VECTOR_SEARCH_INDEX: il percorso dell'indice di Vector Search
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: ID profilo applicazione Bigtable
• BIGTABLE_READ_INSTANCE_ID: l'ID dell'istanza Bigtable di origine
• BIGTABLE_READ_TABLE_ID: l'ID 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);
  }
}