Data engineering for AI: RAG pipelines and feature stores

AI products are only as good as the data behind them. A chatbot that answers from last year's policy, a recommendation model trained on duplicated rows, a search that shows documents a user is not allowed to read: these are data engineering failures, not model failures.

This lesson shows the two pipelines data engineers build most often for AI: RAG pipelines, which feed documents to language models, and feature stores, which feed consistent inputs to machine learning models. You will learn the stages, the classic pitfalls and a few runnable snippets. It assumes you know batch pipelines (*Batch pipelines with Python*) and quality checks (*Data quality, testing and observability*). For how RAG works from the application side, see *Retrieval-augmented generation (RAG) end to end* in the AI Engineering track, and for vector search see *Vector databases: pgvector, Chroma and friends*.

A language model does not know your company's documents. Retrieval-augmented generation (RAG) fixes that: when a user asks a question, the system searches your documents for relevant passages and gives them to the model along with the question. The model then answers from those passages.

The search part is only as good as the indexed data. Building and maintaining that index is a pipeline, and it looks very much like ETL.

Analogy: a librarian who must answer questions from a shelf of books. If the books are old, torn into random pieces, or include a staff-only file, the answers will be bad however clever the librarian is.

Chunking is the step you will most often tune. This function splits text on paragraph boundaries into chunks of about 500 characters, carrying a little overlap forward so context is not lost at the edges. It also computes a content hash per chunk. The hash is the key to cheap updates: if the hash is unchanged, you skip re-embedding.

import hashlib


def chunk_text(text: str, max_chars: int = 500, overlap: int = 80) -> list[str]:
    """Split text on paragraph boundaries into chunks of about max_chars, with overlap."""
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    chunks, current = [], ""
    for p in paragraphs:
        if current and len(current) + len(p) + 2 > max_chars:
            chunks.append(current)
            current = current[-overlap:] + "\n\n" + p  # carry a little context forward
        else:
            current = (current + "\n\n" + p) if current else p
    if current:
        chunks.append(current)
    return chunks


def make_records(doc_id: str, text: str, source_url: str) -> list[dict]:
    records = []
    for i, chunk in enumerate(chunk_text(text)):
        records.append({
            "chunk_id": f"{doc_id}:{i}",
            "doc_id": doc_id,
            "chunk_index": i,
            "content": chunk,
            "content_hash": hashlib.sha256(chunk.encode()).hexdigest(),
            "source_url": source_url,
        })
    return records


if __name__ == "__main__":
    policy = (
        "Refund policy\n\nOrders can be cancelled before they reach the preparing status. "
        "Once preparing starts, we cannot refund the order.\n\n"
        "Pickup\n\nOrders are held for 30 minutes after the ready status. " * 6
    )
    recs = make_records("policy-refunds", policy, "https://example.com/policies/refunds")
    print(len(recs), "chunks")
    print(recs[0]["chunk_id"], recs[0]["content_hash"][:12], len(recs[0]["content"]))

Pure Python, no packages needed. Run it with python rag_chunks.py. The real size and overlap depend on your documents and embedding model; characters are used here only to keep the example simple.

Metadata is as important as the text. Store doc_id, chunk_index, source_url, title, updated_at, the language, and who is allowed to read the document. Answers can then cite sources, results can be filtered, and stale chunks can be deleted.

The next step would be to call an embedding model on each chunk's content and store the vector. The embedding model is a design choice (a hosted API or a local open-source model), so we leave it out here. What matters for the pipeline is that the vector dimension is fixed by the model, and that changing the model means re-embedding everything. Store the model name next to each vector.

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE chunks (
    chunk_id     TEXT PRIMARY KEY,         -- doc_id:chunk_index
    doc_id       TEXT NOT NULL,
    content      TEXT NOT NULL,
    content_hash TEXT NOT NULL,
    source_url   TEXT,
    allowed_role TEXT NOT NULL DEFAULT 'public',
    embed_model  TEXT NOT NULL,
    embedding    vector(384),              -- dimension depends on your embedding model
    updated_at   TIMESTAMPTZ NOT NULL DEFAULT now()
);

CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops);

-- Retrieval: the 5 closest chunks the user is allowed to see.
-- $1 is the question's embedding, $2 is the user's role.
SELECT chunk_id, content, source_url
FROM chunks
WHERE allowed_role = ANY($2)
ORDER BY embedding <=> $1
LIMIT 5;

<=> is cosine distance in pgvector. The WHERE clause is the important part: permissions are checked in the same query that retrieves, so a user can never be shown a chunk their role cannot read. $2 is an array of allowed roles.

A RAG index is a living copy of your documents. Four jobs keep it honest:

