Plantilla de flujos de cambios de Bigtable para la búsqueda vectorial

Esta plantilla crea un flujo de procesamiento en streaming para transmitir registros de cambios de datos de Bigtable y escribirlos en Vertex AI Vector Search mediante Dataflow Runner V2.

Requisitos del flujo de procesamiento

• La instancia de origen de Bigtable debe existir.
• La tabla de origen de Bigtable debe existir y tener habilitados los flujos de cambios.
• El perfil de aplicación de Bigtable debe existir.
• La ruta del índice de búsqueda de vectores debe existir.

Parámetros de plantilla

Parámetros obligatorios

• embeddingColumn nombre de columna completo en el que se almacenan las inserciones. En el formato cf:col.
• embeddingByteSize el tamaño en bytes de cada entrada de la matriz de incrustaciones. Usa 4 para Float y 8 para Double. El valor predeterminado es 4.
• vectorSearchIndex el índice de búsqueda vectorial en el que se transmitirán los cambios, con el formato "projects/{projectID}/locations/{region}/indexes/{indexID}" (sin espacios al principio ni al final). Por ejemplo, projects/123/locations/us-east1/indexes/456.
• bigtableChangeStreamAppProfile el ID del perfil de aplicación de Bigtable. El perfil de la aplicación debe usar el enrutamiento de un solo clúster y permitir transacciones de una sola fila.
• bigtableReadInstanceId el ID de la instancia de Bigtable de origen.
• bigtableReadTableId el ID de la tabla de Bigtable de origen.

Parámetros opcionales

• bigtableMetadataTableTableId ID de la tabla que se usa para crear la tabla de metadatos.
• crowdingTagColumn nombre de columna completo en el que se almacena la etiqueta de aglomeración. En el formato cf:col.
• allowRestrictsMappings los nombres de columna completos separados por comas de las columnas que se deben usar como allow restricciones, con su alias. En el formato cf:col->alias.
• denyRestrictsMappings los nombres de columna completos separados por comas de las columnas que se deben usar como restricciones de deny, con su alias. En el formato cf:col->alias.
• intNumericRestrictsMappings los nombres de columna completos separados por comas de las columnas que se deben usar como numeric_restricts entero, con su alias. En el formato cf:col->alias.
• floatNumericRestrictsMappings los nombres de columna completos separados por comas de las columnas que se deben usar como float (4 bytes) numeric_restricts, con su alias. En el formato cf:col->alias.
• doubleNumericRestrictsMappings los nombres de columna completos separados por comas que se deben usar como numeric_restricts de doble precisión (8 bytes), con su alias. En el formato cf:col->alias.
• upsertMaxBatchSize número máximo de inserciones para almacenar en el búfer antes de insertar el lote en el índice de búsqueda vectorial. Los lotes se enviarán cuando haya registros upsertBatchSize listos o cuando haya transcurrido el tiempo upsertBatchDelay desde que se creó un registro. Por ejemplo, 10. El valor predeterminado es 10.
• upsertMaxBufferDuration el retraso máximo antes de que se envíe un lote de inserciones y actualizaciones a la búsqueda vectorial.Los lotes se enviarán cuando haya registros upsertBatchSize listos o cuando haya transcurrido el tiempo upsertBatchDelay desde que se creó un registro. Los formatos permitidos son: Ns (para segundos, por ejemplo, 5s), Nm (para minutos, por ejemplo, 12m) y Nh (para horas, por ejemplo, 2h). Por ejemplo, 10s. El valor predeterminado es 10 s.
• deleteMaxBatchSize: número máximo de eliminaciones que se almacenarán en el búfer antes de eliminar el lote del índice de búsqueda vectorial. Los lotes se enviarán cuando haya registros deleteBatchSize listos o cuando haya transcurrido el tiempo deleteBatchDelay desde que se añadió un registro. Por ejemplo, 10. El valor predeterminado es 10.
• deleteMaxBufferDuration el tiempo máximo que debe transcurrir antes de que se envíe un lote de eliminaciones a la búsqueda vectorial.Los lotes se enviarán cuando haya deleteBatchSize registros listos o cuando haya transcurrido el tiempo deleteBatchDelay desde que se añadió un registro. Los formatos permitidos son: Ns (para segundos, por ejemplo, 5s), Nm (para minutos, por ejemplo, 12m) y Nh (para horas, por ejemplo, 2h). Por ejemplo, 10s. El valor predeterminado es 10 s.
• dlqDirectory la ruta para almacenar los registros sin procesar con el motivo por el que no se han podido procesar. El valor predeterminado es un directorio de la ubicación temporal de la tarea de Dataflow. El valor predeterminado es suficiente en la mayoría de las condiciones.
• bigtableChangeStreamMetadataInstanceId el ID de instancia de metadatos de los flujos de cambios de Bigtable. El valor predeterminado es una cadena vacía.
• bigtableChangeStreamMetadataTableTableId el ID de la tabla de metadatos del conector de flujos de cambios de Bigtable. Si no se proporciona, se crea automáticamente una tabla de metadatos del conector de flujos de cambios de Bigtable durante la ejecución de la canalización. El valor predeterminado es una cadena vacía.
• bigtableChangeStreamCharset nombre del conjunto de caracteres de los flujos de cambios de Bigtable. El valor predeterminado es UTF-8.
• bigtableChangeStreamStartTimestamp marca de tiempo inicial (https://tools.ietf.org/html/rfc3339) (inclusive) que se usará para leer los flujos de cambios. Por ejemplo, 2022-05-05T07:59:59Z. El valor predeterminado es la marca de tiempo de la hora de inicio de la canalización.
• bigtableChangeStreamIgnoreColumnFamilies lista de cambios en los nombres de las familias de columnas separados por comas que se deben ignorar. El valor predeterminado es una cadena vacía.
• bigtableChangeStreamIgnoreColumns lista de cambios de nombres de columnas separados por comas que se deben ignorar. Ejemplo: "cf1:col1,cf2:col2". El valor predeterminado es una cadena vacía.
• bigtableChangeStreamName nombre único de la canalización del cliente. Te permite reanudar el procesamiento desde el punto en el que se detuvo una canalización que se estaba ejecutando. El valor predeterminado es un nombre generado automáticamente. Consulte los registros de trabajos de Dataflow para ver el valor utilizado.
• bigtableChangeStreamResume si se define como true, una nueva canalización reanuda el procesamiento desde el punto en el que se detuvo una canalización que se estaba ejecutando anteriormente con el mismo valor de bigtableChangeStreamName. Si la canalización con el valor bigtableChangeStreamName proporcionado no se ha ejecutado nunca, no se iniciará ninguna canalización. Si se define como false, se inicia una nueva canalización. Si ya se ha ejecutado una canalización con el mismo valor de bigtableChangeStreamName para la fuente en cuestión, no se iniciará una nueva canalización. El valor predeterminado es false.
• bigtableReadChangeStreamTimeoutMs tiempo de espera de las solicitudes ReadChangeStream de Bigtable en milisegundos.
• bigtableReadProjectId el ID del proyecto de Bigtable. El valor predeterminado es el proyecto de la tarea de Dataflow.

Ejecutar la plantilla

Consola

1. Ve a la página Crear tarea a partir de plantilla de Dataflow.
Ir a Crear tarea a partir de plantilla
2. En el campo Nombre de la tarea, introduce un nombre único.
3. Opcional: En Endpoint regional, seleccione un valor en el menú desplegable. La región predeterminada es us-central1.

   Para ver una lista de las regiones en las que puedes ejecutar una tarea de Dataflow, consulta Ubicaciones de Dataflow.

4. En el menú desplegable Plantilla de flujo de datos, seleccione the Bigtable Change Streams to Vector Search template.
5. En los campos de parámetros proporcionados, introduzca los valores de los parámetros.
6. Haz clic en Ejecutar trabajo.

CLI de gcloud

Nota: Para usar la CLI de Google Cloud y ejecutar plantillas de flex, debes tener la versión CLI de Google Cloud 284.0.0 o una posterior.

En tu shell o terminal, ejecuta la plantilla:

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

Haz los cambios siguientes:

• JOB_NAME: un nombre de trabajo único que elijas
• VERSION: la versión de la plantilla que quieres usar

  Puedes usar los siguientes valores:

  • latest para usar la última versión de la plantilla, que está disponible en la carpeta principal sin fecha del contenedor: gs://dataflow-templates-REGION_NAME/latest/
  • el nombre de la versión, como 2023-09-12-00_RC00, para usar una versión específica de la plantilla, que se encuentra anidada en la carpeta principal correspondiente con la fecha en el bucket: gs://dataflow-templates-REGION_NAME/
  Precaución: La última versión de las plantillas puede actualizarse con cambios importantes. Tus entornos de producción deben usar plantillas que se encuentren en la carpeta principal con fecha más reciente para evitar que estos cambios afecten a tus flujos de trabajo de producción.
• REGION_NAME: la región en la que quieras desplegar tu trabajo de Dataflow. Por ejemplo, us-central1
• EMBEDDING_COLUMN: la columna Embedding
• EMBEDDING_BYTE_SIZE: tamaño en bytes de la matriz de inserciones. Puede ser 4 u 8.
• VECTOR_SEARCH_INDEX: la ruta del índice de búsqueda vectorial
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: el ID del perfil de aplicación de Bigtable
• BIGTABLE_READ_INSTANCE_ID: el ID de instancia de Bigtable de origen
• BIGTABLE_READ_TABLE_ID: el ID de la tabla de Bigtable de origen

API

Para ejecutar la plantilla mediante la API REST, envía una solicitud HTTP POST. Para obtener más información sobre la API y sus ámbitos de autorización, 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" }
  }
}

