From 89861aefe23863946ad382b50a3549134c50e471 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Mon, 12 Oct 2026 13:42:41 +0530 Subject: [PATCH] feat: Migrate from TurboPuffer --- .github/workflows/pr-workflow.yaml | 2 + README.md | 46 ++- cmd/migrate_from_turbopuffer.go | 379 ++++++++++++++++++ go.mod | 2 + go.sum | 5 + .../migrate_from_turbopuffer_test.go | 202 ++++++++++ pkg/commons/config.go | 6 + 7 files changed, 639 insertions(+), 3 deletions(-) create mode 100644 cmd/migrate_from_turbopuffer.go create mode 100644 integration_tests/migrate_from_turbopuffer_test.go diff --git a/.github/workflows/pr-workflow.yaml b/.github/workflows/pr-workflow.yaml index 64ba8f0..4694f45 100644 --- a/.github/workflows/pr-workflow.yaml +++ b/.github/workflows/pr-workflow.yaml @@ -101,6 +101,8 @@ jobs: AWS_DYNAMODB_VECTOR_INDEX: ${{ secrets.AWS_DYNAMODB_VECTOR_INDEX }} AZURE_SEARCH_ENDPOINT: ${{ secrets.AZURE_SEARCH_ENDPOINT }} AZURE_SEARCH_API_KEY: ${{ secrets.AZURE_SEARCH_API_KEY }} + TURBOPUFFER_API_KEY: ${{ secrets.TURBOPUFFER_API_KEY }} + TURBOPUFFER_REGION: ${{ secrets.TURBOPUFFER_REGION }} run: | make test_unit diff --git a/README.md b/README.md index 7f63e8a..0d28977 100644 --- a/README.md +++ b/README.md @@ -14,8 +14,9 @@ CLI tool for migrating data to [Qdrant](http://qdrant.tech) with support for res * Elasticsearch * Azure AI Search * Postgres (pgvector) -* S3 Vectors +* Amazon S3 Vectors * Amazon DynamoDB +* Turbopuffer * FAISS * Apache Solr * Another Qdrant instance @@ -515,9 +516,9 @@ docker run --net=host --rm -it registry.cloud.qdrant.io/library/qdrant-migration
-

From S3 Vectors

 +

From Amazon S3 Vectors

 -Migrate data from an **S3 Vectors** index to **Qdrant**: +Migrate data from an **Amazon S3 Vectors** index to **Qdrant**: ### 📥 Example @@ -611,6 +612,45 @@ docker run --net=host --rm -it \
+
+

From Turbopuffer

 + +Migrate vectors and attributes from a **Turbopuffer** namespace to **Qdrant**: + +### 📥 Example + +```bash +docker run --net=host --rm -it registry.cloud.qdrant.io/library/qdrant-migration turbopuffer \ + --turbopuffer.namespace 'your-namespace' \ + --turbopuffer.api-key 'your-turbopuffer-api-key' \ + --turbopuffer.region 'gcp-us-central1' \ + --qdrant.url 'http://target-hostname:6334' \ + --qdrant.api-key 'optional-qdrant-api-key' \ + --qdrant.collection 'target-collection' \ + --migration.batch-size 64 +``` + +#### Turbopuffer Options + +| Flag | Description | +| ------------------------- | -------------------------------------------------------------- | +| `--turbopuffer.namespace` | Turbopuffer namespace to migrate (required) | +| `--turbopuffer.api-key` | Turbopuffer API key (required) | +| `--turbopuffer.region` | Turbopuffer region, e.g. `gcp-us-central1` (required) | + +#### Qdrant Options + +| Flag | Description | +| --------------------- | ------------------------------------------------------------------ | +| `--qdrant.url` | Qdrant gRPC URL. Default: `"http://localhost:6334"` | +| `--qdrant.collection` | Target collection name | +| `--qdrant.api-key` | Qdrant API key (optional) | +| `--qdrant.id-field` | Payload field storing the Turbopuffer document ID. Default: `"__id__"` | + +* See [Shared Migration Options](#shared-migration-options) for common migration parameters. + +
+

From FAISS

 diff --git a/cmd/migrate_from_turbopuffer.go b/cmd/migrate_from_turbopuffer.go new file mode 100644 index 0000000..2a5bae1 --- /dev/null +++ b/cmd/migrate_from_turbopuffer.go @@ -0,0 +1,379 @@ +//go:build !no_turbopuffer + +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "os/signal" + "regexp" + "strconv" + "strings" + "syscall" + "time" + + "github.com/pterm/pterm" + "github.com/turbopuffer/turbopuffer-go" + "github.com/turbopuffer/turbopuffer-go/option" + + "github.com/qdrant/go-client/qdrant" + + "github.com/qdrant/migration/pkg/commons" +) + +var turbopufferVectorType = regexp.MustCompile(`^\[(\d+)\]f(?:16|32)$`) + +type turbopufferSource struct { + vectors map[string]*qdrant.VectorParams + attributeTypes map[string]string +} + +func init() { + registerCommand[MigrateFromTurbopufferCmd]("turbopuffer", "Migrate data from a turbopuffer namespace to Qdrant.") +} + +type MigrateFromTurbopufferCmd struct { + Turbopuffer commons.TurbopufferConfig `embed:"" prefix:"turbopuffer."` + Qdrant commons.QdrantConfig `embed:"" prefix:"qdrant."` + Migration commons.MigrationConfig `embed:"" prefix:"migration."` + IdField string `prefix:"qdrant." help:"Field storing turbopuffer IDs in Qdrant." default:"__id__"` + + targetHost string + targetPort int + targetTLS bool +} + +func (r *MigrateFromTurbopufferCmd) Parse() error { + var err error + r.targetHost, r.targetPort, r.targetTLS, err = parseQdrantUrl(r.Qdrant.Url) + if err != nil { + return fmt.Errorf("failed to parse target URL: %w", err) + } + return nil +} + +func (r *MigrateFromTurbopufferCmd) Validate() error { + return validateBatchSize(r.Migration.BatchSize) +} + +func (r *MigrateFromTurbopufferCmd) Run(globals *Globals) error { + pterm.DefaultHeader.WithFullWidth().Println("turbopuffer to Qdrant Data Migration") + + if err := r.Parse(); err != nil { + return fmt.Errorf("failed to parse input: %w", err) + } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + sourceClient := turbopuffer.NewClient(option.WithAPIKey(r.Turbopuffer.APIKey), option.WithRegion(r.Turbopuffer.Region)) + namespace := sourceClient.Namespace(r.Turbopuffer.Namespace) + + targetClient, err := connectToQdrant(globals, r.targetHost, r.targetPort, r.Qdrant.APIKey, r.targetTLS, 0) + if err != nil { + return fmt.Errorf("failed to connect to Qdrant target: %w", err) + } + defer targetClient.Close() + + if err := commons.PrepareOffsetsCollection(ctx, r.Migration.OffsetsCollection, targetClient); err != nil { + return fmt.Errorf("failed to prepare migration marker collection: %w", err) + } + + source, err := r.describeNamespace(ctx, namespace) + if err != nil { + return err + } + sourcePointCount, err := r.countTurbopufferRows(ctx, namespace) + if err != nil { + return fmt.Errorf("failed to count points in source: %w", err) + } + if err := r.prepareTargetCollection(ctx, source, targetClient); err != nil { + return fmt.Errorf("error preparing target collection: %w", err) + } + + displayMigrationStart("turbopuffer", r.Turbopuffer.Namespace, r.Qdrant.Collection) + + if err := r.migrateData(ctx, namespace, targetClient, source, sourcePointCount); err != nil { + return fmt.Errorf("failed to migrate data: %w", err) + } + + targetPointCount, err := targetClient.Count(ctx, &qdrant.CountPoints{ + CollectionName: r.Qdrant.Collection, + Exact: qdrant.PtrOf(true), + }) + if err != nil { + return fmt.Errorf("failed to count points in target: %w", err) + } + pterm.Info.Printfln("Target collection has %d points\n", targetPointCount) + + if err := commons.DeleteOffsetsCollection(ctx, r.Migration.OffsetsCollection, targetClient); err != nil { + return fmt.Errorf("failed to delete migration marker collection: %w", err) + } + return nil +} + +func (r *MigrateFromTurbopufferCmd) describeNamespace(ctx context.Context, namespace turbopuffer.Namespace) (*turbopufferSource, error) { + metadata, err := namespace.Metadata(ctx, turbopuffer.NamespaceMetadataParams{}) + if err != nil { + return nil, fmt.Errorf("failed to get turbopuffer namespace metadata: %w", err) + } + + source := &turbopufferSource{ + vectors: make(map[string]*qdrant.VectorParams), + attributeTypes: make(map[string]string, len(metadata.Schema)), + } + for name, attribute := range metadata.Schema { + source.attributeTypes[name] = attribute.Type + match := turbopufferVectorType.FindStringSubmatch(attribute.Type) + if match == nil { + continue + } + dimensions, err := strconv.ParseUint(match[1], 10, 64) + if err != nil { + return nil, fmt.Errorf("invalid dimensions for vector attribute %q: %w", name, err) + } + + var distance qdrant.Distance + switch attribute.Ann.DistanceMetric { + case turbopuffer.DistanceMetricCosineDistance: + distance = qdrant.Distance_Cosine + case turbopuffer.DistanceMetricEuclideanSquared: + distance = qdrant.Distance_Euclid + case "": + pterm.Warning.Printfln("Vector attribute %q has no distance metric, defaulting to cosine", name) + distance = qdrant.Distance_Cosine + default: + return nil, fmt.Errorf("unsupported turbopuffer distance metric: %s", attribute.Ann.DistanceMetric) + } + + source.vectors[name] = &qdrant.VectorParams{Size: dimensions, Distance: distance} + } + if len(source.vectors) == 0 { + return nil, fmt.Errorf("turbopuffer namespace %q has no vector attributes", r.Turbopuffer.Namespace) + } + return source, nil +} + +func (r *MigrateFromTurbopufferCmd) countTurbopufferRows(ctx context.Context, namespace turbopuffer.Namespace) (uint64, error) { + var response struct { + Aggregations struct { + Count uint64 `json:"count"` + } `json:"aggregations"` + } + _, err := namespace.Query(ctx, turbopuffer.NamespaceQueryParams{ + AggregateBy: map[string]turbopuffer.AggregateBy{"count": turbopuffer.NewAggregateByCount()}, + }, option.WithResponseBodyInto(&response)) + if err != nil { + return 0, err + } + return response.Aggregations.Count, nil +} + +func (r *MigrateFromTurbopufferCmd) prepareTargetCollection(ctx context.Context, source *turbopufferSource, targetClient *qdrant.Client) error { + if !r.Migration.CreateCollection { + return nil + } + exists, err := targetClient.CollectionExists(ctx, r.Qdrant.Collection) + if err != nil { + return fmt.Errorf("failed to check if collection exists: %w", err) + } + if exists { + pterm.Info.Printfln("Target collection %q already exists. Skipping creation.", r.Qdrant.Collection) + return nil + } + + if err := targetClient.CreateCollection(ctx, &qdrant.CreateCollection{ + CollectionName: r.Qdrant.Collection, + VectorsConfig: qdrant.NewVectorsConfigMap(source.vectors), + }); err != nil { + return fmt.Errorf("failed to create target collection: %w", err) + } + pterm.Success.Printfln("Created target collection %q", r.Qdrant.Collection) + return nil +} + +func (r *MigrateFromTurbopufferCmd) migrateData(ctx context.Context, namespace turbopuffer.Namespace, targetClient *qdrant.Client, source *turbopufferSource, sourcePointCount uint64) error { + offsetKey := "turbopuffer/" + r.Turbopuffer.Namespace + var lastID json.RawMessage + offsetCount := uint64(0) + + if !r.Migration.Restart { + offsetID, count, err := commons.GetStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, offsetKey) + if err != nil { + return fmt.Errorf("failed to get start offset: %w", err) + } + offsetCount = count + if offsetID != nil { + lastID = json.RawMessage(offsetID.GetUuid()) + } + } + + bar, _ := pterm.DefaultProgressbar.WithTotal(int(sourcePointCount)).Start() + defer func() { _, _ = bar.Stop() }() + displayMigrationProgress(bar, offsetCount) + + for { + params := turbopuffer.NamespaceQueryParams{ + RankBy: turbopuffer.NewRankByAttribute("id", turbopuffer.RankByAttributeOrderAsc), + TopK: turbopuffer.Int(int64(r.Migration.BatchSize)), + IncludeAttributes: turbopuffer.IncludeAttributesParam{Bool: turbopuffer.Bool(true)}, + VectorEncoding: turbopuffer.VectorEncodingFloat, + } + if lastID != nil { + params.Filters = turbopuffer.NewFilterGt("id", lastID) + } + + // The raw body keeps full integer precision. + var body []byte + if _, err := namespace.Query(ctx, params, option.WithResponseBodyInto(&body)); err != nil { + return fmt.Errorf("failed to query turbopuffer: %w", err) + } + var response struct { + Rows []map[string]any `json:"rows"` + } + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + if err := decoder.Decode(&response); err != nil { + return fmt.Errorf("failed to decode turbopuffer response: %w", err) + } + if len(response.Rows) == 0 { + break + } + + points := make([]*qdrant.PointStruct, 0, len(response.Rows)) + for _, row := range response.Rows { + point, err := r.turbopufferRowToPoint(row, source) + if err != nil { + return err + } + points = append(points, point) + } + if err := upsertWithRetry(ctx, targetClient, &qdrant.UpsertPoints{ + CollectionName: r.Qdrant.Collection, + Points: points, + Wait: qdrant.PtrOf(true), + }); err != nil { + return err + } + offsetCount += uint64(len(points)) + bar.Add(len(points)) + + checkpoint, err := json.Marshal(response.Rows[len(response.Rows)-1]["id"]) + if err != nil { + return fmt.Errorf("failed to encode checkpoint: %w", err) + } + if err := commons.StoreStartOffset(ctx, r.Migration.OffsetsCollection, targetClient, offsetKey, qdrant.NewID(string(checkpoint)), offsetCount); err != nil { + return fmt.Errorf("failed to store checkpoint: %w", err) + } + lastID = checkpoint + + if len(response.Rows) < r.Migration.BatchSize { + break + } + if r.Migration.BatchDelay > 0 { + time.Sleep(time.Duration(r.Migration.BatchDelay) * time.Millisecond) + } + } + + pterm.Success.Printfln("Data migration finished successfully") + return nil +} + +func (r *MigrateFromTurbopufferCmd) turbopufferRowToPoint(row map[string]any, source *turbopufferSource) (*qdrant.PointStruct, error) { + pointID, err := turbopufferPointID(row["id"]) + if err != nil { + return nil, err + } + + vectors := make(map[string]*qdrant.Vector) + payload := make(map[string]any, len(row)) + for name, value := range row { + if name == "id" { + continue + } + if _, ok := source.vectors[name]; ok { + if value == nil { + continue + } + vector, err := turbopufferVector(value) + if err != nil { + return nil, fmt.Errorf("invalid vector attribute %q: %w", name, err) + } + vectors[name] = qdrant.NewVectorDense(vector) + continue + } + payload[name] = convertTurbopufferValue(value, source.attributeTypes[name]) + } + payload[r.IdField] = convertTurbopufferValue(row["id"], source.attributeTypes["id"]) + + qdrantPayload, err := qdrant.TryValueMap(payload) + if err != nil { + return nil, fmt.Errorf("failed to convert turbopuffer attributes: %w", err) + } + return &qdrant.PointStruct{ + Id: pointID, + Payload: qdrantPayload, + Vectors: qdrant.NewVectorsMap(vectors), + }, nil +} + +func turbopufferPointID(id any) (*qdrant.PointId, error) { + switch v := id.(type) { + case json.Number: + num, err := strconv.ParseUint(v.String(), 10, 64) + if err != nil { + return nil, fmt.Errorf("unsupported turbopuffer ID %q: %w", v, err) + } + return qdrant.NewIDNum(num), nil + case string: + return arbitraryIDToUUID(v), nil + default: + return nil, fmt.Errorf("unsupported turbopuffer ID type %T", id) + } +} + +func turbopufferVector(value any) ([]float32, error) { + elements, ok := value.([]any) + if !ok { + return nil, fmt.Errorf("expected a list of numbers, got %T", value) + } + vector := make([]float32, len(elements)) + for i, element := range elements { + number, ok := element.(json.Number) + if !ok { + return nil, fmt.Errorf("element %d is %T, expected a number", i, element) + } + parsed, err := strconv.ParseFloat(number.String(), 32) + if err != nil { + return nil, fmt.Errorf("element %d has invalid number %q: %w", i, number, err) + } + vector[i] = float32(parsed) + } + return vector, nil +} + +// Uses the schema type so numbers keep one type across rows. +func convertTurbopufferValue(value any, attributeType string) any { + switch v := value.(type) { + case json.Number: + if attributeType == "float" { + f, _ := v.Float64() + return f + } + if i, err := v.Int64(); err == nil { + return i + } + // Qdrant payloads have no unsigned type, so integers beyond int64 become strings. + return v.String() + case []any: + elementType := strings.TrimPrefix(attributeType, "[]") + for i, element := range v { + v[i] = convertTurbopufferValue(element, elementType) + } + return v + default: + return v + } +} diff --git a/go.mod b/go.mod index e0d5d5d..89cf57c 100644 --- a/go.mod +++ b/go.mod @@ -27,6 +27,7 @@ require ( github.com/stretchr/testify v1.11.1 github.com/testcontainers/testcontainers-go v0.44.0 github.com/testcontainers/testcontainers-go/modules/opensearch v0.43.0 + github.com/turbopuffer/turbopuffer-go v1.21.0 github.com/weaviate/weaviate v1.38.1 github.com/weaviate/weaviate-go-client/v5 v5.7.1 go.mongodb.org/mongo-driver v1.17.8 @@ -189,6 +190,7 @@ require ( github.com/tidwall/gjson v1.18.0 // indirect github.com/tidwall/match v1.2.0 // indirect github.com/tidwall/pretty v1.2.1 // indirect + github.com/tidwall/sjson v1.2.5 // indirect github.com/tklauser/go-sysconf v0.4.0 // indirect github.com/tklauser/numcpus v0.12.0 // indirect github.com/tmc/grpc-websocket-proxy v0.0.0-20220101234140-673ab2c3ae75 // indirect diff --git a/go.sum b/go.sum index 0fdc563..15dcafc 100644 --- a/go.sum +++ b/go.sum @@ -539,6 +539,7 @@ github.com/testcontainers/testcontainers-go/modules/chroma v0.36.0 h1:aP1Xifh3Ig github.com/testcontainers/testcontainers-go/modules/chroma v0.36.0/go.mod h1:4VyK3KXTZ6ATn08mKfW6BdCTknMzj9wTd6ANFCkZ1N4= github.com/testcontainers/testcontainers-go/modules/opensearch v0.43.0 h1:a1ipjF7d/VxPX1dgVPIk4F+t6YkgMbE2OtBuRQCHJt8= github.com/testcontainers/testcontainers-go/modules/opensearch v0.43.0/go.mod h1:OWSeUDiGMUy30iMsAltIJIo9uh/CleLv6KyxjYOsgR8= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY= github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= @@ -547,12 +548,16 @@ github.com/tidwall/match v1.2.0/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JT github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4= github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= github.com/tklauser/go-sysconf v0.4.0 h1:7H0uAN+7RkwWRaxhYXDLqa5V3LPrJeV8wmD9dRUgPQU= github.com/tklauser/go-sysconf v0.4.0/go.mod h1:8mTNWyog7H+MpKijp4VmKJAd2bbYQ2zuUwkYRbUArPI= github.com/tklauser/numcpus v0.12.0 h1:NR85qdvHA9pFse3x3weVZ0r0ST8R6l5RHbZrlRaqob4= github.com/tklauser/numcpus v0.12.0/go.mod h1:ABHeXzJnr/qqwguhClkZKT1/8VABcYrsyUiUGobwWJg= github.com/tmc/grpc-websocket-proxy v0.0.0-20220101234140-673ab2c3ae75 h1:6fotK7otjonDflCTK0BCfls4SPy3NcCVb5dqqmbRknE= github.com/tmc/grpc-websocket-proxy v0.0.0-20220101234140-673ab2c3ae75/go.mod h1:KO6IkyS8Y3j8OdNO85qEYBsRPuteD+YciPomcXdrMnk= +github.com/turbopuffer/turbopuffer-go v1.21.0 h1:bEgfVdjKwR/+k0wCaQkXirGqNaPO7YEiUZh1kPZtP8k= +github.com/turbopuffer/turbopuffer-go v1.21.0/go.mod h1:1UhmfUzh4rFtRoONvF1H/X3LinF8t7YyxpGwYFGE6Ok= github.com/twmb/murmur3 v1.1.8 h1:8Yt9taO/WN3l08xErzjeschgZU2QSrwm1kclYq+0aRg= github.com/twmb/murmur3 v1.1.8/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ= github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= diff --git a/integration_tests/migrate_from_turbopuffer_test.go b/integration_tests/migrate_from_turbopuffer_test.go new file mode 100644 index 0000000..fcf6f68 --- /dev/null +++ b/integration_tests/migrate_from_turbopuffer_test.go @@ -0,0 +1,202 @@ +package integrationtests + +import ( + "context" + "encoding/json" + "fmt" + "math" + "math/rand" + "net/http" + "os" + "strconv" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "github.com/turbopuffer/turbopuffer-go" + "github.com/turbopuffer/turbopuffer-go/option" + + "github.com/qdrant/go-client/qdrant" + + "github.com/qdrant/migration/pkg/commons" +) + +const ( + turbopufferVectorName = "vector" + turbopufferOffsetCollection = "_migration_offsets" + turbopufferCheckpoint = 59 +) + +type turbopufferTestCase struct { + name string + idType string + ids []any // In ascending order, as turbopuffer pages through them. + vectorType string + metric turbopuffer.DistanceMetric + tolerance float64 + resume bool +} + +func TestMigrateFromTurbopuffer(t *testing.T) { + ctx := context.Background() + apiKey := os.Getenv("TURBOPUFFER_API_KEY") + if apiKey == "" { + t.Skip("Skipping turbopuffer source test. TURBOPUFFER_API_KEY not set") + } + region := os.Getenv("TURBOPUFFER_REGION") + if region == "" { + t.Skip("Skipping turbopuffer source test. TURBOPUFFER_REGION not set") + } + client := turbopuffer.NewClient(option.WithAPIKey(apiKey), option.WithRegion(region)) + + qdrantContainer := qdrantContainer(ctx, t, qdrantAPIKey) + defer func() { + require.NoError(t, qdrantContainer.Terminate(ctx)) + }() + qdrantHost, err := qdrantContainer.Host(ctx) + require.NoError(t, err) + qdrantPort, err := qdrantContainer.MappedPort(ctx, qdrantGRPCPort) + require.NoError(t, err) + qdrantClient, err := qdrant.NewClient(&qdrant.Config{ + Host: qdrantHost, + Port: int(qdrantPort.Num()), + APIKey: qdrantAPIKey, + SkipCompatibilityCheck: true, + }) + require.NoError(t, err) + defer qdrantClient.Close() + + // Spaced so that about half of the IDs exceed math.MaxInt64. + uintIDs := make([]any, totalEntries) + stringIDs := make([]any, totalEntries) + for i := range totalEntries { + uintIDs[i] = uint64(i) * (math.MaxUint64 / totalEntries) + stringIDs[i] = fmt.Sprintf("doc-%03d", i) + } + + cases := []turbopufferTestCase{ + {"uint-f32-cosine", "uint", uintIDs, fmt.Sprintf("[%d]f32", dimension), turbopuffer.DistanceMetricCosineDistance, 1e-6, true}, + {"string-f16-euclidean", "string", stringIDs, fmt.Sprintf("[%d]f16", dimension), turbopuffer.DistanceMetricEuclideanSquared, 1e-3, false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + namespaceName := fmt.Sprintf("migration-test-%s-%d", tc.name, rand.Int()) + collectionName := testCollectionName + "-" + namespaceName + vectors := writeTurbopufferNamespace(ctx, t, client.Namespace(namespaceName), tc) + + from := 0 + if tc.resume { + checkpoint, err := json.Marshal(tc.ids[turbopufferCheckpoint]) + require.NoError(t, err) + require.NoError(t, commons.PrepareOffsetsCollection(ctx, turbopufferOffsetCollection, qdrantClient)) + require.NoError(t, commons.StoreStartOffset(ctx, turbopufferOffsetCollection, qdrantClient, + "turbopuffer/"+namespaceName, qdrant.NewID(string(checkpoint)), turbopufferCheckpoint+1)) + from = turbopufferCheckpoint + 1 + } + + args := []string{ + "turbopuffer", + fmt.Sprintf("--turbopuffer.namespace=%s", namespaceName), + fmt.Sprintf("--turbopuffer.api-key=%s", apiKey), + fmt.Sprintf("--turbopuffer.region=%s", region), + fmt.Sprintf("--qdrant.url=http://%s:%s", qdrantHost, qdrantPort.Port()), + fmt.Sprintf("--qdrant.collection=%s", collectionName), + fmt.Sprintf("--qdrant.api-key=%s", qdrantAPIKey), + fmt.Sprintf("--qdrant.id-field=%s", idField), + fmt.Sprintf("--migration.offsets-collection=%s", turbopufferOffsetCollection), + "--migration.batch-size=7", + } + runMigrationBinary(t, args) + + points, err := qdrantClient.Scroll(ctx, &qdrant.ScrollPoints{ + CollectionName: collectionName, + Limit: qdrant.PtrOf(uint32(totalEntries + 1)), + WithPayload: qdrant.NewWithPayload(true), + WithVectors: qdrant.NewWithVectors(true), + }) + require.NoError(t, err) + require.Len(t, points, len(tc.ids)-from) + + byID := make(map[string]*qdrant.RetrievedPoint, len(points)) + for _, point := range points { + key, err := commons.EncodePointID(point.GetId()) + require.NoError(t, err) + byID[key] = point + } + for i := from; i < len(tc.ids); i++ { + key, err := commons.EncodePointID(turbopufferPointID(tc.ids[i])) + require.NoError(t, err) + point, ok := byID[key] + require.True(t, ok, "point for turbopuffer ID %v not found", tc.ids[i]) + + require.Equal(t, turbopufferIDPayload(tc.ids[i]), point.Payload[idField].GetKind()) + require.Equal(t, fmt.Sprintf("name-%d", i), point.Payload["name"].GetStringValue()) + require.IsType(t, &qdrant.Value_DoubleValue{}, point.Payload["score"].GetKind()) + require.InDelta(t, float64(i), point.Payload["score"].GetDoubleValue(), 0) + + expected := vectors[i] + if tc.metric == turbopuffer.DistanceMetricCosineDistance { + expected = normalizedVector(expected) + } + vector := point.Vectors.GetVectors().GetVectors()[turbopufferVectorName].GetDenseVector().GetData() + require.InDeltaSlice(t, expected, vector, tc.tolerance) + } + }) + } +} + +func writeTurbopufferNamespace(ctx context.Context, t *testing.T, namespace turbopuffer.Namespace, tc turbopufferTestCase) [][]float32 { + t.Helper() + + // Never write to or delete a namespace this test did not create. + _, err := namespace.Metadata(ctx, turbopuffer.NamespaceMetadataParams{}) + var apiErr *turbopuffer.Error + require.ErrorAs(t, err, &apiErr) + require.Equal(t, http.StatusNotFound, apiErr.StatusCode, "namespace %q already exists", namespace.ID()) + + vectors := make([][]float32, len(tc.ids)) + rows := make([]turbopuffer.RowParam, len(tc.ids)) + for i, id := range tc.ids { + vectors[i] = randFloat32Values(dimension) + rows[i] = turbopuffer.RowParam{ + "id": id, + turbopufferVectorName: vectors[i], + "name": fmt.Sprintf("name-%d", i), + "score": float64(i), + } + } + t.Cleanup(func() { + _, err := namespace.DeleteAll(context.WithoutCancel(ctx), turbopuffer.NamespaceDeleteAllParams{}) + require.NoError(t, err) + }) + _, err = namespace.Write(ctx, turbopuffer.NamespaceWriteParams{ + UpsertRows: rows, + DistanceMetric: tc.metric, + Schema: map[string]turbopuffer.AttributeSchemaConfigParam{ + "id": {Type: tc.idType}, + turbopufferVectorName: {Type: tc.vectorType, Ann: turbopuffer.AttributeSchemaConfigAnnParam{DistanceMetric: tc.metric}}, + "score": {Type: "float"}, + }, + }) + require.NoError(t, err) + return vectors +} + +func turbopufferPointID(id any) *qdrant.PointId { + if num, ok := id.(uint64); ok { + return qdrant.NewIDNum(num) + } + return qdrant.NewIDUUID(uuid.NewSHA1(uuid.NameSpaceURL, []byte(id.(string))).String()) +} + +// Integers beyond int64 are stored as strings. +func turbopufferIDPayload(id any) any { + num, ok := id.(uint64) + if !ok { + return &qdrant.Value_StringValue{StringValue: id.(string)} + } + if num <= math.MaxInt64 { + return &qdrant.Value_IntegerValue{IntegerValue: int64(num)} + } + return &qdrant.Value_StringValue{StringValue: strconv.FormatUint(num, 10)} +} diff --git a/pkg/commons/config.go b/pkg/commons/config.go index d3280e6..6d11023 100644 --- a/pkg/commons/config.go +++ b/pkg/commons/config.go @@ -121,6 +121,12 @@ type DynamoDBConfig struct { Index string `help:"DynamoDB vector index name" required:"true"` } +type TurbopufferConfig struct { + Namespace string `help:"turbopuffer namespace to migrate" required:"true"` + APIKey string `help:"turbopuffer API key" required:"true"` + Region string `help:"turbopuffer region (e.g. 'gcp-us-central1')" required:"true"` +} + type FaissConfig struct { IndexPath string `help:"Path to the FAISS index file" required:"true"` }