Modelo do Spanner para o Vertex AI Vector Search

O modelo para ficheiros do Spanner para a pesquisa vetorial da Vertex AI no Cloud Storage cria um pipeline em lote que exporta dados de incorporações vetoriais de uma tabela do Spanner para o Cloud Storage no formato JSON. Use parâmetros de modelo para especificar a pasta do Cloud Storage para a qual exportar as incorporações vetoriais. A pasta do Cloud Storage contém a lista de ficheiros .json exportados, que representam as incorporações vetoriais num formato suportado pelo índice de pesquisa vetorial do Vertex AI.

Para mais informações, consulte o artigo Formato e estrutura dos dados de entrada.

Requisitos do pipeline

• A base de dados do Spanner tem de existir.
• O contentor do Cloud Storage para a saída de dados tem de existir.
• Além das funções de gestão de identidade e de acesso (IAM) necessárias para executar tarefas do Dataflow, precisa das funções de IAM necessárias para ler os seus dados do Spanner e escrever no seu contentor do Cloud Storage.

Parâmetros de modelos

Parâmetros obrigatórios

• spannerProjectId: o ID do projeto da instância do Spanner.
• spannerInstanceId: o ID da instância do Spanner a partir da qual exportar as incorporações de vetores.
• spannerDatabaseId: o ID da base de dados do Spanner a partir da qual exportar as incorporações vetoriais.
• spannerTable: a tabela do Spanner a partir da qual ler.
• spannerColumnsToExport: uma lista separada por vírgulas das colunas necessárias para o índice do Vertex AI Vector Search. As colunas de ID e incorporação são obrigatórias para o Vector Search. Se os nomes das colunas não corresponderem à estrutura de entrada do índice do Vertex AI Vector Search, crie mapeamentos de colunas através de alias. Se os nomes das colunas não corresponderem ao formato esperado pelo Vertex AI, use a notação de:para. Por exemplo, se tiver colunas denominadas id e my_embedding, especifique id, my_embedding:embedding.
• gcsOutputFolder: a pasta do Cloud Storage para a qual os ficheiros de saída são escritos. O caminho tem de terminar com uma barra. Por exemplo, gs://your-bucket/folder1/.
• gcsOutputFilePrefix: o prefixo do nome do ficheiro para escrever ficheiros de saída. Por exemplo, vector-embeddings.

Parâmetros opcionais

