Merge #149940

149940: vecbench: download dataset as separate files r=drewkimball a=andy-kimball

Update DatasetLoader (used by vecbench and the vecann workload) to download the train, test, and neighbors vector sets as separate files rather than as a single file. Separating the files will allow huge train datasets to be separated into multiple files. For the new files, use the .fbin and .ibin formats:

.fbin
[num_vectors (uint32), vector_dim (uint32), vector_array (float32)]

.ibin
[num_vectors (uint32), num_neighbors_per_vector (uint32), neighbor_array (int32)]

Epic: CRDB-42943
Release note: None

Co-authored-by: Andrew Kimball <[email protected]>

Author: craig[bot]

pkg/cmd/vecbench/main.go

@@ -94,6 +94,8 @@ var flagDBConnStr = flag.String("db", "postgresql://root@localhost:26257",
 //	dbpedia-openai-100k-angular (100K vectors, 1536 dims)
 //	dbpedia-openai-1000k-angular (1M vectors, 1536 dims)
 //	laion-1m-test-ip (1M vectors, 768 dims)
+//	coco-t2i-512-angular (113K vectors, 512 dims)
+//	coco-i2i-512-angular (113K vectors, 512 dims)
 //
 // After download, the datasets are cached in a local temp directory and a
 // vector index is created. The built vector index is also cached in the temp
@@ -442,6 +444,9 @@ func (vb *vectorBench) ensureDataset(ctx context.Context, forSearch bool) {
 			fmt.Printf(Cyan+"\rDownloaded %s / %s (%.2f%%) in %v          "+Reset,
 				humanizeutil.IBytes(downloaded), humanizeutil.IBytes(total),
 				(float64(downloaded)/float64(total))*100, elapsed.Truncate(time.Second))
+			if downloaded >= total {
+				fmt.Println()
+			}
 		},
 	}
 

pkg/workload/vecann/BUILD.bazel

@@ -18,7 +18,6 @@ go_library(
         "//pkg/workload/histogram",
         "//pkg/workload/workloadimpl",
         "@com_github_cockroachdb_errors//:errors",
-        "@com_github_cockroachdb_errors//oserror",
         "@com_github_spf13_pflag//:pflag",
         "@com_google_cloud_go_storage//:storage",
     ],