Prepare documents for RAG with PySpark: chunk Markdown by heading into a Delta table

Module 12 costed the shop's pipeline. The shop now wants an assistant that answers customers' questions from its help articles, and every answer it gives will only be as good as the text it is handed. This module builds that text: a table of chunks, a vector index kept in step with it, and a model called from SQL. This part needs only module 4's delta_session.py. To rebuild from nothing, run the script below.

The documents

Retrieval-augmented generation, RAGRAGA way to answer from your own documents. A search finds the passages that match the question, and they are added to the prompt before the model writes., answers a question in 2 steps: find the few passages most likely to contain the answer, then give them to a model with the question. The first step searches chunks, short passages of the documents, by meaning. A chunk that holds half of 2 answers, or an answer with no sign of what it is about, is never found for the right question, however good the model.

The shop's help articles are Markdown files, as most help centres and wikis can export. This script writes 4 of them into a folder help. Save it as make_help.py and run it:

"""Write the shop's help articles as Markdown files in help/."""
from pathlib import Path

ARTICLES = {
    "returns": """# Returns and refunds
## How to return an order
Start a return from your order page within 30 days of delivery. Print the label we email you and drop the parcel at any post office. Items must be unused and in their original packaging.
## How long a refund takes
We refund your payment within 5 working days of the parcel reaching us. Your bank may take a few more days to show it. You will get an email when the refund is sent.
## Items that cannot be returned
Gift cards, opened food and personalised items cannot be returned, unless they arrived damaged.
""",
    "delivery": """# Delivery
## Delivery times
Standard delivery takes 3 to 5 working days. Express delivery arrives the next working day if you order before 14:00.
## Delivery costs
Standard delivery is free on orders over 40.00. Below that it costs 3.95. Express delivery costs 6.95 on any order.
## A parcel that has not arrived
If your parcel has not arrived 2 working days after its delivery date, contact us with your order number and we will trace it or send a replacement.
""",
    "payments": """# Payments
## Ways to pay
We accept debit and credit cards, PayPal and gift cards. You can pay with 1 gift card and 1 card together.
## When you are charged
Your card is charged when your order is dispatched, not when you place it. A cancelled order is never charged.
""",
    "gift-cards": """# Gift cards
## Using a gift card
Enter the 16-digit code at checkout. Any balance left stays on the card for next time.
## How long a gift card lasts
Gift cards are valid for 2 years from the day they are bought. The expiry date is printed on the card and in its email.
""",
}

Path("help").mkdir(exist_ok=True)
for slug, text in ARTICLES.items():
    Path(f"help/{slug}.md").write_text(text)

Each article is a # title and a few ## sections, and each section answers 1 kind of question. That is the structure to chunk by.

Split at the headings

Databricks' guide to RAG pipelines lists the ways to chunk: a fixed number of characters or tokens, paragraphs, the document's own format, or topics found by a model. It warns that fixed-size chunks rarely work in production, because they cut through sentences and answers. Markdown has its structure written in: a ## heading starts a new subject. Splitting there keeps each answer whole. Save this as chunk_help.py:

"""Split the shop's help articles into 1 chunk per section, with the article, title and heading beside the text."""
from pyspark.sql import functions as F
from delta_session import delta_spark

spark = delta_spark("chunk_help")

articles = spark.read.text("help", wholetext=True).select(
    F.regexp_extract(F.col("_metadata.file_name"), r"^(.+)\.md$", 1).alias("article"),
    F.regexp_replace("value", r"\s+$", "").alias("markdown"),  # without the file's last newline
)
sections = articles.select(
    "article",
    F.regexp_extract("markdown", r"^# (.+)", 1).alias("title"),
    F.posexplode(F.split("markdown", r"\n## ")).alias("position", "section"),
).where("position > 0")  # position 0 is the title line

chunks = sections.select(
    F.concat_ws("-", "article", F.col("position").cast("string")).alias("chunk_id"),
    "article",
    "title",
    F.regexp_extract("section", r"^([^\n]+)", 1).alias("heading"),
    F.trim(F.regexp_replace("section", r"^[^\n]*\n", "")).alias("body"),
)
# The text to embed starts with what the chunk is about.
chunks = chunks.withColumn("text", F.concat_ws("\n", F.concat_ws(": ", "title", "heading"), "body"))
chunks = chunks.withColumn("words", F.size(F.split("text", r"\s+")))

chunks.select("chunk_id", "heading", "words").orderBy("chunk_id").show(truncate=False)
chunks.drop("words").write.format("delta").mode("overwrite").saveAsTable("help_chunks")

What each step does:

• wholetext=True reads each file as 1 row instead of 1 row per line, so a section's lines stay together. _metadata.file_name, as in module 10, gives the file each row came from, and the article's name is the file's name without .md. The file's last newline is removed first, so an article's last section ends like the others. trim removes only spaces, and a section that kept the newline would change its text, and be embedded again, the day the section after it is removed.
• split at \n## cuts each article where a section starts, and posexplode turns the pieces into rows with their position. Position 0 is the # title line, which is kept as a column instead.
• The heading is the section's first line, and the body the rest.
• chunk_id, such as returns-2, comes from the article and the section's place in it, so the same section has the same id on every run.

The table has 1 row per section: 10 chunks from 4 articles. Run it with python chunk_help.py:

+------------+-----------------------------+-----+
|chunk_id    |heading                      |words|
+------------+-----------------------------+-----+
|delivery-1  |Delivery times               |23   |
|delivery-2  |Delivery costs               |23   |
|delivery-3  |A parcel that has not arrived|35   |
|gift-cards-1|Using a gift card            |22   |
|gift-cards-2|How long a gift card lasts   |33   |
|payments-1  |Ways to pay                  |25   |
|payments-2  |When you are charged         |25   |
|returns-1   |How to return an order       |43   |
|returns-2   |How long a refund takes      |42   |
|returns-3   |Items that cannot be returned|22   |
+------------+-----------------------------+-----+

Give each chunk its context

Read alone, the section "How long a refund takes" never says "return". A customer who asks about returning a parcel would share few words with it. So the text to embed starts with the article's title and the section's heading: Returns and refunds: How long a refund takes, then the body. Databricks' guide lists exactly this kind of metadata, the document's name and the section heading, as one that improves retrieval. The separate article, title and heading columns stay too, for filters and for showing the reader where an answer came from.

Check the sizes

An embedding model reads a limited number of tokens. The model part 2 uses on Databricks, databricks-gte-large-en, reads up to 8,192 tokens and turns each chunk into a vector of 1,024 numbers. The shop's chunks are a few dozen words each, far inside that.

Size is still a choice. Smaller chunks match a specific detail but lose what is around it; larger ones keep the context but bring text that does not answer the question, and every chunk handed to the model costs tokens. A help centre's sections are a good size because a person wrote each one to answer 1 thing. For long documents with few headings, split long sections again at paragraphs, with a sentence or 2 of overlap, so an answer that crosses a break is still found whole. There is no rule that fits every collection: part 4 measures retrieval with real questions, which is how a size is chosen.

For PDFs and Word files, the first step is parsing, not chunking. On Databricks, ai_parse_document turns a document's bytes into its text, tables and layout, and from there the chunking is the same. To see how many vectors a collection becomes, and the memory its index needs, use the vector database size calculator; part 2 builds that index.

ToolVector database size calculatorDocuments, articles, tickets or products in; the vectors they make and the memory and disk the index needs out.