How to Query Vector Embeddings with Athena for Real‑Time RAG, Explained Simply

Most engineers think you need a dedicated vector database for similarity search, but Athena’s Iceberg support lets you run nearest‑neighbor queries directly on S3 data. In a few lines of TypeScript you can turn a Parquet file of embeddings into a searchable knowledge base and feed the results to an LLM. The result is a serverless RAG pipeline that costs pennies and scales automatically.

---

Why Athena Makes Sense for Vector Search

When you hear “vector search” you usually picture a separate service that holds millions of high‑dimensional points and answers “which points are closest to this one?” Athena can do the same thing without provisioning another cluster.

• Serverless – you write SQL, Athena runs it on demand, and you pay only for the bytes it scans.
• Zero‑ops – the data lives in S3, so you don’t have to manage backups, scaling, or networking.
• Native integration – the same IAM role that lets your code read S3 can also start Athena queries, keeping the permission surface tiny.

Think of Athena as a giant, on‑demand librarian. The “books” are your Parquet files, and the “catalog” is a Glue table. When you ask for “the five books most similar to this paragraph,” the librarian looks through the indexed pages (the columns) and hands you the results, charging you only for the pages she actually opened.

In plain English: Athena gives you a cheap, managed way to ask “what’s closest to this vector?” as long as you store the vectors in a format it can read (Parquet) and describe them in a Glue table.

Minimal code: create an Athena client

// Import the v3 AWS SDK client for Athena
import { AthenaClient, StartQueryExecutionCommand } from "@aws-sdk/client-athena";

// The client automatically picks up credentials from the environment (e.g. IAM role)
const athena = new AthenaClient({ region: "us-east-1" });

---

Preparing an Iceberg Table with Embedding Parquet Files

What is Iceberg?

Iceberg is an open‑source table format that sits on top of raw files (like Parquet) and adds schema evolution, partition pruning, and metadata caching. For Athena, Iceberg means you can ask for “only the rows where the date partition is today” without scanning the whole bucket.

What is Parquet?

Parquet is a column‑oriented file format. Each column is stored separately, which lets Athena read only the columns it needs—perfect for vector search where you only need the numeric embedding column and the text chunk column.

Step 1 – Write embeddings to a Parquet file

Below is a tiny helper that converts a list of 1536‑dimensional embeddings (produced by OpenAI’s text-embedding-3-large) into a Parquet file using the parquetjs-lite library. The code is deliberately short; in production you would stream large batches instead of loading everything into memory.

import { S3Client, PutObjectCommand } from "@aws-sdk/client-s3";
import { ParquetWriter, ParquetSchema } from "parquetjs-lite";

// Define the shape of each row we will store
const schema = new ParquetSchema({
  chunk_id:   { type: "INT64" },                 // unique identifier for the text piece
  text:       { type: "UTF8" },                  // the raw text chunk
  // Store the vector as a list of doubles – Parquet calls this "FLOAT64"
  embedding:  { type: "FLOAT64", repeated: true } // 1536 numbers
});

/**
 * Upload a batch of embeddings to S3 as a Parquet file.
 * @param bucket   The S3 bucket name.
 * @param key      The object key (path) where the file will live.
 * @param rows     An array of { chunk_id, text, embedding } objects.
 */
async function uploadEmbeddingsAsParquet(
  bucket: string,
  key: string,
  rows: { chunk_id: number; text: string; embedding: number[] }[]
) {
  // 1️⃣ Write rows to an in‑memory buffer
  const writer = await ParquetWriter.openFile(schema, "/tmp/embeddings.parquet");
  for (const row of rows) {
    await writer.appendRow(row);
  }
  await writer.close();

  // 2️⃣ Push the file to S3
  const s3 = new S3Client({ region: "us-east-1" });
  const fileBody = await import("fs").then(fs => fs.promises.readFile("/tmp/embeddings.parquet"));
  await s3.send(
    new PutObjectCommand({
      Bucket: bucket,
      Key: key,
      Body: fileBody,
      // Optional: set a storage class that matches your cost profile
      StorageClass: "STANDARD_IA"
    })
  );
}

Step 2 – Create an Iceberg table that points at the Parquet file

