Java Vector Search: Embeddings, Cosine Similarity & Top-K Retrieval
Your support team has 10,000 past tickets. A customer writes in: "I can't get into my account — the login screen just loops." You want to answer the classic question: "find tickets like this one" so you can reuse a past resolution. You reach for keyword search:
SELECT * FROM tickets WHERE text LIKE '%login%';It returns the ticket titled "login screen loops" — but misses the one titled "reset my password" and the one titled "cannot access account after email change", which share zero keywords with the query yet describe the same problem. Keyword search matches words; you need to match meaning. This is exactly what embeddings buy you — and the machinery underneath is small enough to write yourself, right now, in plain Java.
This post builds similarity search from scratch: what embeddings are, cosine similarity with a hand-worked example, a top-k search you can compile and run, how to get real vectors from an embeddings API, and — critically — what similarity scores do not mean. This is post 4 of the AI for Java Developers track.
What an embedding actually is
An embedding is a fixed-length list of doubles that represents a piece of text — its dense vector. You send text to an embedding model; it returns something like:
[0.13, -0.88, 0.42, ..., -0.07] // hundreds or thousands of numbersTwo properties matter, and both are easy to misunderstand:
- Dimensions are opaque. The number of values (the dimensionality) varies by model — typically hundreds to thousands — and no individual position has a human meaning. There is no dimension 47 that means "anger" or "billing". Treat the vector as one indivisible fingerprint of meaning. (Check your provider's current docs for the exact dimensions; they change as models are replaced.)
- Meaning becomes direction. The useful property is geometric: texts with similar meaning point in similar directions. "reset my password", "I can't get into my account", and "login screen loops" land near each other even though they share no words. "Your invoice is overdue" points somewhere else entirely.
The whole trick of vector search is therefore a geometry question: which stored vector points most nearly in the same direction as the query vector? That question has a one-formula answer.
Cosine similarity, worked by hand
Given two vectors a and b, cosine similarity is the cosine of the angle between them:
cos(a, b) = (a · b) / (|a| × |b|)
a · b = a₁b₁ + a₂b₂ + ... + aₙbₙ (dot product)
|a| = √(a₁² + a₂² + ... + aₙ²) (magnitude)Dividing by the magnitudes is the whole point: it throws away length and keeps only direction, which is where the meaning lives. The result is always between −1 and 1: 1.0 = identical direction, 0.0 = orthogonal (no directional overlap), −1.0 = opposite. (For text embeddings, scores in practice mostly fall between 0 and 1.)
Let's compute one by hand so you trust the code that follows. Take a = [2, 1, 0] and b = [1, 2, 1]:
dot product: 2×1 + 1×2 + 0×1 = 4
|a| = √(4 + 1 + 0) = √5 ≈ 2.2361
|b| = √(1 + 4 + 1) = √6 ≈ 2.4495
cos(a, b) = 4 / (2.2361 × 2.4495) ≈ 4 / 5.4772 ≈ 0.7303Now the Java. Pure JDK, no dependencies — with the edge cases handled loudly rather than silently:
/** Cosine similarity and friends — pure JDK, no dependencies. */
public final class VectorMath {
private VectorMath() {}
public static double dot(double[] a, double[] b) {
requireSameLength(a, b);
double sum = 0.0;
for (int i = 0; i < a.length; i++) {
sum += a[i] * b[i];
}
return sum;
}
public static double norm(double[] v) {
if (v == null) {
throw new IllegalArgumentException("vector must not be null");
}
double sum = 0.0;
for (double x : v) {
sum += x * x;
}
return Math.sqrt(sum);
}
/**
* Cosine similarity: the cosine of the angle between two vectors.
* 1.0 = identical direction, 0.0 = orthogonal, -1.0 = opposite.
*
* @throws IllegalArgumentException if lengths differ or either vector
* is the zero vector (similarity is undefined there)
*/
public static double cosineSimilarity(double[] a, double[] b) {
requireSameLength(a, b);
double na = norm(a);
double nb = norm(b);
if (na == 0.0 || nb == 0.0) {
throw new IllegalArgumentException(
"cosine similarity is undefined for a zero vector; "
+ "filter or fix the embedding first");
}
return dot(a, b) / (na * nb);
}
private static void requireSameLength(double[] a, double[] b) {
if (a == null || b == null) {
throw new IllegalArgumentException("vectors must not be null");
}
if (a.length != b.length) {
throw new IllegalArgumentException(
"vector length mismatch: " + a.length + " vs " + b.length);
}
}
}Two deliberate choices: a zero vector throws instead of returning a made-up 0.0, because cosine similarity is mathematically undefined there — a silent default would corrupt your rankings without a trace. And length mismatch throws: comparing vectors from different models (different dimensions) is meaningless, which is trap #1 in the "when it breaks" section below.
Top-k search: the whole engine in 20 lines
With similarity in hand, "find tickets like this one" is: score every stored vector against the query, sort descending, take the first k:
import java.util.Comparator;
import java.util.List;
/** In-memory top-k similarity search over a small corpus. */
public final class VectorSearch {
private VectorSearch() {}
/** A stored document: an id, a human-readable title, and its embedding. */
public record StoredDoc(String id, String title, double[] vector) {}
/** One ranked search result. */
public record Hit(String id, String title, double score) {}
/**
* Returns up to k corpus documents ranked by cosine similarity
* to the query, most similar first.
*/
public static List<Hit> topK(double[] query, List<StoredDoc> corpus, int k) {
if (query == null || corpus == null) {
throw new IllegalArgumentException("query and corpus must not be null");
}
if (k <= 0) {
throw new IllegalArgumentException("k must be positive, was " + k);
}
return corpus.stream()
.map(d -> new Hit(d.id(), d.title(),
VectorMath.cosineSimilarity(query, d.vector())))
.sorted(Comparator.comparingDouble(Hit::score).reversed())
.limit(k)
.toList();
}
}Running it: the hand computation, verified
This is the part most tutorials skip: prove the code agrees with the by-hand math. The demo below uses small hand-made vectors — explicitly stand-ins, not real embeddings — so every expected value can be checked on paper:
// VecDemo.main — assertions that must hold for the code to be trusted
double[] a = {2, 1, 0};
double[] b = {1, 2, 1};
double sim = VectorMath.cosineSimilarity(a, b); // hand-computed: ~0.7303
assert Math.abs(sim - 0.7302967433) < 1e-9;
// Stand-in unit vectors for four tickets (real ones come from the API below):
// d1 "password reset help" [1, 0, 0]
// d2 "login trouble" [0.8, 0.6, 0]
// d3 "billing question" [0, 1, 0]
// d4 "cannot access account" [0.6, 0.8, 0]
// Query "reset my password" -> [1, 0, 0]:
// sim = 1.0, 0.8, 0.0, 0.6 → top-2 = d1, d2
List<VectorSearch.Hit> top2 = VectorSearch.topK(query, corpus, 2);
assert top2.get(0).id().equals("d1"); // score 1.0
assert top2.get(1).id().equals("d2"); // score 0.8Compiled with javac --release 21 and run — actual output:
cosine([2,1,0],[1,2,1]) = 0.7302967433
ok - hand-computed cosine similarity matches (expected ~0.7303)
ok - identical vectors score 1.0
ok - orthogonal vectors score 0.0
ok - opposite vectors score -1.0
ok - zero vector throws IllegalArgumentException
rank: password reset help score=1.0000
rank: login trouble score=0.8000
ok - top-2 returns exactly 2 hits
ok - rank 1 is d1 with score 1.0
ok - rank 2 is d2 with score 0.8
ALL CHECKS PASSEDThe implementation reproduces the hand computation to 10 decimal places. Notice the ranking did the thing keyword search couldn't: "login trouble" and "cannot access account" — paraphrases of the query with no shared words — scored 0.8 and 0.6, above the unrelated billing ticket at 0.0.
Getting real vectors: calling an embeddings API
The code above is honest about one thing: the demo vectors are invented. In production, every vector comes from an embeddings API. The request shape is nearly universal — POST some JSON, get back a vector:
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* Minimal embeddings API client: sends text, gets back a dense vector.
* Needs a real API key in the EMBEDDINGS_API_KEY environment variable —
* do not hard-code keys.
*/
public final class EmbeddingClient {
private static final ObjectMapper MAPPER = new ObjectMapper();
private static final Duration TIMEOUT = Duration.ofSeconds(30);
private final HttpClient http;
private final String apiKey;
private final String endpoint;
private final String model;
public EmbeddingClient(String endpoint, String model) {
this.http = HttpClient.newBuilder()
.connectTimeout(TIMEOUT)
.build();
String key = System.getenv("EMBEDDINGS_API_KEY");
if (key == null || key.isBlank()) {
throw new IllegalStateException(
"EMBEDDINGS_API_KEY environment variable is not set");
}
this.apiKey = key;
this.endpoint = endpoint;
this.model = model;
}
/** Returns the embedding vector for {@code text}. */
public double[] embed(String text) throws EmbeddingException {
if (text == null || text.isBlank()) {
throw new IllegalArgumentException("text must not be blank");
}
String requestBody;
try {
requestBody = MAPPER.writeValueAsString(
java.util.Map.of("model", model, "input", text));
} catch (Exception e) {
throw new EmbeddingException("failed to build request JSON", e);
}
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(endpoint))
.timeout(TIMEOUT)
.header("Content-Type", "application/json")
.header("Authorization", "Bearer " + apiKey)
.POST(HttpRequest.BodyPublishers.ofString(requestBody))
.build();
HttpResponse<String> response;
try {
response = http.send(request, HttpResponse.BodyHandlers.ofString());
} catch (java.io.IOException | InterruptedException e) {
Thread.currentThread().interrupt();
throw new EmbeddingException("embeddings request failed", e);
}
if (response.statusCode() / 100 != 2) {
throw new EmbeddingException(
"embeddings API returned HTTP " + response.statusCode()
+ ": " + response.body());
}
try {
JsonNode root = MAPPER.readTree(response.body());
JsonNode data = root.path("data");
if (!data.isArray() || data.isEmpty()) {
throw new EmbeddingException("unexpected API response: no data array");
}
JsonNode embedding = data.get(0).path("embedding");
if (!embedding.isArray() || embedding.isEmpty()) {
throw new EmbeddingException("unexpected API response: no embedding array");
}
double[] vector = new double[embedding.size()];
for (int i = 0; i < vector.length; i++) {
vector[i] = embedding.get(i).asDouble();
}
return vector;
} catch (EmbeddingException e) {
throw e;
} catch (Exception e) {
throw new EmbeddingException("failed to parse embeddings response", e);
}
}
/** Checked wrapper so callers must decide how to handle API failures. */
public static final class EmbeddingException extends Exception {
public EmbeddingException(String message) { super(message); }
public EmbeddingException(String message, Throwable cause) { super(message, cause); }
}
}Notes from the trenches:
- Key via
System.getenv, never in code. The client refuses to construct without it. This code compiles as shown (Jacksonjackson-databindon the classpath) but needs a real provider, a real key, and a real endpoint to actually run — no free lunch here. - The model name is a constructor parameter, not a constant. Providers rename and retire models; pinning it in code means your pipeline silently drifts. Keep it in config (an env var or properties file) and version it alongside your data.
- Embed once, store forever. Embedding 10,000 tickets is 10,000 paid API calls — do it once at ingest time and store the vectors next to the ticket id. Queries embed only the one incoming ticket.
Storing vectors: in-memory now, pgvector later
For this lesson, a List<StoredDoc> is the entire database, and that's fine — with a few hundred vectors, a linear scan with cosine similarity is microseconds per comparison. But know where the ceiling is: scanning a million vectors per query doesn't scale, which is why production systems use pgvector, a Postgres extension that stores embeddings in a column and does indexed approximate-nearest-neighbor search in SQL. (Purpose-built vector databases exist too, but if you already run Postgres, pgvector is the least-new-infrastructure path.) The lesson code's contract — text in, ranked hits out — stays identical; only the storage layer changes.
The honesty section: what scores do NOT mean
This is the section that separates a working demo from a trustworthy system. Read it twice.
Similarity scores are not calibrated probabilities. A score of 0.87 does not mean "87% confident this ticket matches". It means the angle between two vectors has cosine 0.87 — nothing more. Two scores from two different queries aren't comparable either: 0.8 on one query might be a perfect match and 0.8 on another might be garbage.
"Nearest" is not "relevant". Top-k always returns k hits, even when nothing in the corpus is actually related. The ranking answers "which stored vectors point most similarly", not "which stored tickets solve this problem". Your code must treat the ranking as a candidate list for a human or a downstream check, never as an answer.
There is no universal cosine threshold. You will read blog posts claiming "0.8 is a good match" or "filter below 0.75". There is no such number. The right cutoff depends on your model, your data, and your tolerance for misses — it must be tuned on your own labeled data: take a few hundred query/ticket pairs, label which are genuinely helpful, and pick the threshold where precision and recall meet your bar. A number you didn't measure is a superstition.
When it breaks: 4 traps
- Comparing embeddings from different models. Vector spaces are not interchangeable. The same sentence embedded by model A and model B produces two vectors that live in different coordinate systems — their cosine similarity is meaningless noise. If you change models, re-embed the entire corpus. The length-mismatch guard in
VectorMathcatches the obvious case (different dimensions); the dangerous case is two models with the same dimensions, where nothing throws and everything is wrong. Pin the model name in config and treat "model changed" as a full re-embed event. - Ignoring normalization. Some providers return unit-length vectors, some don't. If you shortcut cosine similarity to a bare dot product — a common optimization — it is only valid on normalized vectors. On unnormalized vectors the dot product mixes direction with magnitude, and long documents (bigger magnitudes) outrank better matches. The implementation above divides by both norms explicitly, so it's correct either way. If you ever optimize to dot-product-only, normalize first.
- Using top-k=1 blindly.
topK(query, corpus, 1)always returns something — even when the best score is 0.12 and the "match" is nonsense. Never present a single nearest neighbor as the answer without a relevance gate: a minimum score tuned on your data (see the honesty section — not a blog-post number), below which you return "no good match found". Silent wrong answers are worse than honest empty ones. - Embedding huge documents as one vector. An embedding is a fixed-size fingerprint. Stuff a 40-page manual into one vector and the meaning of any single paragraph drowns — the vector becomes an average of everything and near nothing. Chunk first: split documents into paragraphs or small sections, embed each chunk, and store the parent document id alongside each chunk vector. Search matches chunks; you show the user the chunk plus its source document.
The principle
Nearest is a ranking, not a verdict. Embeddings turn "find tickets like this one" from a word-matching problem into a geometry problem, and cosine similarity plus top-k is genuinely all the machinery you need to start. But the ranking is the beginning of a decision — which candidate actually resolves the ticket — not the decision itself. Build the engine, tune your thresholds on real labeled pairs, gate the low scores, and let a human (or a well-tested downstream step) render the verdict.
Continue with the AI for Java Developers track.
Comments
Post a Comment