Haz los cambios siguientes:

• PROJECT_ID: el ID del proyecto Google Cloud en el que quieres ejecutar la tarea de Dataflow
• JOB_NAME: un nombre de trabajo único que elijas
• VERSION: la versión de la plantilla que quieres usar

  Puedes usar los siguientes valores:

  • latest para usar la última versión de la plantilla, que está disponible en la carpeta principal sin fecha del contenedor: gs://dataflow-templates-REGION_NAME/latest/
  • el nombre de la versión, como 2023-09-12-00_RC00, para usar una versión específica de la plantilla, que se encuentra anidada en la carpeta principal correspondiente con la fecha en el bucket: gs://dataflow-templates-REGION_NAME/
  Precaución: La última versión de las plantillas puede actualizarse con cambios importantes. Tus entornos de producción deben usar plantillas que se encuentren en la carpeta principal con fecha más reciente para evitar que estos cambios afecten a tus flujos de trabajo de producción.
• LOCATION: la región en la que quieras desplegar tu trabajo de Dataflow. Por ejemplo, us-central1
• EMBEDDING_COLUMN: la columna Embedding
• EMBEDDING_BYTE_SIZE: tamaño en bytes de la matriz de inserciones. Puede ser 4 u 8.
• VECTOR_SEARCH_INDEX: la ruta del índice de búsqueda vectorial
• BIGTABLE_CHANGE_STREAM_APP_PROFILE: el ID del perfil de aplicación de Bigtable
• BIGTABLE_READ_INSTANCE_ID: el ID de instancia de Bigtable de origen
• BIGTABLE_READ_TABLE_ID: el ID de la tabla de Bigtable de origen

Código fuente de la plantilla

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