diff --git a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java index 3a3b9c8bf..80f6f4107 100644 --- a/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java +++ b/benchmarks-jmh/src/main/java/io/github/jbellis/jvector/bench/CompactorBenchmark.java @@ -21,6 +21,7 @@ import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSetInfo; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; +import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet; import io.github.jbellis.jvector.example.reporting.GitInfo; import io.github.jbellis.jvector.example.reporting.JfrRecorder; import io.github.jbellis.jvector.example.reporting.JsonlWriter; @@ -251,7 +252,7 @@ private static void writeCompletedCount(int count) { private List> queryVectors; private List> baseVectors; private List> groundTruth; - private DataSet ds; + private InMemoryDataSet ds; private VectorSimilarityFunction similarityFunction; private final List graphs = new ArrayList<>(); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java index 24c39ae47..326704e6e 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/AutoBenchYAML.java @@ -22,6 +22,7 @@ import io.github.jbellis.jvector.example.util.CheckpointManager; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; +import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet; import io.github.jbellis.jvector.example.yaml.DatasetCollection; import io.github.jbellis.jvector.example.yaml.MultiConfig; @@ -123,10 +124,10 @@ public static void main(String[] args) throws IOException { logger.info("Loading dataset: {}", datasetName); try { - DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( + InMemoryDataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( () -> new RuntimeException("Dataset " + datasetName + " not found") ).getDataSet(); - logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseVectors().size()); + logger.info("Dataset loaded: {} with {} vectors", datasetName, ds.getBaseRavv().size()); String normalizedDatasetName = datasetName; if (normalizedDatasetName.endsWith(".hdf5")) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java index 15543ebbc..33f06847d 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/CompactionBench.java @@ -19,6 +19,7 @@ import io.github.jbellis.jvector.disk.ReaderSupplier; import io.github.jbellis.jvector.disk.ReaderSupplierFactory; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet; import io.github.jbellis.jvector.example.util.AccuracyMetrics; import io.github.jbellis.jvector.example.util.CompactionPartitionSource; import io.github.jbellis.jvector.example.yaml.TestDataPartition.Distribution; @@ -97,7 +98,7 @@ private CompactionBench() {} * one result per config. A config that fails (e.g. missing partitions) is logged and skipped so * the remaining configs still run. Throws if the dataset has no query vectors or ground truth. */ - public static List run(DataSet ds) throws Exception { + public static List run(InMemoryDataSet ds) throws Exception { var queryVectors = ds.getQueryVectors(); var groundTruth = ds.getGroundTruth(); if (queryVectors == null || queryVectors.isEmpty()) { @@ -118,10 +119,10 @@ public static List run(DataSet ds) throws Exception { return results; } - private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exception { + private static BenchResult runConfig(InMemoryDataSet ds, PartitionConfig cfg) throws Exception { String datasetName = ds.getName(); logger.info("Compaction bench [{}] config {}: {} vectors", - datasetName, cfg.dirName(), ds.getBaseVectors().size()); + datasetName, cfg.dirName(), ds.getBaseRavv().size()); // 1. Fetch pre-built partitions from S3 (cached locally). List partitionPaths = CompactionPartitionSource.ensurePartitions( @@ -135,7 +136,7 @@ private static BenchResult runConfig(DataSet ds, PartitionConfig cfg) throws Exc } } - private static BenchResult compactAndMeasure(DataSet ds, PartitionConfig cfg, + private static BenchResult compactAndMeasure(InMemoryDataSet ds, PartitionConfig cfg, List partitionPaths, Path tempDir) throws Exception { List> baseVectors = ds.getBaseVectors(); int dimension = ds.getDimension(); @@ -253,7 +254,6 @@ private static SearchStats searchCompacted(Path indexPath, DataSet ds, try (var rs = ReaderSupplierFactory.open(indexPath)) { var graph = OnDiskGraphIndex.load(rs); try (var searcher = new GraphSearcher(graph)) { - searcher.usePruning(false); int n = queryVectors.size(); List results = new ArrayList<>(n); long[] latenciesNanos = new long[n]; diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java index 8f45df2a0..af5e54347 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/Grid.java @@ -58,6 +58,7 @@ import io.github.jbellis.jvector.util.ExplicitThreadLocal; import io.github.jbellis.jvector.util.PhysicalCoreExecutor; import io.github.jbellis.jvector.vector.types.VectorFloat; +import me.tongfei.progressbar.ProgressBar; import java.io.FileNotFoundException; import java.io.IOException; @@ -329,7 +330,7 @@ static void runOneGraph(OnDiskGraphIndexCache cache, "Compressor '%s' was provided but failed to encode vectors for dataset '%s'. " + "Aborting to prevent false recall results.", compressor, ds.getName())); } - System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.getBaseVectors().size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS); + System.out.format("%s: %s encoded %d vectors [%.2f MB] in %.2fs%n", ds.getName(), compressor, ds.getBaseRavv().size(), (cv.ramBytesUsed() / 1024f / 1024f), encodingTimeS); } } @@ -409,6 +410,7 @@ private static Map, ImmutableGraphIndex> buildOnDisk(List, ImmutableGraphIndex> buildOnDisk(List { - IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { - writers.forEach((features, writer) -> { - try { - var stateMap = new EnumMap(FeatureId.class); - suppliers.get(features).forEach((featureId, supplier) -> { - stateMap.put(featureId, supplier.apply(node)); - }); - writer.writeInline(node, stateMap); - } catch (IOException e) { - throw new UncheckedIOException(e); - } + try (ProgressBar pb = new ProgressBar("build onDiskGraph", floatVectors.size())) { + IntStream.range(0, floatVectors.size()).parallel().forEach(node -> { + writers.forEach((features, writer) -> { + try { + var stateMap = new EnumMap(FeatureId.class); + suppliers.get(features).forEach((featureId, supplier) -> { + stateMap.put(featureId, supplier.apply(node)); + }); + writer.writeFeaturesInline(node, stateMap); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }); + builder.addGraphNode(node, vv.get().getVector(node)); + pb.step(); }); - builder.addGraphNode(node, vv.get().getVector(node)); - }); + } }).join(); builder.cleanup(); @@ -513,7 +518,8 @@ private static BuilderWithSuppliers builderWithSuppliers(Set features builder.with(new NVQ(nvq)); suppliers.put(FeatureId.NVQ_VECTORS, ordinal -> new NVQ.State(nvq.encode(floatVectors.getVector(ordinal)))); break; - + default: + break; } } return new BuilderWithSuppliers(builder, suppliers); @@ -605,6 +611,7 @@ private static Map, ImmutableGraphIndex> buildInMemory(List runAllAndCollectResults( searchCompressorObj, ds.getName())); } System.out.format("%s: %s encoded %d vectors [%.2f MB] for search%n", - ds.getName(), searchCompressorObj, ds.getBaseVectors().size(), + ds.getName(), searchCompressorObj, ds.getBaseRavv().size(), (cvArg.ramBytesUsed() / 1024f / 1024f)); } } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java index a33a40d31..606a67a68 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSet.java @@ -52,12 +52,6 @@ public interface DataSet { */ VectorSimilarityFunction getSimilarityFunction(); - /** - * The base vectors as a list. - * @return a list of base vectors - */ - List> getBaseVectors(); - /** * The query vectors as a list. * Each major index corresponds to the self-same index from {@link #getGroundTruth()}. @@ -69,7 +63,7 @@ public interface DataSet { /** * The ground truth as a list. * Each major index corresponds to the self-same index from {@link #getQueryVectors()}. - * Each minor index within represents the corresponding ordinal from {@link #getBaseVectors()} and {@link #getBaseRavv()}. + * Each minor index within represents the corresponding ordinal from {@link #getBaseRavv()}. * @return a list of query vectors. */ List> getGroundTruth(); diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java index 7fd38e3ca..a18d63491 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetInfo.java @@ -52,7 +52,7 @@ public class DataSetInfo implements DataSetProperties { private final DataSetFiles dsFiles; private final DataSetProperties baseProperties; - private volatile DataSet cached; + private volatile InMemoryDataSet cached; /// Creates a new dataset info handle. /// @@ -123,7 +123,7 @@ public boolean isDuplicateVectorFree() { /// completes, after which all callers share the same cached instance. /// /// @return the ready-to-use {@link DataSet} - public DataSet getDataSet() { + public InMemoryDataSet getDataSet() { if (cached == null) { synchronized (this) { if (cached == null) { diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetUtils.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetUtils.java index 61dc64652..acf210aa2 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetUtils.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/DataSetUtils.java @@ -43,7 +43,7 @@ public class DataSetUtils { /** * Processes a dataset using the configured load behavior from the dataset metadata. */ - public static DataSet processDataSet(String pathStr, + public static InMemoryDataSet processDataSet(String pathStr, DataSetProperties props, List> baseVectors, List> queryVectors, @@ -76,7 +76,7 @@ public static DataSet getScrubbedDataSet(String pathStr, return legacyScrubDataSet(pathStr, vsf, baseVectors, queryVectors, groundTruth); } - private static DataSet legacyScrubDataSet(String pathStr, + private static InMemoryDataSet legacyScrubDataSet(String pathStr, VectorSimilarityFunction vsf, List> baseVectors, List> queryVectors, diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/InMemoryDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/InMemoryDataSet.java new file mode 100644 index 000000000..7bf5aec17 --- /dev/null +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/InMemoryDataSet.java @@ -0,0 +1,32 @@ +/* + * Copyright DataStax, Inc. + * + * 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 io.github.jbellis.jvector.example.benchmarks.datasets; + +import java.util.List; + +import io.github.jbellis.jvector.vector.types.VectorFloat; + +/** + * A {@link DataSet} which allows for fetching the entire base vector corpus as a list + */ +public interface InMemoryDataSet extends DataSet { + /** + * The base vectors as a list. + * @return a list of base vectors + */ + List> getBaseVectors(); +} diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SimpleDataSet.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SimpleDataSet.java index bf9c69376..a0d912893 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SimpleDataSet.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/benchmarks/datasets/SimpleDataSet.java @@ -23,7 +23,7 @@ import java.util.List; -public class SimpleDataSet implements DataSet { +public class SimpleDataSet implements InMemoryDataSet { private final String name; private final VectorSimilarityFunction similarityFunction; private final List> baseVectors; diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java index fd3b94ce3..05a10ae86 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/reporting/DatasetInfoWriter.java @@ -120,7 +120,7 @@ public static Row fromDataSet(String datasetName, basePath, queryPath, groundTruthPath, - ds.getBaseVectors().size(), + ds.getBaseRavv().size(), ds.getQueryVectors().size(), ds.getGroundTruth().size(), ds.getDimension(), diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java index 1e6a83f40..be8c74d6b 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/example/util/DataSetPartitioner.java @@ -16,7 +16,7 @@ package io.github.jbellis.jvector.example.util; -import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; +import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet; import io.github.jbellis.jvector.example.yaml.TestDataPartition; import io.github.jbellis.jvector.vector.types.VectorFloat; @@ -39,7 +39,7 @@ public PartitionedData(List>> vectors, List sizes) } } - public static PartitionedData partition(DataSet ds, int numParts, TestDataPartition.Distribution distribution) { + public static PartitionedData partition(InMemoryDataSet ds, int numParts, TestDataPartition.Distribution distribution) { return partition(ds.getBaseVectors(), numParts, distribution); } diff --git a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java index f3728234c..6df490f25 100644 --- a/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java +++ b/jvector-examples/src/main/java/io/github/jbellis/jvector/graph/disk/ParallelWriteExample.java @@ -305,7 +305,7 @@ public static void main(String[] args) throws IOException { DataSet ds = DataSets.loadDataSet(datasetName).orElseThrow( () -> new RuntimeException("Dataset " + datasetName + " not found") ).getDataSet(); - System.out.printf("Loaded %d vectors of dimension %d%n", ds.getBaseVectors().size(), ds.getDimension()); + System.out.printf("Loaded %d vectors of dimension %d%n", ds.getBaseRavv().size(), ds.getDimension()); var floatVectors = ds.getBaseRavv(); diff --git a/jvector-tests/src/test/java/io/github/jbellis/jvector/microbench/GraphBuildBench.java b/jvector-tests/src/test/java/io/github/jbellis/jvector/microbench/GraphBuildBench.java index 28127fb34..ce078f4b9 100644 --- a/jvector-tests/src/test/java/io/github/jbellis/jvector/microbench/GraphBuildBench.java +++ b/jvector-tests/src/test/java/io/github/jbellis/jvector/microbench/GraphBuildBench.java @@ -15,9 +15,8 @@ */ package io.github.jbellis.jvector.microbench; - -import io.github.jbellis.jvector.example.benchmarks.datasets.DataSet; import io.github.jbellis.jvector.example.benchmarks.datasets.DataSets; +import io.github.jbellis.jvector.example.benchmarks.datasets.InMemoryDataSet; import io.github.jbellis.jvector.graph.GraphIndexBuilder; import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues; import org.openjdk.jmh.annotations.Benchmark; @@ -31,6 +30,7 @@ import org.openjdk.jmh.annotations.Warmup; import org.openjdk.jmh.infra.Blackhole; +import java.io.IOException; import java.util.concurrent.TimeUnit; @Warmup(iterations = 1, time = 5) @@ -40,7 +40,7 @@ public class GraphBuildBench { @State(Scope.Benchmark) public static class Parameters { - final DataSet ds; + final InMemoryDataSet ds; final ListRandomAccessVectorValues ravv; public Parameters() { @@ -54,10 +54,11 @@ public Parameters() { @Benchmark @BenchmarkMode(Mode.Throughput) @OutputTimeUnit(TimeUnit.SECONDS) - public void testGraphBuild(Blackhole bh, Parameters p) { + public void testGraphBuild(Blackhole bh, Parameters p) throws IOException { long start = System.nanoTime(); - GraphIndexBuilder graphIndexBuilder = new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, false); - graphIndexBuilder.build(p.ravv); + try (GraphIndexBuilder graphIndexBuilder = new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, false)) { + graphIndexBuilder.build(p.ravv); + } System.out.format("Build M=%d ef=%d in %.2fs%n", 32, 600, (System.nanoTime() - start) / 1_000_000_000.0); } @@ -65,10 +66,11 @@ public void testGraphBuild(Blackhole bh, Parameters p) { @Benchmark @BenchmarkMode(Mode.Throughput) @OutputTimeUnit(TimeUnit.SECONDS) - public void testGraphBuildWithHierarchy(Blackhole bh, Parameters p) { + public void testGraphBuildWithHierarchy(Blackhole bh, Parameters p) throws IOException { long start = System.nanoTime(); - GraphIndexBuilder graphIndexBuilder = new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, true); - graphIndexBuilder.build(p.ravv); + try (GraphIndexBuilder graphIndexBuilder = new GraphIndexBuilder(p.ravv, p.ds.getSimilarityFunction(), 8, 60, 1.2f, 1.4f, true)) { + graphIndexBuilder.build(p.ravv); + } System.out.format("Build M=%d ef=%d in %.2fs%n", 32, 600, (System.nanoTime() - start) / 1_000_000_000.0); }