[Databricks] PDF ファイル検索パイプラインを構築する — Volume → ai_parse_document → Embedding 検索

はじめに

Databricks 上で「手元のドキュメントを検索したい」というニーズに対して、外部ツールを使わず SQL と組み込み AI 関数だけでファイル検索パイプラインを構築できます。

本記事では、以下の 4 ステップを Databricks ノートブック上で実装します。

1. Unity Catalog Volume に PDF をアップロード
2. ai_parse_document でテキスト抽出
3. チャンクを Delta テーブルに保存
4. ai_query + コサイン類似度でセマンティック検索

Free Edition の Serverless コンピュートで動作確認済みです。Vector Search Index を使わず、SQL の ai_query で埋め込みベクトルを取得してコサイン類似度を計算する方式のため、追加のエンドポイント作成は不要です。

前提条件

• Databricks ワークスペース（Free Edition 可）
• Unity Catalog が有効であること
• Serverless コンピュート（ai_query / ai_parse_document の実行に必要）

全体像

PDF ファイル
  ↓ Volume にアップロード
Unity Catalog Volume（/Volumes/dev/file_search_demo/documents/）
  ↓ read_files + ai_parse_document
構造化テキスト（VARIANT 型）
  ↓ Python で要素抽出 → チャンク結合
Delta テーブル（dev.file_search_demo.document_chunks）
  ↓ ai_query('databricks-gte-large-en') で埋め込み → コサイン類似度
検索結果（類似度スコア付きランキング）

「チャンク」とは、ドキュメントを検索に適したサイズに分割したテキスト断片のことです。PDF 全文をそのまま埋め込みベクトル化すると精度が下がるため、段落やページ単位で区切ってから検索対象にします。本記事では簡略化のためドキュメント1件＝1チャンクとしていますが、実運用ではページ単位や一定トークン数ごとに分割するのが一般的です。

Step 1: ダミー PDF の作成と Volume へのアップロード

まず、検索対象となるダミー PDF を作成します。実際のユースケースでは既存の PDF を Volume にアップロードするだけですが、デモとして fpdf2 で3つのサンプル文書を生成します。

fpdf2 のインストール

%pip install fpdf2 --quiet
dbutils.library.restartPython()

PDF 作成 + Volume 保存

from fpdf import FPDF
import os

VOLUME_PATH = "/Volumes/dev/file_search_demo/documents"

def create_pdf(filename, title, sections):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.set_font("Helvetica", "B", 16)
    pdf.cell(0, 10, title, new_x="LMARGIN", new_y="NEXT")
    pdf.ln(4)
    for heading, lines in sections:
        pdf.set_font("Helvetica", "B", 13)
        pdf.cell(0, 8, heading, new_x="LMARGIN", new_y="NEXT")
        pdf.set_font("Helvetica", size=11)
        for line in lines:
            pdf.cell(0, 7, line, new_x="LMARGIN", new_y="NEXT")
        pdf.ln(3)
    out_path = f"{VOLUME_PATH}/{filename}"
    pdf.output(out_path)
    print(f"Created: {out_path} ({os.path.getsize(out_path)} bytes)")

spark.sql("CREATE SCHEMA IF NOT EXISTS dev.file_search_demo")
spark.sql("""
    CREATE VOLUME IF NOT EXISTS dev.file_search_demo.documents
""")

create_pdf("sales_report_2024q4.pdf", "Sales Report: FY2024 Q4", [
    ("Overview", [
        "Revenue for Oct-Dec 2024 was 850M JPY, up 12% YoY.",
        "Cloud services drove growth with ARR increasing 15% QoQ.",
    ]),
    ("Revenue by Segment", [
        "Cloud Services: 420M JPY (+25% YoY)",
        "On-Premise: 210M JPY (-5% YoY)",
        "Consulting: 150M JPY (+8% YoY)",
        "Support: 70M JPY (+3% YoY)",
    ]),
    ("Key KPIs", [
        "New customers acquired: 142 (target: 130)",
        "Churn rate: 2.1% (improved from 2.8%)",
        "ARPU: 580K JPY/month (up from 550K)",
        "NPS: +42 (up from +38)",
    ]),
])