Athena’s DDL (Data Definition Language) is just SQL. The table we create tells Athena “the data lives in this S3 prefix, the columns are defined by this schema, and treat it as an Iceberg table so we can partition later.”

import { StartQueryExecutionCommandInput } from "@aws-sdk/client-athena";

// Replace these placeholders with your own values
const DATABASE = "rag_db";
const TABLE    = "embeddings";
const S3_LOCATION = "s3://my-knowledge-base/embeddings/";

// The DDL string – notice the `EXTERNAL` keyword (data stays in S3)
const createTableSQL = `
CREATE EXTERNAL TABLE ${DATABASE}.${TABLE}
WITH (
  format = 'ICEBERG',
  location = '${S3_LOCATION}'
) AS SELECT
  chunk_id,
  text,
  embedding
FROM ${DATABASE}.staging_embeddings; -- a temporary staging table you load first
`;

// Athena expects a StartQueryExecutionCommandInput object
const startCreateCmd: StartQueryExecutionCommandInput = {
  QueryString: createTableSQL,
  ResultConfiguration: {
    OutputLocation: "s3://my-knowledge-base/athena-results/"
  }
} satisfies StartQueryExecutionCommandInput; // <-- type‑check at compile time

// Fire the query (it runs asynchronously)
await athena.send(new StartQueryExecutionCommand(startCreateCmd));

Tip: Athena charges per TB scanned. By using Iceberg and partitioning on a column like date or source, the query can skip whole partitions, keeping cost low.

---

SQL‑Based Cosine Similarity and Nearest‑Neighbor Query

What is cosine similarity?

Cosine similarity measures the angle between two vectors. If you picture each embedding as an arrow pointing somewhere in a high‑dimensional space, the cosine tells you how closely the arrows line up, regardless of their length. The formula is

[
\text{cosine}(a, b) = \frac{a \cdot b}{|a| |b|}
]

where (a \cdot b) is the dot‑product (sum of element‑wise products) and (|a|) is the vector’s length (Euclidean norm).

Why normalize? If you forget to divide by the lengths, Athena has to compute the raw dot‑product for every row, but the result will be biased by vector magnitude. Moreover, without pre‑normalizing you lose the ability to prune with a simple “WHERE similarity > threshold” because the raw numbers are not comparable.

Implementing cosine similarity in Athena SQL

Athena does not ship a built‑in cosine_similarity function, but you can express it with a few WITH clauses. The trick is to store a pre‑computed length column when you first load the embeddings. If you didn’t, you can compute it on the fly, but that forces a full table scan.

Assume we have a column norm that holds the pre‑computed length of each embedding.

WITH query_vec AS (
  -- The vector you are searching for – 1536 numbers.
  SELECT
    CAST(array[0.01, 0.23, … , 0.45] AS ARRAY<double>) AS q_vec,
    -- Pre‑compute its length once
    sqrt(REDUCE(q_vec, 0D, (s, x) -> s + x * x)) AS q_norm
),
scores AS (
  SELECT
    e.chunk_id,
    e.text,
    -- Compute dot‑product between stored embedding and query vector
    REDUCE(
      TRANSFORM(zip(e.embedding, q_vec.q_vec), x -> x[0] * x[1]),
      0D,
      (s, prod) -> s + prod
    ) AS dot,
    e.norm AS e_norm,
    q_norm.q_norm AS q_norm
  FROM ${DATABASE}.${TABLE} AS e
  CROSS JOIN query_vec AS q_norm
)
SELECT
  chunk_id,
  text,
  dot / (e_norm * q_norm) AS cosine_similarity
FROM scores
ORDER BY cosine_similarity DESC
LIMIT 5;

The REDUCE + TRANSFORM pattern iterates over the 1536 numbers without pulling the whole array into a single scalar, keeping the query readable.

Type‑checked Athena command in TypeScript

import {
  StartQueryExecutionCommand,
  StartQueryExecutionCommandInput
} from "@aws-sdk/client-athena";

// The query we just built (you would probably generate the array dynamically)
const query = `WITH ... (the SQL from above) ...`;

// Athena needs to know where to dump the CSV result file
const execInput: StartQueryExecutionCommandInput = {
  QueryString: query,
  ResultConfiguration: {
    OutputLocation: "s3://my-knowledge-base/athena-results/"
  }
} satisfies StartQueryExecutionCommandInput; // compile‑time safety

