DatabricksでS3にあるファイルに対してRAGを構築してみた

DatabricksのAI関数とAI Searchを組み合わせて、S3のPDFに対するRAGパイプラインをSQLで構築する手順を紹介します。パース、チャンク化、埋め込み、検索、回答生成を一通り実装してみました。
2026.08.05
データ事業本部のueharaです。
今回は、DatabricksでS3にあるファイルに対してRAGを構築してみたいと思います。
 はじめに前提として、PDFのような非構造化ドキュメントを対象にしたRAG (Retrieval-Augmented Generation)は、DatabricksではAgent BricksのKnowledge Assistantを使うのが手軽です。
https://docs.databricks.com/aws/ja/agents/agent-bricks/knowledge-assistant
Knowledge Assistantを利用すれば、Unity Catalogで管理されるボリューム内のファイルなどをナレッジソースとして指定するだけで簡単にチャットボットを作れます。
一方で、処理パイプラインの各部（パース、チャンク化、埋め込み、検索、回答生成）を個別に制御できれば、チャンクの粒度や埋め込みモデル、プロンプトを要件に合わせて差し替えられます。
今回は、Knowledge Assistantがブラックボックスとして包んでくれている処理を、SQLとDeltaテーブルの組み合わせとして明示的に組むことを目的としたいと思います。
具体的には、S3に置いたPDF（日本語）に対し、AI関数である ai_parse_document でテキスト化、 ai_prep_search  でチャンク化、Databricks AI Search（旧 Vector Search）で埋め込みとインデックス作成を行い、最後に ai_query  での回答生成までを一通り検証したいと思います。
なお、今回はPDFを例にしていますが、ai_parse_document はPDFに加えてDOCX、PPTX、JPG、PNG、TIFFをサポートしているので、その他ドキュメントにも拡張可能です。
 構成と処理の流れメダリオンアーキテクチャに沿って、Bronze, Silver, Goldの3層でパイプラインを組みます。
各層の役割は次の通りです。
Bronze：S3バケットに配置した生のPDF。Unity Catalogの外部ボリューム経由で参照し、原本として保持する
Silver：ai_parse_document の出力を保存したテーブル。テキスト、表、レイアウト情報を含む中間表現
Gold：ai_prep_search のチャンクを1行1チャンクにフラット化したテーブルと、そこから同期される AI Searchインデックス
 前提条件今回の前提条件は以下の通りです。
Databricks Free Editionのワークスペースを利用します。
S3バケットに対する外部ロケーションが作成済みであることとします。
検証用のPDFはデジタル庁が公開している「デジタル社会の実現に向けた重点計画」を使います（デジタル庁の政策ページからPDFを取得できます）。
 手順0：スキーマと外部ボリュームの作成まず、スキーマとS3上のPDFを指す外部ボリュームを作成します。
CREATE SCHEMA IF NOT EXISTS workspace.rag_demo;

CREATE EXTERNAL VOLUME IF NOT EXISTS workspace.rag_demo.pdf_source
LOCATION 's3://<YOUR-S3-BUCKET>/rag-demo/pdf/';
作成が完了すると、以下のようになっているかと思います。
外部ボリュームを挟む理由は2つあります。
1つ目は、S3のパスが /Volumes/workspace/rag_demo/pdf_source/ というボリュームパスで参照できるようになり、後続のSQLが管理ボリュームの場合と同じ形で書けることです。
2つ目は、読み取り権限をUnity Catalogのボリューム権限として管理できることです。

s3:// というパスを直接 READ_FILES に渡すこともできますが、その場合は外部ロケーションに対する READ FILES 権限が別途必要になります。
 手順1（Bronze）：ファイルをS3にアップロード手順0で設定したS3のパス（プレフィックス）に対し、PDFファイルをS3にアップロードします。
アップロードができたら、外部ボリューム経由でファイルが見えることを確認します。
LIST '/Volumes/workspace/rag_demo/pdf_source/';
 手順2（Silver）：ドキュメントをパースREAD_FILES でバイナリとして読み、ai_parse_document に渡します。
CREATE OR REPLACE TABLE workspace.rag_demo.parsed_documents AS
SELECT
  path,
  ai_parse_document(content, map('version', '2.0')) AS parsed