create_pdf("product_spec_alpha.pdf", "Product Spec: Project Alpha v3.0", [
    ("1. Overview", [
        "Project Alpha is an enterprise self-service BI platform.",
        "Users can query data using natural language without SQL knowledge.",
    ]),
    ("2. Key Features", [
        "Natural Language Query Engine (Japanese/English)",
        "Auto table structure recognition and SQL generation",
        "Benchmark evaluation for SQL accuracy",
        "Auto Dashboard Generation with 10 templates",
    ]),
    ("3. Technical Requirements", [
        "Supported DBs: PostgreSQL 14+, MySQL 8+, Snowflake, BigQuery, Databricks SQL",
        "Min config: 4vCPU / 16GB RAM / 100GB SSD",
        "Browsers: Chrome 100+, Edge 100+, Safari 16+",
    ]),
])

create_pdf("meeting_notes_20250115.pdf", "Executive Meeting Minutes - Jan 15 2025", [
    ("Agenda 1: FY2025 Business Plan", [
        "Revenue target: 4B JPY (18% YoY growth)",
        "Operating margin target: 15% (prev 12%)",
        "New hires: 50 (30 engineers, 15 sales, 5 others)",
    ]),
    ("Agenda 2: Product Roadmap", [
        "Project Alpha v3.0 release: April 2025",
        "NLQ accuracy target: 85%+",
        "Mobile app beta: H2 2025",
    ]),
    ("Agenda 3: Security", [
        "Dec 2024 phishing attack detected and blocked",
        "Company-wide security training in Feb",
        "SOC2 Type II certification target: June 2025",
    ]),
])

Volume の FUSE マウントパス /Volumes/<catalog>/<schema>/<volume>/ に直接書き込んでいます。Serverless コンピュートでは dbutils.fs.cp("file:/...") が使えないため、この方式が必要です。

実行すると、3つの PDF（各 1,500 バイト前後）が Volume に保存されます。

Step 2: ai_parse_document でテキスト抽出

ai_parse_document は PDF・DOCX・PPTX・画像から構造化テキストを抽出する SQL 組み込み関数です。read_files と組み合わせて Volume 上のファイルを一括処理できます。

parsed_df = spark.sql("""
    SELECT
        _metadata.file_name,
        ai_parse_document(content, MAP('version', '2.0')) AS parsed
    FROM read_files(
        '/Volumes/dev/file_search_demo/documents/*.pdf',
        format => 'binaryFile'
    )
""")
display(parsed_df.select("file_name", "parsed"))

戻り値は VARIANT 型で、document.elements[] 配列にテキスト要素が格納されています。各要素には content（テキスト）、element_type（heading / paragraph 等）、page_id が含まれます。

ai_parse_document は .txt ファイルをサポートしていません。テキストファイルを渡すと Unsupported file format: unknown が返ります。対応フォーマットは PDF・DOCX・PPTX・画像ファイルです。

Step 3: テキスト要素の抽出とチャンク化

VARIANT 型の解析結果から Python でテキスト要素を取り出し、ドキュメント単位でチャンク結合して Delta テーブルに保存します。

テキスト要素の抽出

import json

str_df = spark.sql("""
    SELECT
        _metadata.file_name AS file_name,
        CAST(ai_parse_document(content, MAP('version', '2.0')) AS STRING) AS parsed_json
    FROM read_files(
        '/Volumes/dev/file_search_demo/documents/*.pdf',
        format => 'binaryFile'
    )
""")

rows = str_df.collect()
records = []
for row in rows:
    fname = row["file_name"]
    parsed = json.loads(row["parsed_json"])
    for elem in parsed.get("document", {}).get("elements", []):
        content = elem.get("content", "")
        if content:
            records.append({
                "file_name": fname,
                "content": content,
                "element_type": elem.get("element_type", ""),
                "page_id": elem.get("page_id", 0),
            })

chunks_df = spark.createDataFrame(records)
display(chunks_df)
print(f"Total elements: {len(records)}")

VARIANT 型を Python で扱う際のポイントが2つあります。

