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/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 @@ *