FROM READ_FILES('/Volumes/workspace/rag_demo/pdf_source/', format => 'binaryFile');
ここでのポイントは2つあります。
1つ目は format => 'binaryFile' です。

ai_parse_document は入力としてBINARY型（ファイルのバイト列）を要求するため、 READ_FILES をバイナリモードで使って content 列を生成しています。
2つ目は map('version', '2.0') です。

ai_parse_document の出力スキーマはバージョン管理されており、省略するとデフォルトバージョンが使われます。

マイナーバージョンアップは後方互換ですが、メジャーバージョンアップではフィールドの改名や削除があり得るため、パイプラインとして組む場合はバージョンを固定しておくと安心です。
出力はVARIANT型で、テキスト、表（HTML 表現）、ページやヘッダなどのレイアウト要素を含みます。
パースの成否は error_status で確認できます。
SELECT
  path,
  parsed:error_status::STRING AS error_status
FROM workspace.rag_demo.parsed_documents;
error_status が NULL ならパース成功です。
ページ数は次のように数えられます。
SELECT
  path,
  count(*) AS num_pages
FROM workspace.rag_demo.parsed_documents,
     LATERAL variant_explode(parsed:document:pages) AS page
GROUP BY path;
VARIANT型のフィールドアクセスには parsed:document:pages のようなコロン記法を使い、配列の展開には variant_explode を使います。
今回のドキュメントは47ページなので、以下の通り num_pages が 47 になっていることが確認できます。
抽出された要素を以下のクエリで確認してみます。
SELECT
  path,
  element.value:id::INT           AS element_id,
  element.value:type::STRING      AS element_type,
  element.value:content::STRING   AS content
FROM workspace.rag_demo.parsed_documents,
     LATERAL variant_explode(parsed:document:elements) AS element
