From 1ccdbb086fd0367e11f74532ef5da5539df3a9b9 Mon Sep 17 00:00:00 2001 From: Nadav Har'El Date: Wed, 6 May 2026 22:07:40 +0300 Subject: [PATCH 1/2] Add Alternator vector search support Alternator extends the DynamoDB API with vector indexes and vector similarity search. These features are not available on AWS DynamoDB, so the standard AWS SDK for Java has no knowledge of the new parameters - `VectorIndexes`, `VectorSearch`, `FLOAT32VECTOR`, etc. In this patch we introduce these features into the Alternator Java Clien library without modifying the SDK source. Our implementation uses an `ExecutionInterceptor` that intercepts the serialised JSON body of each request before it is sent, and of each response before the SDK parses it into Java objects: * For requests, extra parameters are injected into the JSON body before transmission. * For Responses, extra fields are extracted from the JSON body before the SDK discards them. The Java implementation is in a new package: com.scylladb.alternator.vectorsearch The new VectorSearchInterceptor is registered by AlternatorDynamoDbClient and AlternatorDynamoDbAsyncClient so no manual setup is needed when using the Alternator builders. Besides the implementation, this patch also includes documentation in README.md of how to create vector indexes and use them for vector searches in Alternator. And also tests: - VectorSearchInterceptorTest: 27 unit tests covering JSON serialisation, Float32Vector encoding/decoding, request/response interception, scores extraction, and the List round-trip. No server needed. - VectorSearchIT: integration tests requiring a live Alternator instance with vector store support (enabled via INTEGRATION_TESTS=true). Signed-off-by: Nadav Har'El --- README.md | 237 +++++++ pom.xml | 17 + .../scylladb/alternator/VectorSearchIT.java | 409 +++++++++++ .../AlternatorDynamoDbAsyncClient.java | 21 +- .../alternator/AlternatorDynamoDbClient.java | 21 +- .../vectorsearch/CreateVectorIndexAction.java | 97 +++ .../vectorsearch/DeleteVectorIndexAction.java | 50 ++ .../vectorsearch/Float32Vector.java | 213 ++++++ .../vectorsearch/VectorAttribute.java | 67 ++ .../alternator/vectorsearch/VectorIndex.java | 152 ++++ .../vectorsearch/VectorIndexUpdate.java | 81 +++ .../vectorsearch/VectorQueryResult.java | 67 ++ .../alternator/vectorsearch/VectorSearch.java | 119 ++++ .../vectorsearch/VectorSearchInterceptor.java | 663 ++++++++++++++++++ .../VectorSearchResultHolder.java | 33 + .../vectorsearch/VectorSearchSupport.java | 315 +++++++++ .../VectorSearchInterceptorTest.java | 367 ++++++++++ 17 files changed, 2909 insertions(+), 20 deletions(-) create mode 100644 src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/CreateVectorIndexAction.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/DeleteVectorIndexAction.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/Float32Vector.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorAttribute.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorIndex.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorIndexUpdate.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorQueryResult.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearch.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java create mode 100644 src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java diff --git a/README.md b/README.md index 6e90b12..e0224bf 100644 --- a/README.md +++ b/README.md @@ -779,3 +779,240 @@ Key route affinity may not be beneficial for: - Read-heavy workloads (reads don't use Paxos) - Workloads with uniformly distributed writes across many keys - Async clients (use sync clients only) + +### Vector Search (Alternator extension) + +Alternator extends the DynamoDB API with **vector indexes** and **vector similarity search**. +These features are not available on AWS DynamoDB, so the standard AWS SDK for Java has no +knowledge of the new parameters (`VectorIndexes`, `VectorSearch`, `FLOAT32VECTOR`, etc.). + +This library bridges the gap without modifying the SDK source. It uses an +`ExecutionInterceptor` that intercepts the serialised JSON body of each request before +it is sent, and of each response before the SDK parses it into Java objects: +- **Requests** — extra parameters are injected into the JSON body before transmission. +- **Responses** — extra fields are extracted from the JSON body before the SDK discards them. + +#### Setup + +`VectorSearchInterceptor` is registered automatically by `AlternatorDynamoDbClient` and +`AlternatorDynamoDbAsyncClient`. No extra configuration is needed: + +```java +DynamoDbClient client = AlternatorDynamoDbClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .credentialsProvider(myCredentials) + .build(); +``` + +If you use the plain `DynamoDbClient.builder()` / `DynamoDbAsyncClient.builder()` directly +(without the Alternator builder), register the interceptor manually: + +```java +import com.scylladb.alternator.vectorsearch.VectorSearchInterceptor; + +DynamoDbClient client = DynamoDbClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .credentialsProvider(myCredentials) + .overrideConfiguration(c -> c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE)) + .build(); +``` + +#### CreateTable with a vector index + +```java +import com.scylladb.alternator.vectorsearch.*; + +VectorIndex vi = VectorIndex.builder() + .indexName("embedding-index") + .vectorAttribute(VectorAttribute.builder() + .attributeName("embedding") + .dimensions(128) + .build()) + .similarityFunction("COSINE") // optional; "DOT_PRODUCT" and "EUCLIDEAN" are also supported + .build(); + +CreateTableRequest base = CreateTableRequest.builder() + .tableName("items") + .keySchema(KeySchemaElement.builder().attributeName("id").keyType(KeyType.HASH).build()) + .attributeDefinitions( + AttributeDefinition.builder().attributeName("id").attributeType(ScalarAttributeType.S).build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build(); + +VectorSearchSupport.createTable(client, base, List.of(vi)); +``` + +You can also use the lower-level helper if you want to call `client.createTable()` yourself: + +```java +client.createTable(VectorSearchSupport.withVectorIndexes(base, List.of(vi))); +``` + +#### UpdateTable — adding or removing a vector index + +```java +VectorIndexUpdate addIndex = VectorIndexUpdate.builder() + .create(CreateVectorIndexAction.builder() + .indexName("new-index") + .vectorAttribute(VectorAttribute.builder() + .attributeName("embedding") + .dimensions(64) + .build()) + .build()) + .build(); + +VectorIndexUpdate removeIndex = VectorIndexUpdate.builder() + .delete(DeleteVectorIndexAction.builder() + .indexName("old-index") + .build()) + .build(); + +VectorSearchSupport.updateTable(client, + UpdateTableRequest.builder().tableName("items").build(), + List.of(addIndex, removeIndex)); +``` + +You can also use the lower-level helper if you want to call `client.updateTable()` yourself: + +```java +client.updateTable(VectorSearchSupport.withVectorIndexUpdates( + UpdateTableRequest.builder().tableName("items").build(), + List.of(addIndex, removeIndex))); +``` + +#### DescribeTable — reading vector index metadata + +```java +VectorSearchSupport.DescribeTableWithVectorIndexes result = + VectorSearchSupport.describeTable(client, + DescribeTableRequest.builder().tableName("items").build()); + +for (VectorIndex vi : result.vectorIndexes()) { + System.out.printf("index=%s attribute=%s dimensions=%d status=%s%n", + vi.indexName(), + vi.vectorAttribute().attributeName(), + vi.vectorAttribute().dimensions(), + vi.indexStatus()); +} +``` + +#### Writing items with the optimized FLOAT32VECTOR type + +Alternator stores vectors in a compact binary format called `FLOAT32VECTOR`. This is +significantly more efficient than the standard DynamoDB list-of-numbers encoding (`L`) +and is recommended for use with a vector index. + +Use `Float32Vector.toAttributeValue(float...)` to create the attribute value. The +interceptor automatically converts it to `{"FLOAT32VECTOR": [...]}` in the JSON body: + +```java +import com.scylladb.alternator.vectorsearch.Float32Vector; + +Map item = new HashMap<>(); +item.put("id", AttributeValue.fromS("item-1")); +item.put("embedding", Float32Vector.toAttributeValue(0.1f, 0.2f, 0.3f, ...)); + +client.putItem(PutItemRequest.builder().tableName("items").item(item).build()); +``` + +This works transparently for all write operations: `PutItem`, `UpdateItem` +(`ExpressionAttributeValues`), and `BatchWriteItem`. + +#### Reading FLOAT32VECTOR attributes back + +When Alternator returns a `FLOAT32VECTOR` attribute, the interceptor converts it +transparently to a standard DynamoDB list-of-numbers (`L` type). Access the values +using the normal SDK API: + +```java +AttributeValue av = resp.item().get("embedding"); +List numbers = av.l(); // each element has type N +float[] values = new float[numbers.size()]; +for (int i = 0; i < values.length; i++) { + values[i] = Float.parseFloat(numbers.get(i).n()); +} +``` + +To write the vector back with the efficient `FLOAT32VECTOR` storage format, +re-encode it using the `List` overload of `Float32Vector.toAttributeValue()`: + +```java +AttributeValue reEncoded = Float32Vector.toAttributeValue(av.l()); +``` + +> **Warning — copying items:** When you read an item that contains a `FLOAT32VECTOR` +> attribute, the interceptor converts it to a standard DynamoDB `L` (list-of-numbers) +> `AttributeValue`. If you then pass that `AttributeValue` straight back in a write +> (e.g. `putItem(...item(response.item())...)`), the vector will be stored as a plain +> `L` list — **not** as `FLOAT32VECTOR`. To preserve the compact storage format you +> must explicitly re-encode every vector attribute before writing: +> +> ```java +> Map item = new HashMap<>(response.item()); +> item.put("embedding", Float32Vector.toAttributeValue(item.get("embedding").l())); +> client.putItem(PutItemRequest.builder().tableName("items").item(item).build()); +> ``` + +#### Vector similarity search (Query) + +```java +VectorSearch vs = VectorSearch.builder() + .queryVector(0.1f, 0.2f, 0.3f, ...) // sent as FLOAT32VECTOR + .returnScores(true) // optional; include similarity scores + .build(); + +VectorQueryResult result = VectorSearchSupport.query(client, + QueryRequest.builder() + .tableName("items") + .indexName("embedding-index") + .limit(10) + .build(), + vs); + +Standard `QueryRequest` parameters work alongside `VectorSearch`: `filterExpression`, +`select`, `projectionExpression`, `expressionAttributeValues`, and +`keyConditionExpression` are all supported. Note that `limit` is **required** for vector +search queries (it specifies how many nearest neighbours to return), and +`exclusiveStartKey` (pagination) is not supported. + +for (int i = 0; i < result.items().size(); i++) { + System.out.printf("item=%s score=%.4f%n", + result.items().get(i).get("id").s(), + result.scores().get(i)); +} +``` + +For the async client: + +```java +CompletableFuture future = + VectorSearchSupport.queryAsync(asyncClient, queryRequest, vs); +``` + +You can also pass the query vector as an `AttributeValue` (standard DynamoDB list type) +instead of a `float[]`, which is useful when the items were stored without the +`FLOAT32VECTOR` optimisation: + +```java +VectorSearch vs = VectorSearch.builder() + .queryVector(AttributeValue.fromL(Arrays.asList( + AttributeValue.fromN("0.1"), + AttributeValue.fromN("0.2"), + AttributeValue.fromN("0.3")))) + .build(); +``` + +#### Summary of the `vectorsearch` package + +| Class | Purpose | +|---|---| +| `VectorSearchInterceptor` | Core interceptor — register once on the client | +| `VectorSearchSupport` | Convenience facade for all operations | +| `VectorIndex` | Descriptor for a vector index (create/describe) | +| `VectorAttribute` | Attribute name + dimensionality for a `VectorIndex` | +| `VectorSearch` | Query parameters: query vector + optional score request | +| `VectorQueryResult` | Wraps `QueryResponse` and adds `scores()` | +| `VectorIndexUpdate` | A single create/delete change for `UpdateTable` | +| `CreateVectorIndexAction` | Create action inside a `VectorIndexUpdate` | +| `DeleteVectorIndexAction` | Delete action inside a `VectorIndexUpdate` | +| `Float32Vector` | Encode/decode `float...` as the compact `FLOAT32VECTOR` type | diff --git a/pom.xml b/pom.xml index 044e73b..93d10a4 100644 --- a/pom.xml +++ b/pom.xml @@ -27,6 +27,14 @@ pom import + + + com.fasterxml.jackson.core + jackson-databind + 2.15.4 + @@ -51,6 +59,15 @@ aws-crt-client provided + + + com.fasterxml.jackson.core + jackson-databind + provided + net.sourceforge.argparse4j argparse4j diff --git a/src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java b/src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java new file mode 100644 index 0000000..8dcf36a --- /dev/null +++ b/src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java @@ -0,0 +1,409 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator; + +import static org.junit.Assert.*; +import static org.junit.Assume.*; + +import com.scylladb.alternator.vectorsearch.*; +import java.net.URI; +import java.util.*; +import java.util.concurrent.TimeUnit; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.*; + +/** + * Integration tests for Alternator's vector search extension. + * + *

These tests verify that {@link VectorSearchInterceptor} and {@link VectorSearchSupport} + * correctly inject vector-search parameters into requests and extract them from responses. + * + *

Tests require a running Alternator cluster. Enable with {@code INTEGRATION_TESTS=true}. + */ +@RunWith(Parameterized.class) +public class VectorSearchIT { + + private final URI seedUri; + private DynamoDbClient client; + private String tableName; + + public VectorSearchIT(String scheme, URI seedUri) { + this.seedUri = seedUri; + } + + @Parameters(name = "{0}") + public static Collection data() { + return IntegrationTestConfig.httpAndHttpsEndpoints(); + } + + @Before + public void setUp() { + assumeTrue( + "Integration tests disabled. Set INTEGRATION_TESTS=true to enable.", + IntegrationTestConfig.ENABLED); + + client = + AlternatorDynamoDbClient.builder() + .endpointOverride(seedUri) + .credentialsProvider(IntegrationTestConfig.CREDENTIALS) + .build(); + + tableName = "vs_it_" + UUID.randomUUID().toString().replace("-", "").substring(0, 8); + } + + @After + public void tearDown() { + if (client != null) { + try { + client.deleteTable(DeleteTableRequest.builder().tableName(tableName).build()); + } catch (ResourceNotFoundException ignored) { + } + client.close(); + } + } + + // ------------------------------------------------------------------------- + // Helper: create a plain hash-only table with a single vector index + // ------------------------------------------------------------------------- + + private static final int DIMENSIONS = 4; + + /** + * Skips the test if the server does not support vector indexes. Must be called after + * {@link #createTableWithVectorIndex()} so that the table exists. + * + *

Two cases are detected: + * + *

    + *
  1. The server predates vector-store support entirely: it ignores {@code VectorIndexes} in + * {@code CreateTable}, so {@code DescribeTable} returns an empty vector-indexes list. + *
  2. The server has vector-store code but it is disabled: it acknowledges the index in + * {@code DescribeTable} but a {@code VectorSearch} query returns "Vector Store is + * disabled". + *
+ * + *

Mirrors the {@code needs_vector_store} fixture in Scylla's Python tests. + */ + private void assumeVectorStoreEnabled() { + // Case 1: server does not know about vector indexes at all. + VectorSearchSupport.DescribeTableWithVectorIndexes desc = + VectorSearchSupport.describeTable( + client, DescribeTableRequest.builder().tableName(tableName).build()); + assumeTrue( + "Skipping: server does not support vector indexes (VectorIndexes absent in DescribeTable)", + !desc.vectorIndexes().isEmpty()); + + // Case 2: server has vector-store code but it is disabled. + try { + VectorSearchSupport.query( + client, + QueryRequest.builder().tableName(tableName).indexName("vi1").limit(1).build(), + VectorSearch.builder().queryVector(new float[] {0f, 0f, 0f, 0f}).build()); + } catch (DynamoDbException e) { + if (e.getMessage() != null && e.getMessage().contains("Vector Store is disabled")) { + assumeTrue("Skipping: Vector Store is disabled on this server", false); + } + // Any other error (e.g. index not yet active) means the vector store is reachable. + } + } + + /** Polls until the table is ACTIVE and the named vector index is ACTIVE (up to 60 seconds). */ + private void waitForVectorIndexActive(String indexName) throws InterruptedException { + for (int i = 0; i < 120; i++) { + VectorSearchSupport.DescribeTableWithVectorIndexes desc = + VectorSearchSupport.describeTable( + client, DescribeTableRequest.builder().tableName(tableName).build()); + if (desc.response().table().tableStatus() != TableStatus.ACTIVE) { + TimeUnit.MILLISECONDS.sleep(500); + continue; + } + boolean indexActive = desc.vectorIndexes().stream() + .anyMatch(vi -> indexName.equals(vi.indexName()) && "ACTIVE".equals(vi.indexStatus())); + if (indexActive) { + return; + } + TimeUnit.MILLISECONDS.sleep(500); + } + fail("Timed out waiting for vector index '" + indexName + "' to become ACTIVE on table '" + tableName + "'"); + } + + private void createTableWithVectorIndex() { + VectorIndex vi = + VectorIndex.builder() + .indexName("vi1") + .vectorAttribute( + VectorAttribute.builder() + .attributeName("embedding") + .dimensions(DIMENSIONS) + .build()) + .build(); + + CreateTableRequest base = + CreateTableRequest.builder() + .tableName(tableName) + .keySchema( + KeySchemaElement.builder().attributeName("pk").keyType(KeyType.HASH).build()) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("pk") + .attributeType(ScalarAttributeType.S) + .build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build(); + + VectorSearchSupport.createTable(client, base, Collections.singletonList(vi)); + } + + // ------------------------------------------------------------------------- + // Tests + // ------------------------------------------------------------------------- + + /** + * Verify that CreateTable with a VectorIndexes parameter is accepted by Alternator without + * throwing an exception. + */ + @Test + public void testCreateTableWithVectorIndex() throws Exception { + createTableWithVectorIndex(); + // Wait until ACTIVE — the vector index build may take a moment. + TableStatus status = TableStatus.UNKNOWN_TO_SDK_VERSION; + for (int i = 0; i < 20; i++) { + DescribeTableResponse desc = + client.describeTable(DescribeTableRequest.builder().tableName(tableName).build()); + status = desc.table().tableStatus(); + if (status == TableStatus.ACTIVE) { + break; + } + TimeUnit.MILLISECONDS.sleep(500); + } + assertEquals(TableStatus.ACTIVE, status); + } + + /** + * Verify that DescribeTable response VectorIndexes can be parsed back via + * VectorSearchSupport.describeTable. + */ + @Test + public void testDescribeTableWithVectorIndexes() throws Exception { + createTableWithVectorIndex(); + assumeVectorStoreEnabled(); + waitForVectorIndexActive("vi1"); + + VectorSearchSupport.DescribeTableWithVectorIndexes result = + VectorSearchSupport.describeTable( + client, DescribeTableRequest.builder().tableName(tableName).build()); + + assertNotNull(result.response()); + List vis = result.vectorIndexes(); + assertFalse("Expected at least one vector index in the response", vis.isEmpty()); + + VectorIndex vi = vis.get(0); + assertEquals("vi1", vi.indexName()); + assertEquals("embedding", vi.vectorAttribute().attributeName()); + assertEquals(DIMENSIONS, vi.vectorAttribute().dimensions()); + } + + /** + * Verify that a basic vector similarity Query (without ReturnScores) does not throw and returns + * a non-null result. + */ + @Test + public void testQueryWithVectorSearch() throws Exception { + createTableWithVectorIndex(); + assumeVectorStoreEnabled(); + waitForVectorIndexActive("vi1"); + for (int i = 0; i < 3; i++) { + Map item = new HashMap<>(); + item.put("pk", AttributeValue.fromS("item-" + i)); + // Store embedding as a regular DynamoDB list + List vec = new ArrayList<>(); + for (int d = 0; d < DIMENSIONS; d++) { + vec.add(AttributeValue.fromN(String.valueOf((float) (i + d)))); + } + item.put("embedding", AttributeValue.fromL(vec)); + client.putItem(PutItemRequest.builder().tableName(tableName).item(item).build()); + } + + // Query vector (closest to item-0's embedding) + VectorSearch vs = + VectorSearch.builder() + .queryVector(new float[] {0.0f, 1.0f, 2.0f, 3.0f}) + .build(); + + VectorQueryResult result = + VectorSearchSupport.query( + client, + QueryRequest.builder().tableName(tableName).indexName("vi1").limit(3).build(), + vs); + assertNotNull(result); + assertNotNull(result.items()); + } + + /** + * Verify that ReturnScores causes the server to include per-item scores in the response and that + * VectorSearchSupport.query parses them correctly. + */ + @Test + public void testQueryWithReturnScores() throws Exception { + createTableWithVectorIndex(); + assumeVectorStoreEnabled(); + waitForVectorIndexActive("vi1"); + + // Put an item + Map item = new HashMap<>(); + item.put("pk", AttributeValue.fromS("item-0")); + List vec = new ArrayList<>(); + for (int d = 0; d < DIMENSIONS; d++) { + vec.add(AttributeValue.fromN(String.valueOf((float) d))); + } + item.put("embedding", AttributeValue.fromL(vec)); + client.putItem(PutItemRequest.builder().tableName(tableName).item(item).build()); + + VectorSearch vs = + VectorSearch.builder() + .queryVector(new float[]{0.0f, 1.0f, 2.0f, 3.0f}) + .returnScores(true) + .build(); + + VectorQueryResult result = + VectorSearchSupport.query( + client, + QueryRequest.builder().tableName(tableName).indexName("vi1").limit(10).build(), + vs); + + assertNotNull(result); + assertNotNull(result.items()); + // If items were returned and ReturnScores was true, scores list should be non-empty + if (!result.items().isEmpty()) { + assertFalse("Expected scores to be returned when ReturnScores=true", result.scores().isEmpty()); + assertEquals(result.items().size(), result.scores().size()); + for (double score : result.scores()) { + // Scores should be finite, positive numbers + assertTrue("Score should be finite", Double.isFinite(score)); + } + } + } + + /** + * Verify that withVectorIndexes preserves any existing overrideConfiguration on the request. + */ + @Test + public void testWithVectorIndexesPreservesExistingOverrideConfiguration() { + // Just test that the helper method doesn't throw and returns a non-null request + VectorIndex vi = + VectorIndex.builder() + .indexName("idx") + .vectorAttribute( + VectorAttribute.builder().attributeName("v").dimensions(2).build()) + .build(); + + CreateTableRequest base = + CreateTableRequest.builder() + .tableName("dummy") + .keySchema( + KeySchemaElement.builder().attributeName("pk").keyType(KeyType.HASH).build()) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("pk") + .attributeType(ScalarAttributeType.S) + .build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build(); + + CreateTableRequest enriched = + VectorSearchSupport.withVectorIndexes(base, Collections.singletonList(vi)); + assertNotNull(enriched); + assertTrue(enriched.overrideConfiguration().isPresent()); + } + + /** + * Verify that Alternator accepts the uppercase similarity function values "COSINE", + * "DOT_PRODUCT", and "EUCLIDEAN". Each valid value must result in a successfully created table. + */ + @Test + public void testValidSimilarityFunctions() throws Exception { + assumeTrue( + "Integration tests disabled. Set INTEGRATION_TESTS=true to enable.", + IntegrationTestConfig.ENABLED); + + for (String sf : Arrays.asList("COSINE", "DOT_PRODUCT", "EUCLIDEAN")) { + String tbl = tableName + "_" + sf.toLowerCase(); + VectorIndex vi = + VectorIndex.builder() + .indexName("vi-" + sf.toLowerCase()) + .vectorAttribute( + VectorAttribute.builder().attributeName("embedding").dimensions(4).build()) + .similarityFunction(sf) + .build(); + + CreateTableRequest base = + CreateTableRequest.builder() + .tableName(tbl) + .keySchema( + KeySchemaElement.builder().attributeName("pk").keyType(KeyType.HASH).build()) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("pk") + .attributeType(ScalarAttributeType.S) + .build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build(); + + // Must not throw — Alternator should accept this similarity function value. + VectorSearchSupport.createTable(client, base, Collections.singletonList(vi)); + + // Clean up immediately so tearDown doesn't have to know about these extra tables. + client.deleteTable(DeleteTableRequest.builder().tableName(tbl).build()); + } + } + + /** + * Verify that Alternator rejects lowercase similarity function values such as "cosine". + * The server requires uppercase: "COSINE", "DOT_PRODUCT", "EUCLIDEAN". + */ + @Test + public void testLowercaseSimilarityFunctionIsRejected() { + assumeTrue( + "Integration tests disabled. Set INTEGRATION_TESTS=true to enable.", + IntegrationTestConfig.ENABLED); + + VectorIndex vi = + VectorIndex.builder() + .indexName("vi1") + .vectorAttribute( + VectorAttribute.builder().attributeName("embedding").dimensions(4).build()) + .similarityFunction("cosine") // intentionally wrong casing + .build(); + + CreateTableRequest base = + CreateTableRequest.builder() + .tableName(tableName) + .keySchema( + KeySchemaElement.builder().attributeName("pk").keyType(KeyType.HASH).build()) + .attributeDefinitions( + AttributeDefinition.builder() + .attributeName("pk") + .attributeType(ScalarAttributeType.S) + .build()) + .billingMode(BillingMode.PAY_PER_REQUEST) + .build(); + + try { + VectorSearchSupport.createTable(client, base, Collections.singletonList(vi)); + // Server did not reject lowercase — this Alternator version may not enforce the validation. + assumeTrue( + "Skipping: this Alternator version does not reject lowercase similarity functions", + false); + } catch (DynamoDbException e) { + // Expected — Alternator rejects invalid SimilarityFunction values. + } + } +} diff --git a/src/main/java/com/scylladb/alternator/AlternatorDynamoDbAsyncClient.java b/src/main/java/com/scylladb/alternator/AlternatorDynamoDbAsyncClient.java index 9a5b609..c827f73 100644 --- a/src/main/java/com/scylladb/alternator/AlternatorDynamoDbAsyncClient.java +++ b/src/main/java/com/scylladb/alternator/AlternatorDynamoDbAsyncClient.java @@ -10,6 +10,7 @@ import com.scylladb.alternator.queryplan.AffinityQueryPlanInterceptor; import com.scylladb.alternator.queryplan.BasicQueryPlanInterceptor; import com.scylladb.alternator.routing.RoutingScope; +import com.scylladb.alternator.vectorsearch.VectorSearchInterceptor; import java.net.URI; import java.util.Collection; import java.util.Objects; @@ -585,16 +586,6 @@ public AlternatorDynamoDbAsyncClientWrapper buildWithAlternatorAPI() { AlternatorConfig alternatorConfig = configBuilder.build(); - if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { - ClientOverrideConfiguration.Builder overrideBuilder = - delegate.overrideConfiguration() != null - ? delegate.overrideConfiguration().toBuilder() - : ClientOverrideConfiguration.builder(); - overrideBuilder.addExecutionInterceptor( - new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); - delegate.overrideConfiguration(overrideBuilder.build()); - } - TlsConfig tlsConfig = alternatorConfig.getTlsConfig(); boolean optimizeHeaders = alternatorConfig.isOptimizeHeaders(); @@ -622,6 +613,16 @@ public AlternatorDynamoDbAsyncClientWrapper buildWithAlternatorAPI() { ? delegate.overrideConfiguration().toBuilder() : ClientOverrideConfiguration.builder(); + // Interceptor registration order matters for request bodies: + // 1. VectorSearchInterceptor — injects VectorIndexes/VectorSearch JSON and converts + // FLOAT32VECTOR markers; must see the original JSON body. + // 2. GzipRequestInterceptor — compresses the (already vector-modified) body. + // 3. QueryPlanInterceptor — selects the target node; body-independent. + overrideBuilder.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE); + if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { + overrideBuilder.addExecutionInterceptor( + new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); + } KeyRouteAffinityConfig keyAffinityConfig = alternatorConfig.getKeyRouteAffinityConfig(); if (keyAffinityConfig != null && keyAffinityConfig.getType() != null diff --git a/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java b/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java index ca6d730..3dace2b 100644 --- a/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java +++ b/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java @@ -9,6 +9,7 @@ import com.scylladb.alternator.queryplan.AffinityQueryPlanInterceptor; import com.scylladb.alternator.queryplan.BasicQueryPlanInterceptor; import com.scylladb.alternator.routing.RoutingScope; +import com.scylladb.alternator.vectorsearch.VectorSearchInterceptor; import java.net.URI; import java.util.Collection; import java.util.Objects; @@ -606,16 +607,6 @@ public AlternatorDynamoDbClientWrapper buildWithAlternatorAPI() { AlternatorConfig alternatorConfig = configBuilder.build(); - if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { - ClientOverrideConfiguration.Builder overrideBuilder = - delegate.overrideConfiguration() != null - ? delegate.overrideConfiguration().toBuilder() - : ClientOverrideConfiguration.builder(); - overrideBuilder.addExecutionInterceptor( - new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); - delegate.overrideConfiguration(overrideBuilder.build()); - } - TlsConfig tlsConfig = alternatorConfig.getTlsConfig(); boolean optimizeHeaders = alternatorConfig.isOptimizeHeaders(); @@ -645,6 +636,16 @@ public AlternatorDynamoDbClientWrapper buildWithAlternatorAPI() { ? delegate.overrideConfiguration().toBuilder() : ClientOverrideConfiguration.builder(); + // Interceptor registration order matters for request bodies: + // 1. VectorSearchInterceptor — injects VectorIndexes/VectorSearch JSON and converts + // FLOAT32VECTOR markers; must see the original JSON body. + // 2. GzipRequestInterceptor — compresses the (already vector-modified) body. + // 3. QueryPlanInterceptor — selects the target node; body-independent. + overrideBuilder.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE); + if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { + overrideBuilder.addExecutionInterceptor( + new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); + } AffinityQueryPlanInterceptor affinityInterceptor = null; KeyRouteAffinityConfig keyAffinityConfig = alternatorConfig.getKeyRouteAffinityConfig(); if (keyAffinityConfig != null diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/CreateVectorIndexAction.java b/src/main/java/com/scylladb/alternator/vectorsearch/CreateVectorIndexAction.java new file mode 100644 index 0000000..65259f8 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/CreateVectorIndexAction.java @@ -0,0 +1,97 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import software.amazon.awssdk.services.dynamodb.model.Projection; + +/** + * Describes a new vector index to create via {@code UpdateTable}. + * + * @see VectorIndexUpdate + */ +public final class CreateVectorIndexAction { + + private final String indexName; + private final VectorAttribute vectorAttribute; + private final Projection projection; + private final String similarityFunction; + + private CreateVectorIndexAction(Builder builder) { + this.indexName = builder.indexName; + this.vectorAttribute = builder.vectorAttribute; + this.projection = builder.projection; + this.similarityFunction = builder.similarityFunction; + } + + /** Returns the name of the vector index to create. */ + public String indexName() { + return indexName; + } + + /** Returns the vector attribute specification. */ + public VectorAttribute vectorAttribute() { + return vectorAttribute; + } + + /** Returns the projection, or {@code null} for the server default (ALL). */ + public Projection projection() { + return projection; + } + + /** Returns the similarity function, or {@code null} for the server default. */ + public String similarityFunction() { + return similarityFunction; + } + + /** Returns a new builder for {@link CreateVectorIndexAction}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link CreateVectorIndexAction}. */ + public static final class Builder { + private String indexName; + private VectorAttribute vectorAttribute; + private Projection projection; + private String similarityFunction; + + private Builder() {} + + /** Sets the index name. */ + public Builder indexName(String indexName) { + this.indexName = indexName; + return this; + } + + /** Sets the vector attribute specification. */ + public Builder vectorAttribute(VectorAttribute vectorAttribute) { + this.vectorAttribute = vectorAttribute; + return this; + } + + /** Sets the projection. */ + public Builder projection(Projection projection) { + this.projection = projection; + return this; + } + + /** Sets the similarity function. */ + public Builder similarityFunction(String similarityFunction) { + this.similarityFunction = similarityFunction; + return this; + } + + /** Builds the {@link CreateVectorIndexAction}. */ + public CreateVectorIndexAction build() { + if (indexName == null) { + throw new IllegalStateException("indexName must be set"); + } + if (vectorAttribute == null) { + throw new IllegalStateException("vectorAttribute must be set"); + } + return new CreateVectorIndexAction(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/DeleteVectorIndexAction.java b/src/main/java/com/scylladb/alternator/vectorsearch/DeleteVectorIndexAction.java new file mode 100644 index 0000000..160882b --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/DeleteVectorIndexAction.java @@ -0,0 +1,50 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +/** + * Describes a vector index to delete via {@code UpdateTable}. + * + * @see VectorIndexUpdate + */ +public final class DeleteVectorIndexAction { + + private final String indexName; + + private DeleteVectorIndexAction(Builder builder) { + this.indexName = builder.indexName; + } + + /** Returns the name of the vector index to delete. */ + public String indexName() { + return indexName; + } + + /** Returns a new builder for {@link DeleteVectorIndexAction}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link DeleteVectorIndexAction}. */ + public static final class Builder { + private String indexName; + + private Builder() {} + + /** Sets the index name. */ + public Builder indexName(String indexName) { + this.indexName = indexName; + return this; + } + + /** Builds the {@link DeleteVectorIndexAction}. */ + public DeleteVectorIndexAction build() { + if (indexName == null) { + throw new IllegalStateException("indexName must be set"); + } + return new DeleteVectorIndexAction(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/Float32Vector.java b/src/main/java/com/scylladb/alternator/vectorsearch/Float32Vector.java new file mode 100644 index 0000000..ab4bfc2 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/Float32Vector.java @@ -0,0 +1,213 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.List; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +/** + * Utility class for the Alternator {@code FLOAT32VECTOR} attribute type. + * + *

Alternator stores vector attributes in a compact binary format on disk ({@code FLOAT32VECTOR}) + * rather than the standard DynamoDB list-of-numbers encoding ({@code L}). The standard AWS SDK for + * Java has no knowledge of this type, so this class provides a marker-based encoding that the + * {@link VectorSearchInterceptor} recognises and converts automatically: + * + *

    + *
  • Writes — Call {@link #toAttributeValue(float[])} to create a {@code Binary} ({@code + * B}) {@link AttributeValue} that embeds a magic prefix followed by the raw IEEE-754 + * big-endian float bytes. The interceptor detects the magic prefix in the serialised JSON and + * replaces the attribute with {@code {"FLOAT32VECTOR": [...]}} before transmission. This + * works transparently for {@code PutItem}, {@code UpdateItem} (in {@code + * ExpressionAttributeValues}), {@code BatchWriteItem}, and any other operation that carries + * {@link AttributeValue}s in its request body. + *
  • Reads — When Alternator returns {@code {"FLOAT32VECTOR": [...]}} in a response, the + * interceptor converts it transparently to a standard DynamoDB list-of-numbers ({@code L}) + * {@link AttributeValue}. The float values are accessible via {@link AttributeValue#l()} — + * each element is an {@code N}-typed {@link AttributeValue}. + *
+ * + *

Example — writing a vector item

+ * + * 
{@code
+ * Map item = new HashMap<>();
+ * item.put("pk", AttributeValue.fromS("item-1"));
+ * item.put("embedding", Float32Vector.toAttributeValue(new float[]{0.1f, 0.2f, 0.3f}));
+ * client.putItem(PutItemRequest.builder().tableName("t").item(item).build());
+ * }
+ * + *

Example — reading a vector back

+ * + * 
{@code
+ * GetItemResponse resp = client.getItem(...);
+ * AttributeValue embedding = resp.item().get("embedding");
+ * List numbers = embedding.l(); // each element has type N
+ * float[] values = new float[numbers.size()];
+ * for (int i = 0; i < values.length; i++) {
+ *     values[i] = Float.parseFloat(numbers.get(i).n());
+ * }
+ * }
+ * + *

Requirement

+ * + *

{@link VectorSearchInterceptor#INSTANCE} must be registered on the DynamoDB client for the + * automatic conversion to take effect. Without it, the magic-{@code B} attribute is sent as plain + * binary data, which Alternator will not recognise as a vector. + */ +public final class Float32Vector { + + /** + * 8-byte magic prefix that marks a Binary {@link AttributeValue} as a Float32Vector placeholder. + * + *

Chosen to be highly unlikely to appear at the start of legitimate binary data (probability + * of random collision ≈ 1/2^64). + */ + static final byte[] MAGIC = { + (byte) 0xF2, (byte) 0xF3, (byte) 0x2F, (byte) 0xEC, + (byte) 0x4A, (byte) 0x7B, (byte) 0x19, (byte) 0xD3 + }; + + /** + * The guaranteed base64 prefix of any Float32Vector-encoded {@link AttributeValue}'s {@code B} + * field in the serialised DynamoDB JSON, derived from the first 6 bytes of {@link #MAGIC} (two + * complete 3-byte base64 groups → 8 base64 characters). + * + *

Used internally by {@link VectorSearchInterceptor} for a fast substring scan of serialised + * request bodies before committing to a full JSON parse. + */ + static final String BASE64_PREFIX = "8vMv7Ep7"; + + private Float32Vector() {} + + /** + * Creates a DynamoDB Binary ({@code B}) {@link AttributeValue} that encodes {@code values} in the + * Alternator {@code FLOAT32VECTOR} wire format. + * + *

When used in a write request on a client that has {@link VectorSearchInterceptor} registered + * (via {@code .overrideConfiguration(c -> + * c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE))}), this attribute value is + * automatically converted to {@code {"FLOAT32VECTOR": [...]}} in the JSON body, enabling compact + * on-disk storage. + * + * @param values the float array to encode; must not be {@code null} + * @return a {@code B}-typed {@link AttributeValue} that the interceptor converts to {@code + * FLOAT32VECTOR} + */ + public static AttributeValue toAttributeValue(float... values) { + ByteBuffer buf = + ByteBuffer.allocate(MAGIC.length + values.length * Float.BYTES).order(ByteOrder.BIG_ENDIAN); + buf.put(MAGIC); + for (float f : values) { + buf.putFloat(f); + } + buf.flip(); + return AttributeValue.fromB(SdkBytes.fromByteBuffer(buf)); + } + + /** + * Creates a DynamoDB Binary ({@code B}) {@link AttributeValue} that encodes the numbers in {@code + * values} in the Alternator {@code FLOAT32VECTOR} wire format. + * + *

This overload is convenient for re-encoding a vector that was read back from Alternator as + * an {@code L}-typed {@link AttributeValue} (each element is an {@code N}-typed {@link + * AttributeValue}): + * + * 

{@code
+   * AttributeValue read = resp.item().get("embedding"); // L type after round-trip
+   * AttributeValue reEncoded = Float32Vector.toAttributeValue(read.l());
+   * }
+ * + * @param values a list of {@code N}-typed {@link AttributeValue}s; must not be {@code null} + * @return a {@code B}-typed {@link AttributeValue} that the interceptor converts to {@code + * FLOAT32VECTOR} + * @throws NumberFormatException if any element's {@code n()} string is not a valid float + */ + public static AttributeValue toAttributeValue(List values) { + ByteBuffer buf = + ByteBuffer.allocate(MAGIC.length + values.size() * Float.BYTES).order(ByteOrder.BIG_ENDIAN); + buf.put(MAGIC); + for (AttributeValue av : values) { + buf.putFloat(Float.parseFloat(av.n())); + } + buf.flip(); + return AttributeValue.fromB(SdkBytes.fromByteBuffer(buf)); + } + + /** + * Returns {@code true} if {@code av} is a Float32Vector-encoded {@link AttributeValue} — i.e., a + * Binary ({@code B}) attribute whose bytes start with the Float32Vector magic prefix. + * + *

This identifies values created by {@link #toAttributeValue(float...)} or {@link + * #toAttributeValue(List)} on the write path, before the interceptor has converted them + * to {@code FLOAT32VECTOR} JSON. It does not apply to values read back from Alternator — + * those arrive as standard {@code L}-typed {@link AttributeValue}s. + * + * @param av the {@link AttributeValue} to test; must not be {@code null} + * @return {@code true} if {@code av} encodes a {@code FLOAT32VECTOR} + */ + public static boolean isFloat32Vector(AttributeValue av) { + return av.b() != null && hasFloat32VectorMagic(av.b().asByteArray()); + } + + /** + * Extracts the float array from a Float32Vector-encoded {@link AttributeValue}. + * + * @param av an {@link AttributeValue} satisfying {@link #isFloat32Vector(AttributeValue)} + * @return the decoded float array + * @throws IllegalArgumentException if {@code av} is not a Float32Vector + */ + public static float[] toFloats(AttributeValue av) { + if (!isFloat32Vector(av)) { + throw new IllegalArgumentException( + "AttributeValue is not a Float32Vector " + + "(expected a B attribute with the Float32Vector magic prefix)"); + } + return bytesToFloats(av.b().asByteArray()); + } + + // ------------------------------------------------------------------------- + // Package-private helpers used by VectorSearchInterceptor + // ------------------------------------------------------------------------- + + /** Returns {@code true} if {@code bytes} starts with the Float32Vector magic prefix. */ + static boolean hasFloat32VectorMagic(byte[] bytes) { + if (bytes.length < MAGIC.length) { + return false; + } + for (int i = 0; i < MAGIC.length; i++) { + if (bytes[i] != MAGIC[i]) { + return false; + } + } + return true; + } + + /** + * Decodes a float array from magic-prefixed bytes. The caller must have already verified the + * magic prefix. + */ + static float[] bytesToFloats(byte[] bytes) { + int payload = bytes.length - MAGIC.length; + if (payload < 0 || payload % Float.BYTES != 0) { + throw new IllegalArgumentException( + "Invalid Float32Vector payload length: " + + bytes.length + + " bytes (expected MAGIC.length + N * " + + Float.BYTES + + ")"); + } + int floatCount = payload / Float.BYTES; + float[] result = new float[floatCount]; + ByteBuffer buf = + ByteBuffer.wrap(bytes, MAGIC.length, floatCount * Float.BYTES).order(ByteOrder.BIG_ENDIAN); + for (int i = 0; i < floatCount; i++) { + result[i] = buf.getFloat(); + } + return result; + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorAttribute.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorAttribute.java new file mode 100644 index 0000000..7c445d5 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorAttribute.java @@ -0,0 +1,67 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +/** + * Describes the vector attribute for a {@link VectorIndex}. + * + *

Specifies the item attribute that holds vector data and its dimensionality. + */ +public final class VectorAttribute { + + private final String attributeName; + private final int dimensions; + + private VectorAttribute(Builder builder) { + this.attributeName = builder.attributeName; + this.dimensions = builder.dimensions; + } + + /** Returns the name of the item attribute that stores vector data. */ + public String attributeName() { + return attributeName; + } + + /** Returns the number of dimensions in the vector. */ + public int dimensions() { + return dimensions; + } + + /** Returns a new builder for {@link VectorAttribute}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link VectorAttribute}. */ + public static final class Builder { + private String attributeName; + private int dimensions; + + private Builder() {} + + /** Sets the name of the attribute that stores vector data. */ + public Builder attributeName(String attributeName) { + this.attributeName = attributeName; + return this; + } + + /** Sets the number of dimensions in the vector. */ + public Builder dimensions(int dimensions) { + this.dimensions = dimensions; + return this; + } + + /** Builds the {@link VectorAttribute}. */ + public VectorAttribute build() { + if (attributeName == null) { + throw new IllegalStateException("attributeName must be set"); + } + if (dimensions <= 0) { + throw new IllegalStateException("dimensions must be a positive integer"); + } + return new VectorAttribute(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndex.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndex.java new file mode 100644 index 0000000..a6f726e --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndex.java @@ -0,0 +1,152 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import software.amazon.awssdk.services.dynamodb.model.Projection; + +/** + * Describes a vector index for Alternator's vector search feature. + * + *

Used in {@code CreateTable} and {@code UpdateTable} requests to define a vector index, and + * returned in {@code DescribeTable} and {@code CreateTable} responses. + * + *

Example: + * + * 

{@code
+ * VectorIndex vi = VectorIndex.builder()
+ *     .indexName("my-vector-index")
+ *     .vectorAttribute(VectorAttribute.builder()
+ *         .attributeName("embedding")
+ *         .dimensions(128)
+ *         .build())
+ *     .similarityFunction("COSINE")
+ *     .build();
+ * }
+ */ +public final class VectorIndex { + + private final String indexName; + private final VectorAttribute vectorAttribute; + private final Projection projection; + private final String similarityFunction; + // Response-only fields: + private final String indexStatus; + private final Boolean backfilling; + + private VectorIndex(Builder builder) { + this.indexName = builder.indexName; + this.vectorAttribute = builder.vectorAttribute; + this.projection = builder.projection; + this.similarityFunction = builder.similarityFunction; + this.indexStatus = builder.indexStatus; + this.backfilling = builder.backfilling; + } + + /** Returns the name of this vector index. */ + public String indexName() { + return indexName; + } + + /** Returns the vector attribute specification. */ + public VectorAttribute vectorAttribute() { + return vectorAttribute; + } + + /** Returns the projection for this index, or {@code null} if using the default (ALL). */ + public Projection projection() { + return projection; + } + + /** + * Returns the similarity function (e.g., {@code "COSINE"}, {@code "DOT_PRODUCT"}, {@code + * "EUCLIDEAN"}), or {@code null} to use the server default. + */ + public String similarityFunction() { + return similarityFunction; + } + + /** + * Returns the index status as reported by the server (e.g., {@code "CREATING"}, {@code + * "ACTIVE"}). Only populated in responses from {@code DescribeTable} or {@code CreateTable}. + */ + public String indexStatus() { + return indexStatus; + } + + /** + * Returns whether the index is backfilling, or {@code null} if not reported by the server. Only + * populated in responses from {@code DescribeTable} or {@code CreateTable}. + */ + public Boolean backfilling() { + return backfilling; + } + + /** Returns a new builder for {@link VectorIndex}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link VectorIndex}. */ + public static final class Builder { + private String indexName; + private VectorAttribute vectorAttribute; + private Projection projection; + private String similarityFunction; + private String indexStatus; + private Boolean backfilling; + + private Builder() {} + + /** Sets the name of the vector index. */ + public Builder indexName(String indexName) { + this.indexName = indexName; + return this; + } + + /** Sets the vector attribute specification. */ + public Builder vectorAttribute(VectorAttribute vectorAttribute) { + this.vectorAttribute = vectorAttribute; + return this; + } + + /** Sets the projection for this index. If not set, the server defaults to ALL. */ + public Builder projection(Projection projection) { + this.projection = projection; + return this; + } + + /** + * Sets the similarity function (e.g., {@code "COSINE"}, {@code "DOT_PRODUCT"}, {@code + * "EUCLIDEAN"}). If not set, the server uses its default. + */ + public Builder similarityFunction(String similarityFunction) { + this.similarityFunction = similarityFunction; + return this; + } + + /** Sets the index status (populated from server responses). */ + public Builder indexStatus(String indexStatus) { + this.indexStatus = indexStatus; + return this; + } + + /** Sets the backfilling flag (populated from server responses). */ + public Builder backfilling(Boolean backfilling) { + this.backfilling = backfilling; + return this; + } + + /** Builds the {@link VectorIndex}. */ + public VectorIndex build() { + if (indexName == null) { + throw new IllegalStateException("indexName must be set"); + } + if (vectorAttribute == null) { + throw new IllegalStateException("vectorAttribute must be set"); + } + return new VectorIndex(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndexUpdate.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndexUpdate.java new file mode 100644 index 0000000..e716d5b --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorIndexUpdate.java @@ -0,0 +1,81 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +/** + * Represents a single vector index change in an {@code UpdateTable} request. + * + *

Exactly one of {@link #create()} or {@link #delete()} must be set. + * + *

Example: + * + * 

{@code
+ * VectorIndexUpdate update = VectorIndexUpdate.builder()
+ *     .create(CreateVectorIndexAction.builder()
+ *         .indexName("my-index")
+ *         .vectorAttribute(VectorAttribute.builder()
+ *             .attributeName("embedding")
+ *             .dimensions(128)
+ *             .build())
+ *         .build())
+ *     .build();
+ * }
+ */ +public final class VectorIndexUpdate { + + private final CreateVectorIndexAction create; + private final DeleteVectorIndexAction delete; + + private VectorIndexUpdate(Builder builder) { + this.create = builder.create; + this.delete = builder.delete; + } + + /** Returns the create action, or {@code null} if this is a delete update. */ + public CreateVectorIndexAction create() { + return create; + } + + /** Returns the delete action, or {@code null} if this is a create update. */ + public DeleteVectorIndexAction delete() { + return delete; + } + + /** Returns a new builder for {@link VectorIndexUpdate}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link VectorIndexUpdate}. */ + public static final class Builder { + private CreateVectorIndexAction create; + private DeleteVectorIndexAction delete; + + private Builder() {} + + /** Sets the create action. */ + public Builder create(CreateVectorIndexAction create) { + this.create = create; + return this; + } + + /** Sets the delete action. */ + public Builder delete(DeleteVectorIndexAction delete) { + this.delete = delete; + return this; + } + + /** Builds the {@link VectorIndexUpdate}. */ + public VectorIndexUpdate build() { + if (create == null && delete == null) { + throw new IllegalStateException("exactly one of create or delete must be set"); + } + if (create != null && delete != null) { + throw new IllegalStateException("exactly one of create or delete must be set, not both"); + } + return new VectorIndexUpdate(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorQueryResult.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorQueryResult.java new file mode 100644 index 0000000..426aac9 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorQueryResult.java @@ -0,0 +1,67 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; + +/** + * The result of a vector search {@code Query} request. + * + *

Wraps the standard {@link QueryResponse} and adds access to per-item similarity scores + * returned by Alternator when {@link VectorSearch#returnScores()} is {@code true}. + * + *

Similarity scores are in the same order as the items returned by {@link #items()}. + */ +public final class VectorQueryResult { + + private final QueryResponse response; + private final List scores; + + public VectorQueryResult(QueryResponse response, List scores) { + this.response = response; + this.scores = scores != null ? Collections.unmodifiableList(scores) : Collections.emptyList(); + } + + /** Returns the underlying {@link QueryResponse}. */ + public QueryResponse response() { + return response; + } + + /** + * Returns the list of items returned by the query. + * + *

Convenience delegate for {@link QueryResponse#items()}. + */ + public List> items() { + return response.items(); + } + + /** + * Returns the number of items that matched the query before any {@code Limit} was applied. + * + *

Convenience delegate for {@link QueryResponse#count()}. + */ + public int count() { + return response.count(); + } + + /** + * Returns per-item similarity scores in the same order as {@link #items()}, or an empty list if + * scores were not requested or the server did not return them. + */ + public List scores() { + return scores; + } + + /** Returns the consumed capacity, or {@code null} if not requested. */ + public ConsumedCapacity consumedCapacity() { + return response.consumedCapacity(); + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearch.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearch.java new file mode 100644 index 0000000..2ce2d21 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearch.java @@ -0,0 +1,119 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +/** + * Parameters for a vector similarity search in a {@code Query} request. + * + *

Carries the query vector and optional flags that control the search behavior. Set this on a + * {@code QueryRequest} via {@link VectorSearchSupport#query}. + * + *

The query vector may be supplied in two forms: + * + *

    + *
  • {@code float[]} — serialized as the compact {@code FLOAT32VECTOR} wire format. + *
  • {@link AttributeValue} — serialized as the standard DynamoDB JSON format; useful when the + * attribute was stored without the {@code FLOAT32VECTOR} optimization. + *
+ * + *

Example: + * + * 

{@code
+ * VectorSearch vs = VectorSearch.builder()
+ *     .queryVector(new float[]{0.1f, 0.2f, 0.3f})
+ *     .returnScores(true)
+ *     .build();
+ * }
+ */ +public final class VectorSearch { + + private final float[] queryVectorFloats; + private final AttributeValue queryVectorAttributeValue; + private final boolean returnScores; + + private VectorSearch(Builder builder) { + this.queryVectorFloats = builder.queryVectorFloats; + this.queryVectorAttributeValue = builder.queryVectorAttributeValue; + this.returnScores = builder.returnScores; + } + + /** + * Returns the query vector as a {@code float[]}, or {@code null} if the vector was provided as an + * {@link AttributeValue}. + */ + public float[] queryVectorFloats() { + return queryVectorFloats; + } + + /** + * Returns the query vector as an {@link AttributeValue}, or {@code null} if the vector was + * provided as a {@code float[]}. + */ + public AttributeValue queryVectorAttributeValue() { + return queryVectorAttributeValue; + } + + /** + * Returns whether the server should return per-item similarity scores alongside the results. + * Scores are accessible via {@link VectorQueryResult#scores()}. + */ + public boolean returnScores() { + return returnScores; + } + + /** Returns a new builder for {@link VectorSearch}. */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link VectorSearch}. */ + public static final class Builder { + private float[] queryVectorFloats; + private AttributeValue queryVectorAttributeValue; + private boolean returnScores; + + private Builder() {} + + /** + * Sets the query vector as a float array. It will be sent to Alternator using the compact + * {@code FLOAT32VECTOR} wire encoding. + */ + public Builder queryVector(float... queryVector) { + this.queryVectorFloats = queryVector; + this.queryVectorAttributeValue = null; + return this; + } + + /** + * Sets the query vector as an {@link AttributeValue}. Use this when the vectors in the table + * were stored using the standard DynamoDB list type rather than the {@code FLOAT32VECTOR} + * format. + */ + public Builder queryVector(AttributeValue queryVector) { + this.queryVectorAttributeValue = queryVector; + this.queryVectorFloats = null; + return this; + } + + /** + * When {@code true}, asks the server to include per-item similarity scores in the response. + * Access them via {@link VectorQueryResult#scores()}. + */ + public Builder returnScores(boolean returnScores) { + this.returnScores = returnScores; + return this; + } + + /** Builds the {@link VectorSearch}. */ + public VectorSearch build() { + if (queryVectorFloats == null && queryVectorAttributeValue == null) { + throw new IllegalStateException("queryVector must be set"); + } + return new VectorSearch(this); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java new file mode 100644 index 0000000..9a3547e --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java @@ -0,0 +1,663 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.ExecutionInterceptor; +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.http.SdkHttpRequest; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.Projection; +import software.amazon.awssdk.services.dynamodb.model.ProjectionType; + +/** + * An {@link ExecutionInterceptor} that enables Alternator's vector search extension. + * + *

Alternator extends the DynamoDB API with vector indexes ({@code VectorIndexes} in {@code + * CreateTable}/{@code UpdateTable}/{@code DescribeTable}) and vector similarity search ({@code + * VectorSearch} in {@code Query}). Because the standard AWS SDK for Java does not know about these + * new fields, this interceptor bridges the gap by injecting them into the raw JSON request bodies + * and extracting them from the raw JSON response bodies at the HTTP layer. + * + *

Usage

+ * + *

Register the interceptor once when building the DynamoDB client: + * + * 

{@code
+ * DynamoDbClient client = DynamoDbClient.builder()
+ *     .overrideConfiguration(c -> c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE))
+ *     ...
+ *     .build();
+ * }
+ * + *

Then use {@link VectorSearchSupport} to attach vector parameters to individual requests. + * + *

Implementation notes

+ * + *
    + *
  • For vector search params (CreateTable/UpdateTable/Query): attached as {@link + * ExecutionAttribute}s via {@link VectorSearchSupport}; injected into the JSON body in {@link + * #modifyHttpContent(Context.ModifyHttpRequest, ExecutionAttributes)}. + *
  • For {@code FLOAT32VECTOR} item attributes on writes: the user creates a marker {@link + * AttributeValue} via {@link Float32Vector#toAttributeValue(float[])}; the interceptor + * detects the magic binary prefix in the serialised JSON and replaces it with {@code + * {"FLOAT32VECTOR": [...]}} before transmission. Works for any operation that carries {@link + * AttributeValue}s ({@code PutItem}, {@code UpdateItem}, {@code BatchWriteItem}, etc.). + *
  • For {@code FLOAT32VECTOR} item attributes on reads: {@link + * #modifyHttpResponseContent(Context.ModifyHttpResponse, ExecutionAttributes)} scans all + * response bodies for {@code FLOAT32VECTOR} attribute values and converts them transparently + * to standard DynamoDB list-of-numbers ({@code L}) {@link AttributeValue}s, accessible via + * {@link AttributeValue#l()}. + *
+ */ +public final class VectorSearchInterceptor implements ExecutionInterceptor { + + /** Singleton instance — stateless, safe to share across clients. */ + public static final VectorSearchInterceptor INSTANCE = new VectorSearchInterceptor(); + + // ------------------------------------------------------------------------- + // ExecutionAttribute keys + // ------------------------------------------------------------------------- + + /** + * List of vector indexes to add to a {@code CreateTable} request. Set via {@link + * VectorSearchSupport#withVectorIndexes}. + */ + public static final ExecutionAttribute> VECTOR_INDEXES = + new ExecutionAttribute<>("AlternatorVectorIndexes"); + + /** + * List of vector index updates to add to an {@code UpdateTable} request. Set via {@link + * VectorSearchSupport#withVectorIndexUpdates}. + */ + public static final ExecutionAttribute> VECTOR_INDEX_UPDATES = + new ExecutionAttribute<>("AlternatorVectorIndexUpdates"); + + /** + * Vector search parameters to add to a {@code Query} request. Set via {@link + * VectorSearchSupport#query} or {@link VectorSearchSupport#withVectorSearch}. + */ + public static final ExecutionAttribute VECTOR_SEARCH = + new ExecutionAttribute<>("AlternatorVectorSearch"); + + /** + * Per-request holder for extra response fields (scores, returned vector indexes). Set internally + * by {@link VectorSearchSupport}. + */ + static final ExecutionAttribute RESULT_HOLDER = + new ExecutionAttribute<>("AlternatorVectorSearchResultHolder"); + + /** + * Cache for the modified request body bytes, set by {@link + * #modifyHttpContent(Context.ModifyHttpRequest, ExecutionAttributes)} and read by {@link + * #modifyHttpRequest(Context.ModifyHttpRequest, ExecutionAttributes)}. + * + *

The SDK calls {@code modifyHttpContent} before {@code modifyHttpRequest} for each + * interceptor (in a single {@code modifyHttpRequestAndHttpContent} pass), and passes the + * same {@link ExecutionAttributes} instance to both. Caching here eliminates the second + * full body read and JSON parse+rewrite that would otherwise be needed to compute the {@code + * Content-Length} in {@code modifyHttpRequest}. + */ + private static final ExecutionAttribute MODIFIED_BODY_CACHE = + new ExecutionAttribute<>("AlternatorModifiedBodyCache"); + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + // DynamoDB X-Amz-Target suffixes + private static final String TARGET_CREATE_TABLE = "DynamoDB_20120810.CreateTable"; + private static final String TARGET_UPDATE_TABLE = "DynamoDB_20120810.UpdateTable"; + private static final String TARGET_DESCRIBE_TABLE = "DynamoDB_20120810.DescribeTable"; + private static final String TARGET_QUERY = "DynamoDB_20120810.Query"; + + private VectorSearchInterceptor() {} + + // ------------------------------------------------------------------------- + // Request interception — inject extra JSON fields + // ------------------------------------------------------------------------- + + /** + * Updates the {@code Content-Length} header to match the body produced by {@link + * #modifyHttpContent(Context.ModifyHttpRequest, ExecutionAttributes)}, which runs first. + * + *

The AWS SDK sets {@code Content-Length} from the original (pre-interceptor) body length. If + * we only change the body in {@code modifyHttpContent}, the server receives a stale {@code + * Content-Length}. We therefore read the cached modified bytes (written by {@code + * modifyHttpContent}) here and update the header to match. + * + *

Call order: The SDK calls {@code modifyHttpContent} before {@code modifyHttpRequest} + * for each interceptor (verified from {@code ExecutionInterceptorChain} bytecode), passing the + * same {@link ExecutionAttributes} instance to both, so the cache is always populated by the time + * this method runs. + */ + @Override + public SdkHttpRequest modifyHttpRequest( + Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + byte[] modifiedBytes = executionAttributes.getAttribute(MODIFIED_BODY_CACHE); + if (modifiedBytes == null) { + return context.httpRequest(); + } + return context.httpRequest().toBuilder() + .putHeader("Content-Length", String.valueOf(modifiedBytes.length)) + .build(); + } + + /** + * Computes the modified body (injecting vector search parameters and converting {@code + * FLOAT32VECTOR} markers), caches the result in {@link ExecutionAttributes} for {@link + * #modifyHttpRequest}, and returns the new body. + */ + @Override + public Optional modifyHttpContent( + Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) { + try { + byte[] modifiedBytes = computeModifiedBodyBytes(context, executionAttributes); + if (modifiedBytes == null) { + return context.requestBody(); + } + executionAttributes.putAttribute(MODIFIED_BODY_CACHE, modifiedBytes); + return Optional.of(RequestBody.fromBytes(modifiedBytes)); + } catch (IOException e) { + throw new RuntimeException("Failed to process vector search parameters in request body", e); + } + } + + /** + * Computes the modified body bytes for the request, or returns {@code null} if no modification is + * needed. Only called from {@link #modifyHttpContent}; the result is cached in {@link + * ExecutionAttributes} and reused by {@link #modifyHttpRequest} to set {@code Content-Length}. + */ + private static byte[] computeModifiedBodyBytes( + Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) + throws IOException { + List vectorIndexes = executionAttributes.getAttribute(VECTOR_INDEXES); + List vectorIndexUpdates = + executionAttributes.getAttribute(VECTOR_INDEX_UPDATES); + VectorSearch vectorSearch = executionAttributes.getAttribute(VECTOR_SEARCH); + + byte[] originalBytes = readBytes(context.requestBody()); + if (originalBytes == null) { + return null; + } + + boolean hasVectorParams = + vectorIndexes != null || vectorIndexUpdates != null || vectorSearch != null; + boolean hasFloat32Vectors = containsAsciiSubstring(originalBytes, Float32Vector.BASE64_PREFIX); + + if (!hasVectorParams && !hasFloat32Vectors) { + return null; + } + + ObjectNode json = (ObjectNode) MAPPER.readTree(originalBytes); + boolean modified = false; + + if (hasVectorParams) { + String target = getTarget(context); + if (TARGET_CREATE_TABLE.equals(target) && vectorIndexes != null) { + json.set("VectorIndexes", vectorIndexesToJson(vectorIndexes)); + modified = true; + } else if (TARGET_UPDATE_TABLE.equals(target) && vectorIndexUpdates != null) { + json.set("VectorIndexUpdates", vectorIndexUpdatesToJson(vectorIndexUpdates)); + modified = true; + } else if (TARGET_QUERY.equals(target) && vectorSearch != null) { + json.set("VectorSearch", vectorSearchToJson(vectorSearch)); + modified = true; + } + } + + if (hasFloat32Vectors) { + modified |= replaceFloat32VectorInRequest(json); + } + + if (!modified) { + return null; + } + + byte[] outBytes = MAPPER.writeValueAsBytes(json); + return outBytes; + } + + // ------------------------------------------------------------------------- + // Response interception — extract extra JSON fields + // ------------------------------------------------------------------------- + + /** + * Intercepts raw response bodies to: + * + *

    + *
  1. Convert any {@code {"FLOAT32VECTOR": [...]}} attribute values to standard DynamoDB + * list-of-numbers ({@code L}) {@link AttributeValue}s so the SDK can parse the item + * normally and callers can access the floats via {@link AttributeValue#l()}. A fast + * substring scan is used to skip parsing when no {@code FLOAT32VECTOR} field is present. + *
  2. Extract Alternator-specific response fields ({@code Scores} from Query, {@code + * VectorIndexes} from CreateTable/DescribeTable) into the per-request {@link + * VectorSearchResultHolder} when one was set by {@link VectorSearchSupport}. + *
+ */ + @Override + public Optional modifyHttpResponseContent( + Context.ModifyHttpResponse context, ExecutionAttributes executionAttributes) { + + VectorSearchResultHolder holder = executionAttributes.getAttribute(RESULT_HOLDER); + + Optional bodyOpt = context.responseBody(); + if (!bodyOpt.isPresent()) { + return bodyOpt; + } + + try { + byte[] bytes = readAllBytes(bodyOpt.get()); + + // Quick scan: skip JSON parsing entirely when neither FLOAT32VECTOR nor a result holder + // needing extraction are in play. + boolean hasFloat32Vector = containsAsciiSubstring(bytes, "FLOAT32VECTOR"); + if (!hasFloat32Vector && holder == null) { + return Optional.of(new ByteArrayInputStream(bytes)); + } + + JsonNode json = MAPPER.readTree(bytes); + boolean modified = false; + + // Convert FLOAT32VECTOR → standard L (list-of-numbers) so the SDK can unmarshal items + // normally. + if (hasFloat32Vector) { + modified = replaceFloat32VectorInResponse(json); + } + + // Extract Scores / VectorIndexes for callers using VectorSearchSupport. + if (holder != null) { + String target = getTarget(context.httpRequest()); + if (TARGET_QUERY.equals(target)) { + extractScores(json, holder); + } else if (TARGET_CREATE_TABLE.equals(target) || TARGET_DESCRIBE_TABLE.equals(target)) { + extractTableDescriptionVectorIndexes(json, holder); + } + } + + byte[] outBytes = modified ? MAPPER.writeValueAsBytes(json) : bytes; + return Optional.of(new ByteArrayInputStream(outBytes)); + + } catch (IOException e) { + throw new RuntimeException("Failed to process vector search fields in response body", e); + } + } + + // ------------------------------------------------------------------------- + // JSON serialisation helpers + // ------------------------------------------------------------------------- + + private static ArrayNode vectorIndexesToJson(List indexes) { + ArrayNode arr = MAPPER.createArrayNode(); + for (VectorIndex vi : indexes) { + arr.add(vectorIndexToJson(vi)); + } + return arr; + } + + private static ObjectNode vectorIndexToJson(VectorIndex vi) { + ObjectNode node = MAPPER.createObjectNode(); + node.put("IndexName", vi.indexName()); + node.set("VectorAttribute", vectorAttributeToJson(vi.vectorAttribute())); + if (vi.projection() != null) { + node.set("Projection", projectionToJson(vi.projection())); + } + if (vi.similarityFunction() != null) { + node.put("SimilarityFunction", vi.similarityFunction()); + } + return node; + } + + private static ObjectNode vectorAttributeToJson(VectorAttribute va) { + ObjectNode node = MAPPER.createObjectNode(); + node.put("AttributeName", va.attributeName()); + node.put("Dimensions", va.dimensions()); + return node; + } + + private static ObjectNode projectionToJson(Projection projection) { + ObjectNode node = MAPPER.createObjectNode(); + if (projection.projectionType() != null) { + node.put("ProjectionType", projection.projectionTypeAsString()); + } + if (projection.nonKeyAttributes() != null && !projection.nonKeyAttributes().isEmpty()) { + ArrayNode nka = node.putArray("NonKeyAttributes"); + projection.nonKeyAttributes().forEach(nka::add); + } + return node; + } + + private static ArrayNode vectorIndexUpdatesToJson(List updates) { + ArrayNode arr = MAPPER.createArrayNode(); + for (VectorIndexUpdate u : updates) { + ObjectNode node = MAPPER.createObjectNode(); + if (u.create() != null) { + node.set("Create", createVectorIndexActionToJson(u.create())); + } + if (u.delete() != null) { + ObjectNode del = MAPPER.createObjectNode(); + del.put("IndexName", u.delete().indexName()); + node.set("Delete", del); + } + arr.add(node); + } + return arr; + } + + private static ObjectNode createVectorIndexActionToJson(CreateVectorIndexAction action) { + ObjectNode node = MAPPER.createObjectNode(); + node.put("IndexName", action.indexName()); + node.set("VectorAttribute", vectorAttributeToJson(action.vectorAttribute())); + if (action.projection() != null) { + node.set("Projection", projectionToJson(action.projection())); + } + if (action.similarityFunction() != null) { + node.put("SimilarityFunction", action.similarityFunction()); + } + return node; + } + + private static ObjectNode vectorSearchToJson(VectorSearch vs) { + ObjectNode node = MAPPER.createObjectNode(); + if (vs.queryVectorFloats() != null) { + // Compact wire format: {"FLOAT32VECTOR": [1.0, 2.0, ...]} + ObjectNode qv = MAPPER.createObjectNode(); + ArrayNode floats = qv.putArray("FLOAT32VECTOR"); + for (float f : vs.queryVectorFloats()) { + floats.add(f); + } + node.set("QueryVector", qv); + } else { + node.set("QueryVector", attributeValueToJson(vs.queryVectorAttributeValue())); + } + if (vs.returnScores()) { + node.put("ReturnScores", "SIMILARITY"); + } + return node; + } + + /** + * Converts an {@link AttributeValue} to its DynamoDB JSON representation (e.g., {@code {"N": + * "42"}}, {@code {"S": "hello"}}, {@code {"L": [...]}}). + * + *

This method is {@code public} so that callers can build raw DynamoDB-style JSON payloads + * when working directly with the low-level HTTP API. + */ + public static ObjectNode attributeValueToJson(AttributeValue av) { + ObjectNode node = MAPPER.createObjectNode(); + if (av.s() != null) { + node.put("S", av.s()); + } else if (av.n() != null) { + node.put("N", av.n()); + } else if (Boolean.TRUE.equals(av.bool())) { + node.put("BOOL", true); + } else if (Boolean.FALSE.equals(av.bool())) { + node.put("BOOL", false); + } else if (Boolean.TRUE.equals(av.nul())) { + node.put("NULL", true); + } else if (av.b() != null) { + // If this is a Float32Vector marker, emit the compact wire format directly. + if (Float32Vector.hasFloat32VectorMagic(av.b().asByteArray())) { + float[] floats = Float32Vector.bytesToFloats(av.b().asByteArray()); + ArrayNode f32v = node.putArray("FLOAT32VECTOR"); + for (float f : floats) { + f32v.add(f); + } + } else { + node.put("B", Base64.getEncoder().encodeToString(av.b().asByteArray())); + } + } else if (av.hasSs()) { + ArrayNode arr = node.putArray("SS"); + av.ss().forEach(arr::add); + } else if (av.hasNs()) { + ArrayNode arr = node.putArray("NS"); + av.ns().forEach(arr::add); + } else if (av.hasBs()) { + ArrayNode arr = node.putArray("BS"); + av.bs().forEach(b -> arr.add(b.asByteArray())); + } else if (av.hasL()) { + ArrayNode arr = node.putArray("L"); + av.l().forEach(elem -> arr.add(attributeValueToJson(elem))); + } else if (av.hasM()) { + ObjectNode map = node.putObject("M"); + for (Map.Entry entry : av.m().entrySet()) { + map.set(entry.getKey(), attributeValueToJson(entry.getValue())); + } + } + return node; + } + + // ------------------------------------------------------------------------- + // JSON deserialisation helpers + // ------------------------------------------------------------------------- + + private static void extractScores(JsonNode root, VectorSearchResultHolder holder) { + JsonNode scoresNode = root.get("Scores"); + if (scoresNode != null && scoresNode.isArray()) { + List scores = new ArrayList<>(scoresNode.size()); + for (JsonNode n : scoresNode) { + scores.add(n.asDouble()); + } + holder.setScores(scores); + } + } + + private static void extractTableDescriptionVectorIndexes( + JsonNode root, VectorSearchResultHolder holder) { + // CreateTable response wraps the table description under "TableDescription" + JsonNode tableDesc = root.get("TableDescription"); + // DescribeTable response also uses "Table" in some SDK versions; try both + if (tableDesc == null) { + tableDesc = root.get("Table"); + } + if (tableDesc == null) { + // Might be a flat response (rare) + tableDesc = root; + } + JsonNode viNode = tableDesc.get("VectorIndexes"); + if (viNode != null && viNode.isArray()) { + List indexes = new ArrayList<>(viNode.size()); + for (JsonNode n : viNode) { + indexes.add(parseVectorIndex((ObjectNode) n)); + } + holder.setVectorIndexes(indexes); + } + } + + private static VectorIndex parseVectorIndex(ObjectNode node) { + String indexName = node.get("IndexName").asText(); + ObjectNode vaNode = (ObjectNode) node.get("VectorAttribute"); + VectorAttribute va = + VectorAttribute.builder() + .attributeName(vaNode.get("AttributeName").asText()) + .dimensions(vaNode.get("Dimensions").asInt()) + .build(); + + Projection projection = null; + if (node.has("Projection")) { + ObjectNode projNode = (ObjectNode) node.get("Projection"); + Projection.Builder projBuilder = Projection.builder(); + if (projNode.has("ProjectionType")) { + projBuilder.projectionType( + ProjectionType.fromValue(projNode.get("ProjectionType").asText())); + } + if (projNode.has("NonKeyAttributes")) { + List nka = new ArrayList<>(); + projNode.get("NonKeyAttributes").forEach(n -> nka.add(n.asText())); + projBuilder.nonKeyAttributes(nka); + } + projection = projBuilder.build(); + } + + String similarityFunction = + node.has("SimilarityFunction") ? node.get("SimilarityFunction").asText() : null; + String indexStatus = node.has("IndexStatus") ? node.get("IndexStatus").asText() : null; + Boolean backfilling = node.has("Backfilling") ? node.get("Backfilling").asBoolean() : null; + + return VectorIndex.builder() + .indexName(indexName) + .vectorAttribute(va) + .projection(projection) + .similarityFunction(similarityFunction) + .indexStatus(indexStatus) + .backfilling(backfilling) + .build(); + } + + // ------------------------------------------------------------------------- + // FLOAT32VECTOR JSON replacement helpers + // ------------------------------------------------------------------------- + + /** + * Recursively scans {@code node} for {@code {"B": "..."}} objects whose base64 value decodes to + * bytes starting with the Float32Vector magic prefix, and replaces them in-place with {@code + * {"FLOAT32VECTOR": [...]}}. + * + * @return {@code true} if any replacement was made + */ + private static boolean replaceFloat32VectorInRequest(JsonNode node) { + if (node.isObject()) { + ObjectNode obj = (ObjectNode) node; + JsonNode bField = obj.get("B"); + if (bField != null && bField.isTextual() && obj.size() == 1) { + byte[] decoded; + try { + decoded = Base64.getDecoder().decode(bField.asText()); + } catch (IllegalArgumentException ignored) { + decoded = new byte[0]; + } + if (Float32Vector.hasFloat32VectorMagic(decoded)) { + float[] floats = Float32Vector.bytesToFloats(decoded); + obj.remove("B"); + ArrayNode arr = obj.putArray("FLOAT32VECTOR"); + for (float f : floats) { + arr.add(f); + } + return true; + } + } + // Not a Float32Vector marker node — recurse into children. + boolean modified = false; + for (JsonNode child : obj) { + modified |= replaceFloat32VectorInRequest(child); + } + return modified; + } else if (node.isArray()) { + boolean modified = false; + for (JsonNode child : node) { + modified |= replaceFloat32VectorInRequest(child); + } + return modified; + } + return false; + } + + /** + * Recursively scans {@code node} for {@code {"FLOAT32VECTOR": [...]}} objects and replaces them + * in-place with a standard DynamoDB list-of-numbers ({@code {"L": [{"N": "x"}, ...]}}) so the SDK + * returns a plain {@code L}-typed {@link AttributeValue} that callers can use directly via {@link + * software.amazon.awssdk.services.dynamodb.model.AttributeValue#l()}. + * + * @return {@code true} if any replacement was made + */ + private static boolean replaceFloat32VectorInResponse(JsonNode node) { + if (node.isObject()) { + ObjectNode obj = (ObjectNode) node; + JsonNode f32vField = obj.get("FLOAT32VECTOR"); + if (f32vField != null && f32vField.isArray() && obj.size() == 1) { + ArrayNode lArray = MAPPER.createArrayNode(); + for (JsonNode numNode : f32vField) { + ObjectNode n = MAPPER.createObjectNode(); + n.put("N", Float.toString((float) numNode.asDouble())); + lArray.add(n); + } + obj.remove("FLOAT32VECTOR"); + obj.set("L", lArray); + return true; + } + boolean modified = false; + for (JsonNode child : obj) { + modified |= replaceFloat32VectorInResponse(child); + } + return modified; + } else if (node.isArray()) { + boolean modified = false; + for (JsonNode child : node) { + modified |= replaceFloat32VectorInResponse(child); + } + return modified; + } + return false; + } + + /** + * Returns {@code true} if {@code haystack} contains {@code needle} as an ASCII substring. Used + * for fast pre-screening of JSON bodies before committing to a full parse. + */ + private static boolean containsAsciiSubstring(byte[] haystack, String needle) { + byte[] needleBytes = needle.getBytes(StandardCharsets.US_ASCII); + outer: + for (int i = 0; i <= haystack.length - needleBytes.length; i++) { + for (int j = 0; j < needleBytes.length; j++) { + if (haystack[i + j] != needleBytes[j]) { + continue outer; + } + } + return true; + } + return false; + } + + // ------------------------------------------------------------------------- + // I/O helpers + // ------------------------------------------------------------------------- + + private static String getTarget(Context.ModifyHttpRequest context) { + return getTarget(context.httpRequest()); + } + + private static String getTarget(SdkHttpRequest httpRequest) { + List targets = httpRequest.headers().get("X-Amz-Target"); + if (targets == null || targets.isEmpty()) { + return null; + } + return targets.get(0); + } + + private static byte[] readBytes(Optional requestBodyOpt) throws IOException { + if (!requestBodyOpt.isPresent()) { + return null; + } + return readAllBytes(requestBodyOpt.get().contentStreamProvider().newStream()); + } + + private static byte[] readAllBytes(InputStream in) throws IOException { + try { + byte[] buf = new byte[4096]; + ByteArrayOutputStream out = new ByteArrayOutputStream(); + int n; + while ((n = in.read(buf)) != -1) { + out.write(buf, 0, n); + } + return out.toByteArray(); + } finally { + in.close(); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java new file mode 100644 index 0000000..aedccc4 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java @@ -0,0 +1,33 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.util.List; + +/** + * Mutable holder used by {@link VectorSearchInterceptor} to pass extra response fields back to + * {@link VectorSearchSupport} without exposing internal state to callers. + */ +final class VectorSearchResultHolder { + + private List scores; + private List vectorIndexes; + + List getScores() { + return scores; + } + + void setScores(List scores) { + this.scores = scores; + } + + List getVectorIndexes() { + return vectorIndexes; + } + + void setVectorIndexes(List vectorIndexes) { + this.vectorIndexes = vectorIndexes; + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java new file mode 100644 index 0000000..8777767 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java @@ -0,0 +1,315 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.CreateTableResponse; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; +import software.amazon.awssdk.services.dynamodb.model.UpdateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.UpdateTableResponse; + +/** + * Utility facade for Alternator's vector search extension. + * + *

Alternator extends the DynamoDB API with vector indexes and vector similarity search. Because + * the standard AWS SDK for Java does not know about these extensions, this class provides + * convenience methods that attach the extra parameters to standard SDK requests via {@link + * VectorSearchInterceptor}. + * + *

Setup

+ * + *

Register {@link VectorSearchInterceptor#INSTANCE} when building the client once: + * + * 

{@code
+ * DynamoDbClient client = DynamoDbClient.builder()
+ *     .overrideConfiguration(c ->
+ *         c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE))
+ *     .endpointOverride(URI.create("http://localhost:8000"))
+ *     .credentialsProvider(...)
+ *     .build();
+ * }
+ * + *

CreateTable with a vector index

+ * + * 
{@code
+ * VectorIndex vi = VectorIndex.builder()
+ *     .indexName("embedding-index")
+ *     .vectorAttribute(VectorAttribute.builder()
+ *         .attributeName("embedding")
+ *         .dimensions(128)
+ *         .build())
+ *     .similarityFunction("COSINE")
+ *     .build();
+ *
+ * CreateTableRequest base = CreateTableRequest.builder()
+ *     .tableName("items")
+ *     .keySchema(KeySchemaElement.builder().attributeName("id").keyType(KeyType.HASH).build())
+ *     .attributeDefinitions(
+ *         AttributeDefinition.builder().attributeName("id").attributeType(ScalarAttributeType.S).build())
+ *     .billingMode(BillingMode.PAY_PER_REQUEST)
+ *     .build();
+ *
+ * CreateTableResponse resp = client.createTable(
+ *     VectorSearchSupport.withVectorIndexes(base, List.of(vi)));
+ * }
+ * + *

Query with vector similarity search

+ * + * 
{@code
+ * VectorSearch vs = VectorSearch.builder()
+ *     .queryVector(new float[]{0.1f, 0.2f, 0.3f, ...})
+ *     .returnScores(true)
+ *     .build();
+ *
+ * QueryRequest qr = QueryRequest.builder()
+ *     .tableName("items")
+ *     .indexName("embedding-index")
+ *     .limit(10)
+ *     .build();
+ *
+ * VectorQueryResult result = VectorSearchSupport.query(client, qr, vs);
+ * result.items().forEach(item -> System.out.println(item));
+ * result.scores().forEach(score -> System.out.println("score: " + score));
+ * }
+ */ +public final class VectorSearchSupport { + + private VectorSearchSupport() {} + + // ------------------------------------------------------------------------- + // Request enrichment helpers (for use with standard client.operation()) + // ------------------------------------------------------------------------- + + /** + * Returns a copy of {@code request} with the given vector indexes attached so that {@link + * VectorSearchInterceptor} will inject them as the {@code VectorIndexes} field in the {@code + * CreateTable} JSON body. + * + *

Any existing {@code overrideConfiguration} on the request is preserved. + */ + public static CreateTableRequest withVectorIndexes( + CreateTableRequest request, List vectorIndexes) { + return request.toBuilder() + .overrideConfiguration( + mergeExecutionAttribute( + request.overrideConfiguration().orElse(null), + VectorSearchInterceptor.VECTOR_INDEXES, + vectorIndexes)) + .build(); + } + + /** + * Returns a copy of {@code request} with the given vector index updates attached so that {@link + * VectorSearchInterceptor} will inject them as the {@code VectorIndexUpdates} field in the {@code + * UpdateTable} JSON body. + * + *

Any existing {@code overrideConfiguration} on the request is preserved. + */ + public static UpdateTableRequest withVectorIndexUpdates( + UpdateTableRequest request, List vectorIndexUpdates) { + return request.toBuilder() + .overrideConfiguration( + mergeExecutionAttribute( + request.overrideConfiguration().orElse(null), + VectorSearchInterceptor.VECTOR_INDEX_UPDATES, + vectorIndexUpdates)) + .build(); + } + + /** + * Returns a copy of {@code request} with the given vector search parameters attached so that + * {@link VectorSearchInterceptor} will inject the {@code VectorSearch} field in the {@code Query} + * JSON body and capture the {@code Scores} field in the response. + * + *

Use this when you need the raw {@link QueryResponse} and will retrieve scores separately via + * a {@link VectorSearchResultHolder}. Prefer {@link #query(DynamoDbClient, QueryRequest, + * VectorSearch)} for a more convenient API that bundles the response and scores. + */ + static QueryRequest withVectorSearch( + QueryRequest request, VectorSearch vectorSearch, VectorSearchResultHolder resultHolder) { + AwsRequestOverrideConfiguration base = request.overrideConfiguration().orElse(null); + AwsRequestOverrideConfiguration config = + mergeExecutionAttributes( + base, + VectorSearchInterceptor.VECTOR_SEARCH, + vectorSearch, + VectorSearchInterceptor.RESULT_HOLDER, + resultHolder); + return request.toBuilder().overrideConfiguration(config).build(); + } + + // ------------------------------------------------------------------------- + // Convenience methods that bundle request + response + // ------------------------------------------------------------------------- + + /** + * Executes a vector similarity {@code Query} and returns the items together with any per-item + * similarity scores. + * + *

The client must have {@link VectorSearchInterceptor#INSTANCE} registered (see class + * javadoc). + * + * @param client the DynamoDB client + * @param request the base {@code QueryRequest}; the {@code VectorSearch} parameter will be + * injected automatically + * @param vectorSearch the vector search parameters + * @return a {@link VectorQueryResult} wrapping the response and scores + */ + public static VectorQueryResult query( + DynamoDbClient client, QueryRequest request, VectorSearch vectorSearch) { + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + QueryRequest enriched = withVectorSearch(request, vectorSearch, holder); + QueryResponse response = client.query(enriched); + return new VectorQueryResult(response, holder.getScores()); + } + + /** + * Asynchronously executes a vector similarity {@code Query} and returns a future that resolves to + * the items and per-item similarity scores. + * + *

The client must have {@link VectorSearchInterceptor#INSTANCE} registered (see class + * javadoc). + * + * @param client the async DynamoDB client + * @param request the base {@code QueryRequest} + * @param vectorSearch the vector search parameters + * @return a future resolving to a {@link VectorQueryResult} + */ + public static CompletableFuture queryAsync( + DynamoDbAsyncClient client, QueryRequest request, VectorSearch vectorSearch) { + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + QueryRequest enriched = withVectorSearch(request, vectorSearch, holder); + return client + .query(enriched) + .thenApply(resp -> new VectorQueryResult(resp, holder.getScores())); + } + + /** + * Executes a {@code CreateTable} request with the given vector indexes and returns the response. + * + *

The returned vector indexes (as reported by the server, with status fields populated) are + * available in the table description via {@link DescribeTableResponse} or via a subsequent {@link + * #describeTable(DynamoDbClient, DescribeTableRequest)}. + * + * @param client the DynamoDB client + * @param request the base {@code CreateTableRequest} + * @param vectorIndexes the vector indexes to create together with the table + * @return the {@code CreateTableResponse} + */ + public static CreateTableResponse createTable( + DynamoDbClient client, CreateTableRequest request, List vectorIndexes) { + return client.createTable(withVectorIndexes(request, vectorIndexes)); + } + + /** + * Executes an {@code UpdateTable} request that adds or removes vector indexes. + * + * @param client the DynamoDB client + * @param request the base {@code UpdateTableRequest} + * @param vectorIndexUpdates the vector index changes to apply + * @return the {@code UpdateTableResponse} + */ + public static UpdateTableResponse updateTable( + DynamoDbClient client, + UpdateTableRequest request, + List vectorIndexUpdates) { + return client.updateTable(withVectorIndexUpdates(request, vectorIndexUpdates)); + } + + /** + * Executes a {@code DescribeTable} request and returns the standard response alongside any vector + * indexes defined on the table. + * + *

The {@link VectorSearchInterceptor} must be registered on the client. + * + * @param client the DynamoDB client + * @param request the {@code DescribeTableRequest} + * @return a pair of the standard response and the list of vector indexes (may be empty) + */ + public static DescribeTableWithVectorIndexes describeTable( + DynamoDbClient client, DescribeTableRequest request) { + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + DescribeTableRequest enriched = + request.toBuilder() + .overrideConfiguration( + mergeExecutionAttribute( + request.overrideConfiguration().orElse(null), + VectorSearchInterceptor.RESULT_HOLDER, + holder)) + .build(); + DescribeTableResponse response = client.describeTable(enriched); + List indexes = holder.getVectorIndexes(); + return new DescribeTableWithVectorIndexes( + response, indexes != null ? indexes : Collections.emptyList()); + } + + // ------------------------------------------------------------------------- + // Private helpers for building overrideConfiguration + // ------------------------------------------------------------------------- + + private static AwsRequestOverrideConfiguration mergeExecutionAttribute( + AwsRequestOverrideConfiguration existing, ExecutionAttribute key, T value) { + AwsRequestOverrideConfiguration.Builder builder = + existing != null ? existing.toBuilder() : AwsRequestOverrideConfiguration.builder(); + builder.putExecutionAttribute(key, value); + return builder.build(); + } + + @SuppressWarnings("unchecked") + private static AwsRequestOverrideConfiguration mergeExecutionAttributes( + AwsRequestOverrideConfiguration existing, + ExecutionAttribute keyA, + A valueA, + ExecutionAttribute keyB, + B valueB) { + AwsRequestOverrideConfiguration.Builder builder = + existing != null ? existing.toBuilder() : AwsRequestOverrideConfiguration.builder(); + builder.putExecutionAttribute(keyA, valueA); + builder.putExecutionAttribute(keyB, valueB); + return builder.build(); + } + + // ------------------------------------------------------------------------- + // Inner result class for DescribeTable + // ------------------------------------------------------------------------- + + /** + * Holds the result of {@link #describeTable(DynamoDbClient, DescribeTableRequest)}: the standard + * SDK response together with the vector indexes parsed from the raw JSON response. + */ + public static final class DescribeTableWithVectorIndexes { + private final DescribeTableResponse response; + private final List vectorIndexes; + + DescribeTableWithVectorIndexes( + DescribeTableResponse response, List vectorIndexes) { + this.response = response; + this.vectorIndexes = Collections.unmodifiableList(vectorIndexes); + } + + /** Returns the standard {@link DescribeTableResponse}. */ + public DescribeTableResponse response() { + return response; + } + + /** + * Returns the vector indexes defined on the table, or an empty list if none are defined or the + * interceptor was not registered. + */ + public List vectorIndexes() { + return vectorIndexes; + } + } +} diff --git a/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java b/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java new file mode 100644 index 0000000..06daac7 --- /dev/null +++ b/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java @@ -0,0 +1,367 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator; + +import static org.junit.Assert.*; + +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.scylladb.alternator.vectorsearch.*; +import java.util.*; +import org.junit.Test; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +/** + * Unit tests for {@link VectorSearchInterceptor} and {@link VectorSearchSupport}. These tests do + * not require a running Alternator instance and verify the JSON serialisation logic. + */ +public class VectorSearchInterceptorTest { + + // ------------------------------------------------------------------------- + // VectorAttribute / VectorIndex serialisation + // ------------------------------------------------------------------------- + + @Test + public void testVectorAttributeBuilder() { + VectorAttribute va = + VectorAttribute.builder().attributeName("embedding").dimensions(128).build(); + assertEquals("embedding", va.attributeName()); + assertEquals(128, va.dimensions()); + } + + @Test + public void testVectorIndexBuilder() { + VectorAttribute va = VectorAttribute.builder().attributeName("v").dimensions(4).build(); + VectorIndex vi = + VectorIndex.builder() + .indexName("idx") + .vectorAttribute(va) + .similarityFunction("COSINE") + .build(); + + assertEquals("idx", vi.indexName()); + assertEquals("v", vi.vectorAttribute().attributeName()); + assertEquals(4, vi.vectorAttribute().dimensions()); + assertEquals("COSINE", vi.similarityFunction()); + assertNull(vi.projection()); + assertNull(vi.indexStatus()); + assertNull(vi.backfilling()); + } + + // ------------------------------------------------------------------------- + // VectorSearch builder + // ------------------------------------------------------------------------- + + @Test + public void testVectorSearchBuilderWithFloatArray() { + float[] floats = {0.1f, 0.2f, 0.3f}; + VectorSearch vs = VectorSearch.builder().queryVector(floats).returnScores(true).build(); + + assertArrayEquals(floats, vs.queryVectorFloats(), 1e-6f); + assertNull(vs.queryVectorAttributeValue()); + assertTrue(vs.returnScores()); + } + + @Test + public void testVectorSearchBuilderWithAttributeValue() { + AttributeValue av = AttributeValue.fromN("42"); + VectorSearch vs = VectorSearch.builder().queryVector(av).build(); + + assertNull(vs.queryVectorFloats()); + assertEquals(av, vs.queryVectorAttributeValue()); + assertFalse(vs.returnScores()); + } + + @Test(expected = IllegalStateException.class) + public void testVectorSearchBuilderRequiresQueryVector() { + VectorSearch.builder().returnScores(true).build(); + } + + // ------------------------------------------------------------------------- + // attributeValueToJson + // ------------------------------------------------------------------------- + + @Test + public void testAttributeValueToJsonString() throws Exception { + AttributeValue av = AttributeValue.fromS("hello"); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertEquals("hello", node.get("S").asText()); + assertEquals(1, node.size()); + } + + @Test + public void testAttributeValueToJsonNumber() throws Exception { + AttributeValue av = AttributeValue.fromN("42.5"); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertEquals("42.5", node.get("N").asText()); + assertEquals(1, node.size()); + } + + @Test + public void testAttributeValueToJsonList() throws Exception { + AttributeValue av = + AttributeValue.fromL( + Arrays.asList( + AttributeValue.fromN("1.0"), + AttributeValue.fromN("2.0"), + AttributeValue.fromN("3.0"))); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertNotNull(node.get("L")); + assertEquals(3, node.get("L").size()); + assertEquals("1.0", node.get("L").get(0).get("N").asText()); + assertEquals("2.0", node.get("L").get(1).get("N").asText()); + assertEquals("3.0", node.get("L").get(2).get("N").asText()); + } + + @Test + public void testAttributeValueToJsonBool() throws Exception { + ObjectNode trueNode = + VectorSearchInterceptor.attributeValueToJson(AttributeValue.fromBool(true)); + assertTrue(trueNode.get("BOOL").asBoolean()); + + ObjectNode falseNode = + VectorSearchInterceptor.attributeValueToJson(AttributeValue.fromBool(false)); + assertFalse(falseNode.get("BOOL").asBoolean()); + } + + @Test + public void testAttributeValueToJsonNull() throws Exception { + AttributeValue av = AttributeValue.fromNul(true); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertTrue(node.get("NULL").asBoolean()); + } + + // ------------------------------------------------------------------------- + // VectorSearchSupport helpers + // ------------------------------------------------------------------------- + + @Test + public void testWithVectorIndexesSetsOverrideConfiguration() { + VectorIndex vi = + VectorIndex.builder() + .indexName("idx") + .vectorAttribute(VectorAttribute.builder().attributeName("v").dimensions(2).build()) + .build(); + + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest base = + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest.builder() + .tableName("test") + .keySchema( + software.amazon.awssdk.services.dynamodb.model.KeySchemaElement.builder() + .attributeName("pk") + .keyType(software.amazon.awssdk.services.dynamodb.model.KeyType.HASH) + .build()) + .attributeDefinitions( + software.amazon.awssdk.services.dynamodb.model.AttributeDefinition.builder() + .attributeName("pk") + .attributeType( + software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType.S) + .build()) + .billingMode(software.amazon.awssdk.services.dynamodb.model.BillingMode.PAY_PER_REQUEST) + .build(); + + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest enriched = + VectorSearchSupport.withVectorIndexes(base, Collections.singletonList(vi)); + + assertNotNull(enriched); + assertTrue(enriched.overrideConfiguration().isPresent()); + // The VECTOR_INDEXES attribute must be set + software.amazon.awssdk.core.interceptor.ExecutionAttributes ea = + enriched.overrideConfiguration().get().executionAttributes(); + List vis = ea.getAttribute(VectorSearchInterceptor.VECTOR_INDEXES); + assertNotNull(vis); + assertEquals(1, vis.size()); + assertEquals("idx", vis.get(0).indexName()); + } + + @Test + public void testVectorIndexUpdateBuilder() { + CreateVectorIndexAction createAction = + CreateVectorIndexAction.builder() + .indexName("new-idx") + .vectorAttribute(VectorAttribute.builder().attributeName("v").dimensions(3).build()) + .similarityFunction("DOT_PRODUCT") + .build(); + VectorIndexUpdate update = VectorIndexUpdate.builder().create(createAction).build(); + + assertNotNull(update.create()); + assertNull(update.delete()); + assertEquals("new-idx", update.create().indexName()); + assertEquals("DOT_PRODUCT", update.create().similarityFunction()); + + DeleteVectorIndexAction deleteAction = + DeleteVectorIndexAction.builder().indexName("old-idx").build(); + VectorIndexUpdate deleteUpdate = VectorIndexUpdate.builder().delete(deleteAction).build(); + + assertNull(deleteUpdate.create()); + assertNotNull(deleteUpdate.delete()); + assertEquals("old-idx", deleteUpdate.delete().indexName()); + } + + @Test + public void testVectorQueryResultEmptyScores() { + software.amazon.awssdk.services.dynamodb.model.QueryResponse response = + software.amazon.awssdk.services.dynamodb.model.QueryResponse.builder().build(); + VectorQueryResult result = new VectorQueryResult(response, null); + assertNotNull(result.scores()); + assertTrue(result.scores().isEmpty()); + } + + @Test + public void testVectorQueryResultWithScores() { + software.amazon.awssdk.services.dynamodb.model.QueryResponse response = + software.amazon.awssdk.services.dynamodb.model.QueryResponse.builder().build(); + List scores = Arrays.asList(0.9, 0.8, 0.7); + VectorQueryResult result = new VectorQueryResult(response, scores); + assertEquals(3, result.scores().size()); + assertEquals(0.9, result.scores().get(0), 1e-9); + } + + // ------------------------------------------------------------------------- + // Float32Vector — encoding / decoding + // ------------------------------------------------------------------------- + + @Test + public void testFloat32VectorRoundTrip() { + float[] original = {1.0f, -2.5f, 0.0f, Float.MAX_VALUE, Float.MIN_VALUE}; + AttributeValue av = Float32Vector.toAttributeValue(original); + + assertTrue("toAttributeValue must produce a B attribute", av.b() != null); + assertTrue("isFloat32Vector must return true", Float32Vector.isFloat32Vector(av)); + + float[] decoded = Float32Vector.toFloats(av); + assertArrayEquals("round-trip must be lossless", original, decoded, 0.0f); + } + + @Test + public void testFloat32VectorEmptyArray() { + float[] empty = {}; + AttributeValue av = Float32Vector.toAttributeValue(empty); + assertTrue(Float32Vector.isFloat32Vector(av)); + assertArrayEquals(empty, Float32Vector.toFloats(av), 0.0f); + } + + @Test + public void testFloat32VectorIsNotStringAttribute() { + assertFalse(Float32Vector.isFloat32Vector(AttributeValue.fromS("hello"))); + } + + @Test + public void testFloat32VectorIsNotNumberAttribute() { + assertFalse(Float32Vector.isFloat32Vector(AttributeValue.fromN("42"))); + } + + @Test + public void testFloat32VectorIsNotArbitraryBinaryAttribute() { + // A B attribute without the magic prefix must not be considered a Float32Vector. + assertFalse( + Float32Vector.isFloat32Vector( + AttributeValue.fromB(SdkBytes.fromByteArray(new byte[] {1, 2, 3, 4, 5, 6, 7, 8})))); + } + + @Test + public void testFloat32VectorBinaryTooShortForMagic() { + assertFalse( + Float32Vector.isFloat32Vector( + AttributeValue.fromB(SdkBytes.fromByteArray(new byte[] {(byte) 0xF2})))); + } + + @Test(expected = IllegalArgumentException.class) + public void testFloat32VectorToFloatsRequiresMagic() { + Float32Vector.toFloats(AttributeValue.fromS("not a vector")); + } + + @Test + public void testFloat32VectorBase64PrefixIsCorrect() { + // Build a Float32Vector and verify the B field in its JSON representation + // starts with a known fixed prefix (derived from the 8-byte magic). + float[] values = {1.0f, 2.0f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + String b64 = java.util.Base64.getEncoder().encodeToString(av.b().asByteArray()); + // The first 8 base64 chars are fully determined by the first 6 magic bytes (two groups of 3). + // We verify stability of the encoding rather than hard-coding the string here. + AttributeValue av2 = Float32Vector.toAttributeValue(new float[] {9.0f, 8.0f}); + String b64_2 = java.util.Base64.getEncoder().encodeToString(av2.b().asByteArray()); + assertEquals( + "All Float32Vectors must share the same 8-char base64 prefix", + b64.substring(0, 8), + b64_2.substring(0, 8)); + } + + // ------------------------------------------------------------------------- + // Float32Vector — JSON-level request replacement (B → FLOAT32VECTOR) + // ------------------------------------------------------------------------- + + @Test + public void testAttributeValueToJsonFloat32VectorUsesCompactFormat() throws Exception { + // When Float32Vector.toAttributeValue() is used, attributeValueToJson should emit + // {"FLOAT32VECTOR": [...]} rather than the generic {"B": "..."}. + float[] values = {0.5f, -1.5f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + ObjectNode json = VectorSearchInterceptor.attributeValueToJson(av); + + assertNull("B key must be absent for Float32Vector", json.get("B")); + assertNotNull("FLOAT32VECTOR key must be present", json.get("FLOAT32VECTOR")); + assertEquals(2, json.get("FLOAT32VECTOR").size()); + assertEquals(0.5f, (float) json.get("FLOAT32VECTOR").get(0).asDouble(), 1e-6f); + assertEquals(-1.5f, (float) json.get("FLOAT32VECTOR").get(1).asDouble(), 1e-6f); + } + + @Test + public void testAttributeValueToJsonOrdinaryBinaryIsBase64() throws Exception { + // A regular B attribute (no magic) must be emitted as base64 string. + byte[] rawBytes = {1, 2, 3, 4}; + AttributeValue av = AttributeValue.fromB(SdkBytes.fromByteArray(rawBytes)); + ObjectNode json = VectorSearchInterceptor.attributeValueToJson(av); + + assertNotNull("B key must be present for plain binary", json.get("B")); + assertNull("FLOAT32VECTOR key must be absent for plain binary", json.get("FLOAT32VECTOR")); + } + + // ------------------------------------------------------------------------- + // Float32Vector — JSON-level response replacement (FLOAT32VECTOR → L) + // ------------------------------------------------------------------------- + + @Test + public void testFloat32VectorWritePathProducesCompactFormat() throws Exception { + // Verify that encoding float[] → toAttributeValue → attributeValueToJson produces + // {"FLOAT32VECTOR": [...]} (compact wire format for writes). + float[] original = {1.0f, 2.0f, 3.0f}; + + AttributeValue av = Float32Vector.toAttributeValue(original); + ObjectNode asJson = VectorSearchInterceptor.attributeValueToJson(av); + assertNotNull( + "FLOAT32VECTOR key must be present in compact wire format", asJson.get("FLOAT32VECTOR")); + assertNull("B key must be absent for Float32Vector", asJson.get("B")); + assertEquals(3, asJson.get("FLOAT32VECTOR").size()); + } + + @Test + public void testFloat32VectorBytesRoundTrip() { + // Round-trip through toAttributeValue → isFloat32Vector → toFloats + float[] values = {-1.0f, 0.0f, 3.14159f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + assertTrue(Float32Vector.isFloat32Vector(av)); + float[] back = Float32Vector.toFloats(av); + assertArrayEquals(values, back, 1e-6f); + } + + @Test + public void testFloat32VectorFromListRoundTrip() { + // Simulate the read-back path: interceptor returns an L-typed AttributeValue; + // re-encoding it via toAttributeValue(List) must produce the same result as + // toAttributeValue(float...) with the original values. + float[] original = {1.5f, -2.5f, 0.0f}; + AttributeValue fromFloats = Float32Vector.toAttributeValue(original); + + // Build an L-typed AttributeValue as the interceptor would return + List numbers = + java.util.Arrays.asList( + AttributeValue.fromN("1.5"), AttributeValue.fromN("-2.5"), AttributeValue.fromN("0.0")); + AttributeValue fromList = Float32Vector.toAttributeValue(numbers); + + assertArrayEquals(fromFloats.b().asByteArray(), fromList.b().asByteArray()); + } +} From 70a32f23015784da96137aa2d4024d464ba6984c Mon Sep 17 00:00:00 2001 From: Nadav Har'El Date: Tue, 16 Jun 2026 15:46:35 +0300 Subject: [PATCH 2/2] tests: allow disabling HTTPS variants by setting ALTERNATOR_HTTPS_PORT= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When ALTERNATOR_HTTPS_PORT is set to an empty string, httpAndHttpsEndpoints() returns only the HTTP entry, so parameterized test classes run only the HTTP variant. This avoids spurious connection-refused errors when testing against an Alternator node that is only reachable over HTTP. Example — run integration tests against an HTTP-only node: INTEGRATION_TESTS=true \ ALTERNATOR_HOST=127.8.209.117 \ ALTERNATOR_PORT=8000 \ ALTERNATOR_HTTPS_PORT= \ mvn test -Dtest=VectorSearchIT --- .../alternator/IntegrationTestConfig.java | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/src/test/java/com/scylladb/alternator/IntegrationTestConfig.java b/src/test/java/com/scylladb/alternator/IntegrationTestConfig.java index 2263e2f..46d4648 100644 --- a/src/test/java/com/scylladb/alternator/IntegrationTestConfig.java +++ b/src/test/java/com/scylladb/alternator/IntegrationTestConfig.java @@ -3,8 +3,9 @@ import java.net.URI; import java.net.URISyntaxException; import java.nio.file.Path; -import java.util.Arrays; +import java.util.ArrayList; import java.util.Collection; +import java.util.List; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; @@ -17,7 +18,8 @@ *