1. SQL 側で CAST(... AS STRING) してから collect() する必要があります。VARIANT 型のまま collect すると VariantVal オブジェクトが返り、json.loads に渡せません
2. SQL の : 演算子（parsed:document:elements）は LATERAL VIEW EXPLODE の中では構文エラーになります。Python 側で json.loads → ループが安定する方法です

ドキュメント単位でチャンク結合 + Delta テーブル保存

from pyspark.sql import functions as F

doc_chunks = (
    chunks_df
    .groupBy("file_name")
    .agg(
        F.concat_ws("\n", F.collect_list("content")).alias("chunk_text"),
        F.first("page_id").alias("page_id"),
    )
    .withColumn("chunk_id", F.monotonically_increasing_id())
)

CATALOG = "dev"
SCHEMA = "file_search_demo"
TABLE = "document_chunks"
FULL_TABLE = f"{CATALOG}.{SCHEMA}.{TABLE}"

spark.sql(f"CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}")

doc_chunks.write.mode("overwrite").option("overwriteSchema", "true").saveAsTable(FULL_TABLE)

saved = spark.sql(f"SELECT * FROM {FULL_TABLE}")
display(saved)
print(f"Saved {saved.count()} chunks to {FULL_TABLE}")

この例ではドキュメント全体を1チャンクとして結合しています。実運用ではページ単位や段落単位で分割するのが一般的です。

Step 4: Embedding ベースのセマンティック検索

Delta テーブルに保存したチャンクに対して、ai_query で埋め込みベクトルを生成し、コサイン類似度で検索します。

results = spark.sql("""
    WITH query AS (
        SELECT ai_query(
            'databricks-gte-large-en',
            'What is the churn rate and customer acquisition?'
        ) AS query_embedding
    ),
    doc_embeddings AS (
        SELECT
            file_name,
            chunk_text,
            ai_query('databricks-gte-large-en', chunk_text) AS doc_embedding
        FROM dev.file_search_demo.document_chunks
    )
    SELECT
        d.file_name,
        LEFT(d.chunk_text, 100) AS chunk_preview,
        aggregate(
            zip_with(
                d.doc_embedding,
                q.query_embedding,
                (x, y) -> x * y
            ),
            CAST(0.0 AS DOUBLE),
            (acc, x) -> acc + x
        ) AS similarity
    FROM doc_embeddings d
    CROSS JOIN query q
    ORDER BY similarity DESC
""")

display(results)
print("Search complete")

検索結果

file_name	chunk_preview	similarity
sales_report_2024q4.pdf	Sales Report: FY2024 Q4 Overview Revenue for Oct-Dec 2024 was 850M JPY...	318.71
meeting_notes_20250115.pdf	Executive Meeting Minutes - Jan 15 2025 Agenda 1: FY2025 Business Plan...	255.50
product_spec_alpha.pdf	Product Spec: Project Alpha v3.0 1. Overview Project Alpha is an enterprise...	247.11

「churn rate（解約率）」と「customer acquisition（顧客獲得）」を含む売上レポートが最も高いスコアで返っています。会議議事録は事業計画（売上目標等）を含むため2位、製品仕様書は直接関連しないため3位です。

コサイン類似度の計算方法

ai_query('databricks-gte-large-en', text) は ARRAY<DOUBLE> 型の埋め込みベクトルを直接返します（from_json は不要）。

コサイン類似度は SQL の高階関数で計算しています。

aggregate(
    zip_with(vec1, vec2, (x, y) -> x * y),  -- 要素ごとの積
    CAST(0.0 AS DOUBLE),
    (acc, x) -> acc + x                      -- 総和（= 内積）
) AS similarity

databricks-gte-large-en は正規化済みベクトルを返すため、内積がそのままコサイン類似度になります。

Vector Search Index との比較

本記事では ai_query + SQL によるオンデマンド計算方式を採用しました。Databricks には Vector Search Index（マネージド近似最近傍検索）も用意されていますが、以下の違いがあります。

観点	ai_query + SQL（本記事）	Vector Search Index
セットアップ	不要（SQL のみ）	エンドポイント作成が必要
Free Edition	利用可	利用不可（Pro 以上）
検索速度	ドキュメント数に比例（全件スキャン）	大規模でも高速（ANN）
適したケース	検証・小規模（数百〜数千件）	本番・大規模（数万件〜）

