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"`
}