const execResponse = await athena.send(new StartQueryExecutionCommand(execInput));
const queryExecutionId = execResponse.QueryExecutionId!;

Key takeaway: By normalizing vectors ahead of time and expressing cosine similarity with plain SQL functions, you avoid a full scan and keep Athena’s cost model happy.

---

Wiring the Athena Query into a Node.js RAG Service

What is RAG?

Retrieval‑Augmented Generation (RAG) is a pattern where you first retrieve relevant documents (or text chunks) and then generate a response using a language model, feeding the retrieved text as part of the prompt. It improves factuality and reduces hallucinations.

End‑to‑end flow

1. Receive a user question → compute its embedding via OpenAI (or any provider).
2. Run the Athena nearest‑neighbor query → get the top 5 most similar chunks.
3. Concatenate the chunks into a prompt.
4. Call an LLM (e.g., Claude on Bedrock) → return the answer.

Below is a compact but complete TypeScript snippet that glues those steps together. It uses the @aws-sdk/client-bedrock-runtime package only for illustration; you can swap it for any HTTP client that talks to your LLM.

import { S3Client, GetObjectCommand } from "@aws-sdk/client-s3";
import {
  GetQueryResultsCommand,
  GetQueryResultsCommandInput
} from "@aws-sdk/client-athena";
import { BedrockRuntimeClient, InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";

// Helper: wait for Athena query to finish (polling)
async function waitForResult(queryId: string): Promise<string> {
  const poll = async (attempt = 0): Promise<string> => {
    const resp = await athena.send(
      new GetQueryResultsCommand({ QueryExecutionId: queryId })
    );
    if (resp.ResultSet?.Rows?.length) return queryId;
    if (attempt > 30) throw new Error("Athena query timed out");
    await new Promise(r => setTimeout(r, 1000 * Math.pow(2, attempt))); // exponential back‑off
    return poll(attempt + 1);
  };
  return poll();
}

// 1️⃣ Compute query embedding (placeholder – you would call OpenAI’s API)
async function getQueryEmbedding(question: string): Promise<number[]> {
  // In a real app, replace this with a fetch to https://api.openai.com/v1/embeddings
  return Array(1536).fill(0).map(() => Math.random()); // dummy vector
}

// 2️⃣ Build the Athena SQL with the embedding we just got
function buildSimilaritySQL(embedding: number[]): string {
  const vecLiteral = `array[${embedding.map(v => v.toFixed(8)).join(", ")}]`;
  return `
    WITH query_vec AS (
      SELECT CAST(${vecLiteral} AS ARRAY<double>) AS q_vec,
             sqrt(REDUCE(CAST(${vecLiteral} AS ARRAY<double>), 0D, (s, x) -> s + x * x)) AS q_norm
    ),
    scores AS (
      SELECT
        e.chunk_id,
        e.text,
        REDUCE(
          TRANSFORM(zip(e.embedding, q_vec.q_vec), x -> x[0] * x[1]),
          0D,
          (s, prod) -> s + prod
        ) AS dot,
        e.norm,
        q_vec.q_norm
      FROM rag_db.embeddings e
      CROSS JOIN query_vec q_vec
    )
    SELECT chunk_id, text, dot / (norm * q_norm) AS similarity
    FROM scores
    ORDER BY similarity DESC
    LIMIT 5;
  `;
}

// 3️⃣ Run the query and fetch the rows
async function fetchTopChunks(question: string): Promise<string[]> {
  const embedding = await getQueryEmbedding(question);
  const sql = buildSimilaritySQL(embedding);

  // Start execution
  const startResp = await athena.send(
    new StartQueryExecutionCommand({
      QueryString: sql,
      ResultConfiguration: { OutputLocation: "s3://my-knowledge-base/athena-results/" }
    })
  );
  const execId = startResp.QueryExecutionId!;
  await waitForResult(execId);

  // Pull results (CSV rows)
  const resultResp = await athena.send(
    new GetQueryResultsCommand({ QueryExecutionId: execId })
  );

  // Rows[0] is the header; slice it off
  const rows = resultResp.ResultSet?.Rows?.slice(1) ?? [];
  return rows.map(r => r.Data?.[1].VarCharValue ?? ""); // column 1 is `text`
}

// 4️⃣ Send the concatenated chunks to Claude (Bedrock)
async function answerWithClaude(question: string): Promise<string> {
  const chunks = await fetchTopChunks(question);
  const prompt = `
You are a helpful assistant. Use the following retrieved passages to answer the user's question.

Passages:
${chunks.map((c, i) => `${i + 1}. ${c}`).join("\n")}

Question: ${question}
Answer:
`;

  const bedrock = new BedrockRuntimeClient({ region: "us-east-1" });
  const invokeCmd = new InvokeModelCommand({
    ModelId: "anthropic.claude-v2",
    ContentType: "application/json",
    Accept: "application/json",
    Body: JSON.stringify({ prompt, max_tokens_to_sample: 500 })
  });

  const resp = await bedrock.send(invokeCmd);
  const payload = JSON.parse(Buffer.from(resp.Body as Uint8Array).toString("utf-8"));
  return payload.completion;
}

// Example usage
(async () => {
  const userQuestion = "How does Athena charge for queries?";
  const answer = await answerWithClaude(userQuestion);
  console.log("🧠 Answer:", answer);
})();

Helpful tip: Athena writes query results to the S3 prefix you specify. Schedule a nightly Lambda (or a simple aws s3 rm command) to delete old CSV files, otherwise storage costs will creep up.

---

Performance Gotchas and How to Avoid Them

Gotcha	Why it hurts	Simple fix
No pre‑normalized vectors	Athena must compute norm for every row, which forces a full scan and adds CPU time → higher latency and cost.	Store an extra column norm when you first load the embeddings (norm = sqrt(sum(x_i^2))).
Unpartitioned Iceberg table	Without partitions Athena reads all Parquet files even if you only need recent data.	Partition on a low‑cardinality column like source or ingest_date. Iceberg automatically skips irrelevant files.
Result location left untouched	Each query creates a CSV in the OutputLocation bucket; these accumulate forever.	After you finish processing, call s3.deleteObject on the result key, or enable Athena’s ResultReuseConfiguration.
Glue catalog lag	After you run CREATE TABLE … Athena may still see the old schema for a few minutes, causing “column not found” errors.	Wait ~30 seconds or run a dummy SHOW TABLES before the real query; this forces Glue to refresh.
Large scans exceed the 10 GB limit	Athena aborts queries that try to read more than 10 GB in a single execution.	Keep each Parquet file under a few hundred MB and use Iceberg’s metadata.refresh_interval to keep the manifest size small.
S3 eventual consistency on list	If you upload a new Parquet file and immediately run a query, the file may not appear yet, leading to missing results.	Use s3.waiter('objectExists') or add a small delay (e.g., 2 seconds) after the upload before starting the query.

In plain English: The biggest performance win comes from pre‑computing what you can (vector lengths) and letting Iceberg do the heavy lifting (partition pruning). Clean up query results and give Glue a moment to catch up, and you’ll stay in the “pennies” cost zone.

---

The Takeaway

• Athena can act as a serverless vector store when you keep embeddings in Parquet and describe them with an Iceberg table.
• Cosine similarity is expressible in plain SQL; just remember to store the vector length (norm) ahead of time.
• A few TypeScript lines—upload, table creation, query execution, and result handling—turn a static S3 bucket into a live RAG knowledge base.
• Cost stays low because Iceberg lets Athena skip whole partitions, and you only pay for the bytes actually read.
• Watch out for un‑normalized vectors, missing partitions, lingering result files, and Glue‑catalog refresh delays; each can blow up latency and price.
• The whole pipeline stays zero‑ops: no managed vector DB, no servers to patch, and scaling is handled automatically by Athena and S3.

Give it a try on a small dataset, watch the query cost in the Athena console, and then scale up by adding partitions. You’ll see that a fully managed vector database isn’t the only path to real‑time RAG—Athena can get you there with just a handful of lines of TypeScript. Happy querying!

---

Transparency notice

This article was written with the help of an AI system — Groq (GPT OSS 120B).

Published: 2026-09-07 · Primary focus: Athena

All code blocks are intended to be correct and runnable, but please verify them
against the official docs for the tools mentioned before using in production.

Find an error? Drop a comment — corrections are always welcome.