mirror of
https://github.com/nosqlbench/nosqlbench.git
synced 2026-10-12 09:24:52 -05:00
Merge pull request #2299 from nosqlbench/rehome_vectordata
Rehome vectordata
This commit is contained in:
58 files changed
+2787
-55
No files matched your search
@@ -25,7 +25,7 @@
|
||||
<packaging>pom</packaging>
|
||||
|
||||
<properties>
|
||||
<revision>5.25.12-SNAPSHOT</revision>
|
||||
<revision>5.25.13-SNAPSHOT</revision>
|
||||
<!-- Set this level to override the logging level for tests during build -->
|
||||
<project.testlevel>INFO</project.testlevel>
|
||||
<!-- Set this level to override the logging level for tests logging configuration during build -->
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
# nb-vectordata
|
||||
|
||||
`nb-vectordata` is the standalone Java programmatic-access implementation of
|
||||
the vectordata dataset format. It is intentionally limited to catalog,
|
||||
manifest, vector-reader, transport, cache, and integrity-verification APIs;
|
||||
workload generators and legacy command-line tooling are not part of this
|
||||
module.
|
||||
|
||||
The compatibility contract is [RUST_COMPATIBILITY.md](RUST_COMPATIBILITY.md).
|
||||
|
||||
## Minimal use
|
||||
|
||||
```java
|
||||
Catalog catalog = Catalog.of(CatalogSources.defaults());
|
||||
TestDataView view = catalog.open("my-dataset", "default");
|
||||
VectorReader<float[]> base = view.baseVectors();
|
||||
float[] vector = base.get(0);
|
||||
```
|
||||
|
||||
Local paths and HTTP(S) URLs are both accepted for manifests, catalogs, and
|
||||
facet sources. HTTP data is range-cached and uses a `.mref` sidecar whenever a
|
||||
server provides one.
|
||||
|
||||
For a release canary against a Rust-hosted dataset, run Maven with
|
||||
`-Dvectordata.canary.catalog=<catalog-url>` and
|
||||
`-Dvectordata.canary.dataset=<dataset>`; optionally set
|
||||
`-Dvectordata.canary.profile=<profile>`.
|
||||
@@ -0,0 +1,56 @@
|
||||
# vectordata-rs compatibility contract
|
||||
|
||||
This module implements the programmatic dataset-access surface of
|
||||
`vectordata-rs` (baseline: crate 1.7.1, source commit `1249310078785dbb59444f1c9bac14247767c286`).
|
||||
When this document and an older Java implementation disagree, the Rust
|
||||
implementation's current behavior and its format tests are normative.
|
||||
|
||||
## Supported data model
|
||||
|
||||
- Catalogs: JSON or YAML `catalog.json`/`catalog.yaml` entry lists and legacy
|
||||
flat `knn_entries.yaml`/`datasets.yaml` entries. A directory resolves in
|
||||
Rust order: `catalog.json`, `catalog.yaml`, then `knn_entries.yaml`.
|
||||
- Dataset manifests: a `dataset.yaml`-shaped document (identified by its
|
||||
top-level `profiles` field) or a legacy flat entries document, regardless of
|
||||
an explicit YAML/YML filename. A dataset directory or URL resolves in Rust
|
||||
order: `dataset.yaml`, then `knn_entries.yaml`.
|
||||
- Data: scalar, fixed-dimension xvec records, and variable-dimension vvec
|
||||
records with `IDXFOR__<file>.i32` or `.i64` index sidecars.
|
||||
- Types: signed and unsigned 8/16/32/64-bit integers plus f16, f32, and f64.
|
||||
- Sources: local files, `file:` URIs, and HTTP(S) URLs.
|
||||
- Remote caching: sparse range caching, `.mref` SHA-256 Merkle verification,
|
||||
cache promotion to a local mapped reader, and full-transfer fallback when a
|
||||
server cannot service byte ranges.
|
||||
- Cache setup: `cache_dir:` in `settings.yaml` wins. With
|
||||
`$VECTORDATA_HOME`, the default is `$VECTORDATA_HOME/cache` and no settings
|
||||
file is created. Otherwise the client uses an absolute `$XDG_CACHE_HOME` or
|
||||
`~/.cache` candidate, auto-persists `cache_dir:` only when the home
|
||||
filesystem is the largest writable mount (or mount discovery is unavailable),
|
||||
and requires an explicit setting when a larger writable mount is elsewhere.
|
||||
|
||||
## Deliberate compatibility decisions
|
||||
|
||||
`metadata_results` is canonical; `metadata_indices` and `predicate_results`
|
||||
are accepted aliases. `prefiltered` is canonical and `filtered` is accepted as
|
||||
a legacy alias. A facet source is resolved relative to its manifest, while a
|
||||
catalog entry's `path` is resolved relative to its catalog. Explicit client
|
||||
settings win over environment and settings-file values.
|
||||
|
||||
Malformed manifests, unsupported extensions, invalid windows, malformed index
|
||||
sidecars, invalid Merkle references, inconsistent data lengths, unavailable
|
||||
profiles, and failed integrity checks fail with a `VectorDataException`; they
|
||||
are never silently repaired.
|
||||
|
||||
The dispatch behavior above is covered by deterministic local tests and an
|
||||
embedded HTTP-server integration test, so Rust precedence is stable across
|
||||
releases instead of depending on a filename convention.
|
||||
|
||||
Automatic cache setup is tested exclusively with temporary, explicit settings
|
||||
roots. The test suite never writes a developer's real `~/.config/vectordata`.
|
||||
|
||||
## Known source-level clarifications
|
||||
|
||||
The Rust storage code permits a no-`.mref` HTTP fallback even where older
|
||||
documentation only described full transfer. This Java implementation supports
|
||||
that fallback. It also preserves custom manifest facets and supports float
|
||||
typed reads, which are part of the advertised API contract.
|
||||
@@ -0,0 +1,81 @@
|
||||
<?xml version="1.0" encoding="UTF-8"?>
|
||||
<!--
|
||||
~ Copyright (c) 2026 The NoSQLBench Authors.
|
||||
~
|
||||
~ 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.
|
||||
-->
|
||||
<project xmlns="http://maven.apache.org/POM/4.0.0"
|
||||
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
|
||||
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
|
||||
<modelVersion>4.0.0</modelVersion>
|
||||
<parent>
|
||||
<groupId>io.nosqlbench</groupId>
|
||||
<artifactId>mvn-defaults</artifactId>
|
||||
<version>${revision}</version>
|
||||
<relativePath>../mvn-defaults</relativePath>
|
||||
</parent>
|
||||
|
||||
<artifactId>nb-vectordata</artifactId>
|
||||
<packaging>jar</packaging>
|
||||
<name>nb-vectordata</name>
|
||||
<description>Programmatic Java access to vectordata-rs datasets.</description>
|
||||
|
||||
<properties>
|
||||
<maven.compiler.release>17</maven.compiler.release>
|
||||
<automatic.module.name>io.nosqlbench.vectordata</automatic.module.name>
|
||||
</properties>
|
||||
|
||||
<dependencies>
|
||||
<dependency>
|
||||
<groupId>org.snakeyaml</groupId>
|
||||
<artifactId>snakeyaml-engine</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>com.squareup.okhttp3</groupId>
|
||||
<artifactId>okhttp</artifactId>
|
||||
</dependency>
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
<!-- Vector fixtures are binary and must never pass through Maven filtering. -->
|
||||
<testResources>
|
||||
<testResource>
|
||||
<directory>src/test/resources</directory>
|
||||
<filtering>false</filtering>
|
||||
</testResource>
|
||||
</testResources>
|
||||
<plugins>
|
||||
<plugin>
|
||||
<groupId>org.jacoco</groupId>
|
||||
<artifactId>jacoco-maven-plugin</artifactId>
|
||||
<executions>
|
||||
<execution>
|
||||
<id>jacoco-check</id>
|
||||
<configuration>
|
||||
<dataFile>${project.build.directory}/jacoco-data.exec</dataFile>
|
||||
<rules>
|
||||
<rule>
|
||||
<element>BUNDLE</element>
|
||||
<limits>
|
||||
<limit><counter>INSTRUCTION</counter><value>COVEREDRATIO</value><minimum>0.65</minimum></limit>
|
||||
<limit><counter>BRANCH</counter><value>COVEREDRATIO</value><minimum>0.45</minimum></limit>
|
||||
</limits>
|
||||
</rule>
|
||||
</rules>
|
||||
</configuration>
|
||||
</execution>
|
||||
</executions>
|
||||
</plugin>
|
||||
</plugins>
|
||||
</build>
|
||||
</project>
|
||||
@@ -0,0 +1,19 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Backing mode selected for a source. */
|
||||
public enum AccessMode { LOCAL, MERKLE_HASHED, MERKLE_CHUNKED, FULL_TRANSFER }
|
||||
@@ -0,0 +1,20 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Immutable observation of a reader's backing cache. */
|
||||
public record CacheStats(AccessMode accessMode, long totalBytes, long cachedBytes,
|
||||
long hits, long misses, boolean complete) { }
|
||||
@@ -0,0 +1,110 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import io.nosqlbench.vectordata.internal.HttpTransport;
|
||||
import io.nosqlbench.vectordata.internal.YamlData;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.URI;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.ArrayList;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
/** Catalog resolver for canonical `catalog.yaml` and legacy `knn_entries.yaml` forms. */
|
||||
public final class Catalog {
|
||||
private final Map<String, CatalogEntry> entries; private final VectorDataSettings settings;
|
||||
private Catalog(Map<String, CatalogEntry> entries, VectorDataSettings settings) { this.entries = Map.copyOf(entries); this.settings = settings; }
|
||||
public static Catalog of(CatalogSources sources) { return of(sources, VectorDataSettings.defaults()); }
|
||||
public static Catalog of(CatalogSources sources, VectorDataSettings settings) {
|
||||
Map<String, CatalogEntry> entries = new LinkedHashMap<>();
|
||||
for (URI source : sources.locations()) load(source, settings, entries);
|
||||
return new Catalog(entries, settings);
|
||||
}
|
||||
public Map<String, CatalogEntry> entries() { return entries; }
|
||||
public TestDataView open(String dataset, String profile) {
|
||||
CatalogEntry entry = entries.get(dataset);
|
||||
if (entry == null) throw new VectorDataException("Dataset is not in configured catalogs: " + dataset);
|
||||
if ("knn_entries.yaml".equals(entry.datasetType())) {
|
||||
@SuppressWarnings("unchecked") Map<String, Map<String, Object>> profiles = (Map<String, Map<String, Object>>) entry.attributes().get("profiles");
|
||||
return TestDataGroup.fromLegacyEntries(dataset, entry.manifest(), profiles, settings).profile(profile);
|
||||
}
|
||||
return TestDataGroup.load(entry.manifest(), settings).profile(profile);
|
||||
}
|
||||
public TestDataView openProfile(String datasetAndProfile) {
|
||||
int separator = datasetAndProfile.lastIndexOf(':');
|
||||
return separator < 0 ? open(datasetAndProfile, null) : open(datasetAndProfile.substring(0, separator), datasetAndProfile.substring(separator + 1));
|
||||
}
|
||||
private static void load(URI source, VectorDataSettings settings, Map<String, CatalogEntry> output) {
|
||||
if (!isCatalogFile(source)) { loadDirectory(source, settings, output); return; }
|
||||
parse(source, read(source, settings), output);
|
||||
}
|
||||
private static void loadDirectory(URI source, VectorDataSettings settings, Map<String, CatalogEntry> output) {
|
||||
for (String candidate : List.of("catalog.json", "catalog.yaml", "knn_entries.yaml")) {
|
||||
URI file = child(source, candidate);
|
||||
try { parse(file, read(file, settings), output); return; }
|
||||
catch (VectorDataException ignored) { }
|
||||
}
|
||||
if ("file".equalsIgnoreCase(source.getScheme())) return;
|
||||
throw new VectorDataException("No catalog.json, catalog.yaml, or knn_entries.yaml at " + source);
|
||||
}
|
||||
private static void parse(URI source, String text, Map<String, CatalogEntry> output) {
|
||||
Object document = YamlData.parseValue(text, source.toString());
|
||||
if (document instanceof List<?> list) { canonical(source, list, output); return; }
|
||||
Map<String, Object> root = YamlData.map(document, source.toString());
|
||||
Object candidate = root.get("datasets"); if (candidate == null) candidate = root.get("entries");
|
||||
if (candidate instanceof List<?> list) canonical(source, list, output); else legacy(source, root, output);
|
||||
}
|
||||
private static void canonical(URI source, List<?> list, Map<String, CatalogEntry> output) {
|
||||
for (Object item : list) {
|
||||
Map<String, Object> entry = YamlData.map(item, "catalog entry");
|
||||
String name = YamlData.string(entry.get("name"), "catalog entry name");
|
||||
String path = YamlData.string(entry.get("path"), "catalog entry path");
|
||||
output.putIfAbsent(name, new CatalogEntry(name, source.resolve(path), YamlData.optionalString(entry.get("dataset_type")), Map.copyOf(entry)));
|
||||
}
|
||||
}
|
||||
private static void legacy(URI source, Map<String, Object> root, Map<String, CatalogEntry> output) {
|
||||
String base = root.get("_defaults") instanceof Map<?, ?> map ? YamlData.optionalString(YamlData.map(map, "_defaults").get("base_url")) : null;
|
||||
URI baseUri = base == null ? source.resolve(".") : source.resolve(base.endsWith("/") ? base : base + "/");
|
||||
Map<String, Map<String, Map<String, Object>>> datasets = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Object> entry : root.entrySet()) {
|
||||
if (entry.getKey().startsWith("_")) continue;
|
||||
String dataset = entry.getKey(); int colon = dataset.indexOf(':'); String profile = colon >= 0 ? dataset.substring(colon + 1) : "default"; if (colon >= 0) dataset = dataset.substring(0, colon);
|
||||
Object value = entry.getValue();
|
||||
Map<String, Object> map = value instanceof Map<?, ?> ? YamlData.map(value, "legacy catalog entry") : Map.of();
|
||||
Map<String, Object> resolved = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Object> facet : map.entrySet()) {
|
||||
if (facet.getValue() instanceof String path) resolved.put(facet.getKey(), baseUri.resolve(path).toString()); else resolved.put(facet.getKey(), facet.getValue());
|
||||
}
|
||||
datasets.computeIfAbsent(dataset, ignored -> new LinkedHashMap<>()).put(profile, resolved);
|
||||
}
|
||||
for (Map.Entry<String, Map<String, Map<String, Object>>> dataset : datasets.entrySet()) {
|
||||
output.putIfAbsent(dataset.getKey(), new CatalogEntry(dataset.getKey(), source, "knn_entries.yaml", Map.of("profiles", Map.copyOf(dataset.getValue()))));
|
||||
}
|
||||
}
|
||||
private static String read(URI source, VectorDataSettings settings) {
|
||||
try {
|
||||
if ("file".equalsIgnoreCase(source.getScheme())) return Files.readString(Path.of(source));
|
||||
return new String(new HttpTransport(settings).get(source), StandardCharsets.UTF_8);
|
||||
} catch (IOException e) { throw new VectorDataException("Cannot read catalog " + source, e); }
|
||||
}
|
||||
private static boolean isCatalogFile(URI source) { String path = source.getPath() == null ? "" : source.getPath().toLowerCase(); return path.endsWith(".json") || path.endsWith(".yaml") || path.endsWith(".yml"); }
|
||||
private static URI child(URI source, String name) { String text = source.toString(); return URI.create(text.endsWith("/") ? text + name : text + "/" + name); }
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.net.URI;
|
||||
import java.util.Map;
|
||||
|
||||
/** A dataset entry resolved from a catalog source. */
|
||||
public record CatalogEntry(String name, URI manifest, String datasetType, Map<String, Object> attributes) { }
|
||||
@@ -0,0 +1,53 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.net.URI;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
import org.snakeyaml.engine.v2.api.Load;
|
||||
import org.snakeyaml.engine.v2.api.LoadSettings;
|
||||
|
||||
/** Ordered catalog locations. Earlier sources win for a duplicate dataset name. */
|
||||
public final class CatalogSources {
|
||||
private final List<URI> locations;
|
||||
private CatalogSources(List<URI> locations) { this.locations = List.copyOf(locations); }
|
||||
public List<URI> locations() { return locations; }
|
||||
public static CatalogSources of(URI... locations) { return new CatalogSources(List.of(locations)); }
|
||||
public static CatalogSources defaults() {
|
||||
String property = System.getProperty("vectordata.catalog");
|
||||
if (property != null && !property.isBlank()) return of(toUri(property));
|
||||
String environment = System.getenv("VECTORDATA_CATALOG");
|
||||
if (environment != null && !environment.isBlank()) return of(toUri(environment));
|
||||
Path home = Optional.ofNullable(System.getenv("VECTORDATA_HOME")).filter(s -> !s.isBlank()).map(Path::of)
|
||||
.orElse(Path.of(System.getProperty("user.home"), ".config", "vectordata"));
|
||||
Path configured = home.resolve("catalogs.yaml");
|
||||
if (Files.isRegularFile(configured)) {
|
||||
try {
|
||||
Object parsed = new Load(LoadSettings.builder().setLabel(configured.toString()).build()).loadFromString(Files.readString(configured));
|
||||
if (parsed instanceof Map<?, ?> sources) {
|
||||
List<URI> locations = sources.values().stream().filter(String.class::isInstance).map(String.class::cast).map(CatalogSources::toUri).toList();
|
||||
if (!locations.isEmpty()) return new CatalogSources(locations);
|
||||
}
|
||||
} catch (Exception e) { throw new VectorDataException("Cannot read vectordata catalogs " + configured, e); }
|
||||
}
|
||||
return of(home.resolve("catalog.yaml").toUri());
|
||||
}
|
||||
private static URI toUri(String value) { return value.contains("://") || value.startsWith("file:") ? URI.create(value) : Path.of(value).toAbsolutePath().toUri(); }
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.util.Locale;
|
||||
|
||||
/** Scalar element encodings understood by vectordata payload readers. */
|
||||
public enum ElementType {
|
||||
U8(1, false, false), I8(1, false, true),
|
||||
U16(2, false, false), I16(2, false, true),
|
||||
U32(4, false, false), I32(4, false, true),
|
||||
U64(8, false, false), I64(8, false, true),
|
||||
F16(2, true, true), F32(4, true, true), F64(8, true, true);
|
||||
|
||||
private final int width;
|
||||
private final boolean floating;
|
||||
private final boolean signed;
|
||||
|
||||
ElementType(int width, boolean floating, boolean signed) {
|
||||
this.width = width;
|
||||
this.floating = floating;
|
||||
this.signed = signed;
|
||||
}
|
||||
|
||||
public int width() { return width; }
|
||||
public boolean floating() { return floating; }
|
||||
public boolean signed() { return signed; }
|
||||
|
||||
public static ElementType forExtension(String source) {
|
||||
String value = source.toLowerCase(Locale.ROOT);
|
||||
int slash = Math.max(value.lastIndexOf('/'), value.lastIndexOf('\\'));
|
||||
if (slash >= 0) value = value.substring(slash + 1);
|
||||
int query = value.indexOf('?');
|
||||
if (query >= 0) value = value.substring(0, query);
|
||||
int dot = value.lastIndexOf('.');
|
||||
if (dot < 0) throw new VectorDataException("Vector source has no extension: " + source);
|
||||
String extension = value.substring(dot + 1);
|
||||
return switch (extension) {
|
||||
case "bvec", "bvecs", "u8", "u8vec", "u8vecs" -> U8;
|
||||
case "i8", "i8vec", "i8vecs" -> I8;
|
||||
case "u16", "u16vec", "u16vecs" -> U16;
|
||||
case "i16", "i16vec", "i16vecs" -> I16;
|
||||
case "u32", "u32vec", "u32vecs" -> U32;
|
||||
case "ivec", "ivecs", "i32", "i32vec", "i32vecs" -> I32;
|
||||
case "u64", "u64vec", "u64vecs" -> U64;
|
||||
case "i64", "i64vec", "i64vecs" -> I64;
|
||||
case "mvec", "mvecs", "f16", "f16vec", "f16vecs" -> F16;
|
||||
case "fvec", "fvecs", "f32", "f32vec", "f32vecs" -> F32;
|
||||
case "dvec", "dvecs", "f64", "f64vec", "f64vecs" -> F64;
|
||||
default -> throw new VectorDataException("Unsupported vector extension '." + extension + "' in " + source);
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.net.URI;
|
||||
import java.util.Map;
|
||||
|
||||
/** Resolved manifest facet, including custom application-defined fields. */
|
||||
public record FacetDescriptor(String name, URI source, String window, Map<String, Object> attributes) { }
|
||||
@@ -0,0 +1,21 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Raised when a downloaded byte range does not match its advertised hash. */
|
||||
public final class IntegrityException extends VectorDataException {
|
||||
public IntegrityException(String message) { super(message); }
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Receives monotonic cache-fill notifications during {@link VectorReader#prebuffer}. */
|
||||
@FunctionalInterface
|
||||
public interface PrebufferProgress {
|
||||
void onProgress(long cachedBytes, long totalBytes);
|
||||
|
||||
PrebufferProgress NONE = (cached, total) -> { };
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import io.nosqlbench.vectordata.internal.HttpTransport;
|
||||
import io.nosqlbench.vectordata.internal.ManifestView;
|
||||
import io.nosqlbench.vectordata.internal.YamlData;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.URI;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
|
||||
/** Dataset manifest and its selectable profiles. */
|
||||
public final class TestDataGroup {
|
||||
private static final Set<String> NON_FACETS = Set.of("extends", "base_count", "query_count", "maxk", "partition", "name", "description", "tags");
|
||||
private final String name; private final URI manifest; private final Map<String, Map<String, Object>> profiles; private final VectorDataSettings settings;
|
||||
private TestDataGroup(String name, URI manifest, Map<String, Map<String, Object>> profiles, VectorDataSettings settings) {
|
||||
this.name = name; this.manifest = manifest; this.profiles = profiles; this.settings = settings;
|
||||
}
|
||||
public static TestDataGroup load(String source) { return load(uri(source), VectorDataSettings.defaults()); }
|
||||
/** Builds a group from a legacy knn_entries layout whose sources are already resolved. */
|
||||
public static TestDataGroup fromLegacyEntries(String name, URI origin, Map<String, Map<String, Object>> entries, VectorDataSettings settings) {
|
||||
Map<String, Map<String, Object>> profiles = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Map<String, Object>> entry : entries.entrySet()) {
|
||||
Map<String, Object> profile = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Object> facet : entry.getValue().entrySet()) {
|
||||
String key = switch (facet.getKey()) { case "base" -> "base"; case "query" -> "query"; case "gt" -> "neighbor_indices"; default -> facet.getKey(); };
|
||||
profile.put(key, facet.getValue() instanceof Map<?, ?> ? facet.getValue() : Map.of("source", facet.getValue()));
|
||||
}
|
||||
profiles.put(entry.getKey(), profile);
|
||||
}
|
||||
if (profiles.isEmpty()) throw new VectorDataException("Legacy catalog has no profiles for " + name);
|
||||
return new TestDataGroup(name, origin, profiles, settings);
|
||||
}
|
||||
public static TestDataGroup load(URI source, VectorDataSettings settings) {
|
||||
if (!isYaml(source)) {
|
||||
URI canonical = child(source, "dataset.yaml");
|
||||
try { return load(canonical, settings); }
|
||||
catch (VectorDataException ignored) { return load(child(source, "knn_entries.yaml"), settings); }
|
||||
}
|
||||
URI manifest = source;
|
||||
Map<String, Object> root = YamlData.parse(read(manifest, settings), manifest.toString());
|
||||
if (!(root.get("profiles") instanceof Map<?, ?>)) return legacy(root, manifest, settings, directoryName(manifest));
|
||||
String name = YamlData.optionalString(root.get("name")); if (name == null) name = basename(manifest);
|
||||
Object rawProfiles = root.get("profiles");
|
||||
Map<String, Map<String, Object>> profiles = new LinkedHashMap<>();
|
||||
if (rawProfiles instanceof Map<?, ?> map) {
|
||||
map.forEach((key, value) -> profiles.put(String.valueOf(key), YamlData.map(value, "profile " + key)));
|
||||
} else {
|
||||
Map<String, Object> fallback = new LinkedHashMap<>(root); fallback.remove("name"); fallback.remove("profiles"); profiles.put("default", fallback);
|
||||
}
|
||||
if (profiles.isEmpty()) throw new VectorDataException("Manifest contains no profiles: " + manifest);
|
||||
return new TestDataGroup(name, manifest, profiles, settings);
|
||||
}
|
||||
private static TestDataGroup legacy(Map<String, Object> root, URI manifest, VectorDataSettings settings, String preferred) {
|
||||
String configuredBase = root.get("_defaults") instanceof Map<?, ?> map ? YamlData.optionalString(YamlData.map(map, "_defaults").get("base_url")) : null;
|
||||
URI base = configuredBase == null ? manifest.resolve(".") : manifest.resolve(configuredBase.endsWith("/") ? configuredBase : configuredBase + "/");
|
||||
Map<String, Map<String, Map<String, Object>>> grouped = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Object> item : root.entrySet()) {
|
||||
if (item.getKey().startsWith("_")) continue;
|
||||
int colon = item.getKey().indexOf(':'); String dataset = colon < 0 ? item.getKey() : item.getKey().substring(0, colon); String profile = colon < 0 ? "default" : item.getKey().substring(colon + 1);
|
||||
Map<String, Object> values = YamlData.map(item.getValue(), "legacy entry " + item.getKey()); Map<String, Object> resolved = new LinkedHashMap<>();
|
||||
for (Map.Entry<String, Object> facet : values.entrySet()) resolved.put(facet.getKey(), facet.getValue() instanceof String path ? base.resolve(path).toString() : facet.getValue());
|
||||
grouped.computeIfAbsent(dataset, ignored -> new LinkedHashMap<>()).put(profile, resolved);
|
||||
}
|
||||
if (grouped.isEmpty()) throw new VectorDataException("Legacy entries document has no datasets: " + manifest);
|
||||
String selected = grouped.containsKey(preferred) ? preferred : grouped.keySet().iterator().next();
|
||||
return fromLegacyEntries(selected, manifest, grouped.get(selected), settings);
|
||||
}
|
||||
public String name() { return name; }
|
||||
public Map<String, Map<String, Object>> profiles() { return Map.copyOf(profiles); }
|
||||
public TestDataView profile(String profile) {
|
||||
String selected = profile == null || profile.isBlank() ? (profiles.containsKey("default") ? "default" : profiles.keySet().iterator().next()) : profile;
|
||||
return new ManifestView(name, selected, facets(selected, new LinkedHashMap<>()), settings);
|
||||
}
|
||||
private Map<String, FacetDescriptor> facets(String profile, Map<String, FacetDescriptor> inherited) {
|
||||
Map<String, Object> definition = profiles.get(profile);
|
||||
if (definition == null) throw new VectorDataException("Dataset " + name + " has no profile " + profile);
|
||||
String parent = YamlData.optionalString(definition.get("extends"));
|
||||
if (parent != null) {
|
||||
if (parent.equals(profile)) throw new VectorDataException("Profile cannot extend itself: " + profile);
|
||||
inherited.putAll(facets(parent, new LinkedHashMap<>()));
|
||||
}
|
||||
for (Map.Entry<String, Object> entry : definition.entrySet()) {
|
||||
if (NON_FACETS.contains(entry.getKey())) continue;
|
||||
String facetName = canonical(entry.getKey());
|
||||
Map<String, Object> values = entry.getValue() instanceof Map<?, ?> ? YamlData.map(entry.getValue(), "facet " + entry.getKey()) : Map.of("source", entry.getValue());
|
||||
String source = YamlData.optionalString(values.get("source"));
|
||||
if (source == null) source = YamlData.optionalString(values.get("path"));
|
||||
if (source == null) continue;
|
||||
String window = YamlData.optionalString(values.get("window"));
|
||||
inherited.put(facetName, new FacetDescriptor(facetName, manifest.resolve(source), window, Map.copyOf(values)));
|
||||
}
|
||||
return inherited;
|
||||
}
|
||||
private static boolean isYaml(URI source) { String path = source.getPath() == null ? "" : source.getPath().toLowerCase(); return path.endsWith(".yaml") || path.endsWith(".yml"); }
|
||||
private static URI child(URI source, String name) { String text = source.toString(); return URI.create(text.endsWith("/") ? text + name : text + "/" + name); }
|
||||
private static String directoryName(URI source) { String path = source.getPath(); int end = path.lastIndexOf('/'); String parent = path.substring(0, end); int start = parent.lastIndexOf('/'); return parent.substring(start + 1); }
|
||||
private static URI uri(String source) { return source.contains("://") || source.startsWith("file:") ? URI.create(source) : Path.of(source).toAbsolutePath().toUri(); }
|
||||
private static String read(URI source, VectorDataSettings settings) {
|
||||
try {
|
||||
if ("file".equalsIgnoreCase(source.getScheme())) return Files.readString(Path.of(source));
|
||||
return new String(new HttpTransport(settings).get(source), java.nio.charset.StandardCharsets.UTF_8);
|
||||
} catch (IOException e) { throw new VectorDataException("Cannot read manifest " + source, e); }
|
||||
}
|
||||
private static String basename(URI uri) { String path = uri.getPath(); String parent = path.substring(0, path.lastIndexOf('/')); return parent.substring(parent.lastIndexOf('/') + 1); }
|
||||
private static String canonical(String value) { return switch (value) { case "metadata_indices", "predicate_results" -> "metadata_results"; case "filtered" -> "prefiltered"; default -> value; }; }
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
|
||||
/** A selected dataset profile with lazily opened standard and custom facets. */
|
||||
public interface TestDataView {
|
||||
String dataset();
|
||||
String profile();
|
||||
Map<String, FacetDescriptor> facets();
|
||||
Optional<FacetDescriptor> facet(String name);
|
||||
VectorReader<float[]> baseVectors();
|
||||
VectorReader<float[]> queryVectors();
|
||||
VectorReader<int[]> neighborIndices();
|
||||
VectorReader<float[]> neighborDistances();
|
||||
VectorReader<int[]> metadataResults();
|
||||
VectorReader<?> openFacet(String name);
|
||||
VvecReader<?> openVariableFacet(String name);
|
||||
void prebuffer(PrebufferProgress progress);
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Numeric projection of a scalar or fixed-vector source. */
|
||||
public interface TypedReader {
|
||||
long count();
|
||||
int dimension();
|
||||
ElementType elementType();
|
||||
Number scalar(long index);
|
||||
Number[] vector(long index);
|
||||
void prebuffer(PrebufferProgress progress);
|
||||
CacheStats cacheStats();
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import io.nosqlbench.vectordata.internal.ElementCodec;
|
||||
|
||||
import java.lang.reflect.Array;
|
||||
import java.net.URI;
|
||||
|
||||
/** Opens numeric projections for scalar and fixed-vector payloads. */
|
||||
public final class TypedReaders {
|
||||
private TypedReaders() { }
|
||||
public static TypedReader open(URI source, VectorDataSettings settings, String identity) {
|
||||
return new Projection(VectorReaders.open(source, settings, identity));
|
||||
}
|
||||
private record Projection(VectorReader<?> reader) implements TypedReader {
|
||||
@Override public long count() { return reader.count(); }
|
||||
@Override public int dimension() { return reader.dimension(); }
|
||||
@Override public ElementType elementType() { return reader.elementType(); }
|
||||
@Override public Number scalar(long index) {
|
||||
if (reader.dimension() != 1 || !(reader.get(index) instanceof Number number))
|
||||
throw new VectorDataException("Source is not a packed scalar payload");
|
||||
return number;
|
||||
}
|
||||
@Override public Number[] vector(long index) {
|
||||
Object values = reader.get(index);
|
||||
if (values instanceof Number number) return new Number[] {number};
|
||||
int length = Array.getLength(values); Number[] result = new Number[length];
|
||||
for (int i = 0; i < length; i++) result[i] = number(values, i, reader.elementType());
|
||||
return result;
|
||||
}
|
||||
@Override public void prebuffer(PrebufferProgress progress) { reader.prebuffer(progress); }
|
||||
@Override public CacheStats cacheStats() { return reader.cacheStats(); }
|
||||
private static Number number(Object values, int index, ElementType type) {
|
||||
return switch (type) {
|
||||
case U8 -> Byte.toUnsignedInt(((byte[]) values)[index]); case I8 -> ((byte[]) values)[index];
|
||||
case U16 -> Short.toUnsignedInt(((short[]) values)[index]); case I16 -> ((short[]) values)[index];
|
||||
case U32 -> Integer.toUnsignedLong(((int[]) values)[index]); case I32 -> ((int[]) values)[index];
|
||||
case U64 -> UnsignedLong.ofBits(((long[]) values)[index]); case I64 -> ((long[]) values)[index]; case F16 -> ElementCodec.halfToFloat(((short[]) values)[index]);
|
||||
case F32 -> ((float[]) values)[index]; case F64 -> ((double[]) values)[index];
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.math.BigInteger;
|
||||
|
||||
/** Exact unsigned 64-bit integer value, including values above {@link Long#MAX_VALUE}. */
|
||||
public final class UnsignedLong extends Number implements Comparable<UnsignedLong> {
|
||||
private final long bits;
|
||||
private UnsignedLong(long bits) { this.bits = bits; }
|
||||
public static UnsignedLong ofBits(long bits) { return new UnsignedLong(bits); }
|
||||
public BigInteger toBigInteger() {
|
||||
BigInteger magnitude = BigInteger.valueOf(bits & Long.MAX_VALUE);
|
||||
return bits < 0 ? magnitude.setBit(63) : magnitude;
|
||||
}
|
||||
@Override public int intValue() { return (int) bits; }
|
||||
@Override public long longValue() { return bits; }
|
||||
@Override public float floatValue() { return toBigInteger().floatValue(); }
|
||||
@Override public double doubleValue() { return toBigInteger().doubleValue(); }
|
||||
@Override public int compareTo(UnsignedLong other) { return Long.compareUnsigned(bits, other.bits); }
|
||||
@Override public String toString() { return Long.toUnsignedString(bits); }
|
||||
@Override public boolean equals(Object value) { return value instanceof UnsignedLong other && bits == other.bits; }
|
||||
@Override public int hashCode() { return Long.hashCode(bits); }
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Raised when a vectordata catalog, manifest, source, or payload is invalid. */
|
||||
public class VectorDataException extends RuntimeException {
|
||||
public VectorDataException(String message) { super(message); }
|
||||
public VectorDataException(String message, Throwable cause) { super(message, cause); }
|
||||
}
|
||||
@@ -0,0 +1,244 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.snakeyaml.engine.v2.api.Load;
|
||||
import org.snakeyaml.engine.v2.api.LoadSettings;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.FileStore;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
|
||||
/** Client configuration. Builder values override environment and settings files. */
|
||||
public final class VectorDataSettings {
|
||||
private static final int DEFAULT_CHUNK_SIZE = 1 << 20;
|
||||
private final Path cacheDirectory;
|
||||
private final String bearerToken;
|
||||
private final Duration timeout;
|
||||
private final int chunkSize;
|
||||
|
||||
private VectorDataSettings(Builder builder) {
|
||||
// An explicit per-client cache is an isolation boundary. Do not let a
|
||||
// builder used by a test or embedding application auto-write settings
|
||||
// merely to obtain a cache value that will be discarded.
|
||||
SettingsFile file = SettingsFile.load(builder.cacheDirectory == null);
|
||||
this.cacheDirectory = builder.cacheDirectory != null ? builder.cacheDirectory : file.cacheDirectory();
|
||||
this.bearerToken = builder.bearerToken != null ? builder.bearerToken : file.token();
|
||||
this.timeout = builder.timeout != null ? builder.timeout : file.timeout();
|
||||
this.chunkSize = builder.chunkSize != null ? builder.chunkSize : file.chunkSize();
|
||||
if (chunkSize <= 0) throw new VectorDataException("chunkSize must be positive");
|
||||
}
|
||||
|
||||
public Path cacheDirectory() { return cacheDirectory; }
|
||||
public Optional<String> bearerToken() { return Optional.ofNullable(bearerToken).filter(s -> !s.isBlank()); }
|
||||
public Duration timeout() { return timeout; }
|
||||
public int chunkSize() { return chunkSize; }
|
||||
public static Builder builder() { return new Builder(); }
|
||||
public static VectorDataSettings defaults() { return builder().build(); }
|
||||
|
||||
public static final class Builder {
|
||||
private Path cacheDirectory;
|
||||
private String bearerToken;
|
||||
private Duration timeout;
|
||||
private Integer chunkSize;
|
||||
public Builder cacheDirectory(Path value) { this.cacheDirectory = value; return this; }
|
||||
public Builder bearerToken(String value) { this.bearerToken = value; return this; }
|
||||
public Builder timeout(Duration value) { this.timeout = value; return this; }
|
||||
public Builder chunkSize(int value) { this.chunkSize = value; return this; }
|
||||
public VectorDataSettings build() { return new VectorDataSettings(this); }
|
||||
}
|
||||
|
||||
/** Test hook: resolve an isolated settings root without reading process configuration. */
|
||||
static Path resolveCacheDirectory(Path settings, Path userHome, Path vectordataHome, Path xdgCacheHome, boolean bootstrap) {
|
||||
return SettingsFile.load(settings, userHome, vectordataHome, xdgCacheHome, bootstrap).cacheDirectory();
|
||||
}
|
||||
|
||||
/** Test hook for the persistence half of the Rust-compatible bootstrap protocol. */
|
||||
static void writeCacheDirectory(Path settings, Path cacheDirectory, boolean force) {
|
||||
SettingsFile.writeCacheDirectory(settings, cacheDirectory, force);
|
||||
}
|
||||
|
||||
private enum AutoCacheReason { HOME_IS_LARGEST_MOUNT, DIFFERENT_MOUNT_IS_LARGEST, MOUNTS_UNAVAILABLE }
|
||||
private record AutoCache(Path path, AutoCacheReason reason) { }
|
||||
|
||||
private record SettingsFile(Path cacheDirectory, String token, Duration timeout, int chunkSize) {
|
||||
private static SettingsFile load(boolean bootstrap) {
|
||||
Path vectordataHome = Optional.ofNullable(System.getenv("VECTORDATA_HOME"))
|
||||
.filter(s -> !s.isBlank()).map(Path::of).orElse(null);
|
||||
Path userHome = Optional.ofNullable(System.getenv("HOME")).filter(s -> !s.isBlank())
|
||||
.map(Path::of).orElseGet(() -> Path.of(System.getProperty("user.home")));
|
||||
Path settings = vectordataHome != null ? vectordataHome.resolve("settings.yaml") : userHome.resolve(".config/vectordata/settings.yaml");
|
||||
return load(settings, userHome, vectordataHome, xdgCacheHome(userHome), bootstrap);
|
||||
}
|
||||
|
||||
private static SettingsFile load(Path settings, Path userHome, Path vectordataHome, Path xdgCacheHome, boolean bootstrap) {
|
||||
String token = System.getenv("VECTORDATA_TOKEN");
|
||||
Duration timeout = Duration.ofSeconds(30);
|
||||
int chunkSize = DEFAULT_CHUNK_SIZE;
|
||||
Path configuredCache = null;
|
||||
boolean settingsExists = Files.isRegularFile(settings);
|
||||
if (settingsExists) {
|
||||
try {
|
||||
Object loaded = new Load(LoadSettings.builder().setLabel(settings.toString()).build())
|
||||
.loadFromString(Files.readString(settings));
|
||||
if (loaded instanceof Map<?, ?> map) {
|
||||
configuredCache = configuredPath(map.get("cache_dir"));
|
||||
token = text(map.get("token"), token);
|
||||
timeout = duration(map.get("timeout_seconds"), timeout);
|
||||
chunkSize = integer(map.get("chunk_size"), chunkSize);
|
||||
}
|
||||
} catch (IOException | RuntimeException e) {
|
||||
throw new VectorDataException("Cannot read vectordata settings " + settings, e);
|
||||
}
|
||||
}
|
||||
if (configuredCache != null) return new SettingsFile(configuredCache, token, timeout, chunkSize);
|
||||
|
||||
// This is a complete configuration/cache isolation boundary, matching
|
||||
// vectordata-rs: it does not create settings.yaml under this root.
|
||||
if (vectordataHome != null) return new SettingsFile(vectordataHome.resolve("cache"), token, timeout, chunkSize);
|
||||
if (!bootstrap) return new SettingsFile(xdgCacheHome.resolve("vectordata"), token, timeout, chunkSize);
|
||||
|
||||
AutoCache auto = autoCache(userHome, xdgCacheHome);
|
||||
if (auto.reason() == AutoCacheReason.DIFFERENT_MOUNT_IS_LARGEST) {
|
||||
throw new VectorDataException("No cache_dir is configured in " + settings
|
||||
+ "; the largest writable mount is outside the home filesystem. Configure cache_dir explicitly.");
|
||||
}
|
||||
try {
|
||||
writeCacheDirectory(settings, auto.path(), false);
|
||||
System.getLogger(VectorDataSettings.class.getName()).log(System.Logger.Level.WARNING,
|
||||
"No vectordata cache_dir was configured; auto-set it to {0}", auto.path());
|
||||
} catch (VectorDataException e) {
|
||||
throw new VectorDataException("No cache_dir is configured in " + settings + "; configure cache_dir explicitly.", e);
|
||||
}
|
||||
return new SettingsFile(auto.path(), token, timeout, chunkSize);
|
||||
}
|
||||
|
||||
private static Path configuredPath(Object value) {
|
||||
String text = text(value, null);
|
||||
return text == null || text.isBlank() ? null : Path.of(text);
|
||||
}
|
||||
|
||||
private static Path xdgCacheHome(Path userHome) {
|
||||
String configured = System.getenv("XDG_CACHE_HOME");
|
||||
if (configured != null && !configured.isBlank()) {
|
||||
Path candidate = Path.of(configured);
|
||||
if (candidate.isAbsolute()) return candidate;
|
||||
}
|
||||
return userHome.resolve(".cache");
|
||||
}
|
||||
|
||||
/** Mirrors Rust: never silently place a large cache on a different, larger mount. */
|
||||
private static AutoCache autoCache(Path userHome, Path xdgCacheHome) {
|
||||
try {
|
||||
FileStore homeStore = Files.getFileStore(userHome);
|
||||
List<Path> mounts = writableMounts();
|
||||
if (mounts.isEmpty()) return new AutoCache(xdgCacheHome.resolve("vectordata"), AutoCacheReason.MOUNTS_UNAVAILABLE);
|
||||
mounts.sort(Comparator.comparingLong(SettingsFile::usableSpace).reversed());
|
||||
for (Path mount : mounts) {
|
||||
FileStore store;
|
||||
try { store = Files.getFileStore(mount); } catch (IOException ignored) { continue; }
|
||||
if (homeStore.equals(store)) return new AutoCache(xdgCacheHome.resolve("vectordata"), AutoCacheReason.HOME_IS_LARGEST_MOUNT);
|
||||
if (Files.isWritable(mount)) return new AutoCache(mount.resolve("vectordata-cache"), AutoCacheReason.DIFFERENT_MOUNT_IS_LARGEST);
|
||||
}
|
||||
} catch (IOException ignored) {
|
||||
// An unusable mount table must not block the XDG fallback.
|
||||
}
|
||||
return new AutoCache(xdgCacheHome.resolve("vectordata"), AutoCacheReason.MOUNTS_UNAVAILABLE);
|
||||
}
|
||||
|
||||
private static List<Path> writableMounts() {
|
||||
Path mounts = Path.of("/proc/mounts");
|
||||
if (!Files.isRegularFile(mounts)) return List.of();
|
||||
try {
|
||||
List<Path> result = new ArrayList<>();
|
||||
for (String line : Files.readAllLines(mounts)) {
|
||||
String[] fields = line.split("\\s+");
|
||||
if (fields.length >= 4 && fields[3].contains("rw")) {
|
||||
Path mount = Path.of(unescapeMount(fields[1]));
|
||||
if (Files.isDirectory(mount) && Files.isWritable(mount)) result.add(mount);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
} catch (IOException ignored) { return List.of(); }
|
||||
}
|
||||
|
||||
private static long usableSpace(Path path) {
|
||||
try { return Files.getFileStore(path).getUsableSpace(); }
|
||||
catch (IOException ignored) { return -1L; }
|
||||
}
|
||||
|
||||
private static String unescapeMount(String value) {
|
||||
return value.replace("\\040", " ").replace("\\011", "\t").replace("\\012", "\n").replace("\\134", "\\");
|
||||
}
|
||||
|
||||
private static void writeCacheDirectory(Path settings, Path cacheDirectory, boolean force) {
|
||||
try {
|
||||
if (Files.isRegularFile(settings)) {
|
||||
Path existing = existingCacheDirectory(settings);
|
||||
if (cacheDirectory.equals(existing)) return;
|
||||
if (!force) throw new VectorDataException("Refusing to overwrite existing vectordata settings " + settings);
|
||||
}
|
||||
if (settings.getParent() != null) Files.createDirectories(settings.getParent());
|
||||
Files.createDirectories(cacheDirectory);
|
||||
String content = Files.isRegularFile(settings)
|
||||
? rewriteCacheDirectory(Files.readString(settings), cacheDirectory)
|
||||
: "cache_dir: " + cacheDirectory + "\nprotect_settings: true\n";
|
||||
Files.writeString(settings, content);
|
||||
} catch (IOException e) {
|
||||
throw new VectorDataException("Cannot write vectordata settings " + settings, e);
|
||||
}
|
||||
}
|
||||
|
||||
private static Path existingCacheDirectory(Path settings) {
|
||||
try {
|
||||
Object loaded = new Load(LoadSettings.builder().setLabel(settings.toString()).build()).loadFromString(Files.readString(settings));
|
||||
return loaded instanceof Map<?, ?> map ? configuredPath(map.get("cache_dir")) : null;
|
||||
} catch (IOException | RuntimeException e) { throw new VectorDataException("Cannot read vectordata settings " + settings, e); }
|
||||
}
|
||||
|
||||
private static String rewriteCacheDirectory(String existing, Path value) {
|
||||
StringBuilder result = new StringBuilder(); boolean replaced = false;
|
||||
for (String line : existing.split("(?<=\\n)", -1)) {
|
||||
String trimmed = line.stripLeading();
|
||||
if (!replaced && !trimmed.startsWith("#") && trimmed.startsWith("cache_dir:")) {
|
||||
result.append("cache_dir: ").append(value).append('\n'); replaced = true;
|
||||
} else result.append(line);
|
||||
}
|
||||
if (!replaced) {
|
||||
if (!existing.isEmpty() && !existing.endsWith("\n")) result.append('\n');
|
||||
result.append("cache_dir: ").append(value).append('\n');
|
||||
}
|
||||
return result.toString();
|
||||
}
|
||||
|
||||
private static String text(Object value, String fallback) { return value == null ? fallback : String.valueOf(value); }
|
||||
private static int integer(Object value, int fallback) {
|
||||
if (value instanceof Number n) return n.intValue();
|
||||
try { return value == null ? fallback : Integer.parseInt(String.valueOf(value)); }
|
||||
catch (NumberFormatException ignored) { return fallback; }
|
||||
}
|
||||
private static Duration duration(Object value, Duration fallback) {
|
||||
return Duration.ofSeconds(integer(value, Math.toIntExact(fallback.toSeconds())));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Random access to a fixed-dimension vector or scalar payload. */
|
||||
public interface VectorReader<A> {
|
||||
long count();
|
||||
int dimension();
|
||||
A get(long index);
|
||||
default void get(long index, A target) { throw new UnsupportedOperationException("Target-array reads are not supported"); }
|
||||
void prebuffer(PrebufferProgress progress);
|
||||
default void prebuffer() { prebuffer(PrebufferProgress.NONE); }
|
||||
boolean isComplete();
|
||||
CacheStats cacheStats();
|
||||
ElementType elementType();
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import io.nosqlbench.vectordata.internal.ByteStorage;
|
||||
import io.nosqlbench.vectordata.internal.FixedVectorReader;
|
||||
import io.nosqlbench.vectordata.internal.ScalarReader;
|
||||
import io.nosqlbench.vectordata.internal.StorageFactory;
|
||||
import io.nosqlbench.vectordata.internal.VariableVectorReader;
|
||||
|
||||
import java.net.URI;
|
||||
import java.nio.file.Path;
|
||||
import java.util.Locale;
|
||||
|
||||
/** Factories for source readers. */
|
||||
public final class VectorReaders {
|
||||
private VectorReaders() { }
|
||||
public static VectorReader<?> open(String source) { return open(toUri(source), VectorDataSettings.defaults(), "direct"); }
|
||||
public static VectorReader<?> open(URI source, VectorDataSettings settings, String identity) {
|
||||
ElementType type = ElementType.forExtension(source.toString());
|
||||
ByteStorage storage = new StorageFactory(settings).open(source, identity);
|
||||
return isScalar(source.toString()) ? new ScalarReader(storage, type) : new FixedVectorReader<>(storage, type);
|
||||
}
|
||||
public static VvecReader<?> openVvec(URI source, VectorDataSettings settings, String identity) {
|
||||
String text = source.toString().toLowerCase(Locale.ROOT);
|
||||
ElementType type = ElementType.forExtension(text);
|
||||
if (!isVector(text)) throw new VectorDataException("vvec source must use a vector extension: " + source);
|
||||
ByteStorage data = new StorageFactory(settings).open(source, identity);
|
||||
URI i32 = URI.create(source.toString().replaceAll("[^/]+$", "IDXFOR__" + filename(source) + ".i32"));
|
||||
try { return new VariableVectorReader<>(data, new StorageFactory(settings).open(i32, identity), type, false); }
|
||||
catch (VectorDataException missingI32) {
|
||||
URI i64 = URI.create(source.toString().replaceAll("[^/]+$", "IDXFOR__" + filename(source) + ".i64"));
|
||||
return new VariableVectorReader<>(data, new StorageFactory(settings).open(i64, identity), type, true);
|
||||
}
|
||||
}
|
||||
@SuppressWarnings("unchecked") public static VectorReader<float[]> f32(URI source, VectorDataSettings settings, String identity) {
|
||||
VectorReader<?> reader = open(source, settings, identity); if (reader.elementType() != ElementType.F32 || !(reader instanceof FixedVectorReader<?>)) throw new VectorDataException("Expected f32 fixed vectors: " + source); return (VectorReader<float[]>) reader;
|
||||
}
|
||||
@SuppressWarnings("unchecked") public static VectorReader<int[]> i32(URI source, VectorDataSettings settings, String identity) {
|
||||
VectorReader<?> reader = open(source, settings, identity); if (reader.elementType() != ElementType.I32 || !(reader instanceof FixedVectorReader<?>)) throw new VectorDataException("Expected i32 fixed vectors: " + source); return (VectorReader<int[]>) reader;
|
||||
}
|
||||
static boolean isScalar(String source) { String ext = extension(source); return ext.matches("u8|i8|u16|i16|u32|i32|u64|i64"); }
|
||||
static boolean isVector(String source) { return !isScalar(source); }
|
||||
private static String extension(String value) { int dot = value.lastIndexOf('.'); int end = value.indexOf('?', dot); return value.substring(dot + 1, end < 0 ? value.length() : end).toLowerCase(Locale.ROOT); }
|
||||
private static URI toUri(String source) { return source.contains("://") || source.startsWith("file:") ? URI.create(source) : Path.of(source).toAbsolutePath().toUri(); }
|
||||
private static String filename(URI source) { String path = source.getPath(); return path.substring(path.lastIndexOf('/') + 1); }
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
/** Random access to variable-dimension vectors backed by an IDXFOR sidecar. */
|
||||
public interface VvecReader<A> {
|
||||
long count();
|
||||
int dimensionAt(long index);
|
||||
A get(long index);
|
||||
void prebuffer(PrebufferProgress progress);
|
||||
default void prebuffer() { prebuffer(PrebufferProgress.NONE); }
|
||||
boolean isComplete();
|
||||
CacheStats cacheStats();
|
||||
ElementType elementType();
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
|
||||
import java.nio.ByteBuffer;
|
||||
|
||||
/** Thread-safe source of arbitrary byte ranges. */
|
||||
public interface ByteStorage extends AutoCloseable {
|
||||
long size();
|
||||
ByteBuffer read(long offset, int length);
|
||||
void prebuffer(PrebufferProgress progress);
|
||||
boolean isComplete();
|
||||
CacheStats stats();
|
||||
@Override default void close() { }
|
||||
}
|
||||
@@ -0,0 +1,75 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.ElementType;
|
||||
import io.nosqlbench.vectordata.UnsignedLong;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
import java.util.Arrays;
|
||||
|
||||
/** Little-endian primitive decoding shared by all vector formats. */
|
||||
public final class ElementCodec {
|
||||
private ElementCodec() { }
|
||||
public static Object decode(ByteBuffer source, ElementType type, int count) {
|
||||
ByteBuffer bytes = source.order(ByteOrder.LITTLE_ENDIAN);
|
||||
return switch (type) {
|
||||
case U8, I8 -> bytes.get(new byte[count]).array();
|
||||
case U16, I16, F16 -> { short[] result = new short[count]; for (int i = 0; i < count; i++) result[i] = bytes.getShort(); yield result; }
|
||||
case U32, I32 -> { int[] result = new int[count]; for (int i = 0; i < count; i++) result[i] = bytes.getInt(); yield result; }
|
||||
case U64, I64 -> { long[] result = new long[count]; for (int i = 0; i < count; i++) result[i] = bytes.getLong(); yield result; }
|
||||
case F32 -> { float[] result = new float[count]; for (int i = 0; i < count; i++) result[i] = bytes.getFloat(); yield result; }
|
||||
case F64 -> { double[] result = new double[count]; for (int i = 0; i < count; i++) result[i] = bytes.getDouble(); yield result; }
|
||||
};
|
||||
}
|
||||
public static Number number(ByteBuffer source, ElementType type) {
|
||||
ByteBuffer bytes = source.order(ByteOrder.LITTLE_ENDIAN);
|
||||
return switch (type) {
|
||||
case U8 -> Byte.toUnsignedInt(bytes.get()); case I8 -> bytes.get();
|
||||
case U16 -> Short.toUnsignedInt(bytes.getShort()); case I16 -> bytes.getShort();
|
||||
case U32 -> Integer.toUnsignedLong(bytes.getInt()); case I32 -> bytes.getInt();
|
||||
case U64 -> UnsignedLong.ofBits(bytes.getLong()); case I64 -> bytes.getLong();
|
||||
case F16 -> halfToFloat(bytes.getShort()); case F32 -> bytes.getFloat(); case F64 -> bytes.getDouble();
|
||||
};
|
||||
}
|
||||
public static float halfToFloat(short half) {
|
||||
int h = half & 0xffff, sign = (h & 0x8000) << 16, exponent = (h >>> 10) & 0x1f, fraction = h & 0x3ff;
|
||||
int bits;
|
||||
if (exponent == 0) {
|
||||
if (fraction == 0) bits = sign;
|
||||
else { while ((fraction & 0x400) == 0) { fraction <<= 1; exponent--; } exponent++; fraction &= ~0x400; bits = sign | ((exponent + 112) << 23) | (fraction << 13); }
|
||||
} else if (exponent == 31) bits = sign | 0x7f800000 | (fraction << 13);
|
||||
else bits = sign | ((exponent + 112) << 23) | (fraction << 13);
|
||||
return Float.intBitsToFloat(bits);
|
||||
}
|
||||
public static void copy(Object source, Object target) {
|
||||
if (source instanceof byte[] a && target instanceof byte[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else if (source instanceof short[] a && target instanceof short[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else if (source instanceof int[] a && target instanceof int[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else if (source instanceof long[] a && target instanceof long[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else if (source instanceof float[] a && target instanceof float[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else if (source instanceof double[] a && target instanceof double[] b) System.arraycopy(a, 0, b, 0, a.length);
|
||||
else throw new VectorDataException("Reader target array has incompatible element type");
|
||||
}
|
||||
public static int arrayLength(Object value) {
|
||||
if (value instanceof byte[] a) return a.length; if (value instanceof short[] a) return a.length;
|
||||
if (value instanceof int[] a) return a.length; if (value instanceof long[] a) return a.length;
|
||||
if (value instanceof float[] a) return a.length; if (value instanceof double[] a) return a.length;
|
||||
throw new VectorDataException("Not a primitive vector array: " + value);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.ElementType;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
|
||||
/** Reader for dimension-prefixed, fixed-width xvec records. */
|
||||
public final class FixedVectorReader<A> implements VectorReader<A> {
|
||||
private final ByteStorage storage;
|
||||
private final ElementType type;
|
||||
private final int dimension;
|
||||
private final int recordBytes;
|
||||
private final long count;
|
||||
public FixedVectorReader(ByteStorage storage, ElementType type) {
|
||||
this.storage = storage; this.type = type;
|
||||
if (storage.size() < 4) throw new VectorDataException("xvec payload is too short for a dimension prefix");
|
||||
dimension = storage.read(0, 4).order(java.nio.ByteOrder.LITTLE_ENDIAN).getInt();
|
||||
if (dimension < 0) throw new VectorDataException("Negative xvec dimension: " + dimension);
|
||||
try { recordBytes = Math.addExact(4, Math.multiplyExact(dimension, type.width())); }
|
||||
catch (ArithmeticException e) { throw new VectorDataException("xvec record is too large", e); }
|
||||
if (recordBytes == 0 || storage.size() % recordBytes != 0) throw new VectorDataException("xvec file length is not a whole number of records");
|
||||
count = storage.size() / recordBytes;
|
||||
}
|
||||
@Override public long count() { return count; }
|
||||
@Override public int dimension() { return dimension; }
|
||||
@SuppressWarnings("unchecked") @Override public A get(long index) {
|
||||
check(index);
|
||||
long offset = index * (long) recordBytes;
|
||||
int current = storage.read(offset, 4).order(java.nio.ByteOrder.LITTLE_ENDIAN).getInt();
|
||||
if (current != dimension) throw new VectorDataException("Inconsistent xvec dimension at record " + index + ": " + current + " != " + dimension);
|
||||
return (A) ElementCodec.decode(storage.read(offset + 4, recordBytes - 4), type, dimension);
|
||||
}
|
||||
@Override public void get(long index, A target) { Object value = get(index); if (ElementCodec.arrayLength(target) != dimension) throw new VectorDataException("Target dimension differs from " + dimension); ElementCodec.copy(value, target); }
|
||||
private void check(long index) { if (index < 0 || index >= count) throw new IndexOutOfBoundsException("Vector index " + index + " outside 0.." + (count - 1)); }
|
||||
@Override public void prebuffer(PrebufferProgress progress) { storage.prebuffer(progress); }
|
||||
@Override public boolean isComplete() { return storage.isComplete(); }
|
||||
@Override public CacheStats cacheStats() { return storage.stats(); }
|
||||
@Override public ElementType elementType() { return type; }
|
||||
}
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
|
||||
/** HTTP failure retaining the response status for transport fallbacks. */
|
||||
final class HttpStatusException extends VectorDataException {
|
||||
private final int status;
|
||||
HttpStatusException(String message, int status) { super(message); this.status = status; }
|
||||
int status() { return status; }
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import okhttp3.OkHttpClient;
|
||||
import okhttp3.Request;
|
||||
import okhttp3.Response;
|
||||
import okhttp3.ResponseBody;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.URI;
|
||||
import java.time.Duration;
|
||||
import java.util.Optional;
|
||||
|
||||
/** Small authenticated HTTP helper used by manifests and binary sources. */
|
||||
public final class HttpTransport {
|
||||
private static final int MAX_ATTEMPTS = 3;
|
||||
private final OkHttpClient client;
|
||||
private final Optional<String> token;
|
||||
|
||||
public HttpTransport(VectorDataSettings settings) {
|
||||
Duration timeout = settings.timeout();
|
||||
client = new OkHttpClient.Builder().callTimeout(timeout).connectTimeout(timeout).readTimeout(timeout).build();
|
||||
token = settings.bearerToken();
|
||||
}
|
||||
|
||||
public byte[] get(URI uri) { return execute(uri, null, "GET"); }
|
||||
public byte[] range(URI uri, long from, long to) {
|
||||
if (from < 0 || to < from) throw new VectorDataException("Invalid HTTP range " + from + "-" + to);
|
||||
return execute(uri, "bytes=" + from + "-" + to, "GET");
|
||||
}
|
||||
public Probe probe(URI uri) {
|
||||
Request request = request(uri).head().build();
|
||||
try (Response response = client.newCall(request).execute()) {
|
||||
if (response.isSuccessful()) {
|
||||
long length = response.header("Content-Length", "-1").equals("-1") ? -1L : Long.parseLong(response.header("Content-Length"));
|
||||
return new Probe(length, "bytes".equalsIgnoreCase(response.header("Accept-Ranges")));
|
||||
}
|
||||
} catch (IOException ignored) { }
|
||||
Request rangeRequest = request(uri).header("Range", "bytes=0-0").get().build();
|
||||
try (Response response = client.newCall(rangeRequest).execute()) {
|
||||
if (!response.isSuccessful()) throw failure(uri, response);
|
||||
String contentRange = response.header("Content-Range");
|
||||
if (contentRange != null && contentRange.contains("/")) {
|
||||
return new Probe(Long.parseLong(contentRange.substring(contentRange.lastIndexOf('/') + 1)), true);
|
||||
}
|
||||
long length = response.body() == null ? -1L : response.body().contentLength();
|
||||
return new Probe(length, false);
|
||||
} catch (IOException e) { throw new VectorDataException("Cannot probe " + uri, e); }
|
||||
}
|
||||
private byte[] execute(URI uri, String range, String method) {
|
||||
Request.Builder builder = request(uri);
|
||||
if (range != null) builder.header("Range", range);
|
||||
if ("GET".equals(method)) builder.get();
|
||||
VectorDataException last = null;
|
||||
for (int attempt = 1; attempt <= MAX_ATTEMPTS; attempt++) {
|
||||
try (Response response = client.newCall(builder.build()).execute()) {
|
||||
if (!response.isSuccessful()) {
|
||||
HttpStatusException failure = failure(uri, response);
|
||||
if (!transientStatus(response.code()) || attempt == MAX_ATTEMPTS) throw failure;
|
||||
last = failure;
|
||||
} else {
|
||||
ResponseBody body = response.body();
|
||||
if (body == null) throw new VectorDataException("Empty HTTP response for " + uri);
|
||||
return body.bytes();
|
||||
}
|
||||
} catch (IOException e) {
|
||||
last = new VectorDataException("HTTP request failed for " + uri, e);
|
||||
if (attempt == MAX_ATTEMPTS) throw last;
|
||||
}
|
||||
try { Thread.sleep(25L * attempt); }
|
||||
catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new VectorDataException("HTTP retry interrupted for " + uri, e); }
|
||||
}
|
||||
throw last == null ? new VectorDataException("HTTP request failed for " + uri) : last;
|
||||
}
|
||||
private Request.Builder request(URI uri) {
|
||||
Request.Builder request = new Request.Builder().url(uri.toString());
|
||||
token.ifPresent(value -> request.header("Authorization", "Bearer " + value));
|
||||
return request;
|
||||
}
|
||||
private static HttpStatusException failure(URI uri, Response response) {
|
||||
return new HttpStatusException("HTTP " + response.code() + " for " + uri, response.code());
|
||||
}
|
||||
private static boolean transientStatus(int status) { return status == 408 || status == 429 || status >= 500; }
|
||||
public record Probe(long contentLength, boolean ranges) { }
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.FacetDescriptor;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.TestDataView;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import io.nosqlbench.vectordata.VectorReaders;
|
||||
import io.nosqlbench.vectordata.VvecReader;
|
||||
|
||||
import java.net.URI;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
import java.util.Optional;
|
||||
|
||||
/** Default TestDataView implementation backed by resolved manifest descriptors. */
|
||||
public final class ManifestView implements TestDataView {
|
||||
private final String dataset, profile; private final Map<String, FacetDescriptor> facets; private final VectorDataSettings settings;
|
||||
public ManifestView(String dataset, String profile, Map<String, FacetDescriptor> facets, VectorDataSettings settings) {
|
||||
this.dataset = dataset; this.profile = profile; this.facets = Map.copyOf(new LinkedHashMap<>(facets)); this.settings = settings;
|
||||
}
|
||||
@Override public String dataset() { return dataset; }
|
||||
@Override public String profile() { return profile; }
|
||||
@Override public Map<String, FacetDescriptor> facets() { return facets; }
|
||||
@Override public Optional<FacetDescriptor> facet(String name) { return Optional.ofNullable(facets.get(canonical(name))); }
|
||||
@Override public VectorReader<float[]> baseVectors() { return fixed("base", VectorReaders::f32); }
|
||||
@Override public VectorReader<float[]> queryVectors() { return fixed("query", VectorReaders::f32); }
|
||||
@Override public VectorReader<int[]> neighborIndices() { return fixed("neighbor_indices", VectorReaders::i32); }
|
||||
@Override public VectorReader<float[]> neighborDistances() { return fixed("neighbor_distances", VectorReaders::f32); }
|
||||
@Override public VectorReader<int[]> metadataResults() { return fixed("metadata_results", VectorReaders::i32); }
|
||||
@Override public VectorReader<?> openFacet(String name) { FacetDescriptor facet = require(name); return window(VectorReaders.open(facet.source(), settings, dataset), facet); }
|
||||
@Override public VvecReader<?> openVariableFacet(String name) { FacetDescriptor facet = require(name); return VectorReaders.openVvec(facet.source(), settings, dataset); }
|
||||
@Override public void prebuffer(PrebufferProgress progress) { for (String name : facets.keySet()) openFacet(name).prebuffer(progress); }
|
||||
private FacetDescriptor require(String name) { return facet(name).orElseThrow(() -> new VectorDataException("Profile " + profile + " lacks facet " + name)); }
|
||||
private <T> VectorReader<T> fixed(String name, FixedFactory<T> factory) { FacetDescriptor facet = require(name); return window(factory.open(facet.source(), settings, dataset), facet); }
|
||||
private static <T> VectorReader<T> window(VectorReader<T> reader, FacetDescriptor facet) { return facet.window() == null || facet.window().isBlank() ? reader : new WindowedVectorReader<>(reader, facet.window()); }
|
||||
private static String canonical(String name) {
|
||||
return switch (name) { case "metadata_indices", "predicate_results" -> "metadata_results"; case "filtered" -> "prefiltered"; default -> name; };
|
||||
}
|
||||
@FunctionalInterface private interface FixedFactory<T> { VectorReader<T> open(URI source, VectorDataSettings settings, String identity); }
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.AccessMode;
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.channels.FileChannel;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.StandardOpenOption;
|
||||
|
||||
/** Read-only memory mapped local file storage. */
|
||||
public final class MappedStorage implements ByteStorage {
|
||||
private final FileChannel channel;
|
||||
private final ByteBuffer data;
|
||||
private final long size;
|
||||
|
||||
public MappedStorage(Path path) {
|
||||
try {
|
||||
channel = FileChannel.open(path, StandardOpenOption.READ);
|
||||
size = channel.size();
|
||||
data = channel.map(FileChannel.MapMode.READ_ONLY, 0, size);
|
||||
} catch (IOException e) {
|
||||
throw new VectorDataException("Cannot map vector source " + path, e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override public long size() { return size; }
|
||||
@Override public ByteBuffer read(long offset, int length) {
|
||||
bounds(offset, length);
|
||||
ByteBuffer copy = data.asReadOnlyBuffer();
|
||||
copy.position(Math.toIntExact(offset)).limit(Math.toIntExact(offset + length));
|
||||
return copy.slice().asReadOnlyBuffer();
|
||||
}
|
||||
private void bounds(long offset, int length) {
|
||||
if (offset < 0 || length < 0 || offset > size - length)
|
||||
throw new VectorDataException("Byte range outside source: offset=" + offset + ", length=" + length + ", size=" + size);
|
||||
}
|
||||
@Override public void prebuffer(PrebufferProgress progress) { progress.onProgress(size, size); }
|
||||
@Override public boolean isComplete() { return true; }
|
||||
@Override public CacheStats stats() { return new CacheStats(AccessMode.LOCAL, size, size, 0, 0, true); }
|
||||
@Override public void close() { try { channel.close(); } catch (IOException ignored) { } }
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
import java.util.Arrays;
|
||||
|
||||
/** Parsed vectordata `.mref` footer and leaf hashes. */
|
||||
public final class MerkleReference {
|
||||
private final byte[][] leaves;
|
||||
private final long chunkSize;
|
||||
private final long contentSize;
|
||||
|
||||
private MerkleReference(byte[][] leaves, long chunkSize, long contentSize) {
|
||||
this.leaves = leaves;
|
||||
this.chunkSize = chunkSize;
|
||||
this.contentSize = contentSize;
|
||||
}
|
||||
public long chunkSize() { return chunkSize; }
|
||||
public long contentSize() { return contentSize; }
|
||||
public int leafCount() { return leaves.length; }
|
||||
public byte[] leaf(int index) { return leaves[index].clone(); }
|
||||
|
||||
public static MerkleReference parse(byte[] raw, String source) {
|
||||
if (raw.length < 41) throw new VectorDataException("Malformed .mref for " + source + ": too short");
|
||||
int footerLength = raw[raw.length - 1] & 0xff;
|
||||
if (footerLength != 41 && footerLength != 45)
|
||||
throw new VectorDataException("Malformed .mref for " + source + ": unsupported footer length " + footerLength);
|
||||
if (raw.length < footerLength) throw new VectorDataException("Malformed .mref for " + source);
|
||||
int footerOffset = raw.length - footerLength;
|
||||
ByteBuffer footer = ByteBuffer.wrap(raw, footerOffset, footerLength).order(ByteOrder.BIG_ENDIAN);
|
||||
long chunkSize = footer.getLong();
|
||||
long contentSize = footer.getLong();
|
||||
int totalChunks = footer.getInt();
|
||||
int leafCount = footer.getInt();
|
||||
int capacityLeaves = footer.getInt();
|
||||
int nodeCount = footer.getInt();
|
||||
footer.getInt(); // root offset
|
||||
int internalNodeCount = footer.getInt();
|
||||
if (footerLength == 45) footer.getInt(); // valid-bitset size
|
||||
if (chunkSize <= 0 || contentSize < 0 || totalChunks < 0 || leafCount < totalChunks || capacityLeaves < leafCount || nodeCount < leafCount)
|
||||
throw new VectorDataException("Malformed .mref footer for " + source);
|
||||
long hashBytes = (long) nodeCount * 32L;
|
||||
if (hashBytes != footerOffset) throw new VectorDataException("Malformed .mref node section for " + source);
|
||||
long expectedChunks = contentSize == 0 ? 0 : (contentSize + chunkSize - 1) / chunkSize;
|
||||
if (expectedChunks != totalChunks) throw new VectorDataException("Malformed .mref chunk count for " + source);
|
||||
if (internalNodeCount < 0 || internalNodeCount + leafCount != nodeCount)
|
||||
throw new VectorDataException("Malformed .mref tree shape for " + source);
|
||||
byte[][] leaves = new byte[leafCount][];
|
||||
// vectordata stores leaves after all internal nodes. Padded leaves protect no content.
|
||||
for (int i = 0; i < totalChunks; i++) leaves[i] = Arrays.copyOfRange(raw, (internalNodeCount + i) * 32, (internalNodeCount + i + 1) * 32);
|
||||
return new MerkleReference(Arrays.copyOf(leaves, totalChunks), chunkSize, contentSize);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.AccessMode;
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.IntegrityException;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.URI;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.channels.FileChannel;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.nio.file.StandardCopyOption;
|
||||
import java.nio.file.StandardOpenOption;
|
||||
import java.security.MessageDigest;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.BitSet;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
|
||||
/** Persistent sparse HTTP range cache with optional `.mref` leaf verification. */
|
||||
public final class RemoteStorage implements ByteStorage {
|
||||
private static final byte[] STATE_MAGIC = new byte[] {'N', 'B', 'V', 'D', '1'};
|
||||
private final URI uri;
|
||||
private final HttpTransport transport;
|
||||
private final long size;
|
||||
private final int chunkSize;
|
||||
private final MerkleReference reference;
|
||||
private final boolean ranges;
|
||||
private final Path dataPath;
|
||||
private final Path statePath;
|
||||
private final BitSet present;
|
||||
private final FileChannel data;
|
||||
private final Object lock = new Object();
|
||||
private final AtomicLong hits = new AtomicLong();
|
||||
private final AtomicLong misses = new AtomicLong();
|
||||
private volatile boolean complete;
|
||||
|
||||
public RemoteStorage(URI uri, VectorDataSettings settings, Path cachePath) {
|
||||
this.uri = uri;
|
||||
this.transport = new HttpTransport(settings);
|
||||
MerkleReference found = loadReference(uri);
|
||||
HttpTransport.Probe probe = transport.probe(uri);
|
||||
this.reference = found;
|
||||
this.size = found == null ? probe.contentLength() : found.contentSize();
|
||||
if (size < 0) throw new VectorDataException("Server did not provide a content length for " + uri);
|
||||
long configured = found == null ? settings.chunkSize() : found.chunkSize();
|
||||
if (configured > Integer.MAX_VALUE) throw new VectorDataException("Chunk size is too large for Java reader: " + configured);
|
||||
this.chunkSize = (int) configured;
|
||||
this.ranges = probe.ranges();
|
||||
this.dataPath = cachePath.resolveSibling(cachePath.getFileName() + ".data");
|
||||
this.statePath = cachePath.resolveSibling(cachePath.getFileName() + ".mrkl");
|
||||
try {
|
||||
Files.createDirectories(dataPath.getParent());
|
||||
data = FileChannel.open(dataPath, StandardOpenOption.CREATE, StandardOpenOption.READ, StandardOpenOption.WRITE);
|
||||
present = loadState();
|
||||
complete = present.cardinality() == chunks();
|
||||
} catch (IOException e) { throw new VectorDataException("Cannot open cache for " + uri, e); }
|
||||
}
|
||||
|
||||
private MerkleReference loadReference(URI source) {
|
||||
try { return MerkleReference.parse(transport.get(URI.create(source + ".mref")), source.toString()); }
|
||||
catch (HttpStatusException missing) {
|
||||
// Public S3 datasets may expose data while restricting optional sidecars.
|
||||
// Rust falls back to the non-Merkle range path in that situation.
|
||||
if (missing.status() == 401 || missing.status() == 403 || missing.status() == 404) return null;
|
||||
throw missing;
|
||||
}
|
||||
}
|
||||
@Override public long size() { return size; }
|
||||
@Override public ByteBuffer read(long offset, int length) {
|
||||
checkBounds(offset, length);
|
||||
if (length == 0) return ByteBuffer.allocate(0).asReadOnlyBuffer();
|
||||
ensure(offset / chunkSize, (offset + length - 1) / chunkSize);
|
||||
ByteBuffer result = ByteBuffer.allocate(length);
|
||||
synchronized (lock) {
|
||||
try { while (result.hasRemaining()) if (data.read(result, offset + result.position()) < 0) throw new IOException("Unexpected cache EOF"); }
|
||||
catch (IOException e) { throw new VectorDataException("Cannot read cache for " + uri, e); }
|
||||
}
|
||||
return result.flip().asReadOnlyBuffer();
|
||||
}
|
||||
private void ensure(long first, long last) {
|
||||
synchronized (lock) {
|
||||
for (long chunk = first; chunk <= last; chunk++) {
|
||||
if (present.get(Math.toIntExact(chunk))) hits.incrementAndGet(); else fetch(chunk);
|
||||
}
|
||||
}
|
||||
}
|
||||
private void fetch(long chunk) {
|
||||
long offset = chunk * (long) chunkSize;
|
||||
int length = Math.toIntExact(Math.min(chunkSize, size - offset));
|
||||
byte[] bytes;
|
||||
if (ranges) {
|
||||
bytes = transport.range(uri, offset, offset + length - 1);
|
||||
if (bytes.length != length) throw new VectorDataException("Range response length mismatch for " + uri + " chunk " + chunk);
|
||||
} else {
|
||||
bytes = transport.get(uri);
|
||||
if (bytes.length != size) throw new VectorDataException("Full HTTP response length mismatch for " + uri);
|
||||
write(0, bytes);
|
||||
for (int i = 0; i < chunks(); i++) verifyAndMark(i);
|
||||
saveState();
|
||||
return;
|
||||
}
|
||||
if (reference != null && !MessageDigest.isEqual(hash(bytes), reference.leaf(Math.toIntExact(chunk))))
|
||||
throw new IntegrityException("Merkle leaf verification failed for " + uri + " chunk " + chunk);
|
||||
write(offset, bytes);
|
||||
present.set(Math.toIntExact(chunk));
|
||||
misses.incrementAndGet();
|
||||
complete = present.cardinality() == chunks();
|
||||
saveState();
|
||||
}
|
||||
private void verifyAndMark(int chunk) {
|
||||
long offset = chunk * (long) chunkSize;
|
||||
int length = Math.toIntExact(Math.min(chunkSize, size - offset));
|
||||
if (reference != null && !MessageDigest.isEqual(hash(readCached(offset, length)), reference.leaf(chunk)))
|
||||
throw new IntegrityException("Merkle leaf verification failed for " + uri + " chunk " + chunk);
|
||||
present.set(chunk);
|
||||
misses.incrementAndGet();
|
||||
complete = present.cardinality() == chunks();
|
||||
}
|
||||
private byte[] readCached(long offset, int length) {
|
||||
ByteBuffer bytes = ByteBuffer.allocate(length);
|
||||
try { while (bytes.hasRemaining()) data.read(bytes, offset + bytes.position()); }
|
||||
catch (IOException e) { throw new VectorDataException("Cannot verify cache for " + uri, e); }
|
||||
return bytes.array();
|
||||
}
|
||||
private void write(long offset, byte[] bytes) {
|
||||
ByteBuffer buffer = ByteBuffer.wrap(bytes);
|
||||
try { while (buffer.hasRemaining()) data.write(buffer, offset + buffer.position()); data.force(false); }
|
||||
catch (IOException e) { throw new VectorDataException("Cannot write cache for " + uri, e); }
|
||||
}
|
||||
@Override public void prebuffer(PrebufferProgress progress) {
|
||||
for (int i = 0; i < chunks(); i++) { ensure(i, i); progress.onProgress(Math.min(size, (long) present.cardinality() * chunkSize), size); }
|
||||
}
|
||||
@Override public boolean isComplete() { return complete; }
|
||||
@Override public CacheStats stats() {
|
||||
long cached = Math.min(size, (long) present.cardinality() * chunkSize);
|
||||
AccessMode mode = reference == null ? (ranges ? AccessMode.MERKLE_CHUNKED : AccessMode.FULL_TRANSFER) : AccessMode.MERKLE_HASHED;
|
||||
return new CacheStats(mode, size, cached, hits.get(), misses.get(), complete);
|
||||
}
|
||||
private int chunks() { return size == 0 ? 0 : Math.toIntExact((size + chunkSize - 1) / chunkSize); }
|
||||
private void checkBounds(long offset, int length) {
|
||||
if (offset < 0 || length < 0 || offset > size - length) throw new VectorDataException("Byte range outside remote source " + uri);
|
||||
}
|
||||
private BitSet loadState() throws IOException {
|
||||
if (!Files.isRegularFile(statePath)) return new BitSet(chunks());
|
||||
byte[] raw = Files.readAllBytes(statePath);
|
||||
if (raw.length < STATE_MAGIC.length + 16) return new BitSet(chunks());
|
||||
for (int i = 0; i < STATE_MAGIC.length; i++) if (raw[i] != STATE_MAGIC[i]) return new BitSet(chunks());
|
||||
ByteBuffer header = ByteBuffer.wrap(raw).position(STATE_MAGIC.length).slice();
|
||||
long storedSize = header.getLong(); int storedChunkSize = header.getInt(); int words = header.getInt();
|
||||
if (storedSize != size || storedChunkSize != chunkSize || words < 0 || raw.length != STATE_MAGIC.length + 16 + words * 8) return new BitSet(chunks());
|
||||
long[] values = new long[words]; for (int i = 0; i < words; i++) values[i] = header.getLong();
|
||||
return BitSet.valueOf(values);
|
||||
}
|
||||
private void saveState() {
|
||||
long[] bits = present.toLongArray();
|
||||
ByteBuffer output = ByteBuffer.allocate(STATE_MAGIC.length + 16 + bits.length * 8);
|
||||
output.put(STATE_MAGIC).putLong(size).putInt(chunkSize).putInt(bits.length); for (long bit : bits) output.putLong(bit);
|
||||
Path temporary = statePath.resolveSibling(statePath.getFileName() + ".tmp");
|
||||
try {
|
||||
Files.write(temporary, output.array());
|
||||
try { Files.move(temporary, statePath, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING); }
|
||||
catch (IOException ignored) { Files.move(temporary, statePath, StandardCopyOption.REPLACE_EXISTING); }
|
||||
} catch (IOException e) { throw new VectorDataException("Cannot save cache state for " + uri, e); }
|
||||
}
|
||||
private static byte[] hash(byte[] bytes) {
|
||||
try { return MessageDigest.getInstance("SHA-256").digest(bytes); }
|
||||
catch (NoSuchAlgorithmException e) { throw new IllegalStateException(e); }
|
||||
}
|
||||
@Override public void close() { try { data.close(); } catch (IOException ignored) { } }
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.ElementType;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
|
||||
/** Reader for packed scalar payloads. */
|
||||
public final class ScalarReader implements VectorReader<Number> {
|
||||
private final ByteStorage storage; private final ElementType type; private final long count;
|
||||
public ScalarReader(ByteStorage storage, ElementType type) {
|
||||
this.storage = storage; this.type = type;
|
||||
if (storage.size() % type.width() != 0) throw new VectorDataException("Scalar payload is not aligned to " + type.width() + " bytes");
|
||||
count = storage.size() / type.width();
|
||||
}
|
||||
@Override public long count() { return count; }
|
||||
@Override public int dimension() { return 1; }
|
||||
@Override public Number get(long index) { if (index < 0 || index >= count) throw new IndexOutOfBoundsException(); return ElementCodec.number(storage.read(index * type.width(), type.width()), type); }
|
||||
@Override public void prebuffer(PrebufferProgress progress) { storage.prebuffer(progress); }
|
||||
@Override public boolean isComplete() { return storage.isComplete(); }
|
||||
@Override public CacheStats cacheStats() { return storage.stats(); }
|
||||
@Override public ElementType elementType() { return type; }
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
|
||||
import java.lang.ref.WeakReference;
|
||||
import java.net.URI;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.nio.file.Path;
|
||||
import java.security.MessageDigest;
|
||||
import java.security.NoSuchAlgorithmException;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
/** Shares byte storage by canonical source URI within a JVM. */
|
||||
public final class StorageFactory {
|
||||
private static final ConcurrentHashMap<String, WeakReference<ByteStorage>> OPEN = new ConcurrentHashMap<>();
|
||||
private final VectorDataSettings settings;
|
||||
public StorageFactory(VectorDataSettings settings) { this.settings = settings; }
|
||||
|
||||
public ByteStorage open(URI source, String datasetIdentity) {
|
||||
String key = source.normalize().toString();
|
||||
synchronized (OPEN) {
|
||||
ByteStorage existing = OPEN.containsKey(key) ? OPEN.get(key).get() : null;
|
||||
if (existing != null) return existing;
|
||||
ByteStorage created;
|
||||
if ("file".equalsIgnoreCase(source.getScheme()) || source.getScheme() == null) {
|
||||
created = new MappedStorage(Path.of(source));
|
||||
} else if ("http".equalsIgnoreCase(source.getScheme()) || "https".equalsIgnoreCase(source.getScheme())) {
|
||||
Path cache = settings.cacheDirectory().resolve(safe(datasetIdentity)).resolve(digest(key));
|
||||
created = new RemoteStorage(source, settings, cache);
|
||||
} else throw new VectorDataException("Unsupported source URI scheme: " + source);
|
||||
OPEN.put(key, new WeakReference<>(created));
|
||||
return created;
|
||||
}
|
||||
}
|
||||
private static String safe(String identity) { return identity.replaceAll("[^A-Za-z0-9._-]", "_"); }
|
||||
private static String digest(String value) {
|
||||
try {
|
||||
byte[] hashed = MessageDigest.getInstance("SHA-256").digest(value.getBytes(StandardCharsets.UTF_8));
|
||||
StringBuilder result = new StringBuilder(); for (byte b : hashed) result.append(String.format("%02x", b));
|
||||
return result.toString();
|
||||
} catch (NoSuchAlgorithmException e) { throw new IllegalStateException(e); }
|
||||
}
|
||||
}
|
||||
+57
@@ -0,0 +1,57 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.ElementType;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VvecReader;
|
||||
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
|
||||
/** Reader for variable vectors, indexed by byte offsets in an IDXFOR sidecar. */
|
||||
public final class VariableVectorReader<A> implements VvecReader<A> {
|
||||
private final ByteStorage data; private final ByteStorage index; private final ElementType type; private final int width; private final long count;
|
||||
public VariableVectorReader(ByteStorage data, ByteStorage index, ElementType type, boolean index64) {
|
||||
this.data = data; this.index = index; this.type = type; this.width = index64 ? 8 : 4;
|
||||
if (index.size() % width != 0 || index.size() < width * 2L) throw new VectorDataException("Invalid IDXFOR sidecar length");
|
||||
count = index.size() / width - 1;
|
||||
long end = offset(count); if (end != data.size()) throw new VectorDataException("IDXFOR final offset does not equal vvec data size");
|
||||
}
|
||||
@Override public long count() { return count; }
|
||||
@Override public int dimensionAt(long value) { return record(value).order(ByteOrder.LITTLE_ENDIAN).getInt(); }
|
||||
@Override @SuppressWarnings("unchecked") public A get(long value) {
|
||||
ByteBuffer record = record(value).order(ByteOrder.LITTLE_ENDIAN); int dimension = record.getInt();
|
||||
if (dimension < 0 || record.remaining() != Math.multiplyExact(dimension, type.width())) throw new VectorDataException("Malformed vvec record " + value);
|
||||
return (A) ElementCodec.decode(record.slice(), type, dimension);
|
||||
}
|
||||
private ByteBuffer record(long value) {
|
||||
if (value < 0 || value >= count) throw new IndexOutOfBoundsException("vvec index " + value);
|
||||
long begin = offset(value), end = offset(value + 1);
|
||||
if (begin < 0 || end < begin || end > data.size() || end - begin > Integer.MAX_VALUE) throw new VectorDataException("Invalid IDXFOR range at " + value);
|
||||
return data.read(begin, Math.toIntExact(end - begin));
|
||||
}
|
||||
private long offset(long value) {
|
||||
ByteBuffer bytes = index.read(value * width, width).order(ByteOrder.LITTLE_ENDIAN);
|
||||
return width == 4 ? Integer.toUnsignedLong(bytes.getInt()) : bytes.getLong();
|
||||
}
|
||||
@Override public void prebuffer(PrebufferProgress progress) { data.prebuffer(progress); index.prebuffer(PrebufferProgress.NONE); }
|
||||
@Override public boolean isComplete() { return data.isComplete() && index.isComplete(); }
|
||||
@Override public CacheStats cacheStats() { return data.stats(); }
|
||||
@Override public ElementType elementType() { return type; }
|
||||
}
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.CacheStats;
|
||||
import io.nosqlbench.vectordata.ElementType;
|
||||
import io.nosqlbench.vectordata.PrebufferProgress;
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
|
||||
/** Contiguous logical index window over a fixed vector reader. */
|
||||
public final class WindowedVectorReader<A> implements VectorReader<A> {
|
||||
private final VectorReader<A> delegate; private final long begin, count;
|
||||
public WindowedVectorReader(VectorReader<A> delegate, String expression) {
|
||||
this.delegate = delegate;
|
||||
String[] parts = expression.trim().replace("[", "").replace("]", "").split("\\.\\.", -1);
|
||||
if (parts.length != 2) throw new VectorDataException("Window must be a contiguous start..end range: " + expression);
|
||||
try { begin = parts[0].isBlank() ? 0 : Long.parseLong(parts[0].trim()); long end = parts[1].isBlank() ? delegate.count() : Long.parseLong(parts[1].trim()); count = end - begin; }
|
||||
catch (NumberFormatException e) { throw new VectorDataException("Invalid window " + expression, e); }
|
||||
if (begin < 0 || count < 0 || begin > delegate.count() - count) throw new VectorDataException("Window lies outside source: " + expression);
|
||||
}
|
||||
@Override public long count() { return count; }
|
||||
@Override public int dimension() { return delegate.dimension(); }
|
||||
@Override public A get(long index) { check(index); return delegate.get(begin + index); }
|
||||
@Override public void get(long index, A target) { check(index); delegate.get(begin + index, target); }
|
||||
private void check(long index) { if (index < 0 || index >= count) throw new IndexOutOfBoundsException(); }
|
||||
@Override public void prebuffer(PrebufferProgress progress) { delegate.prebuffer(progress); }
|
||||
@Override public boolean isComplete() { return delegate.isComplete(); }
|
||||
@Override public CacheStats cacheStats() { return delegate.cacheStats(); }
|
||||
@Override public ElementType elementType() { return delegate.elementType(); }
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata.internal;
|
||||
|
||||
import io.nosqlbench.vectordata.VectorDataException;
|
||||
import org.snakeyaml.engine.v2.api.Load;
|
||||
import org.snakeyaml.engine.v2.api.LoadSettings;
|
||||
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.Map;
|
||||
|
||||
/** Checked conversion helpers for permissive YAML object trees. */
|
||||
public final class YamlData {
|
||||
private YamlData() { }
|
||||
public static Map<String, Object> map(Object value, String label) {
|
||||
if (!(value instanceof Map<?, ?> raw)) throw new VectorDataException(label + " must be a YAML mapping");
|
||||
Map<String, Object> result = new LinkedHashMap<>(); raw.forEach((key, item) -> result.put(String.valueOf(key), item)); return result;
|
||||
}
|
||||
public static Object parseValue(String text, String label) { return new Load(LoadSettings.builder().setLabel(label).build()).loadFromString(text); }
|
||||
public static Map<String, Object> parse(String text, String label) { return map(parseValue(text, label), label); }
|
||||
public static String string(Object value, String label) { if (value == null) throw new VectorDataException("Missing " + label); return String.valueOf(value); }
|
||||
public static String optionalString(Object value) { return value == null ? null : String.valueOf(value); }
|
||||
public static int integer(Object value, String label) {
|
||||
if (value instanceof Number number) return number.intValue();
|
||||
try { return Integer.parseInt(string(value, label)); } catch (NumberFormatException e) { throw new VectorDataException("Invalid integer " + label, e); }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
class CatalogManifestTest {
|
||||
@TempDir Path temporary;
|
||||
@Test void resolvesCatalogProfilesInheritanceAliasesAndCustomFacets() throws Exception {
|
||||
Path dataset = Files.createDirectories(temporary.resolve("sample"));
|
||||
FixtureSupport.fvec(dataset, "base.fvec", new float[][] {{1f, 2f}});
|
||||
FixtureSupport.fvec(dataset, "query.fvec", new float[][] {{3f, 4f}});
|
||||
FixtureSupport.ivec(dataset, "neighbors.ivecs", new int[][] {{9, 10}});
|
||||
FixtureSupport.ivec(dataset, "metadata.ivecs", new int[][] {{11, 12}});
|
||||
Files.writeString(dataset.resolve("dataset.yaml"), """
|
||||
name: sample
|
||||
profiles:
|
||||
parent:
|
||||
base: base.fvec
|
||||
custom_scores: {source: metadata.ivecs}
|
||||
demo:
|
||||
extends: parent
|
||||
query: {source: query.fvec, window: 0..1}
|
||||
neighbor_indices: neighbors.ivecs
|
||||
metadata_indices: metadata.ivecs
|
||||
""");
|
||||
Files.writeString(temporary.resolve("catalog.yaml"), """
|
||||
datasets:
|
||||
- name: sample
|
||||
path: sample/dataset.yaml
|
||||
dataset_type: dataset.yaml
|
||||
""");
|
||||
Catalog catalog = Catalog.of(CatalogSources.of(temporary.resolve("catalog.yaml").toUri()));
|
||||
TestDataView view = catalog.open("sample", "demo");
|
||||
assertArrayEquals(new float[] {1f, 2f}, view.baseVectors().get(0));
|
||||
assertArrayEquals(new float[] {3f, 4f}, view.queryVectors().get(0));
|
||||
assertArrayEquals(new int[] {9, 10}, view.neighborIndices().get(0));
|
||||
assertArrayEquals(new int[] {11, 12}, view.metadataResults().get(0));
|
||||
assertTrue(view.facet("custom_scores").isPresent());
|
||||
assertEquals(1, view.queryVectors().count());
|
||||
}
|
||||
@Test void opensLegacyKnnEntriesProfilesWithoutDatasetManifest() throws Exception {
|
||||
FixtureSupport.fvec(temporary, "base.fvec", new float[][] {{1f, 2f}});
|
||||
FixtureSupport.fvec(temporary, "query.fvec", new float[][] {{3f, 4f}});
|
||||
FixtureSupport.ivec(temporary, "gt.ivecs", new int[][] {{7, 8}});
|
||||
Path catalogFile = temporary.resolve("knn_entries.yaml");
|
||||
Files.writeString(catalogFile, """
|
||||
"legacy:demo":
|
||||
base: base.fvec
|
||||
query: query.fvec
|
||||
gt: gt.ivecs
|
||||
""");
|
||||
TestDataView view = Catalog.of(CatalogSources.of(catalogFile.toUri())).open("legacy", "demo");
|
||||
assertArrayEquals(new float[] {1f, 2f}, view.baseVectors().get(0));
|
||||
assertArrayEquals(new float[] {3f, 4f}, view.queryVectors().get(0));
|
||||
assertArrayEquals(new int[] {7, 8}, view.neighborIndices().get(0));
|
||||
}
|
||||
@Test void directArbitraryYamlDispatchesLegacyEntriesByShape() throws Exception {
|
||||
FixtureSupport.fvec(temporary, "base.fvec", new float[][] {{1f, 2f}});
|
||||
Path file = temporary.resolve("entries.yaml");
|
||||
Files.writeString(file, "\"shape:default\":\n base: base.fvec\n");
|
||||
assertArrayEquals(new float[] {1f, 2f}, TestDataGroup.load(file.toUri(), VectorDataSettings.defaults()).profile("default").baseVectors().get(0));
|
||||
}
|
||||
@Test void directoryPrefersDatasetYamlThenFallsBackToKnnEntries() throws Exception {
|
||||
Path preferred = Files.createDirectories(temporary.resolve("preferred"));
|
||||
FixtureSupport.fvec(preferred, "base.fvec", new float[][] {{1f, 2f}});
|
||||
Files.writeString(preferred.resolve("dataset.yaml"), "profiles:\n default:\n base: base.fvec\n");
|
||||
Files.writeString(preferred.resolve("knn_entries.yaml"), "\"ignored:default\":\n base: missing.fvec\n");
|
||||
assertArrayEquals(new float[] {1f, 2f}, TestDataGroup.load(preferred.toUri(), VectorDataSettings.defaults()).profile("default").baseVectors().get(0));
|
||||
Path legacy = Files.createDirectories(temporary.resolve("legacy"));
|
||||
FixtureSupport.fvec(legacy, "base.fvec", new float[][] {{3f, 4f}});
|
||||
Files.writeString(legacy.resolve("knn_entries.yaml"), "\"legacy:default\":\n base: base.fvec\n");
|
||||
assertArrayEquals(new float[] {3f, 4f}, TestDataGroup.load(legacy.toUri(), VectorDataSettings.defaults()).profile("default").baseVectors().get(0));
|
||||
}
|
||||
@Test void catalogDirectoryPrefersJsonOverYaml() throws Exception {
|
||||
Path dataset = Files.createDirectories(temporary.resolve("json-dataset"));
|
||||
FixtureSupport.fvec(dataset, "base.fvec", new float[][] {{1f, 2f}}); Files.writeString(dataset.resolve("dataset.yaml"), "profiles:\n default:\n base: base.fvec\n");
|
||||
Files.writeString(temporary.resolve("catalog.json"), "[{\"name\":\"json\",\"path\":\"json-dataset/dataset.yaml\",\"dataset_type\":\"dataset.yaml\"}]");
|
||||
Files.writeString(temporary.resolve("catalog.yaml"), "- name: yaml\n path: absent.yaml\n");
|
||||
assertArrayEquals(new float[] {1f, 2f}, Catalog.of(CatalogSources.of(temporary.toUri())).open("json", "default").baseVectors().get(0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.security.MessageDigest;
|
||||
|
||||
final class FixtureSupport {
|
||||
private FixtureSupport() { }
|
||||
static Path fvec(Path directory, String filename, float[][] values) throws IOException {
|
||||
int dimension = values[0].length; ByteBuffer bytes = ByteBuffer.allocate(values.length * (4 + dimension * 4)).order(ByteOrder.LITTLE_ENDIAN);
|
||||
for (float[] value : values) { bytes.putInt(dimension); for (float element : value) bytes.putFloat(element); }
|
||||
Path result = directory.resolve(filename); Files.write(result, bytes.array()); return result;
|
||||
}
|
||||
static Path ivec(Path directory, String filename, int[][] values) throws IOException {
|
||||
int dimension = values[0].length; ByteBuffer bytes = ByteBuffer.allocate(values.length * (4 + dimension * 4)).order(ByteOrder.LITTLE_ENDIAN);
|
||||
for (int[] value : values) { bytes.putInt(dimension); for (int element : value) bytes.putInt(element); }
|
||||
Path result = directory.resolve(filename); Files.write(result, bytes.array()); return result;
|
||||
}
|
||||
static Path scalarI32(Path directory, String filename, int... values) throws IOException {
|
||||
ByteBuffer bytes = ByteBuffer.allocate(values.length * 4).order(ByteOrder.LITTLE_ENDIAN); for (int value : values) bytes.putInt(value);
|
||||
Path result = directory.resolve(filename); Files.write(result, bytes.array()); return result;
|
||||
}
|
||||
static Path vvec(Path directory, String filename, int[][] values) throws IOException {
|
||||
int size = 0; for (int[] value : values) size += 4 + value.length * 4;
|
||||
ByteBuffer data = ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN); ByteBuffer index = ByteBuffer.allocate((values.length + 1) * 4).order(ByteOrder.LITTLE_ENDIAN);
|
||||
for (int[] value : values) { index.putInt(data.position()); data.putInt(value.length); for (int element : value) data.putInt(element); }
|
||||
index.putInt(data.position()); Path result = directory.resolve(filename); Files.write(result, data.array()); Files.write(directory.resolve("IDXFOR__" + filename + ".i32"), index.array()); return result;
|
||||
}
|
||||
static Path vvec64(Path directory, String filename, int[][] values) throws IOException {
|
||||
int size = 0; for (int[] value : values) size += 4 + value.length * 4;
|
||||
ByteBuffer data = ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN); ByteBuffer index = ByteBuffer.allocate((values.length + 1) * 8).order(ByteOrder.LITTLE_ENDIAN);
|
||||
for (int[] value : values) { index.putLong(data.position()); data.putInt(value.length); for (int element : value) data.putInt(element); }
|
||||
index.putLong(data.position()); Path result = directory.resolve(filename); Files.write(result, data.array()); Files.write(directory.resolve("IDXFOR__" + filename + ".i64"), index.array()); return result;
|
||||
}
|
||||
/** Serializes the vectordata-rs MerkleRef layout (complete, padded binary tree). */
|
||||
static byte[] mref(byte[] data, int chunkSize) throws Exception {
|
||||
int chunks = (data.length + chunkSize - 1) / chunkSize;
|
||||
int leaves = 1; while (leaves < Math.max(1, chunks)) leaves <<= 1;
|
||||
int internals = leaves - 1, nodes = internals + leaves;
|
||||
byte[][] hashes = new byte[nodes][]; MessageDigest digest = MessageDigest.getInstance("SHA-256");
|
||||
for (int i = 0; i < leaves; i++) hashes[internals + i] = digest.digest(i < chunks ? java.util.Arrays.copyOfRange(data, i * chunkSize, Math.min(data.length, (i + 1) * chunkSize)) : new byte[0]);
|
||||
for (int i = internals - 1; i >= 0; i--) { digest.reset(); digest.update(hashes[i * 2 + 1]); hashes[i] = digest.digest(hashes[i * 2 + 2]); }
|
||||
ByteBuffer result = ByteBuffer.allocate(nodes * 32 + 45).order(ByteOrder.BIG_ENDIAN);
|
||||
for (byte[] hash : hashes) result.put(hash);
|
||||
result.putLong(chunkSize).putLong(data.length).putInt(chunks).putInt(leaves).putInt(leaves).putInt(nodes).putInt(internals).putInt(internals).putInt(0).put((byte) 45);
|
||||
return result.array();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.net.URI;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
|
||||
|
||||
/** Verifies the checked-in catalog and binary xvec format fixture through the public API. */
|
||||
class FormatCompatibilityTest {
|
||||
@Test void readsCheckedInFormatFixtureThroughCatalogApi() throws Exception {
|
||||
URI catalogUri = getClass().getClassLoader().getResource("rust-v1/catalog.yaml").toURI();
|
||||
TestDataView view = Catalog.of(CatalogSources.of(catalogUri)).open("rust-v1-demo", "default");
|
||||
assertArrayEquals(new float[] {1f, 2f}, view.baseVectors().get(0));
|
||||
assertArrayEquals(new float[] {3f, 4f}, view.baseVectors().get(1));
|
||||
assertArrayEquals(new float[] {5f, 6f}, view.queryVectors().get(0));
|
||||
assertArrayEquals(new int[] {7, 8}, view.neighborIndices().get(0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,185 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import com.sun.net.httpserver.HttpExchange;
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.net.URI;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
class RemoteCacheIntegrationTest {
|
||||
@TempDir Path temporary;
|
||||
@Test void rangeCachesAndVerifiesMerkleLeaves() throws Exception {
|
||||
Path source = FixtureSupport.fvec(temporary, "remote.fvec", new float[][] {{1f, 2f}, {3f, 4f}});
|
||||
byte[] payload = Files.readAllBytes(source); byte[] mref = FixtureSupport.mref(payload, 8); AtomicInteger ranges = new AtomicInteger();
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/remote.fvec", exchange -> serve(exchange, payload, ranges));
|
||||
server.createContext("/remote.fvec.mref", exchange -> { exchange.getResponseHeaders().add("Content-Length", String.valueOf(mref.length)); exchange.sendResponseHeaders(200, mref.length); exchange.getResponseBody().write(mref); exchange.close(); });
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/remote.fvec");
|
||||
VectorDataSettings settings = VectorDataSettings.builder().cacheDirectory(temporary.resolve("cache")).build();
|
||||
VectorReader<float[]> reader = VectorReaders.f32(uri, settings, "remote-test");
|
||||
assertArrayEquals(new float[] {3f, 4f}, reader.get(1));
|
||||
assertEquals(AccessMode.MERKLE_HASHED, reader.cacheStats().accessMode());
|
||||
int downloaded = ranges.get(); assertTrue(downloaded > 0);
|
||||
assertArrayEquals(new float[] {1f, 2f}, reader.get(0));
|
||||
int afterFirstPass = ranges.get();
|
||||
assertArrayEquals(new float[] {3f, 4f}, reader.get(1));
|
||||
assertEquals(afterFirstPass, ranges.get(), "repeated access must use cached chunks");
|
||||
reader.prebuffer(); assertTrue(reader.isComplete());
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void rejectsMalformedPublishedMerkleReference() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "bad.fvec", new float[][] {{1f, 2f}}));
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/bad.fvec", exchange -> serve(exchange, payload, new AtomicInteger()));
|
||||
server.createContext("/bad.fvec.mref", exchange -> { byte[] invalid = new byte[] {1, 2, 3}; exchange.sendResponseHeaders(200, invalid.length); exchange.getResponseBody().write(invalid); exchange.close(); });
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/bad.fvec");
|
||||
assertThrows(VectorDataException.class, () -> VectorReaders.f32(uri, VectorDataSettings.builder().cacheDirectory(temporary.resolve("bad-cache")).build(), "bad"));
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void fallsBackToOneFullTransferWhenRangesAreUnavailable() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "full.fvec", new float[][] {{1f, 2f}, {3f, 4f}}));
|
||||
AtomicInteger fullRequests = new AtomicInteger();
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/full.fvec", exchange -> {
|
||||
if (exchange.getRequestURI().getPath().endsWith(".mref")) { exchange.sendResponseHeaders(404, -1); exchange.close(); return; }
|
||||
if ("HEAD".equals(exchange.getRequestMethod())) { exchange.getResponseHeaders().add("Content-Length", String.valueOf(payload.length)); exchange.sendResponseHeaders(200, -1); exchange.close(); return; }
|
||||
fullRequests.incrementAndGet(); exchange.sendResponseHeaders(200, payload.length); exchange.getResponseBody().write(payload); exchange.close();
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/full.fvec");
|
||||
VectorReader<float[]> reader = VectorReaders.f32(uri, VectorDataSettings.builder().cacheDirectory(temporary.resolve("full-cache")).build(), "full");
|
||||
assertArrayEquals(new float[] {3f, 4f}, reader.get(1));
|
||||
assertEquals(AccessMode.FULL_TRANSFER, reader.cacheStats().accessMode());
|
||||
assertEquals(1, fullRequests.get());
|
||||
assertArrayEquals(new float[] {1f, 2f}, reader.get(0));
|
||||
assertEquals(1, fullRequests.get());
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void appliesBearerTokenToReferenceAndDataRequests() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "secure.fvec", new float[][] {{1f, 2f}}));
|
||||
byte[] mref = FixtureSupport.mref(payload, payload.length); AtomicInteger authorized = new AtomicInteger();
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/secure.fvec", exchange -> {
|
||||
if (!"Bearer test-token".equals(exchange.getRequestHeaders().getFirst("Authorization"))) { exchange.sendResponseHeaders(401, -1); exchange.close(); return; }
|
||||
authorized.incrementAndGet(); serve(exchange, payload, new AtomicInteger());
|
||||
});
|
||||
server.createContext("/secure.fvec.mref", exchange -> {
|
||||
if (!"Bearer test-token".equals(exchange.getRequestHeaders().getFirst("Authorization"))) { exchange.sendResponseHeaders(401, -1); exchange.close(); return; }
|
||||
authorized.incrementAndGet(); exchange.sendResponseHeaders(200, mref.length); exchange.getResponseBody().write(mref); exchange.close();
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/secure.fvec");
|
||||
VectorDataSettings settings = VectorDataSettings.builder().cacheDirectory(temporary.resolve("secure-cache")).bearerToken("test-token").build();
|
||||
assertArrayEquals(new float[] {1f, 2f}, VectorReaders.f32(uri, settings, "secure").get(0));
|
||||
assertTrue(authorized.get() >= 3, "mref, HEAD probe, and range read must be authorized");
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void retriesTransientRangeFailure() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "retry.fvec", new float[][] {{1f, 2f}}));
|
||||
AtomicInteger attempts = new AtomicInteger();
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/retry.fvec", exchange -> {
|
||||
if (exchange.getRequestURI().getPath().endsWith(".mref")) { exchange.sendResponseHeaders(404, -1); exchange.close(); return; }
|
||||
if ("HEAD".equals(exchange.getRequestMethod())) { exchange.getResponseHeaders().add("Content-Length", String.valueOf(payload.length)); exchange.getResponseHeaders().add("Accept-Ranges", "bytes"); exchange.sendResponseHeaders(200, -1); exchange.close(); return; }
|
||||
if (attempts.getAndIncrement() == 0) { exchange.sendResponseHeaders(503, -1); exchange.close(); return; }
|
||||
serve(exchange, payload, new AtomicInteger());
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/retry.fvec");
|
||||
assertArrayEquals(new float[] {1f, 2f}, VectorReaders.f32(uri, VectorDataSettings.builder().cacheDirectory(temporary.resolve("retry-cache")).build(), "retry").get(0));
|
||||
assertTrue(attempts.get() >= 2);
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void fallsBackWhenOptionalMerkleSidecarIsForbidden() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "forbidden.fvec", new float[][] {{1f, 2f}}));
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/forbidden.fvec", exchange -> {
|
||||
if (exchange.getRequestURI().getPath().endsWith(".mref")) { exchange.sendResponseHeaders(403, -1); exchange.close(); return; }
|
||||
serve(exchange, payload, new AtomicInteger());
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/forbidden.fvec");
|
||||
VectorReader<float[]> reader = VectorReaders.f32(uri, VectorDataSettings.builder().cacheDirectory(temporary.resolve("forbidden-cache")).build(), "forbidden");
|
||||
assertArrayEquals(new float[] {1f, 2f}, reader.get(0));
|
||||
assertEquals(AccessMode.MERKLE_CHUNKED, reader.cacheStats().accessMode());
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void httpDirectoryFallsBackToLegacyEntries() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "http-base.fvec", new float[][] {{9f, 10f}}));
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/remote/knn_entries.yaml", exchange -> {
|
||||
byte[] yaml = "\"remote:default\":\n base: base.fvec\n".getBytes(java.nio.charset.StandardCharsets.UTF_8);
|
||||
exchange.sendResponseHeaders(200, yaml.length); exchange.getResponseBody().write(yaml); exchange.close();
|
||||
});
|
||||
server.createContext("/remote/base.fvec", exchange -> {
|
||||
if (exchange.getRequestURI().getPath().endsWith(".mref")) { exchange.sendResponseHeaders(404, -1); exchange.close(); return; }
|
||||
serve(exchange, payload, new AtomicInteger());
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI root = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/remote");
|
||||
TestDataView view = TestDataGroup.load(root, VectorDataSettings.builder().cacheDirectory(temporary.resolve("http-legacy-cache")).build()).profile("default");
|
||||
assertArrayEquals(new float[] {9f, 10f}, view.baseVectors().get(0));
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
@Test void concurrentReadsDeduplicateTheMissingChunkFetch() throws Exception {
|
||||
byte[] payload = Files.readAllBytes(FixtureSupport.fvec(temporary, "concurrent.fvec", new float[][] {{1f, 2f}}));
|
||||
AtomicInteger ranges = new AtomicInteger();
|
||||
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
|
||||
server.createContext("/concurrent.fvec", exchange -> {
|
||||
if (exchange.getRequestURI().getPath().endsWith(".mref")) { exchange.sendResponseHeaders(404, -1); exchange.close(); return; }
|
||||
serve(exchange, payload, ranges);
|
||||
});
|
||||
server.start();
|
||||
try {
|
||||
URI uri = URI.create("http://127.0.0.1:" + server.getAddress().getPort() + "/concurrent.fvec");
|
||||
VectorReader<float[]> reader = VectorReaders.f32(uri, VectorDataSettings.builder().cacheDirectory(temporary.resolve("concurrent-cache")).build(), "concurrent");
|
||||
java.util.concurrent.ExecutorService pool = java.util.concurrent.Executors.newFixedThreadPool(8);
|
||||
try {
|
||||
var futures = java.util.stream.IntStream.range(0, 16).mapToObj(ignored -> pool.submit(() -> reader.get(0))).toList();
|
||||
for (var future : futures) assertArrayEquals(new float[] {1f, 2f}, future.get());
|
||||
} finally { pool.shutdownNow(); }
|
||||
assertEquals(1, ranges.get(), "all concurrent reads must share the first cached chunk");
|
||||
} finally { server.stop(0); }
|
||||
}
|
||||
private static void serve(HttpExchange exchange, byte[] payload, AtomicInteger ranges) throws IOException {
|
||||
if ("HEAD".equals(exchange.getRequestMethod())) { exchange.getResponseHeaders().add("Content-Length", String.valueOf(payload.length)); exchange.getResponseHeaders().add("Accept-Ranges", "bytes"); exchange.sendResponseHeaders(200, -1); exchange.close(); return; }
|
||||
String range = exchange.getRequestHeaders().getFirst("Range");
|
||||
if (range != null) {
|
||||
ranges.incrementAndGet(); String[] bounds = range.substring("bytes=".length()).split("-"); int from = Integer.parseInt(bounds[0]); int to = Integer.parseInt(bounds[1]);
|
||||
byte[] section = java.util.Arrays.copyOfRange(payload, from, to + 1); exchange.getResponseHeaders().add("Content-Range", "bytes " + from + "-" + to + "/" + payload.length); exchange.sendResponseHeaders(206, section.length); exchange.getResponseBody().write(section);
|
||||
} else { exchange.sendResponseHeaders(200, payload.length); exchange.getResponseBody().write(payload); }
|
||||
exchange.close();
|
||||
}
|
||||
}
|
||||
+40
@@ -0,0 +1,40 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import static org.junit.jupiter.api.Assumptions.assumeTrue;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
/**
|
||||
* Optional release canary for a dataset published by vectordata-rs.
|
||||
* Set {@code vectordata.canary.catalog}, {@code vectordata.canary.dataset},
|
||||
* and optionally {@code vectordata.canary.profile} to enable it.
|
||||
*/
|
||||
class RustHostedCanaryIntegrationTest {
|
||||
@Test void opensConfiguredRustHostedDataset() {
|
||||
String catalog = System.getProperty("vectordata.canary.catalog");
|
||||
String dataset = System.getProperty("vectordata.canary.dataset");
|
||||
assumeTrue(catalog != null && !catalog.isBlank() && dataset != null && !dataset.isBlank(),
|
||||
"Configure vectordata.canary.catalog and vectordata.canary.dataset to enable the Rust-hosted canary");
|
||||
String profile = System.getProperty("vectordata.canary.profile", "default");
|
||||
String cache = System.getProperty("vectordata.canary.cache", System.getProperty("java.io.tmpdir") + "/nb-vectordata-canary-cache");
|
||||
VectorDataSettings settings = VectorDataSettings.builder().cacheDirectory(java.nio.file.Path.of(cache)).build();
|
||||
TestDataView view = Catalog.of(CatalogSources.of(java.net.URI.create(catalog)), settings).open(dataset, profile);
|
||||
assertTrue(view.baseVectors().count() > 0, "canary base vectors must be non-empty");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
/** Settings tests use explicit temporary roots and never inspect or write user configuration. */
|
||||
class VectorDataSettingsTest {
|
||||
@TempDir Path temporary;
|
||||
|
||||
@Test void configuredCacheDirectoryWinsWithoutChangingSettings() throws Exception {
|
||||
Path settings = temporary.resolve("config/settings.yaml");
|
||||
Path configured = temporary.resolve("configured-cache");
|
||||
Files.createDirectories(settings.getParent());
|
||||
Files.writeString(settings, "# retain this comment\ncache_dir: " + configured + "\nchunk_size: 8192\n");
|
||||
|
||||
assertEquals(configured.toAbsolutePath(), VectorDataSettings.resolveCacheDirectory(
|
||||
settings, temporary.resolve("home"), null, temporary.resolve("xdg"), true));
|
||||
assertEquals("# retain this comment\ncache_dir: " + configured + "\nchunk_size: 8192\n", Files.readString(settings));
|
||||
}
|
||||
|
||||
@Test void vectordataHomeProvidesIsolatedDefaultWithoutWritingSettings() {
|
||||
Path root = temporary.resolve("isolated-vectordata-home");
|
||||
Path settings = root.resolve("settings.yaml");
|
||||
|
||||
assertEquals(root.resolve("cache"), VectorDataSettings.resolveCacheDirectory(
|
||||
settings, temporary.resolve("home"), root, temporary.resolve("xdg"), true));
|
||||
assertFalse(Files.exists(settings));
|
||||
assertFalse(Files.exists(root.resolve("cache")), "cache creation remains lazy until a remote reader opens");
|
||||
}
|
||||
|
||||
@Test void autoBootstrapWriterCreatesDirectoryAndRustCompatibleSettingsFile() throws Exception {
|
||||
Path settings = temporary.resolve("config/settings.yaml");
|
||||
Path cache = temporary.resolve("cache/vectordata");
|
||||
|
||||
VectorDataSettings.writeCacheDirectory(settings, cache, false);
|
||||
|
||||
assertTrue(Files.isDirectory(cache));
|
||||
assertEquals("cache_dir: " + cache + "\nprotect_settings: true\n", Files.readString(settings));
|
||||
assertEquals(cache.toAbsolutePath(), VectorDataSettings.resolveCacheDirectory(
|
||||
settings, temporary.resolve("home"), null, temporary.resolve("xdg"), false));
|
||||
}
|
||||
|
||||
@Test void absentConfigurationBootstrapsTheSafeXdgFallback() throws Exception {
|
||||
Path settings = temporary.resolve("config/settings.yaml");
|
||||
Path xdg = temporary.resolve("xdg");
|
||||
Path expected = xdg.resolve("vectordata");
|
||||
|
||||
// A non-existent home makes mount discovery unavailable deterministically;
|
||||
// both clients then use and persist the XDG fallback.
|
||||
assertEquals(expected, VectorDataSettings.resolveCacheDirectory(
|
||||
settings, temporary.resolve("unavailable-home"), null, xdg, true));
|
||||
assertTrue(Files.isDirectory(expected));
|
||||
assertEquals("cache_dir: " + expected + "\nprotect_settings: true\n", Files.readString(settings));
|
||||
}
|
||||
|
||||
@Test void automaticBootstrapRefusesToChangeAnExistingSettingsFileWithoutCacheDir() throws Exception {
|
||||
Path settings = temporary.resolve("config/settings.yaml");
|
||||
Files.createDirectories(settings.getParent());
|
||||
Files.writeString(settings, "token: preserved\n");
|
||||
|
||||
assertThrows(VectorDataException.class, () -> VectorDataSettings.writeCacheDirectory(
|
||||
settings, temporary.resolve("cache"), false));
|
||||
assertEquals("token: preserved\n", Files.readString(settings));
|
||||
}
|
||||
|
||||
@Test void explicitCallerCacheDoesNotRequireBootstrapPersistence() {
|
||||
Path settings = temporary.resolve("config/settings.yaml");
|
||||
Path fallback = temporary.resolve("xdg/vectordata");
|
||||
|
||||
assertEquals(fallback, VectorDataSettings.resolveCacheDirectory(
|
||||
settings, temporary.resolve("home"), null, temporary.resolve("xdg"), false));
|
||||
assertFalse(Files.exists(settings));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
/*
|
||||
* Copyright (c) 2026 The NoSQLBench Authors.
|
||||
*
|
||||
* 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.nosqlbench.vectordata;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
class VectorReadersTest {
|
||||
@TempDir Path temporary;
|
||||
@Test void readsFixedVectorsAndCopiesToCallerArray() throws Exception {
|
||||
Path source = FixtureSupport.fvec(temporary, "base.fvec", new float[][] {{1.5f, 2f}, {3f, 4.5f}});
|
||||
@SuppressWarnings("unchecked") VectorReader<float[]> reader = (VectorReader<float[]>) VectorReaders.open(source.toString());
|
||||
assertEquals(2, reader.count()); assertEquals(2, reader.dimension()); assertArrayEquals(new float[] {3f, 4.5f}, reader.get(1));
|
||||
float[] target = new float[2]; reader.get(0, target); assertArrayEquals(new float[] {1.5f, 2f}, target);
|
||||
assertTrue(reader.isComplete()); assertEquals(AccessMode.LOCAL, reader.cacheStats().accessMode());
|
||||
}
|
||||
@Test void readsPackedScalarsAndVariableVectors() throws Exception {
|
||||
Path scalar = FixtureSupport.scalarI32(temporary, "values.i32", -1, 7);
|
||||
VectorReader<?> scalars = VectorReaders.open(scalar.toString());
|
||||
assertEquals(-1, scalars.get(0)); assertEquals(7, scalars.get(1));
|
||||
TypedReader typed = TypedReaders.open(scalar.toUri(), VectorDataSettings.defaults(), "test");
|
||||
assertEquals(-1, typed.scalar(0)); assertEquals(7, typed.scalar(1));
|
||||
Path variable = FixtureSupport.vvec(temporary, "metadata.ivecs", new int[][] {{1}, {2, 3}});
|
||||
@SuppressWarnings("unchecked") VvecReader<int[]> vectors = (VvecReader<int[]>) VectorReaders.openVvec(variable.toUri(), VectorDataSettings.defaults(), "test");
|
||||
assertEquals(2, vectors.count()); assertEquals(2, vectors.dimensionAt(1)); assertArrayEquals(new int[] {2, 3}, vectors.get(1));
|
||||
Path variable64 = FixtureSupport.vvec64(temporary, "metadata64.ivecs", new int[][] {{4}, {5, 6}});
|
||||
@SuppressWarnings("unchecked") VvecReader<int[]> vectors64 = (VvecReader<int[]>) VectorReaders.openVvec(variable64.toUri(), VectorDataSettings.defaults(), "test");
|
||||
assertArrayEquals(new int[] {5, 6}, vectors64.get(1));
|
||||
}
|
||||
@Test void rejectsMalformedFixedRecord() throws Exception {
|
||||
Path source = FixtureSupport.fvec(temporary, "bad.fvec", new float[][] {{1f, 2f}});
|
||||
java.nio.file.Files.write(source, new byte[] {0}, java.nio.file.StandardOpenOption.APPEND);
|
||||
assertThrows(VectorDataException.class, () -> VectorReaders.open(source.toString()));
|
||||
}
|
||||
@Test void preservesUnsigned64ScalarValues() throws Exception {
|
||||
java.nio.ByteBuffer bytes = java.nio.ByteBuffer.allocate(8).order(java.nio.ByteOrder.LITTLE_ENDIAN).putLong(-1L);
|
||||
Path source = temporary.resolve("largest.u64"); java.nio.file.Files.write(source, bytes.array());
|
||||
assertEquals("18446744073709551615", VectorReaders.open(source.toString()).get(0).toString());
|
||||
assertEquals("18446744073709551615", TypedReaders.open(source.toUri(), VectorDataSettings.defaults(), "test").scalar(0).toString());
|
||||
}
|
||||
@Test void projectsF16AndF64Vectors() throws Exception {
|
||||
java.nio.ByteBuffer half = java.nio.ByteBuffer.allocate(8).order(java.nio.ByteOrder.LITTLE_ENDIAN);
|
||||
half.putInt(2).putShort((short) 0x3c00).putShort((short) 0xc100); // 1.0, -2.5
|
||||
Path f16 = temporary.resolve("half.mvec"); java.nio.file.Files.write(f16, half.array());
|
||||
assertArrayEquals(new Number[] {1.0f, -2.5f}, TypedReaders.open(f16.toUri(), VectorDataSettings.defaults(), "test").vector(0));
|
||||
java.nio.ByteBuffer doubles = java.nio.ByteBuffer.allocate(20).order(java.nio.ByteOrder.LITTLE_ENDIAN);
|
||||
doubles.putInt(2).putDouble(1.25).putDouble(-3.5);
|
||||
Path f64 = temporary.resolve("double.dvec"); java.nio.file.Files.write(f64, doubles.array());
|
||||
assertArrayEquals(new Number[] {1.25d, -3.5d}, TypedReaders.open(f64.toUri(), VectorDataSettings.defaults(), "test").vector(0));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
datasets:
|
||||
- name: rust-v1-demo
|
||||
path: demo/dataset.yaml
|
||||
dataset_type: dataset.yaml
|
||||
Binary file not shown.
@@ -0,0 +1,6 @@
|
||||
name: rust-v1-demo
|
||||
profiles:
|
||||
default:
|
||||
base: base.fvec
|
||||
query: query.fvec
|
||||
neighbor_indices: neighbors.ivecs
|
||||
Binary file not shown.
Binary file not shown.
@@ -55,8 +55,8 @@
|
||||
|
||||
<dependency>
|
||||
<groupId>io.nosqlbench</groupId>
|
||||
<artifactId>datatools-vectordata</artifactId>
|
||||
<version>0.1.25</version>
|
||||
<artifactId>nb-vectordata</artifactId>
|
||||
<version>${revision}</version>
|
||||
</dependency>
|
||||
|
||||
<!-- <dependency>-->
|
||||
|
||||
+13
-16
@@ -22,9 +22,10 @@ import io.nosqlbench.nb.annotations.Service;
|
||||
import io.nosqlbench.nb.api.expr.ExprFunctionProvider;
|
||||
import io.nosqlbench.nb.api.expr.annotations.ExprExample;
|
||||
import io.nosqlbench.nb.api.expr.annotations.ExprFunctionSpec;
|
||||
import io.nosqlbench.vectordata.VectorTestData;
|
||||
import io.nosqlbench.vectordata.discovery.vector.TestDataView;
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.*;
|
||||
import io.nosqlbench.vectordata.Catalog;
|
||||
import io.nosqlbench.vectordata.CatalogSources;
|
||||
import io.nosqlbench.vectordata.TestDataView;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import org.apache.logging.log4j.LogManager;
|
||||
import org.apache.logging.log4j.Logger;
|
||||
|
||||
@@ -43,7 +44,7 @@ public class VectorDataExprs implements ExprFunctionProvider {
|
||||
if (!datasetNameAndProfile.contains(":")) {
|
||||
logger.warn("datasetNameAndProfile missing profile:" + datasetNameAndProfile);
|
||||
}
|
||||
return VectorTestData.catalogs().configure().catalog().profile(datasetNameAndProfile);
|
||||
return Catalog.of(CatalogSources.defaults()).openProfile(datasetNameAndProfile);
|
||||
}
|
||||
@ExprExample(args = {"\"airports:demo\""}, expectNotNull = true)
|
||||
@ExprExample(args = {"\"airports:demo\""}, matches = ".+" )
|
||||
@@ -52,9 +53,8 @@ public class VectorDataExprs implements ExprFunctionProvider {
|
||||
synopsis = "baseVectors(\"dataset:profile\")",
|
||||
description = "Return the BaseVectors associated with the dataset profile."
|
||||
)
|
||||
public BaseVectors baseVectors(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).getBaseVectors().orElseThrow(() ->
|
||||
new RuntimeException("Base vectors are not defined for dataset profile '" + datasetNameAndProfile + "'"));
|
||||
public VectorReader<float[]> baseVectors(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).baseVectors();
|
||||
}
|
||||
|
||||
@ExprExample(args = {"\"airports:demo\""}, expectNotNull = true)
|
||||
@@ -64,9 +64,8 @@ public class VectorDataExprs implements ExprFunctionProvider {
|
||||
synopsis = "queryVectors(\"dataset:profile\")",
|
||||
description = "Return the QueryVectors associated with the dataset profile."
|
||||
)
|
||||
public QueryVectors queryVectors(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).getQueryVectors().orElseThrow(() ->
|
||||
new RuntimeException("Query vectors are not defined for dataset profile '" + datasetNameAndProfile + "'"));
|
||||
public VectorReader<float[]> queryVectors(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).queryVectors();
|
||||
}
|
||||
|
||||
@ExprExample(args = {"\"airports:demo\""}, expectNotNull = true)
|
||||
@@ -76,9 +75,8 @@ public class VectorDataExprs implements ExprFunctionProvider {
|
||||
synopsis = "neighborDistances(\"dataset:profile\")",
|
||||
description = "Return the NeighborDistances associated with the dataset profile."
|
||||
)
|
||||
public NeighborDistances neighborDistances(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).getNeighborDistances().orElseThrow(() ->
|
||||
new RuntimeException("Neighbor distances are not defined for dataset profile '" + datasetNameAndProfile + "'"));
|
||||
public VectorReader<float[]> neighborDistances(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).neighborDistances();
|
||||
}
|
||||
|
||||
@ExprExample(args = {"\"airports:demo\""}, expectNotNull = true)
|
||||
@@ -88,9 +86,8 @@ public class VectorDataExprs implements ExprFunctionProvider {
|
||||
synopsis = "neighborIndices(\"dataset:profile\")",
|
||||
description = "Return the NeighborIndices associated with the dataset profile."
|
||||
)
|
||||
public NeighborIndices neighborIndices(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).getNeighborIndices().orElseThrow(() ->
|
||||
new RuntimeException("Neighbor indices are not defined for dataset profile '" + datasetNameAndProfile + "'"));
|
||||
public VectorReader<int[]> neighborIndices(String datasetNameAndProfile) {
|
||||
return dataset(datasetNameAndProfile).neighborIndices();
|
||||
}
|
||||
|
||||
}
|
||||
+9
-3
@@ -18,7 +18,8 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.FloatVectors;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.virtdata.api.annotations.Categories;
|
||||
import io.nosqlbench.virtdata.api.annotations.Category;
|
||||
import io.nosqlbench.virtdata.api.annotations.Example;
|
||||
@@ -42,9 +43,14 @@ public class BaseVectors extends CoreVectors<float[]> {
|
||||
super (datasetAndProfile, prebuffer);
|
||||
}
|
||||
|
||||
/** Construct with explicit vectordata settings, including an isolated cache location. */
|
||||
public BaseVectors(String datasetAndProfile, boolean prebuffer, VectorDataSettings settings) {
|
||||
super(datasetAndProfile, prebuffer, settings);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected FloatVectors getRandomAccessData() {
|
||||
return super.tdv.getBaseVectors().orElseThrow(()->new RuntimeException("Cannot get base vectors from data view " + super.tdv.getName()));
|
||||
protected VectorReader<float[]> getRandomAccessData() {
|
||||
return super.tdv.baseVectors();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
+14
-16
@@ -18,35 +18,33 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import io.nosqlbench.nbdatatools.api.concurrent.ProgressIndicatingFuture;
|
||||
import io.nosqlbench.vectordata.VectorTestData;
|
||||
import io.nosqlbench.vectordata.discovery.ProfileSelector;
|
||||
import io.nosqlbench.vectordata.discovery.vector.TestDataView;
|
||||
import io.nosqlbench.vectordata.downloader.Catalog;
|
||||
import io.nosqlbench.vectordata.downloader.DatasetEntry;
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.DatasetView;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import io.nosqlbench.vectordata.Catalog;
|
||||
import io.nosqlbench.vectordata.CatalogSources;
|
||||
import io.nosqlbench.vectordata.TestDataView;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import java.util.function.LongFunction;
|
||||
|
||||
public abstract class CoreVectors<T> implements LongFunction<T> {
|
||||
|
||||
protected final TestDataView tdv;
|
||||
protected final DatasetView<T> dataset;
|
||||
protected final VectorReader<T> dataset;
|
||||
|
||||
public CoreVectors(String datasetAndProfile, boolean prebuffer) {
|
||||
Catalog catalog = VectorTestData.catalogs().configure().catalog();
|
||||
tdv = catalog.profile(datasetAndProfile);
|
||||
this(datasetAndProfile, prebuffer, VectorDataSettings.defaults());
|
||||
}
|
||||
|
||||
protected CoreVectors(String datasetAndProfile, boolean prebuffer, VectorDataSettings settings) {
|
||||
Catalog catalog = Catalog.of(CatalogSources.defaults(), settings);
|
||||
tdv = catalog.openProfile(datasetAndProfile);
|
||||
dataset = getRandomAccessData();
|
||||
|
||||
if (prebuffer) {
|
||||
CompletableFuture<Void> pbfuture = dataset.prebuffer();
|
||||
if (pbfuture instanceof ProgressIndicatingFuture<Void> indicator) {
|
||||
indicator.monitorProgress(1000);
|
||||
}
|
||||
dataset.prebuffer();
|
||||
}
|
||||
}
|
||||
|
||||
protected abstract DatasetView<T> getRandomAccessData();
|
||||
protected abstract VectorReader<T> getRandomAccessData();
|
||||
|
||||
@Override
|
||||
public T apply(long value) {
|
||||
|
||||
+9
-3
@@ -18,7 +18,8 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.FloatVectors;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import io.nosqlbench.virtdata.api.annotations.Categories;
|
||||
import io.nosqlbench.virtdata.api.annotations.Category;
|
||||
import io.nosqlbench.virtdata.api.annotations.Example;
|
||||
@@ -39,9 +40,14 @@ public class NeighborDistances extends CoreVectors<float[]> {
|
||||
super (datasetAndProfile, prebuffer);
|
||||
}
|
||||
|
||||
/** Construct with explicit vectordata settings, including an isolated cache location. */
|
||||
public NeighborDistances(String datasetAndProfile, boolean prebuffer, VectorDataSettings settings) {
|
||||
super(datasetAndProfile, prebuffer, settings);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected FloatVectors getRandomAccessData() {
|
||||
return super.tdv.getNeighborDistances().orElseThrow(()->new RuntimeException("Cannot get neighbor distances from data view " + super.tdv.getName()));
|
||||
protected VectorReader<float[]> getRandomAccessData() {
|
||||
return super.tdv.neighborDistances();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
+9
-4
@@ -18,8 +18,8 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.FloatVectors;
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.IntVectors;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import io.nosqlbench.virtdata.api.annotations.Categories;
|
||||
import io.nosqlbench.virtdata.api.annotations.Category;
|
||||
import io.nosqlbench.virtdata.api.annotations.Example;
|
||||
@@ -40,9 +40,14 @@ public class NeighborIndices extends CoreVectors<int[]> {
|
||||
super (datasetAndProfile, prebuffer);
|
||||
}
|
||||
|
||||
/** Construct with explicit vectordata settings, including an isolated cache location. */
|
||||
public NeighborIndices(String datasetAndProfile, boolean prebuffer, VectorDataSettings settings) {
|
||||
super(datasetAndProfile, prebuffer, settings);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected IntVectors getRandomAccessData() {
|
||||
return super.tdv.getNeighborIndices().orElseThrow(()->new RuntimeException("Cannot get neighbor indices from data view " + super.tdv.getName()));
|
||||
protected VectorReader<int[]> getRandomAccessData() {
|
||||
return super.tdv.neighborIndices();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
+9
-3
@@ -18,7 +18,8 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import io.nosqlbench.vectordata.spec.datasets.types.FloatVectors;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import io.nosqlbench.vectordata.VectorReader;
|
||||
import io.nosqlbench.virtdata.api.annotations.Categories;
|
||||
import io.nosqlbench.virtdata.api.annotations.Category;
|
||||
import io.nosqlbench.virtdata.api.annotations.Example;
|
||||
@@ -39,10 +40,15 @@ public class QueryVectors extends CoreVectors<float[]> {
|
||||
super (datasetAndProfile, prebuffer);
|
||||
}
|
||||
|
||||
/** Construct with explicit vectordata settings, including an isolated cache location. */
|
||||
public QueryVectors(String datasetAndProfile, boolean prebuffer, VectorDataSettings settings) {
|
||||
super(datasetAndProfile, prebuffer, settings);
|
||||
}
|
||||
|
||||
|
||||
@Override
|
||||
protected FloatVectors getRandomAccessData() {
|
||||
return super.tdv.getQueryVectors().orElseThrow(()->new RuntimeException("Cannot get query vectors from data view " + super.tdv.getName()));
|
||||
protected VectorReader<float[]> getRandomAccessData() {
|
||||
return super.tdv.queryVectors();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
+34
-7
@@ -18,22 +18,49 @@ package io.nosqlbench.virtdata.lib.vectors.vectordata;
|
||||
*/
|
||||
|
||||
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import io.nosqlbench.vectordata.VectorDataSettings;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Tag;
|
||||
import org.junit.jupiter.api.io.TempDir;
|
||||
|
||||
import java.util.Arrays;
|
||||
import java.nio.ByteBuffer;
|
||||
import java.nio.ByteOrder;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
@Tag("unit")
|
||||
public class BaseVectorsTest {
|
||||
|
||||
@TempDir Path temporary;
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
public void testExampleDataset() {
|
||||
BaseVectors bv1 = new BaseVectors("example1");
|
||||
float[] v23 = bv1.apply(23L);
|
||||
System.out.println("v23:" + Arrays.toString(v23));
|
||||
public void testFixtureDataset() throws Exception {
|
||||
Path dataset = Files.createDirectories(temporary.resolve("example"));
|
||||
ByteBuffer values = ByteBuffer.allocate(24).order(ByteOrder.LITTLE_ENDIAN);
|
||||
values.putInt(2).putFloat(1f).putFloat(2f).putInt(2).putFloat(3f).putFloat(4f);
|
||||
Files.write(dataset.resolve("base.fvec"), values.array());
|
||||
Files.writeString(dataset.resolve("dataset.yaml"), """
|
||||
name: example
|
||||
profiles:
|
||||
demo:
|
||||
base: base.fvec
|
||||
""");
|
||||
Path catalog = temporary.resolve("catalog.yaml");
|
||||
Files.writeString(catalog, """
|
||||
datasets:
|
||||
- name: example
|
||||
path: example/dataset.yaml
|
||||
""");
|
||||
String prior = System.getProperty("vectordata.catalog");
|
||||
System.setProperty("vectordata.catalog", catalog.toString());
|
||||
try {
|
||||
BaseVectors vectors = new BaseVectors("example:demo", true,
|
||||
VectorDataSettings.builder().cacheDirectory(temporary.resolve("cache")).build());
|
||||
assertArrayEquals(new float[] {3f, 4f}, vectors.apply(1));
|
||||
} finally {
|
||||
if (prior == null) System.clearProperty("vectordata.catalog"); else System.setProperty("vectordata.catalog", prior);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user