• spannerHost: o ponto final do Spanner a chamar no modelo. O valor predefinido é https://batch-spanner.googleapis.com. Por exemplo, https://batch-spanner.googleapis.com.
• spannerVersionTime: se estiver definido, especifica a hora em que a versão da base de dados tem de ser obtida. O valor é uma string no formato de data RFC-3339 na hora epoch Unix. Por exemplo: 1990-12-31T23:59:60Z. A data/hora tem de ser no passado e aplica-se a obsolescência máxima da data/hora (https://cloud.google.com/spanner/docs/timestamp-bounds#maximum_timestamp_staleness). Se não estiver definido, é usado um limite forte (https://cloud.google.com/spanner/docs/timestamp-bounds#strong) para ler os dados mais recentes. A predefinição é empty. Por exemplo, 1990-12-31T23:59:60Z.
• spannerDataBoostEnabled: quando definido como true, o modelo usa o cálculo a pedido do Spanner. A tarefa de exportação é executada em recursos de computação independentes que não afetam as cargas de trabalho atuais do Spanner. A utilização desta opção incorre em custos adicionais no Spanner. Para mais informações, consulte a vista geral do Spanner Data Boost (https://cloud.google.com/spanner/docs/databoost/databoost-overview). Predefinição: false.
• spannerPriority: a prioridade do pedido para chamadas do Spanner. Os valores permitidos são HIGH, MEDIUM e LOW. O valor predefinido é MEDIUM.

Execute o modelo

Consola

1. Aceda à página do fluxo de dados 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 as localizações do Dataflow.

4. No menu pendente Modelo do fluxo de dados, selecione the Spanner to Vertex AI Vector Search files on Cloud Storage template.
5. Nos campos de parâmetros fornecidos, introduza os valores dos parâmetros.
6. Clique em Executar tarefa.

gcloud

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

Na shell ou no terminal, execute o modelo:

gcloud dataflow jobs run JOB_NAME \
    --gcs-location=gs://dataflow-templates-REGION_NAME/VERSION/Cloud_Spanner_vectors_to_Cloud_Storage \
    --project=PROJECT_ID \
    --region=REGION_NAME \
    --parameters \
       spannerProjectId=SPANNER_PROJECT_ID,\
       spannerInstanceId=SPANNER_INSTANCE_ID,\
       spannerDatabaseId=SPANNER_DATABASE_ID,\
       spannerTable=SPANNER_TABLE,\
       spannerColumnsToExport=SPANNER_COLUMNS_TO_EXPORT,\
       gcsOutputFolder=GCS_OUTPUT_FOLDER,\
       gcsOutputFilePrefix=GCS_OUTPUT_FILE_PREFIX,\

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
• SPANNER_PROJECT_ID: o ID do projeto do Spanner
• SPANNER_INSTANCE_ID: o ID da instância do Spanner
• SPANNER_DATABASE_ID: o ID da base de dados do Spanner
• SPANNER_TABLE: a tabela do Spanner
• SPANNER_COLUMNS_TO_EXPORT: as colunas a exportar da tabela do Spanner
• GCS_OUTPUT_FOLDER: a pasta do Cloud Storage para a qual os ficheiros são enviados
• GCS_OUTPUT_FILE_PREFIX: o prefixo dos ficheiros de saída no Cloud Storage

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/templates:launch?gcsPath=gs://dataflow-templates-LOCATION/VERSION/Cloud_Spanner_vectors_to_Cloud_Storage
{
   "jobName": "JOB_NAME",
   "parameters": {
     "spannerProjectId": "SPANNER_PROJECT_ID",
     "spannerInstanceId": "SPANNER_INSTANCE_ID",
     "spannerDatabaseId": "SPANNER_DATABASE_ID",
     "spannerTable": "SPANNER_TABLE",
     "spannerColumnsToExport": "SPANNER_COLUMNS_TO_EXPORT",
     "gcsOutputFolder": "GCS_OUTPUT_FOLDER",
     "gcsOutputFilePrefix": "GCS_OUTPUT_FILE_PREFIX",
   },
   "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
• SPANNER_PROJECT_ID: o ID do projeto do Spanner
• SPANNER_INSTANCE_ID: o ID da instância do Spanner
• SPANNER_DATABASE_ID: o ID da base de dados do Spanner
• SPANNER_TABLE: a tabela do Spanner
• SPANNER_COLUMNS_TO_EXPORT: as colunas a exportar da tabela do Spanner
• GCS_OUTPUT_FOLDER: a pasta do Cloud Storage para a qual os ficheiros são enviados
• GCS_OUTPUT_FILE_PREFIX: o prefixo dos ficheiros de saída no Cloud Storage

Código-fonte do modelo

Java

/*
 * Copyright (C) 2023 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.templates;

import com.google.cloud.spanner.Options.RpcPriority;
import com.google.cloud.teleport.metadata.Template;
import com.google.cloud.teleport.metadata.TemplateCategory;
import com.google.cloud.teleport.metadata.TemplateParameter;
import com.google.cloud.teleport.metadata.TemplateParameter.TemplateEnumOption;
import com.google.cloud.teleport.spanner.spannerio.ReadOperation;
import com.google.cloud.teleport.spanner.spannerio.SpannerConfig;
import com.google.cloud.teleport.spanner.spannerio.SpannerIO;
import com.google.cloud.teleport.spanner.spannerio.Transaction;
import com.google.cloud.teleport.templates.SpannerVectorEmbeddingExport.SpannerToVectorEmbeddingJsonOptions;
import com.google.cloud.teleport.templates.common.SpannerConverters;
import com.google.cloud.teleport.templates.common.SpannerConverters.CreateTransactionFnWithTimestamp;
import com.google.cloud.teleport.templates.common.SpannerConverters.VectorSearchStructValidator;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.io.TextIO;
import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.options.ValueProvider;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TypeDescriptors;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Dataflow template which export vector embeddings from Spanner to GCS in json format. It exports a
 * Spanner table using <a
 * href="https://cloud.google.com/spanner/docs/reads#read_data_in_parallel">Batch API</a>, which
 * creates multiple workers in parallel for better performance. The result is written to a JSON file
 * in Google Cloud Storage.
 *
 * <p>Check out <a
 * href="https://github.com/GoogleCloudPlatform/DataflowTemplates/blob/main/v1/README_Cloud_Spanner_to_Vector_Embedding.md">README</a>
 * for instructions on how to use or modify this template.
 */
@Template(
    name = "Cloud_Spanner_vectors_to_Cloud_Storage",
    category = TemplateCategory.BATCH,
    displayName = "Cloud Spanner vectors to Cloud Storage for Vertex Vector Search",
    optionsClass = SpannerToVectorEmbeddingJsonOptions.class,
    description = {
      "The Cloud Spanner to Vector Embeddings on Cloud Storage template is a batch pipeline that exports vector embeddings data from Cloud Spanner's table to Cloud Storage in JSON format. "
          + "Vector embeddings are exported to a Cloud Storage folder specified by the user in the template parameters."
          + " The Cloud Storage folder will contain the list of exported `.json` files representing vector embeddings in a format supported by Vertex AI Vector Search Index.\n",
      "Check <a href=\"https://cloud.google.com/vertex-ai/docs/vector-search/setup/format-structure#json\">Vector Search Format Structure</a> for additional details."
    },
    documentation =
        "https://cloud.google.com/dataflow/docs/guides/templates/provided/cloud-spanner-to-vertex-vector-search",
    contactInformation = "https://cloud.google.com/support",
    requirements = {
      "The Cloud Spanner database must exist.",
      "The output Cloud Storage bucket must exist.",
      "In addition to the Identity and Access Management (IAM) roles necessary to run Dataflow jobs, you must also have the <a href=\"https://cloud.google.com/spanner/docs/export#iam\">appropriate IAM roles</a> for reading your Cloud Spanner data and writing to your Cloud Storage bucket."
    })
@SuppressWarnings("unused")
public class SpannerVectorEmbeddingExport {

  private static final Logger LOG = LoggerFactory.getLogger(SpannerVectorEmbeddingExport.class);

  /** Custom PipelineOptions. */
  public interface SpannerToVectorEmbeddingJsonOptions extends PipelineOptions {
    @TemplateParameter.ProjectId(
        order = 10,
        groupName = "Source",
        description = "Cloud Spanner Project Id",
        helpText = "The project ID of the Spanner instance.")
    ValueProvider<String> getSpannerProjectId();

    void setSpannerProjectId(ValueProvider<String> value);

    @TemplateParameter.Text(
        order = 20,
        groupName = "Source",
        regexes = {"[a-z][a-z0-9\\-]*[a-z0-9]"},
        description = "Cloud Spanner instance ID",
        helpText = "The ID of the Spanner instance to export the vector embeddings from.")
    ValueProvider<String> getSpannerInstanceId();

    void setSpannerInstanceId(ValueProvider<String> spannerInstanceId);

    @TemplateParameter.Text(
        order = 30,
        groupName = "Source",
        regexes = {"[a-z][a-z0-9_\\-]*[a-z0-9]"},
        description = "Cloud Spanner database ID",
        helpText = "The ID of the Spanner database to export the vector embeddings from.")
    ValueProvider<String> getSpannerDatabaseId();

    void setSpannerDatabaseId(ValueProvider<String> spannerDatabaseId);

    @TemplateParameter.Text(
        order = 40,
        groupName = "Source",
        regexes = {"^.+$"},
        description = "Spanner Table",
        helpText = "The Spanner table to read from.")
    ValueProvider<String> getSpannerTable();

    void setSpannerTable(ValueProvider<String> table);

    @TemplateParameter.Text(
        order = 50,
        groupName = "Source",
        description = "Columns to Export from Spanner Table",
        helpText =
            "A comma-separated list of required columns for the Vertex AI Vector Search index. The ID and embedding columns are required by Vector Search. If your column names don't match the Vertex AI Vector Search index input structure, create column mappings by using aliases. If the column names don't match the format expected by Vertex AI, use the notation from:to. For example, if you have columns named id and my_embedding, specify id, my_embedding:embedding.")
    ValueProvider<String> getSpannerColumnsToExport();

    void setSpannerColumnsToExport(ValueProvider<String> value);

    @TemplateParameter.GcsWriteFolder(
        order = 60,
        groupName = "Target",
        description = "Output files folder in Cloud Storage",
        helpText =
            "The Cloud Storage folder to write output files to. The path must end with a slash.",
        example = "gs://your-bucket/folder1/")
    ValueProvider<String> getGcsOutputFolder();

    void setGcsOutputFolder(ValueProvider<String> value);

    @TemplateParameter.Text(
        order = 70,
        groupName = "Target",
        description = "Output files prefix in Cloud Storage",
        helpText = "The filename prefix for writing output files.",
        example = "vector-embeddings")
    ValueProvider<String> getGcsOutputFilePrefix();

    void setGcsOutputFilePrefix(ValueProvider<String> textWritePrefix);

    @TemplateParameter.Text(
        order = 80,
        groupName = "Source",
        optional = true,
        description = "Cloud Spanner Endpoint to call",
        helpText =
            "The Spanner endpoint to call in the template. The default value is https://batch-spanner.googleapis.com.",
        example = "https://batch-spanner.googleapis.com")
    @Default.String("https://batch-spanner.googleapis.com")
    ValueProvider<String> getSpannerHost();

    void setSpannerHost(ValueProvider<String> value);

    @TemplateParameter.Text(
        order = 90,
        groupName = "Source",
        optional = true,
        regexes = {
          "^([0-9]{4})-([0-9]{2})-([0-9]{2})T([0-9]{2}):([0-9]{2}):(([0-9]{2})(\\.[0-9]+)?)Z$"
        },
        description = "Timestamp to read stale data from a version in the past.",
        helpText =
            "If set, specifies the time when the database version must be taken. The value is a string in the RFC-3339 date format in Unix epoch time. For example: `1990-12-31T23:59:60Z`. The timestamp must be in the past, and maximum timestamp staleness (https://cloud.google.com/spanner/docs/timestamp-bounds#maximum_timestamp_staleness) applies. If not set, a strong bound (https://cloud.google.com/spanner/docs/timestamp-bounds#strong) is used to read the latest data. Defaults to `empty`.",
        example = "1990-12-31T23:59:60Z")
    @Default.String(value = "")
    ValueProvider<String> getSpannerVersionTime();

    void setSpannerVersionTime(ValueProvider<String> value);

    @TemplateParameter.Boolean(
        order = 100,
        groupName = "Source",
        optional = true,
        description = "Use independent compute resource (Spanner DataBoost).",
        helpText =
            "When set to `true`, the template uses Spanner on-demand compute. The export job runs on independent compute resources that don't impact current Spanner workloads. Using this option incurs additional charges in Spanner. For more information, see Spanner Data Boost overview (https://cloud.google.com/spanner/docs/databoost/databoost-overview). Defaults to: `false`.")
    @Default.Boolean(false)
    ValueProvider<Boolean> getSpannerDataBoostEnabled();

    void setSpannerDataBoostEnabled(ValueProvider<Boolean> value);

    @TemplateParameter.Enum(
        order = 110,
        groupName = "Source",
        enumOptions = {
          @TemplateEnumOption("LOW"),
          @TemplateEnumOption("MEDIUM"),
          @TemplateEnumOption("HIGH")
        },
        optional = true,
        description = "Priority for Spanner RPC invocations",
        helpText =
            "The request priority for Spanner calls. The allowed values are `HIGH`, `MEDIUM`, and `LOW`. The default value is `MEDIUM`.")
    ValueProvider<RpcPriority> getSpannerPriority();

    void setSpannerPriority(ValueProvider<RpcPriority> value);
  }

  /**
   * Runs a pipeline which reads in vector embeddings records from Spanner, and writes the JSON to
   * TextIO sink.
   *
   * @param args arguments to the pipeline
   */
  public static void main(String[] args) {
    LOG.info("Starting pipeline setup");
    PipelineOptionsFactory.register(SpannerToVectorEmbeddingJsonOptions.class);

    SpannerToVectorEmbeddingJsonOptions options =
        PipelineOptionsFactory.fromArgs(args)
            .withValidation()
            .as(SpannerToVectorEmbeddingJsonOptions.class);

    FileSystems.setDefaultPipelineOptions(options);
    Pipeline pipeline = Pipeline.create(options);

    SpannerConfig spannerConfig =
        SpannerConfig.create()
            .withHost(options.getSpannerHost())
            .withProjectId(options.getSpannerProjectId())
            .withInstanceId(options.getSpannerInstanceId())
            .withDatabaseId(options.getSpannerDatabaseId())
            .withRpcPriority(options.getSpannerPriority())
            .withDataBoostEnabled(options.getSpannerDataBoostEnabled());

    ValueProvider<String> gcsOutputFilePrefix = options.getGcsOutputFilePrefix();

    // Concatenating cloud storage folder with file prefix to get complete path
    ValueProvider<String> gcsOutputFilePathWithPrefix =
        ValueProvider.NestedValueProvider.of(
            options.getGcsOutputFolder(),
            (SerializableFunction<String, String>)
                folder -> {
                  if (!folder.endsWith("/")) {
                    // Appending the slash if not provided by user
                    folder = folder + "/";
                  }
                  return folder + gcsOutputFilePrefix.get();
                });

    PTransform<PBegin, PCollection<ReadOperation>> spannerExport =
        SpannerConverters.ExportTransformFactory.create(
            options.getSpannerTable(),
            spannerConfig,
            gcsOutputFilePathWithPrefix,
            options.getSpannerVersionTime(),
            options.getSpannerColumnsToExport(),
            ValueProvider.StaticValueProvider.of(/* disable_schema_export= */ false));

    /* CreateTransaction and CreateTransactionFn classes in SpannerIO
     * only take a timestamp object for exact staleness which works when
     * parameters are provided during template compile time. They do not work with
     * a Timestamp valueProvider which can take parameters at runtime. Hence a new
     * ParDo class CreateTransactionFnWithTimestamp had to be created for this
     * purpose.
     */
    PCollectionView<Transaction> tx =
        pipeline
            .apply("Setup for Transaction", Create.of(1))
            .apply(
                "Create transaction",
                ParDo.of(
                    new CreateTransactionFnWithTimestamp(
                        spannerConfig, options.getSpannerVersionTime())))
            .apply("As PCollectionView", View.asSingleton());

    PCollection<String> json =
        pipeline
            .apply("Create export", spannerExport)
            // We need to use SpannerIO.readAll() instead of SpannerIO.read()
            // because ValueProvider parameters such as table name required for
            // SpannerIO.read() can be read only inside DoFn but SpannerIO.read() is of
            // type PTransform<PBegin, Struct>, which prevents prepending it with DoFn that reads
            // these parameters at the pipeline execution time.
            .apply(
                "Read all records",
                SpannerIO.readAll().withTransaction(tx).withSpannerConfig(spannerConfig))
            .apply(
                "Struct To JSON",
                MapElements.into(TypeDescriptors.strings())
                    .via(
                        struct ->
                            (new SpannerConverters.StructJSONPrinter(
                                    new VectorSearchStructValidator()))
                                .print(struct)));

    json.apply(
        "Write to storage", TextIO.write().to(gcsOutputFilePathWithPrefix).withSuffix(".json"));

    pipeline.run();
    LOG.info("Completed pipeline setup");
  }
}

O que se segue?

• Saiba mais sobre os modelos do Dataflow.
• Consulte a lista de modelos fornecidos pela Google.