ORDER BY path, element_id
LIMIT 50;
`contentP 列に、日本語本文が入っていることが確認できました。
 手順3（Gold）：チャンク化とフラット化次に ai_prep_search でチャンク化します。
この関数はパース結果を意味のまとまりで分割し、文書タイトルやセクション見出し、ページ参照などの文書レベルの文脈を付与したチャンクを返します。
CREATE OR REPLACE TABLE workspace.rag_demo.rag_chunks AS
WITH prepped AS (
  SELECT
    path,
    ai_prep_search(parsed) AS prepped
  FROM workspace.rag_demo.parsed_documents
)
SELECT
  md5(concat_ws('#', path, chunk.value:chunk_position::STRING)) AS chunk_pk,
  chunk.value:chunk_id::STRING          AS chunk_id,
  chunk.value:chunk_position::INT       AS chunk_position,
  chunk.value:chunk_to_retrieve::STRING AS chunk_to_retrieve,
  chunk.value:chunk_to_embed::STRING    AS chunk_to_embed,
  path                                  AS source_uri
FROM prepped,
     LATERAL variant_explode(prepped:document.contents) AS chunk;
このクエリにもポイントが2つあります。
1つ目は chunk_to_embed と chunk_to_retrieve の使い分けです。
ai_prep_search は各チャンクに2種類の文字列を生成します。
chunk_to_embed は文書タイトルやセクション見出しなどの文脈を付与した埋め込み用の文字列で、検索精度を上げるために構築された1つの文字列です。

chunk_to_retrieve は本文のみの文字列で、検索ヒット後にLLMへ渡す用途に使います。
「埋め込みに使う表現」と「人間やLLMが読む表現」を分けるのがRAGの定石になりますが、この関数はそれを最初から分けて返してくれます。
2つ目は主キー chunk_pk です。
chunk_id は文書内での採番なので、複数のPDFを扱うと重複し得ます。

AI Search インデックスは主キーを要求するため、ソースURIとチャンク位置からハッシュを作って主キーにしています。
最後にChange Data Feedを有効化します。
ALTER TABLE workspace.rag_demo.rag_chunks
SET TBLPROPERTIES (delta.enableChangeDataFeed = true);
これはAI SearchのDelta Sync Indexがソーステーブルの差分を追跡するために必要な設定になります。
 手順4（Gold）：AI Search エンドポイントとインデックスの作成ノートブックでSDKから作成します。
1セル目
%pip install databricks-vectorsearch
dbutils.library.restartPython()
2セル目
from databricks.vector_search.client import VectorSearchClient

vsc = VectorSearchClient()

vsc.create_endpoint(name="rag-demo-endpoint", endpoint_type="STANDARD")

index = vsc.create_delta_sync_index(
    endpoint_name="rag-demo-endpoint",
    index_name="workspace.rag_demo.rag_chunks_index",
    source_table_name="workspace.rag_demo.rag_chunks",
    primary_key="chunk_pk",
    embedding_source_column="chunk_to_embed",
    embedding_model_endpoint_name="databricks-qwen3-embedding-0-6b",
    pipeline_type="TRIGGERED",
)
埋め込みモデルには databricks-qwen3-embedding-0-6b を指定しています。
Databricks管理の埋め込みモデルにはその他のモデルもありますが、 databricks-qwen3-embedding-0-6b は日本語を含む100以上の言語に対応したマルチリンガルモデルなので、今回はこちらを選びます。
embedding_source_column に chunk_to_embed を指定しているのは、手順3で説明した通り、文脈付きの表現をベクトル化したいからです。
pipeline_type="TRIGGERED" は手動またはスケジュールで同期する方式で、2回目以降のソーステーブル更新は index.sync() で明示的に同期します。

※作成時の初期同期は自動で走ります。
作成直後の index.describe() は次のような状態でした。
3セル目
index.describe()
PROVISIONING_ENDPOINT はエンドポイントのプロビジョニング中という意味で、正常な作成途中の状態です。
初回はエンドポイントの起動に数分から数十分程度かかり、その後に初期同期（埋め込み計算）が走ります。
なお describe() の出力には index_subtype: HYBRID とありますが、これはベクトル検索とキーワード検索を組み合わせたハイブリッド検索のインデックスになっているという意味です。
カタログエクスプローラを確認してみると、インデックスが作成されていることが分かります。
しばらく待つと、 Index status は Online になり、
Data Ingest の Update status も Completed になります。
データを確認すると、以下の通り計算されたベクトルが格納されていることが分かります。
 手順5：ベクトル検索を試すSQLの vector_search 関数で類似チャンクを取得します。
SELECT chunk_pk, chunk_to_retrieve, source_uri
FROM vector_search(
  index => 'workspace.rag_demo.rag_chunks_index',
  query_text => 'デジタル社会の実現に向けた日本の課題は？',
  num_results => 5
);
内容に関する質問に対し、関連するチャンクが返ってきました。
ここで返る chunk_to_retrieve が、そのままLLMのコンテキストに使える本文になります。
 手順6：RAGで回答生成を行う検索結果を結合してプロンプトに埋め込み、Foundation Model APIに投げます。
WITH retrieved AS (
  SELECT chunk_to_retrieve
  FROM vector_search(
    index => 'workspace.rag_demo.rag_chunks_index',
    query_text => 'デジタル社会の実現に向けた日本の課題は？',
    num_results => 5
  )
)
SELECT ai_query(
  'databricks-meta-llama-3-3-70b-instruct',
  concat(
    'あなたは社内文書検索のアシスタントです。以下の参考情報だけを根拠に、質問に日本語で簡潔に回答してください。参考情報にない内容は「資料に記載がありません」と答えてください。\n\n# 参考情報\n',
    (SELECT string_agg(chunk_to_retrieve, '\n\n---\n\n') FROM retrieved),
    '\n\n# 質問\nデジタル社会の実現に向けた日本の課題は？'
  )
) AS answer;
ai_query はデプロイ済みの基盤モデルをSQLから呼べる関数で、エンドポイントのプロビジョニングや設定は不要です。

※Free Editionでもpay-per-tokenの基盤モデルを呼べます。
ここで、モデルは databricks-meta-llama-3-3-70b-instruct を選択しています。
結果は以下の通りです。
検索から回答生成まで、SQLだけで一通り動かすことができました。
 最後に今回は、DatabricksでS3にあるファイルに対してRAGを構築してみました。
ai_parse_document と ai_prep_search の組み合わせで、パースからチャンク化までをSQLだけで書けるので非常にお手軽だと感じました。
この記事が同じ構成を試す方の参考になれば幸いです。
 参考文献ai_parse_document function | Databricks on AWS
ai_prep_search function | Databricks on AWS
Databricks AI Search | Databricks on AWS
vector_search function | Databricks on AWS
Use ai_query | Databricks on AWS