小規模なドキュメント検索であれば、本記事の方式で十分実用的です。

ハマりポイントまとめ

本記事の実装中に遭遇した問題と対処法です。

問題	原因	対処
ai_parse_document が Unsupported file format を返す	.txt ファイルは非対応	PDF・DOCX・PPTX・画像のみ使用する
Serverless で dbutils.fs.cp("file:/...") が失敗	Serverless はローカルファイルシステムへのアクセスを制限している	Volume の FUSE マウントパスに直接書き込む
json.loads(row["parsed"]) で TypeError	collect() した VARIANT 値は VariantVal 型であり文字列ではない	SQL 側で CAST(... AS STRING) してから collect する
from_json(embedding, 'ARRAY<DOUBLE>') で型エラー	ai_query は埋め込みモデルに対して ARRAY<DOUBLE> を直接返す（文字列ではない）	from_json を使わず、返り値をそのまま配列として扱う

まとめ

Databricks の組み込み関数だけで、ファイルアップロードからセマンティック検索までの一連のパイプラインを構築できます。

• read_files + ai_parse_document: PDF からの構造化テキスト抽出
• ai_query('databricks-gte-large-en', ...): テキストの埋め込みベクトル化
• SQL 高階関数（aggregate + zip_with）: コサイン類似度の計算

外部の embedding API や Vector DB を用意せずに、ノートブック1つでファイル検索を試せる点が利点です。

参考リンク

• ai_parse_document function | Databricks
• ai_query function | Databricks
• What are Unity Catalog volumes? | Databricks
• Databricks Vector Search | Databricks

全コード（ノートブック一括コピー用）

%pip install fpdf2 --quiet
dbutils.library.restartPython()

from fpdf import FPDF
import os

VOLUME_PATH = "/Volumes/dev/file_search_demo/documents"

def create_pdf(filename, title, sections):
    pdf = FPDF()
    pdf.add_page()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.set_font("Helvetica", "B", 16)
    pdf.cell(0, 10, title, new_x="LMARGIN", new_y="NEXT")
    pdf.ln(4)
    for heading, lines in sections:
        pdf.set_font("Helvetica", "B", 13)
        pdf.cell(0, 8, heading, new_x="LMARGIN", new_y="NEXT")
        pdf.set_font("Helvetica", size=11)
        for line in lines:
            pdf.cell(0, 7, line, new_x="LMARGIN", new_y="NEXT")
        pdf.ln(3)
    out_path = f"{VOLUME_PATH}/{filename}"
    pdf.output(out_path)
    print(f"Created: {out_path} ({os.path.getsize(out_path)} bytes)")

spark.sql("CREATE SCHEMA IF NOT EXISTS dev.file_search_demo")
spark.sql("""
    CREATE VOLUME IF NOT EXISTS dev.file_search_demo.documents
""")

create_pdf("sales_report_2024q4.pdf", "Sales Report: FY2024 Q4", [
    ("Overview", [
        "Revenue for Oct-Dec 2024 was 850M JPY, up 12% YoY.",
        "Cloud services drove growth with ARR increasing 15% QoQ.",
    ]),
    ("Revenue by Segment", [
        "Cloud Services: 420M JPY (+25% YoY)",
        "On-Premise: 210M JPY (-5% YoY)",
        "Consulting: 150M JPY (+8% YoY)",
        "Support: 70M JPY (+3% YoY)",
    ]),
    ("Key KPIs", [
        "New customers acquired: 142 (target: 130)",
        "Churn rate: 2.1% (improved from 2.8%)",
        "ARPU: 580K JPY/month (up from 550K)",
        "NPS: +42 (up from +38)",
    ]),
])

create_pdf("product_spec_alpha.pdf", "Product Spec: Project Alpha v3.0", [
    ("1. Overview", [
        "Project Alpha is an enterprise self-service BI platform.",
        "Users can query data using natural language without SQL knowledge.",
    ]),
    ("2. Key Features", [
        "Natural Language Query Engine (Japanese/English)",
        "Auto table structure recognition and SQL generation",
        "Benchmark evaluation for SQL accuracy",
        "Auto Dashboard Generation with 10 templates",
    ]),
    ("3. Technical Requirements", [
        "Supported DBs: PostgreSQL 14+, MySQL 8+, Snowflake, BigQuery, Databricks SQL",
        "Min config: 4vCPU / 16GB RAM / 100GB SSD",
        "Browsers: Chrome 100+, Edge 100+, Safari 16+",
    ]),
])

