From 56b84da4877b7af5d6001a8ecedba5cc541cb985 Mon Sep 17 00:00:00 2001 From: Dmitry Kropachev Date: Tue, 21 Jul 2026 10:12:26 -0400 Subject: [PATCH] Add Alternator vector search support --- README.md | 237 ++++ pom.xml | 16 + .../scylladb/alternator/VectorSearchIT.java | 488 ++++++++ .../AlternatorDynamoDbAsyncClient.java | 34 +- .../alternator/AlternatorDynamoDbClient.java | 34 +- .../vectorsearch/CreateVectorIndexAction.java | 97 ++ .../vectorsearch/DeleteVectorIndexAction.java | 50 + .../vectorsearch/Float32Vector.java | 213 ++++ .../vectorsearch/VectorAttribute.java | 67 ++ .../alternator/vectorsearch/VectorIndex.java | 152 +++ .../vectorsearch/VectorIndexUpdate.java | 81 ++ .../vectorsearch/VectorQueryResult.java | 67 ++ .../alternator/vectorsearch/VectorSearch.java | 119 ++ .../vectorsearch/VectorSearchInterceptor.java | 1003 +++++++++++++++++ .../VectorSearchResultHolder.java | 33 + .../vectorsearch/VectorSearchSupport.java | 315 ++++++ .../VectorSearchInterceptorTest.java | 367 ++++++ .../VectorSearchHttpInterceptorTest.java | 727 ++++++++++++ 18 files changed, 4070 insertions(+), 30 deletions(-) create mode 100644 src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/CreateVectorIndexAction.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/DeleteVectorIndexAction.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/Float32Vector.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorAttribute.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorIndex.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorIndexUpdate.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorQueryResult.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearch.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java create mode 100644 src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java create mode 100644 src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java create mode 100644 src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java diff --git a/README.md b/README.md index cb71a03..370f71d 100644 --- a/README.md +++ b/README.md @@ -973,3 +973,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 where partition-key names cannot be pre-configured + +### 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 9f810a1..0eea213 100644 --- a/pom.xml +++ b/pom.xml @@ -27,6 +27,14 @@ pom import + + + com.fasterxml.jackson.core + jackson-databind + 2.22.1 + @@ -51,6 +59,14 @@ aws-crt-client provided + + + com.fasterxml.jackson.core + jackson-databind + 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..183aa9c --- /dev/null +++ b/src/integration-test/java/com/scylladb/alternator/VectorSearchIT.java @@ -0,0 +1,488 @@ +// 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() { + assumeVectorStoreEnabled(tableName, "vi1"); + } + + private void assumeVectorStoreEnabled(String targetTableName, String indexName) { + // Case 1: server does not know about vector indexes at all. + VectorSearchSupport.DescribeTableWithVectorIndexes desc = + VectorSearchSupport.describeTable( + client, DescribeTableRequest.builder().tableName(targetTableName).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(targetTableName) + .indexName(indexName) + .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(); + + createTableWithVectorIndex(base, Collections.singletonList(vi)); + } + + private void createTableWithVectorIndex( + CreateTableRequest base, List vectorIndexes) { + try { + VectorSearchSupport.createTable(client, base, vectorIndexes); + } catch (DynamoDbException e) { + skipIfVectorFeatureUnsupported(e); + throw e; + } + } + + private static void skipIfVectorFeatureUnsupported(DynamoDbException e) { + if (isVectorFeatureUnsupported(e)) { + assumeTrue("Skipping: server does not support vector indexes: " + e.getMessage(), false); + } + } + + private static boolean isVectorFeatureUnsupported(DynamoDbException e) { + String message = exceptionText(e); + if (message.contains("vector store is disabled")) { + return true; + } + if ((message.contains("unknown") || message.contains("unrecognized")) + && message.contains("vector")) { + return true; + } + return (message.contains("vectorindexes") + || message.contains("vector indexes") + || message.contains("vector store")) + && (message.contains("unsupported") + || message.contains("not supported") + || message.contains("disabled")); + } + + private static String exceptionText(DynamoDbException e) { + StringBuilder text = new StringBuilder(); + if (e.awsErrorDetails() != null) { + if (e.awsErrorDetails().errorCode() != null) { + text.append(e.awsErrorDetails().errorCode()).append(' '); + } + if (e.awsErrorDetails().errorMessage() != null) { + text.append(e.awsErrorDetails().errorMessage()).append(' '); + } + } + if (e.getMessage() != null) { + text.append(e.getMessage()); + } + return text.toString().toLowerCase(Locale.ROOT); + } + + // ------------------------------------------------------------------------- + // 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); + } + assumeVectorStoreEnabled(); + 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(); + + try { + // Must not throw — Alternator should accept this similarity function value. + createTableWithVectorIndex(base, Collections.singletonList(vi)); + assumeVectorStoreEnabled(tbl, "vi-" + sf.toLowerCase()); + } finally { + try { + client.deleteTable(DeleteTableRequest.builder().tableName(tbl).build()); + } catch (ResourceNotFoundException ignored) { + } + } + } + } + + /** + * 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 { + createTableWithVectorIndex(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 15df7f7..0930e6b 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; @@ -732,21 +733,6 @@ public AlternatorDynamoDbAsyncClientWrapper buildWithAlternatorAPI() { AlternatorConfig alternatorConfig = configBuilder.build(); - ClientOverrideConfiguration.Builder compressionOverrideBuilder = - delegate.overrideConfiguration() != null - ? delegate.overrideConfiguration().toBuilder() - : ClientOverrideConfiguration.builder(); - if (alternatorConfig.isResponseCompressionEnabled()) { - compressionOverrideBuilder.addExecutionInterceptor( - new ResponseCompressionInterceptor( - alternatorConfig.getResponseCompressionAlgorithms())); - } - if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { - compressionOverrideBuilder.addExecutionInterceptor( - new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); - } - delegate.overrideConfiguration(compressionOverrideBuilder.build()); - TlsConfig tlsConfig = alternatorConfig.getTlsConfig(); if (!httpClientSet) { SdkAsyncHttpClient mainClient = @@ -768,6 +754,24 @@ public AlternatorDynamoDbAsyncClientWrapper buildWithAlternatorAPI() { ? delegate.overrideConfiguration().toBuilder() : ClientOverrideConfiguration.builder(); + // Interceptor registration order matters. Request hooks run in the listed order; response + // hooks run in reverse order. + // 1. VectorSearchInterceptor injects VectorIndexes/VectorSearch JSON and converts + // FLOAT32VECTOR markers; it must see the original request body. + // 2. ResponseCompressionInterceptor adds Accept-Encoding and, on responses, decodes before + // vector response parsing. + // 3. GzipRequestInterceptor compresses the already vector-modified request body. + // 4. QueryPlanInterceptor selects the target node; body-independent. + overrideBuilder.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE); + if (alternatorConfig.isResponseCompressionEnabled()) { + overrideBuilder.addExecutionInterceptor( + new ResponseCompressionInterceptor( + alternatorConfig.getResponseCompressionAlgorithms())); + } + if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { + overrideBuilder.addExecutionInterceptor( + new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); + } KeyRouteAffinityConfig keyAffinityConfig = alternatorConfig.getKeyRouteAffinityConfig(); AffinityQueryPlanInterceptor affinityInterceptor = null; if (keyAffinityConfig != null diff --git a/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java b/src/main/java/com/scylladb/alternator/AlternatorDynamoDbClient.java index 632ebe1..fe71895 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; @@ -753,21 +754,6 @@ public AlternatorDynamoDbClientWrapper buildWithAlternatorAPI() { AlternatorConfig alternatorConfig = configBuilder.build(); - ClientOverrideConfiguration.Builder compressionOverrideBuilder = - delegate.overrideConfiguration() != null - ? delegate.overrideConfiguration().toBuilder() - : ClientOverrideConfiguration.builder(); - if (alternatorConfig.isResponseCompressionEnabled()) { - compressionOverrideBuilder.addExecutionInterceptor( - new ResponseCompressionInterceptor( - alternatorConfig.getResponseCompressionAlgorithms())); - } - if (alternatorConfig.getCompressionAlgorithm().isEnabled()) { - compressionOverrideBuilder.addExecutionInterceptor( - new GzipRequestInterceptor(alternatorConfig.getMinCompressionSizeBytes())); - } - delegate.overrideConfiguration(compressionOverrideBuilder.build()); - TlsConfig tlsConfig = alternatorConfig.getTlsConfig(); SdkHttpClient pollingClient = null; if (!httpClientSet) { @@ -790,6 +776,24 @@ public AlternatorDynamoDbClientWrapper buildWithAlternatorAPI() { ? delegate.overrideConfiguration().toBuilder() : ClientOverrideConfiguration.builder(); + // Interceptor registration order matters. Request hooks run in the listed order; response + // hooks run in reverse order. + // 1. VectorSearchInterceptor injects VectorIndexes/VectorSearch JSON and converts + // FLOAT32VECTOR markers; it must see the original request body. + // 2. ResponseCompressionInterceptor adds Accept-Encoding and, on responses, decodes before + // vector response parsing. + // 3. GzipRequestInterceptor compresses the already vector-modified request body. + // 4. QueryPlanInterceptor selects the target node; body-independent. + overrideBuilder.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE); + if (alternatorConfig.isResponseCompressionEnabled()) { + overrideBuilder.addExecutionInterceptor( + new ResponseCompressionInterceptor( + alternatorConfig.getResponseCompressionAlgorithms())); + } + 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..8cca38b --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchInterceptor.java @@ -0,0 +1,1003 @@ +// 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.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Base64; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.zip.GZIPInputStream; +import java.util.zip.InflaterInputStream; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +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.http.SdkHttpResponse; +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)} and {@link + * #modifyAsyncHttpResponseContent(Context.ModifyHttpResponse, ExecutionAttributes)} scan all + * response bodies for {@code FLOAT32VECTOR} attribute values and convert 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 ExecutionAttribute> RESPONSE_CONTENT_ENCODINGS = + new ExecutionAttribute<>("AlternatorResponseContentEncodings"); + private static final ExecutionAttribute PROCESSED_RESPONSE_BODY_CACHE = + new ExecutionAttribute<>("AlternatorProcessedResponseBodyCache"); + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final String CONTENT_ENCODING_HEADER = "Content-Encoding"; + private static final String CONTENT_LENGTH_HEADER = "Content-Length"; + private static final String DYNAMODB_CRC32_HEADER = "x-amz-crc32"; + private static final String DYNAMODB_CRC32C_HEADER = "x-amz-crc32c"; + private static final String AWS_CHECKSUM_HEADER_PREFIX = "x-amz-checksum-"; + + // 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_BATCH_EXECUTE_STATEMENT = + "DynamoDB_20120810.BatchExecuteStatement"; + private static final String TARGET_BATCH_GET_ITEM = "DynamoDB_20120810.BatchGetItem"; + private static final String TARGET_BATCH_WRITE_ITEM = "DynamoDB_20120810.BatchWriteItem"; + private static final String TARGET_DELETE_ITEM = "DynamoDB_20120810.DeleteItem"; + private static final String TARGET_EXECUTE_STATEMENT = "DynamoDB_20120810.ExecuteStatement"; + private static final String TARGET_EXECUTE_TRANSACTION = "DynamoDB_20120810.ExecuteTransaction"; + private static final String TARGET_GET_ITEM = "DynamoDB_20120810.GetItem"; + private static final String TARGET_PUT_ITEM = "DynamoDB_20120810.PutItem"; + private static final String TARGET_QUERY = "DynamoDB_20120810.Query"; + private static final String TARGET_SCAN = "DynamoDB_20120810.Scan"; + private static final String TARGET_TRANSACT_GET_ITEMS = "DynamoDB_20120810.TransactGetItems"; + private static final String TARGET_TRANSACT_WRITE_ITEMS = "DynamoDB_20120810.TransactWriteItems"; + private static final String TARGET_UPDATE_ITEM = "DynamoDB_20120810.UpdateItem"; + + 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 { + ProcessedRequestBody processedBody = + computeProcessedRequestBody(context, executionAttributes); + if (processedBody == null) { + return context.requestBody(); + } + if (processedBody.bodyModified()) { + executionAttributes.putAttribute(MODIFIED_BODY_CACHE, processedBody.bytes()); + } + return Optional.of(RequestBody.fromBytes(processedBody.bytes())); + } catch (IOException e) { + throw new RuntimeException("Failed to process vector search parameters in request body", e); + } + } + + /** + * Computes the request body bytes, including vector rewrites when needed. If the body was read + * but not changed, the original bytes are still returned so later interceptors can replay the + * body safely. + */ + private static ProcessedRequestBody computeProcessedRequestBody( + 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 new ProcessedRequestBody(originalBytes, false); + } + + 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 new ProcessedRequestBody(originalBytes, false); + } + + byte[] outBytes = MAPPER.writeValueAsBytes(json); + return new ProcessedRequestBody(outBytes, true); + } + + private static final class ProcessedRequestBody { + private final byte[] bytes; + private final boolean bodyModified; + + private ProcessedRequestBody(byte[] bytes, boolean bodyModified) { + this.bytes = bytes; + this.bodyModified = bodyModified; + } + + private byte[] bytes() { + return bytes; + } + + private boolean bodyModified() { + return bodyModified; + } + } + + // ------------------------------------------------------------------------- + // 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 SdkHttpResponse modifyHttpResponse( + Context.ModifyHttpResponse context, ExecutionAttributes executionAttributes) { + SdkHttpResponse response = context.httpResponse(); + List contentEncodings = getContentEncodings(response); + executionAttributes.putAttribute(RESPONSE_CONTENT_ENCODINGS, contentEncodings); + + Optional bodyOpt = context.responseBody(); + if (bodyOpt.isPresent()) { + VectorSearchResultHolder holder = executionAttributes.getAttribute(RESULT_HOLDER); + try { + ProcessedResponseBody processedBody = + processResponseBody( + readAllBytes(bodyOpt.get()), context.httpRequest(), holder, contentEncodings); + executionAttributes.putAttribute(PROCESSED_RESPONSE_BODY_CACHE, processedBody); + return processedBody.bodyModified() + ? stripStaleResponseBodyHeaders(response, !contentEncodings.isEmpty()) + : response; + } catch (IOException e) { + throw new RuntimeException("Failed to process vector search fields in response body", e); + } + } + + boolean asyncBodyMayBeModified = + executionAttributes.getAttribute(VECTOR_SEARCH) != null + || responseMayContainAttributeValues(context.httpRequest()); + return contentEncodings.isEmpty() && !asyncBodyMayBeModified + ? response + : stripStaleResponseBodyHeaders(response, !contentEncodings.isEmpty()); + } + + @Override + public Optional modifyHttpResponseContent( + Context.ModifyHttpResponse context, ExecutionAttributes executionAttributes) { + + VectorSearchResultHolder holder = executionAttributes.getAttribute(RESULT_HOLDER); + List contentEncodings = executionAttributes.getAttribute(RESPONSE_CONTENT_ENCODINGS); + + Optional bodyOpt = context.responseBody(); + if (!bodyOpt.isPresent()) { + return bodyOpt; + } + + ProcessedResponseBody cachedBody = + executionAttributes.getAttribute(PROCESSED_RESPONSE_BODY_CACHE); + if (cachedBody != null) { + return Optional.of(new ByteArrayInputStream(cachedBody.bytes())); + } + + try { + byte[] bytes = readAllBytes(bodyOpt.get()); + ProcessedResponseBody processedBody = + processResponseBody(bytes, context.httpRequest(), holder, contentEncodings); + return Optional.of(new ByteArrayInputStream(processedBody.bytes())); + + } catch (IOException e) { + throw new RuntimeException("Failed to process vector search fields in response body", e); + } + } + + @Override + public Optional> modifyAsyncHttpResponseContent( + Context.ModifyHttpResponse context, ExecutionAttributes executionAttributes) { + Optional> publisherOpt = context.responsePublisher(); + if (!publisherOpt.isPresent()) { + return publisherOpt; + } + + VectorSearchResultHolder holder = executionAttributes.getAttribute(RESULT_HOLDER); + List contentEncodings = executionAttributes.getAttribute(RESPONSE_CONTENT_ENCODINGS); + return Optional.of( + new TransformingResponsePublisher( + publisherOpt.get(), context.httpRequest(), holder, contentEncodings)); + } + + private static ProcessedResponseBody processResponseBody( + byte[] bytes, + SdkHttpRequest httpRequest, + VectorSearchResultHolder holder, + List contentEncodings) + throws IOException { + boolean modified = false; + if (contentEncodings != null && !contentEncodings.isEmpty() && bytes.length > 0) { + bytes = decompressResponseBody(bytes, contentEncodings); + modified = true; + } + if (bytes.length == 0) { + return new ProcessedResponseBody(bytes, modified); + } + + // 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 new ProcessedResponseBody(bytes, modified); + } + + JsonNode json = MAPPER.readTree(bytes); + boolean jsonModified = false; + + // Convert FLOAT32VECTOR -> standard L (list-of-numbers) so the SDK can unmarshal items + // normally. + if (hasFloat32Vector) { + jsonModified = replaceFloat32VectorInResponse(json); + } + + // Extract Scores / VectorIndexes for callers using VectorSearchSupport. + if (holder != null) { + String target = getTarget(httpRequest); + if (TARGET_QUERY.equals(target)) { + extractScores(json, holder); + } else if (TARGET_CREATE_TABLE.equals(target) || TARGET_DESCRIBE_TABLE.equals(target)) { + extractTableDescriptionVectorIndexes(json, holder); + } + } + + if (jsonModified) { + bytes = MAPPER.writeValueAsBytes(json); + modified = true; + } + return new ProcessedResponseBody(bytes, modified); + } + + private static final class ProcessedResponseBody { + private final byte[] bytes; + private final boolean bodyModified; + + private ProcessedResponseBody(byte[] bytes, boolean bodyModified) { + this.bytes = bytes; + this.bodyModified = bodyModified; + } + + private byte[] bytes() { + return bytes; + } + + private boolean bodyModified() { + return bodyModified; + } + } + + private static final class TransformingResponsePublisher implements Publisher { + private final Publisher delegate; + private final SdkHttpRequest httpRequest; + private final VectorSearchResultHolder holder; + private final List contentEncodings; + + private TransformingResponsePublisher( + Publisher delegate, + SdkHttpRequest httpRequest, + VectorSearchResultHolder holder, + List contentEncodings) { + this.delegate = delegate; + this.httpRequest = httpRequest; + this.holder = holder; + this.contentEncodings = contentEncodings; + } + + @Override + public void subscribe(Subscriber subscriber) { + delegate.subscribe( + new BufferingResponseSubscriber(subscriber, httpRequest, holder, contentEncodings)); + } + } + + private static final class BufferingResponseSubscriber implements Subscriber { + private final Subscriber downstream; + private final SdkHttpRequest httpRequest; + private final VectorSearchResultHolder holder; + private final List contentEncodings; + private final ByteArrayOutputStream body = new ByteArrayOutputStream(); + private final AtomicBoolean requested = new AtomicBoolean(); + private Subscription upstream; + + private BufferingResponseSubscriber( + Subscriber downstream, + SdkHttpRequest httpRequest, + VectorSearchResultHolder holder, + List contentEncodings) { + this.downstream = downstream; + this.httpRequest = httpRequest; + this.holder = holder; + this.contentEncodings = contentEncodings; + } + + @Override + public void onSubscribe(Subscription subscription) { + this.upstream = subscription; + downstream.onSubscribe( + new Subscription() { + @Override + public void request(long n) { + if (n <= 0) { + upstream.cancel(); + downstream.onError( + new IllegalArgumentException( + "Reactive Streams request amount must be positive")); + return; + } + if (requested.compareAndSet(false, true)) { + upstream.request(Long.MAX_VALUE); + } + } + + @Override + public void cancel() { + upstream.cancel(); + } + }); + } + + @Override + public void onNext(ByteBuffer byteBuffer) { + ByteBuffer copy = byteBuffer.asReadOnlyBuffer(); + byte[] chunk = new byte[copy.remaining()]; + copy.get(chunk); + body.write(chunk, 0, chunk.length); + } + + @Override + public void onError(Throwable throwable) { + downstream.onError(throwable); + } + + @Override + public void onComplete() { + try { + ProcessedResponseBody processedBody = + processResponseBody(body.toByteArray(), httpRequest, holder, contentEncodings); + byte[] outBytes = processedBody.bytes(); + if (outBytes.length > 0) { + downstream.onNext(ByteBuffer.wrap(outBytes)); + } + downstream.onComplete(); + } catch (IOException e) { + downstream.onError( + 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 boolean responseMayContainAttributeValues(SdkHttpRequest httpRequest) { + String target = getTarget(httpRequest); + if (target == null) { + return false; + } + switch (target) { + case TARGET_BATCH_EXECUTE_STATEMENT: + case TARGET_BATCH_GET_ITEM: + case TARGET_BATCH_WRITE_ITEM: + case TARGET_DELETE_ITEM: + case TARGET_EXECUTE_STATEMENT: + case TARGET_EXECUTE_TRANSACTION: + case TARGET_GET_ITEM: + case TARGET_PUT_ITEM: + case TARGET_QUERY: + case TARGET_SCAN: + case TARGET_TRANSACT_GET_ITEMS: + case TARGET_TRANSACT_WRITE_ITEMS: + case TARGET_UPDATE_ITEM: + return true; + default: + return false; + } + } + + private static byte[] readBytes(Optional requestBodyOpt) throws IOException { + if (!requestBodyOpt.isPresent()) { + return null; + } + return readAllBytes(requestBodyOpt.get().contentStreamProvider().newStream()); + } + + private static List getContentEncodings(SdkHttpResponse httpResponse) { + List encodings = new ArrayList<>(); + for (String headerValue : httpResponse.matchingHeaders(CONTENT_ENCODING_HEADER)) { + for (String encoding : headerValue.split(",")) { + String normalized = encoding.trim().toLowerCase(Locale.ROOT); + if (normalized.isEmpty() || "identity".equals(normalized)) { + continue; + } + if (!"gzip".equals(normalized) + && !"x-gzip".equals(normalized) + && !"deflate".equals(normalized)) { + return new ArrayList<>(); + } + encodings.add(normalized); + } + } + return encodings; + } + + private static SdkHttpResponse stripStaleResponseBodyHeaders( + SdkHttpResponse httpResponse, boolean stripContentEncoding) { + SdkHttpResponse.Builder builder = httpResponse.toBuilder(); + for (String headerName : httpResponse.headers().keySet()) { + if (shouldStripResponseBodyHeader(headerName, stripContentEncoding)) { + builder.removeHeader(headerName); + } + } + return builder.build(); + } + + private static boolean shouldStripResponseBodyHeader( + String headerName, boolean stripContentEncoding) { + if (isHeader(headerName, CONTENT_LENGTH_HEADER) + || isHeader(headerName, DYNAMODB_CRC32_HEADER) + || isHeader(headerName, DYNAMODB_CRC32C_HEADER)) { + return true; + } + if (headerName != null + && headerName.toLowerCase(Locale.ROOT).startsWith(AWS_CHECKSUM_HEADER_PREFIX)) { + return true; + } + return stripContentEncoding && isHeader(headerName, CONTENT_ENCODING_HEADER); + } + + private static boolean isHeader(String actual, String expected) { + return actual != null && actual.equalsIgnoreCase(expected); + } + + private static byte[] decompressResponseBody(byte[] bytes, List contentEncodings) + throws IOException { + byte[] decoded = bytes; + for (int i = contentEncodings.size() - 1; i >= 0; i--) { + decoded = decompressResponseBody(decoded, contentEncodings.get(i)); + } + return decoded; + } + + private static byte[] decompressResponseBody(byte[] bytes, String contentEncoding) + throws IOException { + if ("gzip".equals(contentEncoding) || "x-gzip".equals(contentEncoding)) { + return readAllBytes(new GZIPInputStream(new ByteArrayInputStream(bytes))); + } + if ("deflate".equals(contentEncoding)) { + return readAllBytes(new InflaterInputStream(new ByteArrayInputStream(bytes))); + } + return bytes; + } + + private static byte[] readAllBytes(InputStream in) throws IOException { + try { + byte[] buf = new byte[4096]; + ByteArrayOutputStream out = new ByteArrayOutputStream(); + int n; + while ((n = in.read(buf)) != -1) { + out.write(buf, 0, n); + } + return out.toByteArray(); + } finally { + in.close(); + } + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java new file mode 100644 index 0000000..aedccc4 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchResultHolder.java @@ -0,0 +1,33 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.util.List; + +/** + * Mutable holder used by {@link VectorSearchInterceptor} to pass extra response fields back to + * {@link VectorSearchSupport} without exposing internal state to callers. + */ +final class VectorSearchResultHolder { + + private List scores; + private List vectorIndexes; + + List getScores() { + return scores; + } + + void setScores(List scores) { + this.scores = scores; + } + + List getVectorIndexes() { + return vectorIndexes; + } + + void setVectorIndexes(List vectorIndexes) { + this.vectorIndexes = vectorIndexes; + } +} diff --git a/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java new file mode 100644 index 0000000..8777767 --- /dev/null +++ b/src/main/java/com/scylladb/alternator/vectorsearch/VectorSearchSupport.java @@ -0,0 +1,315 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration; +import software.amazon.awssdk.core.interceptor.ExecutionAttribute; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.CreateTableResponse; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableRequest; +import software.amazon.awssdk.services.dynamodb.model.DescribeTableResponse; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; +import software.amazon.awssdk.services.dynamodb.model.QueryResponse; +import software.amazon.awssdk.services.dynamodb.model.UpdateTableRequest; +import software.amazon.awssdk.services.dynamodb.model.UpdateTableResponse; + +/** + * Utility facade for Alternator's vector search extension. + * + *

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