- Incremental updates. On each run, re-parse documents, compute chunk hashes and embed only new or changed chunks. Embedding costs time and money, so do not redo everything daily. - Deletes. When a document is removed or replaced, delete its old chunks. A chatbot quoting a withdrawn refund policy is worse than one that says "I don't know". - Re-indexing. When you change chunk size or embedding model, rebuild into a new table or index, test it, then switch over. Same idea as Write, Audit, Publish. - Backfills. Keep raw parsed text so you can re-chunk without fetching every source again.

You are a data engineer. Design a RAG ingestion pipeline for the following.

Sources: [e.g. 300 PDF policies on a shared drive, a help-centre website, a ticket database]
How often they change: [e.g. daily / weekly]
Who may read what: [e.g. HR documents only for HR]
Stack I know: [e.g. Python, PostgreSQL, Airflow]

Give me: the stages, the table schema for chunks with metadata, how to detect changed documents, how to delete removed ones, how to enforce permissions at retrieval, and three data quality checks for the index. Do not write full code yet.

What to expect: A good plan has incremental updates by hash, delete handling, permission metadata used at retrieval, and checks such as 'no empty chunks' and 'every chunk has a source'.

Treat retrieval like any other data product and test it. Build a small golden set: 20 to 50 real questions, each with the document or chunk that should be retrieved. After every change to chunking, embeddings or sources, measure how often the right chunk appears in the top results. If the score drops, the change made things worse, even if the pipeline ran green. Also add ordinary data quality checks: no empty or very short chunks, no duplicate chunks, every chunk has a source and a timestamp, and the number of chunks per document is plausible.

A feature is a value computed from data and fed to a machine learning model, such as orders_last_30d or avg_basket_value for a customer. Many models in a company need the same features, and they need them in two places:

- Offline, for training: thousands of historical examples, as a big table. - Online, for predictions: the latest value for one customer, in milliseconds.

If the training code and the serving code compute features differently, the model sees different data in production than it learned from. This is called training-serving skew, and it is a classic reason models behave well in the notebook and badly in the app.

A feature store is a system that defines each feature once, computes it, stores it in an offline store (a warehouse or lake) and an online store (a fast key-value database), and serves both consistently. Open-source Feast is a well-known example, and cloud platforms such as Databricks, Amazon SageMaker and Google Vertex AI include feature store services. You do not need one for a first model. A well-tested SQL table can be enough, until several models and teams share features.

The hardest idea in feature engineering is point-in-time correctness. When you build a training set, each example must use only the feature values that were known at the time of the event. Otherwise information from the future leaks into training (data leakage), the model looks brilliant in testing and fails in real life.

Example: you want to predict customer churn. A label is recorded on 3 October for customer C1. The feature orders_last_30d was 2 on 1 October and 6 on 5 October. The correct value for the 3 October label is 2, the latest value before the event, not 6. An ASOF join does this: for each label, take the most recent feature snapshot at or before the label time.

CREATE TABLE feature_snapshots AS SELECT * FROM (VALUES
  ('C1', TIMESTAMP '2026-10-01 00:00:00', 2),
  ('C1', TIMESTAMP '2026-10-05 00:00:00', 6),
  ('C2', TIMESTAMP '2026-10-01 00:00:00', 0)
) t(customer_id, snapshot_ts, orders_last_30d);

CREATE TABLE labels AS SELECT * FROM (VALUES
  ('C1', TIMESTAMP '2026-10-03 12:00:00', 1),
  ('C1', TIMESTAMP '2026-10-06 12:00:00', 0),
  ('C2', TIMESTAMP '2026-10-02 09:00:00', 0)
) t(customer_id, event_ts, churned);

SELECT l.customer_id, l.event_ts, f.orders_last_30d, l.churned
FROM labels l
ASOF LEFT JOIN feature_snapshots f
  ON l.customer_id = f.customer_id AND l.event_ts >= f.snapshot_ts
ORDER BY l.customer_id, l.event_ts;

ASOF JOIN is supported in DuckDB and some other engines such as Snowflake. In others you can get the same result with a window function (row_number over snapshots before the event). Paste it into the DuckDB CLI or Python to try.

The first C1 label gets 2, the value from 1 October, and not the value 6 from 5 October. Feature stores such as Feast do this join for you; knowing how it works lets you check them and build your own when needed.

See *MLOps basics: from notebook to deployed model* in the Data Science track for how these pieces fit around a deployed model.

Expect: "How would you build a pipeline to keep a chatbot's knowledge up to date?" (incremental ingest by hash, deletes, metadata and permissions, re-indexing, monitoring), "What is training-serving skew?" and "What is data leakage and how does a feature store help?" (point-in-time joins, one definition for offline and online). Say what you would measure: retrieval hit rate on a golden set, freshness of the index, feature freshness, and pipeline cost. Honest limits help too: a feature store is not always needed, and RAG quality depends more on clean data and chunking than on the vendor of the vector database.