create_pdf("meeting_notes_20250115.pdf", "Executive Meeting Minutes - Jan 15 2025", [
    ("Agenda 1: FY2025 Business Plan", [
        "Revenue target: 4B JPY (18% YoY growth)",
        "Operating margin target: 15% (prev 12%)",
        "New hires: 50 (30 engineers, 15 sales, 5 others)",
    ]),
    ("Agenda 2: Product Roadmap", [
        "Project Alpha v3.0 release: April 2025",
        "NLQ accuracy target: 85%+",
        "Mobile app beta: H2 2025",
    ]),
    ("Agenda 3: Security", [
        "Dec 2024 phishing attack detected and blocked",
        "Company-wide security training in Feb",
        "SOC2 Type II certification target: June 2025",
    ]),
])

print("All PDFs created in Volume.")

parsed_df = spark.sql("""
    SELECT
        _metadata.file_name,
        ai_parse_document(content, MAP('version', '2.0')) AS parsed
    FROM read_files(
        '/Volumes/dev/file_search_demo/documents/*.pdf',
        format => 'binaryFile'
    )
""")
display(parsed_df.select("file_name", "parsed"))

import json

str_df = spark.sql("""
    SELECT
        _metadata.file_name AS file_name,
        CAST(ai_parse_document(content, MAP('version', '2.0')) AS STRING) AS parsed_json
    FROM read_files(
        '/Volumes/dev/file_search_demo/documents/*.pdf',
        format => 'binaryFile'
    )
""")

rows = str_df.collect()
records = []
for row in rows:
    fname = row["file_name"]
    parsed = json.loads(row["parsed_json"])
    for elem in parsed.get("document", {}).get("elements", []):
        content = elem.get("content", "")
        if content:
            records.append({
                "file_name": fname,
                "content": content,
                "element_type": elem.get("element_type", ""),
                "page_id": elem.get("page_id", 0),
            })

chunks_df = spark.createDataFrame(records)
display(chunks_df)
print(f"Total elements: {len(records)}")

from pyspark.sql import functions as F

doc_chunks = (
    chunks_df
    .groupBy("file_name")
    .agg(
        F.concat_ws("\n", F.collect_list("content")).alias("chunk_text"),
        F.first("page_id").alias("page_id"),
    )
    .withColumn("chunk_id", F.monotonically_increasing_id())
)

CATALOG = "dev"
SCHEMA = "file_search_demo"
TABLE = "document_chunks"
FULL_TABLE = f"{CATALOG}.{SCHEMA}.{TABLE}"

spark.sql(f"CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}")

doc_chunks.write.mode("overwrite").option("overwriteSchema", "true").saveAsTable(FULL_TABLE)

saved = spark.sql(f"SELECT * FROM {FULL_TABLE}")
display(saved)
print(f"Saved {saved.count()} chunks to {FULL_TABLE}")

results = spark.sql("""
    WITH query AS (
        SELECT ai_query(
            'databricks-gte-large-en',
            'What is the churn rate and customer acquisition?'
        ) AS query_embedding
    ),
    doc_embeddings AS (
        SELECT
            file_name,
            chunk_text,
            ai_query('databricks-gte-large-en', chunk_text) AS doc_embedding
        FROM dev.file_search_demo.document_chunks
    )
    SELECT
        d.file_name,
        LEFT(d.chunk_text, 100) AS chunk_preview,
        aggregate(
            zip_with(
                d.doc_embedding,
                q.query_embedding,
                (x, y) -> x * y
            ),
            CAST(0.0 AS DOUBLE),
            (acc, x) -> acc + x
        ) AS similarity
    FROM doc_embeddings d
    CROSS JOIN query q
    ORDER BY similarity DESC
""")

display(results)
print("Search complete")