Setup

+ * + *

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

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

CreateTable with a vector index

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

Query with vector similarity search

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

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

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

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

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

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

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

The {@link VectorSearchInterceptor} must be registered on the client. + * + * @param client the DynamoDB client + * @param request the {@code DescribeTableRequest} + * @return a pair of the standard response and the list of vector indexes (may be empty) + */ + public static DescribeTableWithVectorIndexes describeTable( + DynamoDbClient client, DescribeTableRequest request) { + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + DescribeTableRequest enriched = + request.toBuilder() + .overrideConfiguration( + mergeExecutionAttribute( + request.overrideConfiguration().orElse(null), + VectorSearchInterceptor.RESULT_HOLDER, + holder)) + .build(); + DescribeTableResponse response = client.describeTable(enriched); + List indexes = holder.getVectorIndexes(); + return new DescribeTableWithVectorIndexes( + response, indexes != null ? indexes : Collections.emptyList()); + } + + // ------------------------------------------------------------------------- + // Private helpers for building overrideConfiguration + // ------------------------------------------------------------------------- + + private static AwsRequestOverrideConfiguration mergeExecutionAttribute( + AwsRequestOverrideConfiguration existing, ExecutionAttribute key, T value) { + AwsRequestOverrideConfiguration.Builder builder = + existing != null ? existing.toBuilder() : AwsRequestOverrideConfiguration.builder(); + builder.putExecutionAttribute(key, value); + return builder.build(); + } + + @SuppressWarnings("unchecked") + private static AwsRequestOverrideConfiguration mergeExecutionAttributes( + AwsRequestOverrideConfiguration existing, + ExecutionAttribute keyA, + A valueA, + ExecutionAttribute keyB, + B valueB) { + AwsRequestOverrideConfiguration.Builder builder = + existing != null ? existing.toBuilder() : AwsRequestOverrideConfiguration.builder(); + builder.putExecutionAttribute(keyA, valueA); + builder.putExecutionAttribute(keyB, valueB); + return builder.build(); + } + + // ------------------------------------------------------------------------- + // Inner result class for DescribeTable + // ------------------------------------------------------------------------- + + /** + * Holds the result of {@link #describeTable(DynamoDbClient, DescribeTableRequest)}: the standard + * SDK response together with the vector indexes parsed from the raw JSON response. + */ + public static final class DescribeTableWithVectorIndexes { + private final DescribeTableResponse response; + private final List vectorIndexes; + + DescribeTableWithVectorIndexes( + DescribeTableResponse response, List vectorIndexes) { + this.response = response; + this.vectorIndexes = Collections.unmodifiableList(vectorIndexes); + } + + /** Returns the standard {@link DescribeTableResponse}. */ + public DescribeTableResponse response() { + return response; + } + + /** + * Returns the vector indexes defined on the table, or an empty list if none are defined or the + * interceptor was not registered. + */ + public List vectorIndexes() { + return vectorIndexes; + } + } +} diff --git a/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java b/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java new file mode 100644 index 0000000..06daac7 --- /dev/null +++ b/src/test/java/com/scylladb/alternator/VectorSearchInterceptorTest.java @@ -0,0 +1,367 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator; + +import static org.junit.Assert.*; + +import com.fasterxml.jackson.databind.node.ObjectNode; +import com.scylladb.alternator.vectorsearch.*; +import java.util.*; +import org.junit.Test; +import software.amazon.awssdk.core.SdkBytes; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +/** + * Unit tests for {@link VectorSearchInterceptor} and {@link VectorSearchSupport}. These tests do + * not require a running Alternator instance and verify the JSON serialisation logic. + */ +public class VectorSearchInterceptorTest { + + // ------------------------------------------------------------------------- + // VectorAttribute / VectorIndex serialisation + // ------------------------------------------------------------------------- + + @Test + public void testVectorAttributeBuilder() { + VectorAttribute va = + VectorAttribute.builder().attributeName("embedding").dimensions(128).build(); + assertEquals("embedding", va.attributeName()); + assertEquals(128, va.dimensions()); + } + + @Test + public void testVectorIndexBuilder() { + VectorAttribute va = VectorAttribute.builder().attributeName("v").dimensions(4).build(); + VectorIndex vi = + VectorIndex.builder() + .indexName("idx") + .vectorAttribute(va) + .similarityFunction("COSINE") + .build(); + + assertEquals("idx", vi.indexName()); + assertEquals("v", vi.vectorAttribute().attributeName()); + assertEquals(4, vi.vectorAttribute().dimensions()); + assertEquals("COSINE", vi.similarityFunction()); + assertNull(vi.projection()); + assertNull(vi.indexStatus()); + assertNull(vi.backfilling()); + } + + // ------------------------------------------------------------------------- + // VectorSearch builder + // ------------------------------------------------------------------------- + + @Test + public void testVectorSearchBuilderWithFloatArray() { + float[] floats = {0.1f, 0.2f, 0.3f}; + VectorSearch vs = VectorSearch.builder().queryVector(floats).returnScores(true).build(); + + assertArrayEquals(floats, vs.queryVectorFloats(), 1e-6f); + assertNull(vs.queryVectorAttributeValue()); + assertTrue(vs.returnScores()); + } + + @Test + public void testVectorSearchBuilderWithAttributeValue() { + AttributeValue av = AttributeValue.fromN("42"); + VectorSearch vs = VectorSearch.builder().queryVector(av).build(); + + assertNull(vs.queryVectorFloats()); + assertEquals(av, vs.queryVectorAttributeValue()); + assertFalse(vs.returnScores()); + } + + @Test(expected = IllegalStateException.class) + public void testVectorSearchBuilderRequiresQueryVector() { + VectorSearch.builder().returnScores(true).build(); + } + + // ------------------------------------------------------------------------- + // attributeValueToJson + // ------------------------------------------------------------------------- + + @Test + public void testAttributeValueToJsonString() throws Exception { + AttributeValue av = AttributeValue.fromS("hello"); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertEquals("hello", node.get("S").asText()); + assertEquals(1, node.size()); + } + + @Test + public void testAttributeValueToJsonNumber() throws Exception { + AttributeValue av = AttributeValue.fromN("42.5"); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertEquals("42.5", node.get("N").asText()); + assertEquals(1, node.size()); + } + + @Test + public void testAttributeValueToJsonList() throws Exception { + AttributeValue av = + AttributeValue.fromL( + Arrays.asList( + AttributeValue.fromN("1.0"), + AttributeValue.fromN("2.0"), + AttributeValue.fromN("3.0"))); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertNotNull(node.get("L")); + assertEquals(3, node.get("L").size()); + assertEquals("1.0", node.get("L").get(0).get("N").asText()); + assertEquals("2.0", node.get("L").get(1).get("N").asText()); + assertEquals("3.0", node.get("L").get(2).get("N").asText()); + } + + @Test + public void testAttributeValueToJsonBool() throws Exception { + ObjectNode trueNode = + VectorSearchInterceptor.attributeValueToJson(AttributeValue.fromBool(true)); + assertTrue(trueNode.get("BOOL").asBoolean()); + + ObjectNode falseNode = + VectorSearchInterceptor.attributeValueToJson(AttributeValue.fromBool(false)); + assertFalse(falseNode.get("BOOL").asBoolean()); + } + + @Test + public void testAttributeValueToJsonNull() throws Exception { + AttributeValue av = AttributeValue.fromNul(true); + ObjectNode node = VectorSearchInterceptor.attributeValueToJson(av); + assertTrue(node.get("NULL").asBoolean()); + } + + // ------------------------------------------------------------------------- + // VectorSearchSupport helpers + // ------------------------------------------------------------------------- + + @Test + public void testWithVectorIndexesSetsOverrideConfiguration() { + VectorIndex vi = + VectorIndex.builder() + .indexName("idx") + .vectorAttribute(VectorAttribute.builder().attributeName("v").dimensions(2).build()) + .build(); + + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest base = + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest.builder() + .tableName("test") + .keySchema( + software.amazon.awssdk.services.dynamodb.model.KeySchemaElement.builder() + .attributeName("pk") + .keyType(software.amazon.awssdk.services.dynamodb.model.KeyType.HASH) + .build()) + .attributeDefinitions( + software.amazon.awssdk.services.dynamodb.model.AttributeDefinition.builder() + .attributeName("pk") + .attributeType( + software.amazon.awssdk.services.dynamodb.model.ScalarAttributeType.S) + .build()) + .billingMode(software.amazon.awssdk.services.dynamodb.model.BillingMode.PAY_PER_REQUEST) + .build(); + + software.amazon.awssdk.services.dynamodb.model.CreateTableRequest enriched = + VectorSearchSupport.withVectorIndexes(base, Collections.singletonList(vi)); + + assertNotNull(enriched); + assertTrue(enriched.overrideConfiguration().isPresent()); + // The VECTOR_INDEXES attribute must be set + software.amazon.awssdk.core.interceptor.ExecutionAttributes ea = + enriched.overrideConfiguration().get().executionAttributes(); + List vis = ea.getAttribute(VectorSearchInterceptor.VECTOR_INDEXES); + assertNotNull(vis); + assertEquals(1, vis.size()); + assertEquals("idx", vis.get(0).indexName()); + } + + @Test + public void testVectorIndexUpdateBuilder() { + CreateVectorIndexAction createAction = + CreateVectorIndexAction.builder() + .indexName("new-idx") + .vectorAttribute(VectorAttribute.builder().attributeName("v").dimensions(3).build()) + .similarityFunction("DOT_PRODUCT") + .build(); + VectorIndexUpdate update = VectorIndexUpdate.builder().create(createAction).build(); + + assertNotNull(update.create()); + assertNull(update.delete()); + assertEquals("new-idx", update.create().indexName()); + assertEquals("DOT_PRODUCT", update.create().similarityFunction()); + + DeleteVectorIndexAction deleteAction = + DeleteVectorIndexAction.builder().indexName("old-idx").build(); + VectorIndexUpdate deleteUpdate = VectorIndexUpdate.builder().delete(deleteAction).build(); + + assertNull(deleteUpdate.create()); + assertNotNull(deleteUpdate.delete()); + assertEquals("old-idx", deleteUpdate.delete().indexName()); + } + + @Test + public void testVectorQueryResultEmptyScores() { + software.amazon.awssdk.services.dynamodb.model.QueryResponse response = + software.amazon.awssdk.services.dynamodb.model.QueryResponse.builder().build(); + VectorQueryResult result = new VectorQueryResult(response, null); + assertNotNull(result.scores()); + assertTrue(result.scores().isEmpty()); + } + + @Test + public void testVectorQueryResultWithScores() { + software.amazon.awssdk.services.dynamodb.model.QueryResponse response = + software.amazon.awssdk.services.dynamodb.model.QueryResponse.builder().build(); + List scores = Arrays.asList(0.9, 0.8, 0.7); + VectorQueryResult result = new VectorQueryResult(response, scores); + assertEquals(3, result.scores().size()); + assertEquals(0.9, result.scores().get(0), 1e-9); + } + + // ------------------------------------------------------------------------- + // Float32Vector — encoding / decoding + // ------------------------------------------------------------------------- + + @Test + public void testFloat32VectorRoundTrip() { + float[] original = {1.0f, -2.5f, 0.0f, Float.MAX_VALUE, Float.MIN_VALUE}; + AttributeValue av = Float32Vector.toAttributeValue(original); + + assertTrue("toAttributeValue must produce a B attribute", av.b() != null); + assertTrue("isFloat32Vector must return true", Float32Vector.isFloat32Vector(av)); + + float[] decoded = Float32Vector.toFloats(av); + assertArrayEquals("round-trip must be lossless", original, decoded, 0.0f); + } + + @Test + public void testFloat32VectorEmptyArray() { + float[] empty = {}; + AttributeValue av = Float32Vector.toAttributeValue(empty); + assertTrue(Float32Vector.isFloat32Vector(av)); + assertArrayEquals(empty, Float32Vector.toFloats(av), 0.0f); + } + + @Test + public void testFloat32VectorIsNotStringAttribute() { + assertFalse(Float32Vector.isFloat32Vector(AttributeValue.fromS("hello"))); + } + + @Test + public void testFloat32VectorIsNotNumberAttribute() { + assertFalse(Float32Vector.isFloat32Vector(AttributeValue.fromN("42"))); + } + + @Test + public void testFloat32VectorIsNotArbitraryBinaryAttribute() { + // A B attribute without the magic prefix must not be considered a Float32Vector. + assertFalse( + Float32Vector.isFloat32Vector( + AttributeValue.fromB(SdkBytes.fromByteArray(new byte[] {1, 2, 3, 4, 5, 6, 7, 8})))); + } + + @Test + public void testFloat32VectorBinaryTooShortForMagic() { + assertFalse( + Float32Vector.isFloat32Vector( + AttributeValue.fromB(SdkBytes.fromByteArray(new byte[] {(byte) 0xF2})))); + } + + @Test(expected = IllegalArgumentException.class) + public void testFloat32VectorToFloatsRequiresMagic() { + Float32Vector.toFloats(AttributeValue.fromS("not a vector")); + } + + @Test + public void testFloat32VectorBase64PrefixIsCorrect() { + // Build a Float32Vector and verify the B field in its JSON representation + // starts with a known fixed prefix (derived from the 8-byte magic). + float[] values = {1.0f, 2.0f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + String b64 = java.util.Base64.getEncoder().encodeToString(av.b().asByteArray()); + // The first 8 base64 chars are fully determined by the first 6 magic bytes (two groups of 3). + // We verify stability of the encoding rather than hard-coding the string here. + AttributeValue av2 = Float32Vector.toAttributeValue(new float[] {9.0f, 8.0f}); + String b64_2 = java.util.Base64.getEncoder().encodeToString(av2.b().asByteArray()); + assertEquals( + "All Float32Vectors must share the same 8-char base64 prefix", + b64.substring(0, 8), + b64_2.substring(0, 8)); + } + + // ------------------------------------------------------------------------- + // Float32Vector — JSON-level request replacement (B → FLOAT32VECTOR) + // ------------------------------------------------------------------------- + + @Test + public void testAttributeValueToJsonFloat32VectorUsesCompactFormat() throws Exception { + // When Float32Vector.toAttributeValue() is used, attributeValueToJson should emit + // {"FLOAT32VECTOR": [...]} rather than the generic {"B": "..."}. + float[] values = {0.5f, -1.5f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + ObjectNode json = VectorSearchInterceptor.attributeValueToJson(av); + + assertNull("B key must be absent for Float32Vector", json.get("B")); + assertNotNull("FLOAT32VECTOR key must be present", json.get("FLOAT32VECTOR")); + assertEquals(2, json.get("FLOAT32VECTOR").size()); + assertEquals(0.5f, (float) json.get("FLOAT32VECTOR").get(0).asDouble(), 1e-6f); + assertEquals(-1.5f, (float) json.get("FLOAT32VECTOR").get(1).asDouble(), 1e-6f); + } + + @Test + public void testAttributeValueToJsonOrdinaryBinaryIsBase64() throws Exception { + // A regular B attribute (no magic) must be emitted as base64 string. + byte[] rawBytes = {1, 2, 3, 4}; + AttributeValue av = AttributeValue.fromB(SdkBytes.fromByteArray(rawBytes)); + ObjectNode json = VectorSearchInterceptor.attributeValueToJson(av); + + assertNotNull("B key must be present for plain binary", json.get("B")); + assertNull("FLOAT32VECTOR key must be absent for plain binary", json.get("FLOAT32VECTOR")); + } + + // ------------------------------------------------------------------------- + // Float32Vector — JSON-level response replacement (FLOAT32VECTOR → L) + // ------------------------------------------------------------------------- + + @Test + public void testFloat32VectorWritePathProducesCompactFormat() throws Exception { + // Verify that encoding float[] → toAttributeValue → attributeValueToJson produces + // {"FLOAT32VECTOR": [...]} (compact wire format for writes). + float[] original = {1.0f, 2.0f, 3.0f}; + + AttributeValue av = Float32Vector.toAttributeValue(original); + ObjectNode asJson = VectorSearchInterceptor.attributeValueToJson(av); + assertNotNull( + "FLOAT32VECTOR key must be present in compact wire format", asJson.get("FLOAT32VECTOR")); + assertNull("B key must be absent for Float32Vector", asJson.get("B")); + assertEquals(3, asJson.get("FLOAT32VECTOR").size()); + } + + @Test + public void testFloat32VectorBytesRoundTrip() { + // Round-trip through toAttributeValue → isFloat32Vector → toFloats + float[] values = {-1.0f, 0.0f, 3.14159f}; + AttributeValue av = Float32Vector.toAttributeValue(values); + assertTrue(Float32Vector.isFloat32Vector(av)); + float[] back = Float32Vector.toFloats(av); + assertArrayEquals(values, back, 1e-6f); + } + + @Test + public void testFloat32VectorFromListRoundTrip() { + // Simulate the read-back path: interceptor returns an L-typed AttributeValue; + // re-encoding it via toAttributeValue(List) must produce the same result as + // toAttributeValue(float...) with the original values. + float[] original = {1.5f, -2.5f, 0.0f}; + AttributeValue fromFloats = Float32Vector.toAttributeValue(original); + + // Build an L-typed AttributeValue as the interceptor would return + List numbers = + java.util.Arrays.asList( + AttributeValue.fromN("1.5"), AttributeValue.fromN("-2.5"), AttributeValue.fromN("0.0")); + AttributeValue fromList = Float32Vector.toAttributeValue(numbers); + + assertArrayEquals(fromFloats.b().asByteArray(), fromList.b().asByteArray()); + } +} diff --git a/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java b/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java new file mode 100644 index 0000000..c1e7638 --- /dev/null +++ b/src/test/java/com/scylladb/alternator/vectorsearch/VectorSearchHttpInterceptorTest.java @@ -0,0 +1,727 @@ +// Copyright 2026-present ScyllaDB +// +// SPDX-License-Identifier: LicenseRef-ScyllaDB-Source-Available-1.1 + +package com.scylladb.alternator.vectorsearch; + +import static org.junit.Assert.*; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.net.URI; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; +import java.util.zip.CRC32; +import java.util.zip.DeflaterOutputStream; +import java.util.zip.GZIPOutputStream; +import org.junit.Test; +import org.reactivestreams.Publisher; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; +import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider; +import software.amazon.awssdk.core.interceptor.Context; +import software.amazon.awssdk.core.interceptor.ExecutionAttributes; +import software.amazon.awssdk.core.interceptor.InterceptorContext; +import software.amazon.awssdk.http.SdkHttpMethod; +import software.amazon.awssdk.http.SdkHttpRequest; +import software.amazon.awssdk.http.SdkHttpResponse; +import software.amazon.awssdk.http.async.AsyncExecuteRequest; +import software.amazon.awssdk.http.async.SdkAsyncHttpClient; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.services.dynamodb.model.QueryRequest; + +public class VectorSearchHttpInterceptorTest { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + @Test + public void testAsyncHttpResponseContentExtractsScoresAndConvertsFloat32Vector() + throws Exception { + byte[] responseBody = + bytes( + "{\"Items\":[{\"embedding\":{\"FLOAT32VECTOR\":[1.0,2.5]}}]," + + "\"Scores\":[0.7,0.6]}"); + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + ExecutionAttributes attrs = new ExecutionAttributes(); + attrs.putAttribute(VectorSearchInterceptor.RESULT_HOLDER, holder); + + Context.ModifyHttpResponse context = + responseContext(queryHttpRequest(), singleChunkPublisher(responseBody)); + + Optional> modified = + VectorSearchInterceptor.INSTANCE.modifyAsyncHttpResponseContent(context, attrs); + + assertTrue(modified.isPresent()); + byte[] out = collect(modified.get()).get(5, TimeUnit.SECONDS); + JsonNode json = MAPPER.readTree(out); + + JsonNode embedding = json.get("Items").get(0).get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("1.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("2.5", embedding.get("L").get(1).get("N").asText()); + assertEquals(Arrays.asList(0.7, 0.6), holder.getScores()); + } + + @Test + public void testGzipHttpResponseContentExtractsScoresAndConvertsFloat32Vector() throws Exception { + byte[] responseBody = + bytes( + "{\"Items\":[{\"embedding\":{\"FLOAT32VECTOR\":[1.0,2.5]}}]," + + "\"Scores\":[0.7,0.6]}"); + byte[] compressedBody = gzipCompress(responseBody); + VectorSearchResultHolder holder = new VectorSearchResultHolder(); + ExecutionAttributes attrs = new ExecutionAttributes(); + attrs.putAttribute(VectorSearchInterceptor.RESULT_HOLDER, holder); + SdkHttpResponse compressedResponse = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Encoding", "gzip") + .putHeader("Content-Length", String.valueOf(compressedBody.length)) + .build(); + + Context.ModifyHttpResponse headerContext = + responseContext(queryHttpRequest(), compressedResponse, compressedBody); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + Context.ModifyHttpResponse bodyContext = + responseContext(queryHttpRequest(), modifiedResponse, compressedBody); + + assertFalse(modifiedResponse.firstMatchingHeader("Content-Encoding").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("Content-Length").isPresent()); + + Optional modified = + VectorSearchInterceptor.INSTANCE.modifyHttpResponseContent(bodyContext, attrs); + + assertTrue(modified.isPresent()); + JsonNode json = MAPPER.readTree(readAllBytes(modified.get())); + + JsonNode embedding = json.get("Items").get(0).get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("1.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("2.5", embedding.get("L").get(1).get("N").asText()); + assertEquals(Arrays.asList(0.7, 0.6), holder.getScores()); + } + + @Test + public void testUncompressedHttpResponseStripsChecksumAndLengthHeadersBeforeVectorConversion() + throws Exception { + byte[] responseBody = bytes("{\"Item\":{\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.5]}}}"); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Length", String.valueOf(responseBody.length)) + .putHeader("x-amz-crc32", "123") + .putHeader("X-Amz-Crc32c", "456") + .putHeader("x-amz-checksum-sha256", "abc") + .putHeader("X-Test", "kept") + .build(); + + Context.ModifyHttpResponse headerContext = + responseContext(queryHttpRequest(), response, responseBody); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + + assertFalse(modifiedResponse.firstMatchingHeader("Content-Length").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32c").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-checksum-sha256").isPresent()); + assertEquals("kept", modifiedResponse.firstMatchingHeader("X-Test").get()); + + Context.ModifyHttpResponse bodyContext = + responseContext(queryHttpRequest(), modifiedResponse, responseBody); + Optional modified = + VectorSearchInterceptor.INSTANCE.modifyHttpResponseContent(bodyContext, attrs); + + assertTrue(modified.isPresent()); + JsonNode json = MAPPER.readTree(readAllBytes(modified.get())); + JsonNode embedding = json.get("Item").get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("3.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("4.5", embedding.get("L").get(1).get("N").asText()); + } + + @Test + public void testUnchangedUncompressedHttpResponsePreservesChecksumAndLengthHeaders() + throws Exception { + byte[] responseBody = bytes("{\"TableNames\":[\"items\"]}"); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Length", String.valueOf(responseBody.length)) + .putHeader("x-amz-crc32", "123") + .putHeader("X-Amz-Crc32c", "456") + .putHeader("x-amz-checksum-sha256", "abc") + .putHeader("X-Test", "kept") + .build(); + + Context.ModifyHttpResponse headerContext = + responseContext(listTablesHttpRequest(), response, responseBody); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + + assertEquals( + String.valueOf(responseBody.length), + modifiedResponse.firstMatchingHeader("Content-Length").get()); + assertEquals("123", modifiedResponse.firstMatchingHeader("x-amz-crc32").get()); + assertEquals("456", modifiedResponse.firstMatchingHeader("x-amz-crc32c").get()); + assertEquals("abc", modifiedResponse.firstMatchingHeader("x-amz-checksum-sha256").get()); + assertEquals("kept", modifiedResponse.firstMatchingHeader("X-Test").get()); + + Optional modified = + VectorSearchInterceptor.INSTANCE.modifyHttpResponseContent(headerContext, attrs); + + assertTrue(modified.isPresent()); + assertArrayEquals(responseBody, readAllBytes(modified.get())); + } + + @Test + public void testUnchangedUncompressedAsyncHttpResponsePreservesChecksumAndLengthHeaders() + throws Exception { + byte[] responseBody = bytes("{\"TableNames\":[\"items\"]}"); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Length", String.valueOf(responseBody.length)) + .putHeader("x-amz-crc32", "123") + .putHeader("X-Amz-Crc32c", "456") + .putHeader("x-amz-checksum-sha256", "abc") + .putHeader("X-Test", "kept") + .build(); + + Context.ModifyHttpResponse headerContext = + responseContext(listTablesHttpRequest(), response, singleChunkPublisher(responseBody)); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + + assertEquals( + String.valueOf(responseBody.length), + modifiedResponse.firstMatchingHeader("Content-Length").get()); + assertEquals("123", modifiedResponse.firstMatchingHeader("x-amz-crc32").get()); + assertEquals("456", modifiedResponse.firstMatchingHeader("x-amz-crc32c").get()); + assertEquals("abc", modifiedResponse.firstMatchingHeader("x-amz-checksum-sha256").get()); + assertEquals("kept", modifiedResponse.firstMatchingHeader("X-Test").get()); + + Optional> modified = + VectorSearchInterceptor.INSTANCE.modifyAsyncHttpResponseContent(headerContext, attrs); + + assertTrue(modified.isPresent()); + assertArrayEquals(responseBody, collect(modified.get()).get(5, TimeUnit.SECONDS)); + } + + @Test + public void testUncompressedAsyncGetItemFloat32VectorResponseStripsChecksumAndLengthHeaders() + throws Exception { + byte[] responseBody = bytes("{\"Item\":{\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.5]}}}"); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Length", String.valueOf(responseBody.length)) + .putHeader("x-amz-crc32", "123") + .putHeader("X-Amz-Crc32c", "456") + .putHeader("x-amz-checksum-sha256", "abc") + .putHeader("X-Test", "kept") + .build(); + + Context.ModifyHttpResponse headerContext = responseContext(getItemHttpRequest(), response); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + + assertFalse(modifiedResponse.firstMatchingHeader("Content-Length").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32c").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-checksum-sha256").isPresent()); + assertEquals("kept", modifiedResponse.firstMatchingHeader("X-Test").get()); + + Context.ModifyHttpResponse bodyContext = + responseContext(getItemHttpRequest(), modifiedResponse, singleChunkPublisher(responseBody)); + Optional> modified = + VectorSearchInterceptor.INSTANCE.modifyAsyncHttpResponseContent(bodyContext, attrs); + + assertTrue(modified.isPresent()); + JsonNode json = MAPPER.readTree(collect(modified.get()).get(5, TimeUnit.SECONDS)); + JsonNode embedding = json.get("Item").get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("3.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("4.5", embedding.get("L").get(1).get("N").asText()); + } + + @Test + public void + testUncompressedAsyncBatchWriteItemUnprocessedVectorResponseStripsChecksumAndLengthHeaders() + throws Exception { + byte[] responseBody = + bytes( + "{\"UnprocessedItems\":{\"items\":[{\"PutRequest\":{\"Item\":" + + "{\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.5]}}}}]}}}"); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse response = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Length", String.valueOf(responseBody.length)) + .putHeader("x-amz-crc32", "123") + .putHeader("X-Amz-Crc32c", "456") + .putHeader("x-amz-checksum-sha256", "abc") + .putHeader("X-Test", "kept") + .build(); + + Context.ModifyHttpResponse headerContext = + responseContext(batchWriteItemHttpRequest(), response); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + + assertFalse(modifiedResponse.firstMatchingHeader("Content-Length").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-crc32c").isPresent()); + assertFalse(modifiedResponse.firstMatchingHeader("x-amz-checksum-sha256").isPresent()); + assertEquals("kept", modifiedResponse.firstMatchingHeader("X-Test").get()); + + Context.ModifyHttpResponse bodyContext = + responseContext( + batchWriteItemHttpRequest(), modifiedResponse, singleChunkPublisher(responseBody)); + Optional> modified = + VectorSearchInterceptor.INSTANCE.modifyAsyncHttpResponseContent(bodyContext, attrs); + + assertTrue(modified.isPresent()); + JsonNode json = MAPPER.readTree(collect(modified.get()).get(5, TimeUnit.SECONDS)); + JsonNode embedding = + json.get("UnprocessedItems") + .get("items") + .get(0) + .get("PutRequest") + .get("Item") + .get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("3.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("4.5", embedding.get("L").get(1).get("N").asText()); + } + + @Test + public void testDeflateAsyncHttpResponseContentConvertsFloat32VectorWithoutResultHolder() + throws Exception { + byte[] responseBody = bytes("{\"Item\":{\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.5]}}}"); + byte[] compressedBody = deflateCompress(responseBody); + ExecutionAttributes attrs = new ExecutionAttributes(); + SdkHttpResponse compressedResponse = + SdkHttpResponse.builder().statusCode(200).putHeader("Content-Encoding", "deflate").build(); + + Context.ModifyHttpResponse headerContext = + responseContext( + queryHttpRequest(), compressedResponse, singleChunkPublisher(compressedBody)); + SdkHttpResponse modifiedResponse = + VectorSearchInterceptor.INSTANCE.modifyHttpResponse(headerContext, attrs); + Context.ModifyHttpResponse bodyContext = + responseContext(queryHttpRequest(), modifiedResponse, singleChunkPublisher(compressedBody)); + + assertFalse(modifiedResponse.firstMatchingHeader("Content-Encoding").isPresent()); + + Optional> modified = + VectorSearchInterceptor.INSTANCE.modifyAsyncHttpResponseContent(bodyContext, attrs); + + assertTrue(modified.isPresent()); + JsonNode json = MAPPER.readTree(collect(modified.get()).get(5, TimeUnit.SECONDS)); + + JsonNode embedding = json.get("Item").get("embedding"); + assertNull(embedding.get("FLOAT32VECTOR")); + assertEquals("3.0", embedding.get("L").get(0).get("N").asText()); + assertEquals("4.5", embedding.get("L").get(1).get("N").asText()); + } + + @Test + public void testDynamoDbAsyncClientInjectsVectorSearchAndReadsScores() throws Exception { + RecordingAsyncHttpClient httpClient = + new RecordingAsyncHttpClient( + bytes( + "{\"Items\":[{\"id\":{\"S\":\"item-1\"}," + + "\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.0]}}]," + + "\"Count\":1,\"ScannedCount\":1,\"Scores\":[0.9]}")); + + DynamoDbAsyncClient client = + DynamoDbAsyncClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .region(Region.US_EAST_1) + .httpClient(httpClient) + .overrideConfiguration(c -> c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE)) + .build(); + + try { + VectorQueryResult result = + VectorSearchSupport.queryAsync( + client, + QueryRequest.builder() + .tableName("items") + .indexName("embedding-index") + .limit(1) + .build(), + VectorSearch.builder().queryVector(3.0f, 4.0f).returnScores(true).build()) + .get(5, TimeUnit.SECONDS); + + String requestJson = new String(httpClient.requestBody(), StandardCharsets.UTF_8); + assertTrue(requestJson.contains("\"VectorSearch\"")); + assertTrue(requestJson.contains("\"FLOAT32VECTOR\"")); + assertTrue(requestJson.contains("\"ReturnScores\":\"SIMILARITY\"")); + + assertEquals(1, result.items().size()); + assertEquals("item-1", result.items().get(0).get("id").s()); + assertEquals(0.9, result.scores().get(0), 1e-9); + AttributeValue embedding = result.items().get(0).get("embedding"); + assertEquals("3.0", embedding.l().get(0).n()); + assertEquals("4.0", embedding.l().get(1).n()); + } finally { + client.close(); + } + } + + @Test + public void testDynamoDbAsyncClientReadsGzipVectorSearchResponse() throws Exception { + RecordingAsyncHttpClient httpClient = + new RecordingAsyncHttpClient( + gzipCompress( + bytes( + "{\"Items\":[{\"id\":{\"S\":\"item-1\"}," + + "\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.0]}}]," + + "\"Count\":1,\"ScannedCount\":1,\"Scores\":[0.9]}")), + "gzip"); + + DynamoDbAsyncClient client = + DynamoDbAsyncClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .region(Region.US_EAST_1) + .httpClient(httpClient) + .overrideConfiguration(c -> c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE)) + .build(); + + try { + VectorQueryResult result = + VectorSearchSupport.queryAsync( + client, + QueryRequest.builder() + .tableName("items") + .indexName("embedding-index") + .limit(1) + .build(), + VectorSearch.builder().queryVector(3.0f, 4.0f).returnScores(true).build()) + .get(5, TimeUnit.SECONDS); + + assertEquals(1, result.items().size()); + assertEquals("item-1", result.items().get(0).get("id").s()); + assertEquals(0.9, result.scores().get(0), 1e-9); + AttributeValue embedding = result.items().get(0).get("embedding"); + assertEquals("3.0", embedding.l().get(0).n()); + assertEquals("4.0", embedding.l().get(1).n()); + } finally { + client.close(); + } + } + + @Test + public void testDynamoDbAsyncClientReadsUncompressedVectorResponseWithCrc32Header() + throws Exception { + byte[] responseBody = + bytes( + "{\"Items\":[{\"id\":{\"S\":\"item-1\"}," + + "\"embedding\":{\"FLOAT32VECTOR\":[3.0,4.0]}}]," + + "\"Count\":1,\"ScannedCount\":1,\"Scores\":[0.9]}"); + RecordingAsyncHttpClient httpClient = + new RecordingAsyncHttpClient(responseBody, null, crc32(responseBody)); + + DynamoDbAsyncClient client = + DynamoDbAsyncClient.builder() + .endpointOverride(URI.create("http://localhost:8000")) + .credentialsProvider(AnonymousCredentialsProvider.create()) + .region(Region.US_EAST_1) + .httpClient(httpClient) + .overrideConfiguration(c -> c.addExecutionInterceptor(VectorSearchInterceptor.INSTANCE)) + .build(); + + try { + VectorQueryResult result = + VectorSearchSupport.queryAsync( + client, + QueryRequest.builder() + .tableName("items") + .indexName("embedding-index") + .limit(1) + .build(), + VectorSearch.builder().queryVector(3.0f, 4.0f).returnScores(true).build()) + .get(5, TimeUnit.SECONDS); + + assertEquals(1, result.items().size()); + assertEquals("item-1", result.items().get(0).get("id").s()); + assertEquals(0.9, result.scores().get(0), 1e-9); + AttributeValue embedding = result.items().get(0).get("embedding"); + assertEquals("3.0", embedding.l().get(0).n()); + assertEquals("4.0", embedding.l().get(1).n()); + } finally { + client.close(); + } + } + + @Test(expected = IllegalStateException.class) + public void testDeleteVectorIndexActionRequiresIndexName() { + DeleteVectorIndexAction.builder().build(); + } + + private static SdkHttpRequest queryHttpRequest() { + return SdkHttpRequest.builder() + .protocol("http") + .host("localhost") + .port(8000) + .method(SdkHttpMethod.POST) + .encodedPath("/") + .putHeader("X-Amz-Target", "DynamoDB_20120810.Query") + .putHeader("Content-Type", "application/x-amz-json-1.0") + .build(); + } + + private static SdkHttpRequest getItemHttpRequest() { + return SdkHttpRequest.builder() + .protocol("http") + .host("localhost") + .port(8000) + .method(SdkHttpMethod.POST) + .encodedPath("/") + .putHeader("X-Amz-Target", "DynamoDB_20120810.GetItem") + .putHeader("Content-Type", "application/x-amz-json-1.0") + .build(); + } + + private static SdkHttpRequest batchWriteItemHttpRequest() { + return SdkHttpRequest.builder() + .protocol("http") + .host("localhost") + .port(8000) + .method(SdkHttpMethod.POST) + .encodedPath("/") + .putHeader("X-Amz-Target", "DynamoDB_20120810.BatchWriteItem") + .putHeader("Content-Type", "application/x-amz-json-1.0") + .build(); + } + + private static SdkHttpRequest listTablesHttpRequest() { + return SdkHttpRequest.builder() + .protocol("http") + .host("localhost") + .port(8000) + .method(SdkHttpMethod.POST) + .encodedPath("/") + .putHeader("X-Amz-Target", "DynamoDB_20120810.ListTables") + .putHeader("Content-Type", "application/x-amz-json-1.0") + .build(); + } + + private static Context.ModifyHttpResponse responseContext( + SdkHttpRequest httpRequest, Publisher publisher) { + return responseContext( + httpRequest, SdkHttpResponse.builder().statusCode(200).build(), publisher); + } + + private static Context.ModifyHttpResponse responseContext( + SdkHttpRequest httpRequest, SdkHttpResponse httpResponse, Publisher publisher) { + return InterceptorContext.builder() + .request(QueryRequest.builder().tableName("items").build()) + .httpRequest(httpRequest) + .httpResponse(httpResponse) + .responsePublisher(publisher) + .build(); + } + + private static Context.ModifyHttpResponse responseContext( + SdkHttpRequest httpRequest, SdkHttpResponse httpResponse) { + return InterceptorContext.builder() + .request(QueryRequest.builder().tableName("items").build()) + .httpRequest(httpRequest) + .httpResponse(httpResponse) + .build(); + } + + private static Context.ModifyHttpResponse responseContext( + SdkHttpRequest httpRequest, SdkHttpResponse httpResponse, byte[] responseBody) { + return InterceptorContext.builder() + .request(QueryRequest.builder().tableName("items").build()) + .httpRequest(httpRequest) + .httpResponse(httpResponse) + .responseBody(new ByteArrayInputStream(responseBody)) + .build(); + } + + private static byte[] gzipCompress(byte[] uncompressed) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (GZIPOutputStream gzip = new GZIPOutputStream(out)) { + gzip.write(uncompressed); + } + return out.toByteArray(); + } + + private static String crc32(byte[] bytes) { + CRC32 crc32 = new CRC32(); + crc32.update(bytes, 0, bytes.length); + return Long.toString(crc32.getValue()); + } + + private static byte[] deflateCompress(byte[] uncompressed) throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + try (DeflaterOutputStream deflate = new DeflaterOutputStream(out)) { + deflate.write(uncompressed); + } + return out.toByteArray(); + } + + private static byte[] readAllBytes(java.io.InputStream in) throws IOException { + try { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + byte[] buffer = new byte[4096]; + int len; + while ((len = in.read(buffer)) != -1) { + out.write(buffer, 0, len); + } + return out.toByteArray(); + } finally { + in.close(); + } + } + + private static Publisher singleChunkPublisher(byte[] bytes) { + return subscriber -> + subscriber.onSubscribe( + new Subscription() { + private boolean done; + + @Override + public void request(long n) { + if (done) { + return; + } + if (n <= 0) { + done = true; + subscriber.onError( + new IllegalArgumentException( + "Reactive Streams request amount must be positive")); + return; + } + done = true; + if (bytes.length > 0) { + subscriber.onNext(ByteBuffer.wrap(bytes)); + } + subscriber.onComplete(); + } + + @Override + public void cancel() { + done = true; + } + }); + } + + private static CompletableFuture collect(Publisher publisher) { + CompletableFuture result = new CompletableFuture<>(); + ByteArrayOutputStream out = new ByteArrayOutputStream(); + publisher.subscribe( + new Subscriber() { + @Override + public void onSubscribe(Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer byteBuffer) { + ByteBuffer copy = byteBuffer.asReadOnlyBuffer(); + byte[] bytes = new byte[copy.remaining()]; + copy.get(bytes); + out.write(bytes, 0, bytes.length); + } + + @Override + public void onError(Throwable throwable) { + result.completeExceptionally(throwable); + } + + @Override + public void onComplete() { + result.complete(out.toByteArray()); + } + }); + return result; + } + + private static byte[] bytes(String value) { + return value.getBytes(StandardCharsets.UTF_8); + } + + private static final class RecordingAsyncHttpClient implements SdkAsyncHttpClient { + private final byte[] responseBody; + private final String contentEncoding; + private final String crc32Header; + private volatile byte[] requestBody; + + private RecordingAsyncHttpClient(byte[] responseBody) { + this(responseBody, null); + } + + private RecordingAsyncHttpClient(byte[] responseBody, String contentEncoding) { + this(responseBody, contentEncoding, null); + } + + private RecordingAsyncHttpClient( + byte[] responseBody, String contentEncoding, String crc32Header) { + this.responseBody = responseBody; + this.contentEncoding = contentEncoding; + this.crc32Header = crc32Header; + } + + @Override + public CompletableFuture execute(AsyncExecuteRequest request) { + Publisher requestPublisher = request.requestContentPublisher(); + CompletableFuture requestBodyFuture = + requestPublisher != null + ? collect(requestPublisher) + : CompletableFuture.completedFuture(new byte[0]); + + CompletableFuture responseFuture = + requestBodyFuture.thenAccept( + bytes -> { + requestBody = bytes; + SdkHttpResponse.Builder responseBuilder = + SdkHttpResponse.builder() + .statusCode(200) + .putHeader("Content-Type", "application/x-amz-json-1.0"); + if (contentEncoding != null) { + responseBuilder.putHeader("Content-Encoding", contentEncoding); + } + if (crc32Header != null) { + responseBuilder.putHeader("x-amz-crc32", crc32Header); + } + request.responseHandler().onHeaders(responseBuilder.build()); + request.responseHandler().onStream(singleChunkPublisher(responseBody)); + }); + + responseFuture.whenComplete( + (ignored, error) -> { + if (error != null) { + request.responseHandler().onError(error); + } + }); + return responseFuture; + } + + @Override + public void close() {} + + private byte[] requestBody() { + return requestBody; + } + } +}