work around lastfm dataset containing f64 query vectors

jbellis · Oct 26, 2023 · 39094dc · 39094dc
1 parent 9e02dd5
commit 39094dc
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/Hdf5Loader.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/Hdf5Loader.java
@@ -19,10 +19,14 @@
 import io.github.jbellis.jvector.vector.VectorSimilarityFunction;
 import io.github.jbellis.jvector.vector.VectorUtil;
 import io.jhdf.HdfFile;
+import io.jhdf.api.Dataset;
+import io.jhdf.object.datatype.DataType;
+import io.jhdf.object.datatype.FloatingPoint;
 
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.*;
+import java.util.stream.IntStream;
 
 public class Hdf5Loader {
     public static final String HDF5_DIR = "hdf5/";
@@ -47,7 +51,20 @@ else if (filename.contains("-euclidean")) {
         Path path = Path.of(HDF5_DIR).resolve(filename);
         try (HdfFile hdf = new HdfFile(path)) {
             baseVectors = (float[][]) hdf.getDatasetByPath("train").getData();
-            queryVectors = (float[][]) hdf.getDatasetByPath("test").getData();
+            Dataset queryDataset = hdf.getDatasetByPath("test");
+            if (((FloatingPoint) queryDataset.getDataType()).getBitPrecision() == 64) {
+                // lastfm dataset contains f64 queries but f32 everything else
+                var doubles = ((double[][]) queryDataset.getData());
+                queryVectors = IntStream.range(0, doubles.length).parallel().mapToObj(i -> {
+                    var a = new float[doubles[i].length];
+                    for (int j = 0; j < doubles[i].length; j++) {
+                        a[j] = (float) doubles[i][j];
+                    }
+                    return a;
+                }).toArray(float[][]::new);
+            } else {
+                queryVectors = (float[][]) queryDataset.getData();
+            }
             groundTruth = (int[][]) hdf.getDatasetByPath("neighbors").getData();
         }