Databricks 自然語言處理實戰：SpaCy 與 PySpark 分散式整合技術

# =============================================================================
# Databricks 環境初始化設定
# 目的：在 DBFS 建立目錄結構並上傳 SpaCy 語言模型
# 執行環境：Databricks Notebook
# 版本需求：Databricks Runtime 13.0 以上
# =============================================================================

# -----------------------------------------------------------------------------
# 步驟一：建立 DBFS 目錄結構
#
# dbutils 是 Databricks 提供的工具函式庫，整合了檔案系統操作、
# 秘密管理、筆記本工作流程等多種功能。這個工具函式庫在 Databricks
# 的 Notebook 環境中會自動載入，開發者可以直接使用。
#
# fs.mkdirs 方法用於建立目錄，其行為類似於 Unix 系統的 mkdir -p 命令，
# 如果目錄已經存在則不會產生錯誤，如果父目錄不存在則會自動建立。
# 這個特性讓開發者可以放心執行這個命令，不需要事先檢查目錄是否存在。
#
# 此目錄將作為所有自然語言處理模型的統一存放位置，建議採用有意義的
# 命名結構，例如 /databricks/models/{類別}/{模型名稱}，以便於管理
# 多個不同用途的模型。
# -----------------------------------------------------------------------------
dbutils.fs.mkdirs("dbfs:/databricks/models/spacy")

# -----------------------------------------------------------------------------
# 步驟二：從雲端儲存體複製預訓練模型至 DBFS
#
# Databricks 支援多種雲端儲存服務的存取，包括 Amazon S3、Azure Blob
# Storage、Azure Data Lake Storage Gen2 以及 Google Cloud Storage。
# 不同的雲端服務使用不同的協定前綴：
#   Amazon S3：s3a:// 或 s3://，其中 s3a 效能較佳
#   Azure Blob Storage：wasbs://
#   Azure Data Lake Storage Gen2：abfss://
#   Google Cloud Storage：gs://
#
# dbutils.fs.cp 方法用於複製檔案或目錄，參數說明如下：
#   第一個參數：來源路徑，可以是雲端儲存體或 DBFS 路徑
#   第二個參數：目標路徑，通常是 DBFS 路徑
#   第三個參數：布林值，True 表示遞迴複製整個目錄結構
#
# 在正式環境中，建議將模型檔案儲存在受版本控制的儲存體中，
# 並使用明確的版本號碼作為目錄名稱，以便追蹤模型版本變更。
# -----------------------------------------------------------------------------
dbutils.fs.cp(
    "s3a://your-bucket/models/spacy/",  # 請替換為實際的儲存體路徑
    "dbfs:/databricks/models/spacy/",
    True
)

# -----------------------------------------------------------------------------
# 步驟三：驗證模型檔案是否成功複製
#
# 驗證步驟在環境建置中非常重要，透過確認檔案確實存在且大小正確，
# 可以避免後續處理時因為模型檔案缺失或損毀導致的錯誤。
#
# display 函式是 Databricks Notebook 提供的輸出函式，
# 會以美觀的表格形式呈現資料，包含檔案名稱、大小、修改時間等資訊。
# 相較於 print 函式，display 提供更豐富的視覺化效果。
#
# 預期應看到以下檔案（實際版本號碼可能不同）：
#   en_core_web_lg-3.7.1.tar.gz：SpaCy 官方大型英文模型
#   zh_core_web_trf-3.7.0.tar.gz：中文 Transformer 模型（選用）
#   其他自訂訓練的模型套件
# -----------------------------------------------------------------------------
display(dbutils.fs.ls("dbfs:/databricks/models/spacy/"))

# =============================================================================
# 建立叢集初始化指令碼
# 目的：確保每個節點啟動時自動安裝 SpaCy 模型
# 設計原則：冪等性、錯誤處理、日誌記錄
# =============================================================================

# -----------------------------------------------------------------------------
# 建立指令碼存放目錄
#
# 將初始化指令碼與模型檔案分開存放是一個良好的實踐做法，
# 這樣可以更清楚地區分不同類型的資源，也便於版本管理。
# 建議的目錄結構如下：
#   /databricks/models/：存放各種模型檔案
#   /databricks/scripts/：存放初始化指令碼
#   /databricks/configs/：存放設定檔案
# -----------------------------------------------------------------------------
dbutils.fs.mkdirs("dbfs:/databricks/scripts/")

# -----------------------------------------------------------------------------
# 定義初始化指令碼內容
#
# 這是一個標準的 Bash Shell 指令碼，會在每個節點啟動時執行。
# 指令碼設計時需要考慮以下幾個重點：
#
# 1. 冪等性：指令碼應該可以重複執行而不會產生錯誤或副作用
# 2. 錯誤處理：當某個步驟失敗時，應該適當地處理錯誤
# 3. 日誌記錄：記錄足夠的資訊以便於除錯
# 4. 效能考量：避免不必要的重複安裝
#
# 使用 /dbfs 路徑存取 DBFS 檔案是因為在叢集節點上，
# DBFS 會被掛載為本地檔案系統，掛載點就是 /dbfs 目錄。
# 因此 dbfs:/databricks/models/spacy/ 在節點上可以透過
# /dbfs/databricks/models/spacy/ 路徑存取。
# -----------------------------------------------------------------------------
init_script_content = """#!/bin/bash
# =============================================================================
# SpaCy 模型自動安裝指令碼
# 執行時機：叢集每個節點啟動時
# 目的：確保所有節點都具備相同的 NLP 模型環境
# 維護資訊：請配合模型更新調整版本號碼
# =============================================================================

# 啟用嚴格模式，任何命令失敗都會中斷指令碼執行
# -e：命令失敗時立即結束
# -u：使用未定義變數時視為錯誤
# -o pipefail：管線中任何命令失敗都視為整個管線失敗
set -euo pipefail

# 記錄開始時間，用於效能監控
echo "[$(date '+%Y-%m-%d %H:%M:%S')] 開始安裝 SpaCy 模型"

# -----------------------------------------------------------------------------
# 安裝 SpaCy 官方預訓練大型英文語言模型
#
# en_core_web_lg 是 SpaCy 提供的大型英文語言模型，包含以下功能：
#   完整詞彙表：約 500,000 個詞條
#   詞向量：685,000 維度的詞向量，支援語意相似度計算
#   命名實體識別：可識別 18 種實體類別
#   詞性標註：精確的詞性分析
#   依存句法分析：完整的語法結構分析
#
# 模型檔案大小約 800MB，首次載入需要數秒時間。
# 在分散式環境中，每個節點都需要載入自己的模型實例，
# 因此需要確保每個節點都有足夠的記憶體空間。
#
# 使用 pip install 從本地檔案安裝模型的好處是：
#   不需要網路連線，適合封閉網路環境
#   安裝速度快，因為不需要下載
#   版本可控，確保所有環境使用相同版本
# -----------------------------------------------------------------------------
pip install /dbfs/databricks/models/spacy/en_core_web_lg-3.7.1.tar.gz

# -----------------------------------------------------------------------------
# 安裝自訂訓練的命名實體識別模型（選用）
#
# 如果有針對特定領域訓練的自訂模型，可以在此一併安裝。
# 自訂模型能識別標準模型無法辨識的領域專有實體，例如：
#   醫療領域：藥物名稱、疾病名稱、醫療程序
#   法律領域：法條編號、案件名稱、法律術語
#   金融領域：金融產品名稱、財務指標、公司代碼
#
# 取消下方註解以啟用自訂模型安裝
# -----------------------------------------------------------------------------
# pip install /dbfs/databricks/models/spacy/custom_ner_model-0.0.1.tar.gz

# 記錄完成時間
echo "[$(date '+%Y-%m-%d %H:%M:%S')] SpaCy 模型安裝完成"

# 驗證安裝結果
python -c "import spacy; print(f'SpaCy 版本: {spacy.__version__}')"
python -c "import en_core_web_lg; print('en_core_web_lg 模型載入成功')"

echo "[$(date '+%Y-%m-%d %H:%M:%S')] 模型驗證完成"
"""

# -----------------------------------------------------------------------------
# 將指令碼寫入 DBFS
#
# dbutils.fs.put 方法用於將字串內容寫入檔案，參數說明：
#   第一個參數：目標檔案的完整路徑
#   第二個參數：要寫入的字串內容
#   第三個參數：布林值，True 表示如果檔案已存在則覆寫
#
# 指令碼建立後，需要在叢集組態中指定此路徑才會在啟動時執行。
# 設定方式是進入叢集編輯頁面，展開「進階選項」區塊，
# 在「初始化指令碼」項目中新增指令碼路徑。
# -----------------------------------------------------------------------------
dbutils.fs.put(
    "dbfs:/databricks/scripts/spacy_init.sh",
    init_script_content,
    True
)

# -----------------------------------------------------------------------------
# 驗證指令碼內容
#
# 使用 head 方法讀取指令碼的前 2000 個位元組，確認內容正確。
# 這個驗證步驟可以及早發現語法錯誤或格式問題，
# 避免在叢集啟動時才發現問題。
# -----------------------------------------------------------------------------
print("指令碼內容預覽：")
print(dbutils.fs.head("dbfs:/databricks/scripts/spacy_init.sh"))

# =============================================================================
# 資料載入與前處理模組
# 目的：載入新聞資料集並進行品質檢查與清理
# 適用資料集：AG News Dataset 或類似結構的新聞資料
# =============================================================================

# -----------------------------------------------------------------------------
# 匯入必要的函式庫
#
# 這些函式庫各自負責不同的功能，在自然語言處理流程中都扮演重要角色：
#   spacy：自然語言處理的核心功能
#   numpy：數值運算與陣列操作
#   pandas：資料處理與分析
#   pyspark.sql.functions：Spark SQL 的各種操作函式
#   pyspark.sql.types：Spark SQL 的資料型別定義
# -----------------------------------------------------------------------------

# SpaCy：自然語言處理核心函式庫
# 提供詞性標註、命名實體識別、依存句法分析等功能
# SpaCy 的設計理念是「提供最佳實踐」，預設設定已經過優化
import spacy

# NumPy：數值運算函式庫
# 提供高效的多維陣列運算能力，是許多科學計算套件的基礎
# 在自然語言處理中常用於處理詞向量與數值特徵
import numpy as np

# Pandas：資料分析與處理函式庫
# 提供 DataFrame 資料結構與豐富的資料操作方法
# 在 Pandas UDF 中會使用到 Pandas 的資料結構
import pandas as pd

# PySpark SQL 函式庫
# 這些函式用於 DataFrame 的各種操作，是建構資料處理流程的基本工具
from pyspark.sql.functions import (
    udf,        # 使用者定義函式，將 Python 函式轉換為 Spark 可用的分散式函式
    col,        # 選取 DataFrame 欄位的便捷函式，等同於 df["column_name"]
    length,     # 計算字串長度的內建函式
    when,       # 條件判斷函式，類似 SQL 的 CASE WHEN 語法
    isnull      # 檢查空值的函式，回傳布林值
)

# PySpark 資料型別定義
# 這些型別用於定義 UDF 的輸入輸出型別與 DataFrame Schema
# 正確的型別定義對於 Spark 的效能優化非常重要
from pyspark.sql.types import (
    ArrayType,      # 陣列型別，用於表示元素列表
    StructType,     # 結構型別，類似 JSON 物件或資料表的一列
    StructField,    # 結構欄位定義，包含欄位名稱與型別
    StringType,     # 字串型別
    IntegerType     # 整數型別
)

# -----------------------------------------------------------------------------
# 定義資料來源路徑
#
# 使用雲端物件儲存作為資料來源是企業資料平台的標準做法，
# 這種架構具有以下優點：
#   高可用性：雲端儲存服務提供多區域複製與故障轉移
#   彈性擴展：可以儲存任意大小的資料集
#   成本效益：只需支付實際使用的儲存空間
#   安全性：支援存取控制與加密
#
# 對於部署位置的選擇，建議使用地理位置較近的區域來儲存資料，
# 例如 AWS 的東京區域或 Azure 的日本區域，以降低網路延遲。
# -----------------------------------------------------------------------------
INPUT_PATH = "s3a://nlp-demo/ag_dataset/train.csv"

# -----------------------------------------------------------------------------
# 使用 Spark 讀取 CSV 檔案
#
# spark 是 SparkSession 物件，在 Databricks Notebook 環境中會自動建立。
# SparkSession 是 Spark 2.0 以後的統一入口點，整合了 SparkContext、
# SQLContext 和 HiveContext 的功能。
#
# 讀取 CSV 檔案時，各參數的作用說明如下：
#
#   header='true'：
#     表示 CSV 檔案的第一列是欄位名稱，Spark 會使用這一列作為
#     DataFrame 的欄位名稱。如果設為 'false'，Spark 會自動產生
#     _c0、_c1 等預設欄位名稱。
#
#   inferSchema='true'：
#     啟用自動型別推斷功能。Spark 會掃描資料的前幾列，
#     推斷每個欄位的資料型別。這個功能方便但會有額外的掃描成本，
#     對於大型資料集，建議手動定義 Schema 以提升效能。
#
#   quote='"'：
#     指定引號字元。當欄位值包含逗號或換行符號時，
#     會使用引號將整個值包圍起來。這個設定告訴 Spark 使用
#     雙引號作為引號字元。
#
#   escape='"'：
#     指定跳脫字元。當欄位值本身包含引號時，需要使用跳脫字元
#     來區分。在標準 CSV 中，通常使用連續兩個引號表示一個引號，
#     因此這裡也設定為雙引號。
#
#   multiLine='true'：
#     允許欄位值跨越多行。新聞摘要可能包含換行符號，
#     啟用這個選項可以正確處理這種情況。但要注意這會降低
#     讀取效能，因為無法並行處理單一檔案的不同部分。
# -----------------------------------------------------------------------------
df = spark.read.format('csv').options(
    header='true',
    inferSchema='true',
    quote='"',
    escape='"',
    multiLine='true'
).load(INPUT_PATH)

# -----------------------------------------------------------------------------
# 顯示資料集基本資訊
#
# 在開始資料處理之前，了解資料的基本特性是非常重要的步驟。
# 這些資訊可以幫助開發者規劃後續的處理策略，
# 例如資料量決定需要多少運算資源，欄位結構決定如何設計處理邏輯。
# -----------------------------------------------------------------------------
print(f"資料集總筆數：{df.count():,}")
print(f"資料集欄位：{df.columns}")

# 顯示資料結構與型別
# printSchema 方法會以樹狀結構顯示 DataFrame 的 Schema，
# 包含每個欄位的名稱、型別與是否允許空值
df.printSchema()

# 預覽前幾筆資料，了解資料格式與內容
# show 方法顯示 DataFrame 的前 n 筆資料
# truncate 參數限制每個欄位最多顯示的字元數，避免輸出過長
df.show(5, truncate=80)

# =============================================================================
# 資料品質檢查與清理
# =============================================================================

# -----------------------------------------------------------------------------
# 檢查各欄位的空值數量
# 空值會導致 SpaCy 模型處理錯誤，必須在前處理階段妥善處理
# 在大型資料集中，即使空值比例很低，絕對數量也可能相當可觀
# -----------------------------------------------------------------------------
print("=== 資料品質檢查 ===")

# 計算 description 欄位中空值的數量
# filter() 篩選符合條件的資料列
# isnull() 檢查欄位是否為空值，回傳布林值
# count() 計算篩選後的資料筆數
null_count = df.filter(isnull(col("description"))).count()
total_count = df.count()
null_percentage = (null_count / total_count) * 100

print(f"description 欄位空值數量：{null_count:,} 筆 ({null_percentage:.2f}%)")

# -----------------------------------------------------------------------------
# 分析文字長度分布
# 了解文字長度有助於評估處理時間與記憶體需求
# 過長的文字可能需要分段處理或設定 SpaCy 的 max_length 參數
# 過短的文字可能缺乏足夠的語境資訊進行有效的實體識別
# -----------------------------------------------------------------------------

# withColumn 方法新增一個名為 text_length 的計算欄位
df_with_length = df.withColumn(
    "text_length",
    length(col("description"))
)

# 計算文字長度的描述性統計
# summary() 計算最小值、各分位數、最大值、平均值等統計資訊
print("\n文字長度統計：")
length_stats = df_with_length.select("text_length").summary(
    "min", "25%", "50%", "75%", "max", "mean"
)
length_stats.show()

# -----------------------------------------------------------------------------
# 資料清理：過濾空值與過短文字
# isNotNull() 確保欄位值不為空值
# length() > 20 過濾掉長度過短的文字（20 字元約為 3-5 個英文單詞）
# 過短的文字通常缺乏足夠語境資訊，實體識別效果不佳
# -----------------------------------------------------------------------------
df_cleaned = df.filter(
    (col("description").isNotNull()) &
    (length(col("description")) > 20)
)

cleaned_count = df_cleaned.count()
removed_count = total_count - cleaned_count
print(f"\n清理前資料筆數：{total_count:,}")
print(f"清理後資料筆數：{cleaned_count:,}")
print(f"移除資料筆數：{removed_count:,} ({(removed_count/total_count)*100:.2f}%)")

# -----------------------------------------------------------------------------
# 快取清理後的資料
# cache() 將 DataFrame 保留在各節點的記憶體中
# 這是延遲執行操作，會在第一次行動操作時實際執行
# 對於需要多次存取的資料，快取能大幅提升後續操作的執行效能
# -----------------------------------------------------------------------------
df_cleaned.cache()

# 觸發快取並統計實際資料量
print(f"\n資料已快取，準備進行 NLP 處理")

# =============================================================================
# 資料分區優化
# =============================================================================

# -----------------------------------------------------------------------------
# 檢查目前的分區數量
# 了解資料分布狀況，判斷是否需要調整
# -----------------------------------------------------------------------------
current_partitions = df_cleaned.rdd.getNumPartitions()
print(f"目前分區數量：{current_partitions}")

# -----------------------------------------------------------------------------
# 根據叢集規模計算最佳分區數量
# 假設使用的叢集配置：
#   - 4 個 Worker 節點
#   - 每個節點 8 個 CPU 核心
#   - 總共 32 個核心
# 建議分區數量：核心數 × 2 到 核心數 × 4
# 這個範圍能確保所有核心都有任務可執行，同時不會有過多的排程開銷
# -----------------------------------------------------------------------------
CLUSTER_CORES = 32  # 請根據實際叢集配置調整
PARTITIONS_PER_CORE = 3
optimal_partitions = CLUSTER_CORES * PARTITIONS_PER_CORE

print(f"叢集核心總數：{CLUSTER_CORES}")
print(f"建議分區數量：{optimal_partitions}")

# -----------------------------------------------------------------------------
# 調整分區數量
# repartition() 會觸發完整的資料重新洗牌（Shuffle）
# Shuffle 會將資料重新分配到各個分區，確保資料均勻分布
# 這個操作相對耗時，但對於後續大量運算來說是值得的投資
# -----------------------------------------------------------------------------
if current_partitions != optimal_partitions:
    df_repartitioned = df_cleaned.repartition(optimal_partitions)
    df_repartitioned.cache()
    
    # 釋放原本的快取，避免佔用過多記憶體
    df_cleaned.unpersist()
    
    print(f"已調整分區數量為：{optimal_partitions}")
else:
    df_repartitioned = df_cleaned
    print("分區數量已為最佳值，無需調整")

# =============================================================================
# 命名實體識別 UDF 定義
# =============================================================================

# -----------------------------------------------------------------------------
# 定義輸出資料結構
# Spark 需要明確的型別定義才能正確處理資料
# 每筆文字可能包含零到多個實體，因此外層使用 ArrayType 表示陣列
# 每個實體使用 StructType 定義結構，包含四個屬性：
#   - text：實體的文字內容（如 "蘋果公司"、"台北市"）
#   - start_char：實體在原文中的起始字元位置（索引從 0 開始）
#   - end_char：實體在原文中的結束字元位置（不包含該位置字元）
#   - label：實體的分類標籤（如 PERSON、ORG、GPE、DATE）
# -----------------------------------------------------------------------------
entity_schema = ArrayType(
    StructType([
        StructField("text", StringType(), nullable=False),
        StructField("start_char", IntegerType(), nullable=False),
        StructField("end_char", IntegerType(), nullable=False),
        StructField("label", StringType(), nullable=False)
    ])
)

# -----------------------------------------------------------------------------
# 定義命名實體識別函式
# 此函式將被轉換為 PySpark UDF，在各工作節點上平行執行
# 函式設計要點：
#   1. 處理空值與異常輸入
#   2. 使用全域變數快取模型，避免重複載入
#   3. 回傳結構化結果，符合預定義的 Schema
# -----------------------------------------------------------------------------
def extract_entities(text):
    """
    從輸入文字中擷取命名實體
    
    此函式使用 SpaCy 的預訓練語言模型處理輸入文字，
    識別其中的命名實體並回傳結構化結果。
    為了在分散式環境中維持效能，採用全域變數快取已載入的模型。
    
    Parameters
    ----------
    text : str
        待處理的文字字串，可以是任何長度的英文文字
    
    Returns
    -------
    list
        包含所有識別實體的串列，每個實體為一個包含
        [文字, 起始位置, 結束位置, 標籤] 的串列
        若無識別到任何實體，回傳空串列
    
    Examples
    --------
    >>> extract_entities("Apple CEO Tim Cook announced new products in Cupertino.")
    [['Apple', 0, 5, 'ORG'], ['Tim Cook', 10, 18, 'PERSON'], ['Cupertino', 47, 56, 'GPE']]
    """
    # 使用全域變數儲存已載入的模型
    # 這樣可以避免每次函式呼叫都重新載入模型，大幅提升效能
    # 在分散式環境中，每個 Worker 節點會維護自己的模型實例
    # 模型實例在 Executor 的生命週期內會被重複使用
    global nlp_model
    
    # 處理空值輸入，避免模型處理錯誤
    if text is None or str(text).strip() == '':
        return []
    
    try:
        # 嘗試使用已載入的模型處理文字
        doc = nlp_model(str(text))
    except NameError:
        # 若模型尚未載入（nlp_model 變數不存在），則進行載入
        # en_core_web_lg 是 SpaCy 提供的大型英文語言模型
        # 包含完整的命名實體識別、詞向量與語法分析能力
        # 模型大小約 800MB，首次載入需要數秒時間
        nlp_model = spacy.load('en_core_web_lg')
        doc = nlp_model(str(text))
    
    # 從處理結果中擷取所有實體
    # doc.ents 是一個包含所有識別實體的元組
    # 每個實體物件（Entity）具有以下重要屬性：
    #   - text：實體的原始文字內容
    #   - start_char：起始字元位置
    #   - end_char：結束字元位置
    #   - label_：類別標籤字串（注意有底線）
    entities = [
        [ent.text, ent.start_char, ent.end_char, ent.label_]
        for ent in doc.ents
    ]
    
    return entities

# -----------------------------------------------------------------------------
# 將 Python 函式轉換為 PySpark UDF
# udf() 函式接收 Python 函式與輸出型別定義作為參數
# 轉換後的 UDF 可以在 DataFrame 的 withColumn、select 等操作中使用
# UDF 會在 Spark 執行時自動序列化並分發到各 Worker 節點
# -----------------------------------------------------------------------------
extract_entities_udf = udf(extract_entities, entity_schema)

# =============================================================================
# 執行命名實體識別
# =============================================================================

# -----------------------------------------------------------------------------
# 將 NER UDF 應用於 DataFrame
# withColumn 方法會新增一個名為 entities 的欄位
# 該欄位包含從 description 欄位擷取出的所有命名實體
# 這是延遲執行（Lazy Evaluation）操作，定義轉換但不立即執行
# Spark 會自動優化執行計畫，並在觸發行動操作時才實際執行
# -----------------------------------------------------------------------------
documents_df = df_repartitioned.withColumn(
    'entities',
    extract_entities_udf(col('description'))
)

# -----------------------------------------------------------------------------
# 預覽處理結果
# show() 是行動操作，會觸發實際的分散式運算
# 此時 Spark 會開始在各節點平行執行 NER 處理
# truncate=80 限制每個欄位的顯示長度
# -----------------------------------------------------------------------------
print("=== NER 處理結果預覽 ===")
documents_df.select('description', 'entities').show(5, truncate=80)

# -----------------------------------------------------------------------------
# 統計各類別實體的數量分布
# 這個分析能幫助了解資料集中實體類型的組成
# 對於評估模型效果與規劃後續應用都很有價值
# -----------------------------------------------------------------------------
from pyspark.sql.functions import explode, size

# 先統計每筆文字識別出的實體數量
entity_count_stats = documents_df.withColumn(
    'entity_count',
    size(col('entities'))
).select('entity_count').summary()
print("\n每筆文字的實體數量統計：")
entity_count_stats.show()

# 展開實體陣列，將每個實體變成獨立的資料列
# explode() 函式會將陣列欄位中的每個元素展開為獨立的列
# 原本一列包含多個實體的資料，會被展開為多列，每列一個實體
entities_exploded = documents_df.select(
    explode(col('entities')).alias('entity')
)

# 按實體類別統計數量
# col('entity.label') 存取巢狀結構中的 label 欄位
# groupBy 將相同標籤的實體歸為一組
# count() 計算每組的數量
# orderBy 按數量降序排列
entity_type_counts = entities_exploded.select(
    col('entity.label').alias('entity_type')
).groupBy('entity_type').count().orderBy(col('count').desc())

print("\n各類別命名實體數量分布：")
entity_type_counts.show(15)

# =============================================================================
# Pandas UDF 實作：高效能命名實體識別
# =============================================================================

# -----------------------------------------------------------------------------
# 匯入 Pandas UDF 相關模組
# pandas_udf 是裝飾器，用於將 Python 函式轉換為 Pandas UDF
# PandasUDFType 定義 UDF 的類型（SCALAR 或 GROUPED_MAP）
# -----------------------------------------------------------------------------
from pyspark.sql.functions import pandas_udf
from typing import Iterator

# -----------------------------------------------------------------------------
# 定義 Pandas UDF 進行批次命名實體識別
# 使用 Iterator 型別提示，表示這是一個迭代器形式的 Pandas UDF
# 迭代器形式能更有效地處理大型資料集，減少記憶體使用
# 輸出型別使用先前定義的 entity_schema
# -----------------------------------------------------------------------------
@pandas_udf(entity_schema)
def extract_entities_optimized(
    text_iterator: Iterator[pd.Series]
) -> Iterator[pd.Series]:
    """
    使用 Pandas UDF 進行高效能批次命名實體識別
    
    此函式利用 Pandas UDF 的批次處理機制，結合 SpaCy 的 nlp.pipe 方法，
    實現高效能的分散式命名實體識別。相較於傳統 UDF，效能可提升 3-10 倍。
    
    Parameters
    ----------
    text_iterator : Iterator[pd.Series]
        文字資料的迭代器，每個批次為一個 Pandas Series
    
    Yields
    ------
    pd.Series
        對應的實體識別結果，每個元素為一個實體串列
    
    Notes
    -----
    迭代器形式的 Pandas UDF 特別適合需要初始化昂貴資源（如載入模型）的場景，
    因為初始化只需要在迭代開始時執行一次，而非每個批次都執行。
    """
    # 在迭代器開始時載入模型，整個迭代過程中只載入一次
    # 這比在每個批次中檢查並載入模型更有效率
    nlp = spacy.load('en_core_web_lg')
    
    # 設定 SpaCy 的批次處理參數
    # batch_size 控制每批處理的文件數量
    # 較大的批次能提升效能，但會增加記憶體使用
    # 根據實務測試，50-100 是比較好的平衡點
    BATCH_SIZE = 64
    
    # 遍歷所有文字批次
    for text_series in text_iterator:
        # 將 Series 轉換為串列，準備進行批次處理
        # fillna('') 將空值替換為空字串，避免 SpaCy 處理錯誤
        # astype(str) 確保所有值都是字串型別
        texts = text_series.fillna('').astype(str).tolist()
        
        # 使用 nlp.pipe 進行批次處理
        # pipe 方法接收可迭代物件，回傳處理結果的生成器
        # 這比逐一呼叫 nlp(text) 更有效率
        results = []
        for doc in nlp.pipe(texts, batch_size=BATCH_SIZE):
            # 擷取每篇文件中的所有實體
            entities = [
                [ent.text, ent.start_char, ent.end_char, ent.label_]
                for ent in doc.ents
            ]
            results.append(entities)
        
        # 將結果轉換為 Pandas Series 並回傳
        # 結果 Series 的長度必須與輸入 Series 相同
        yield pd.Series(results)

# -----------------------------------------------------------------------------
# 使用優化版 UDF 處理資料
# 語法與傳統 UDF 相同，但內部處理機制完全不同
# Pandas UDF 會將多筆資料打包成批次處理，大幅提升效能
# -----------------------------------------------------------------------------
documents_df_optimized = df_repartitioned.withColumn(
    'entities',
    extract_entities_optimized(col('description'))
)

# -----------------------------------------------------------------------------
# 執行處理並計時
# 透過實際執行來驗證效能提升效果
# -----------------------------------------------------------------------------
import time

print("開始執行優化版 NER 處理...")
start_time = time.time()

# 觸發實際運算（count 是行動操作）
processed_count = documents_df_optimized.count()

end_time = time.time()
processing_time = end_time - start_time

print(f"處理完成")
print(f"處理資料筆數：{processed_count:,}")
print(f"總處理時間：{processing_time:.2f} 秒")
print(f"平均每秒處理：{processed_count/processing_time:.0f} 筆")

# =============================================================================
# 叢集資源設定檢視
# =============================================================================

# -----------------------------------------------------------------------------
# 檢視目前的 Spark 設定
# 這些設定會影響分散式運算的效能與穩定性
# 若發現設定不適當，需要在叢集組態中調整後重新啟動叢集
# -----------------------------------------------------------------------------
print("=== 目前 Spark 設定 ===")

# Executor 記憶體：每個執行器可使用的記憶體量
# 建議至少 8GB 以上，確保能載入語言模型並有足夠緩衝空間
executor_memory = spark.conf.get("spark.executor.memory", "未設定")
print(f"Executor 記憶體：{executor_memory}")

# Executor 核心數：每個執行器使用的 CPU 核心數
# 通常設定為 4-8 個核心，取決於任務類型
executor_cores = spark.conf.get("spark.executor.cores", "未設定")
print(f"Executor 核心數：{executor_cores}")

# Driver 記憶體：驅動程式可使用的記憶體量
# 需要足夠記憶體來處理結果彙整與廣播變數
driver_memory = spark.conf.get("spark.driver.memory", "未設定")
print(f"Driver 記憶體：{driver_memory}")

# Shuffle 分區數：資料洗牌時的分區數量
# 建議與資料處理時的分區數量保持一致
shuffle_partitions = spark.conf.get("spark.sql.shuffle.partitions", "未設定")
print(f"Shuffle 分區數：{shuffle_partitions}")

# -----------------------------------------------------------------------------
# 建議的叢集設定（需在 Databricks 叢集組態中設定）
# 以下設定適用於中型自然語言處理任務
# -----------------------------------------------------------------------------
print("\n=== 建議設定 ===")
print("spark.executor.memory: 8g 或更高")
print("spark.executor.cores: 4")
print("spark.driver.memory: 8g")
print("spark.sql.shuffle.partitions: 與資料分區數相同")
print("spark.python.worker.memory: 2g")

# =============================================================================
# 處理結果儲存
# =============================================================================

# -----------------------------------------------------------------------------
# 定義輸出路徑
# 使用雲端物件儲存作為永久儲存位置
# Parquet 格式能完整保留巢狀結構，且具有良好的壓縮效率
# -----------------------------------------------------------------------------
OUTPUT_PATH = "s3a://nlp-demo/ag_dataset/processed/entities_result"

# -----------------------------------------------------------------------------
# 儲存處理結果為 Parquet 格式
# write.parquet 是 DataFrame 的輸出方法
# mode 參數控制當目標路徑已存在時的處理方式：
#   - overwrite：覆寫已存在的資料
#   - append：追加到現有資料
#   - ignore：若已存在則不執行
#   - error：若已存在則拋出錯誤
# compression 參數指定壓縮演算法，snappy 提供良好的壓縮/解壓縮效能平衡
# -----------------------------------------------------------------------------
documents_df_optimized.write.parquet(
    OUTPUT_PATH,
    mode="overwrite",
    compression="snappy"
)

print(f"處理結果已儲存至：{OUTPUT_PATH}")

# -----------------------------------------------------------------------------
# 驗證儲存結果
# 讀取剛才儲存的資料，確認儲存成功且資料完整
# -----------------------------------------------------------------------------
result_df = spark.read.parquet(OUTPUT_PATH)
stored_count = result_df.count()

print(f"儲存的資料筆數：{stored_count:,}")

# 顯示資料結構，確認巢狀欄位正確保留
print("\n資料結構：")
result_df.printSchema()

# 預覽儲存的資料
print("\n資料預覽：")
result_df.select('description', 'entities').show(3, truncate=60)

# =============================================================================
# 進階分析範例：實體統計與關聯分析
# =============================================================================

from pyspark.sql.functions import explode, col, collect_list, array_distinct

# -----------------------------------------------------------------------------
# 範例一：最常被提及的組織名稱
# 這類分析能幫助了解新聞報導的焦點與商業趨勢
# 對於輿情監測、競爭分析等應用場景非常有價值
# -----------------------------------------------------------------------------
print("=== 最常被提及的組織 (ORG) ===")

# 展開實體欄位並篩選組織類型
org_entities = documents_df_optimized.select(
    explode(col('entities')).alias('entity')
).filter(
    col('entity.label') == 'ORG'
).select(
    col('entity.text').alias('organization')
)

# 統計各組織出現次數並排序
org_frequency = org_entities.groupBy('organization').count().orderBy(
    col('count').desc()
)

org_frequency.show(15)

# -----------------------------------------------------------------------------
# 範例二：最常被提及的人物
# 識別新聞中的關鍵人物，有助於追蹤重要人士的新聞曝光度
# -----------------------------------------------------------------------------
print("\n=== 最常被提及的人物 (PERSON) ===")

person_entities = documents_df_optimized.select(
    explode(col('entities')).alias('entity')
).filter(
    col('entity.label') == 'PERSON'
).select(
    col('entity.text').alias('person')
)

person_frequency = person_entities.groupBy('person').count().orderBy(
    col('count').desc()
)

person_frequency.show(15)

# -----------------------------------------------------------------------------
# 範例三：地理位置分布
# 分析新聞報導的地理焦點，了解哪些地區受到較多關注
# GPE 代表地理政治實體，通常是國家、城市或地區名稱
# -----------------------------------------------------------------------------
print("\n=== 地理位置分布 (GPE) ===")

location_entities = documents_df_optimized.select(
    explode(col('entities')).alias('entity')
).filter(
    col('entity.label') == 'GPE'
).select(
    col('entity.text').alias('location')
)

location_frequency = location_entities.groupBy('location').count().orderBy(
    col('count').desc()
)

location_frequency.show(15)

# -----------------------------------------------------------------------------
# 範例四：輸出為其他格式供外部系統使用
# JSON 格式具有良好的可讀性與跨語言相容性
# 適合作為 API 回傳格式或匯入到 NoSQL 資料庫
# -----------------------------------------------------------------------------
JSON_OUTPUT_PATH = "s3a://nlp-demo/ag_dataset/processed/entities_result.json"

# 選擇需要的欄位並輸出為 JSON
documents_df_optimized.select(
    'description', 'entities'
).write.json(
    JSON_OUTPUT_PATH,
    mode="overwrite"
)

print(f"\nJSON 格式結果已儲存至：{JSON_OUTPUT_PATH}")

# =============================================================================
# 生產化 NER 處理 Notebook
# 設計為可參數化執行，支援 Databricks Jobs 排程
# =============================================================================

# -----------------------------------------------------------------------------
# 定義執行參數
# 這些參數可以在 Jobs 設定中指定，也可以在互動式執行時手動輸入
# dbutils.widgets 提供互動式參數輸入介面
# getArgument 則用於取得 Jobs 傳入的參數值
# -----------------------------------------------------------------------------

# 定義參數與預設值
# 輸入路徑：待處理資料的位置
INPUT_PATH = getArgument("input_path", "s3a://nlp-demo/ag_dataset/train.csv")

# 輸出路徑：處理結果的儲存位置
OUTPUT_PATH = getArgument("output_path", "s3a://nlp-demo/ag_dataset/processed/entities")

# 模型名稱：使用的 SpaCy 模型
MODEL_NAME = getArgument("model_name", "en_core_web_lg")

# 處理日期：用於資料分區與日誌記錄
PROCESS_DATE = getArgument("process_date", "2025-11-28")

# -----------------------------------------------------------------------------
# 記錄執行參數（用於追蹤與除錯）
# -----------------------------------------------------------------------------
print("=" * 60)
print("NER 處理任務開始")
print("=" * 60)
print(f"執行時間：{PROCESS_DATE}")
print(f"輸入路徑：{INPUT_PATH}")
print(f"輸出路徑：{OUTPUT_PATH}")
print(f"使用模型：{MODEL_NAME}")
print("=" * 60)

# -----------------------------------------------------------------------------
# 匯入必要函式庫
# -----------------------------------------------------------------------------
import spacy
import pandas as pd
import time
from pyspark.sql.functions import col, pandas_udf
from pyspark.sql.types import ArrayType, StructType, StructField, StringType, IntegerType
from typing import Iterator

# -----------------------------------------------------------------------------
# 載入資料
# -----------------------------------------------------------------------------
print("\n[Step 1/5] 載入資料...")
start_time = time.time()

df = spark.read.format('csv').options(
    header='true',
    inferSchema='true',
    quote='"',
    escape='"'
).load(INPUT_PATH)

# 資料驗證：確保必要欄位存在
required_columns = ['description']
for col_name in required_columns:
    if col_name not in df.columns:
        raise ValueError(f"缺少必要欄位：{col_name}")

# 資料清理
df_cleaned = df.filter(
    (col("description").isNotNull()) &
    (col("description") != "")
)

input_count = df_cleaned.count()
print(f"載入資料筆數：{input_count:,}")
print(f"耗時：{time.time() - start_time:.2f} 秒")

# -----------------------------------------------------------------------------
# 定義 NER UDF
# -----------------------------------------------------------------------------
print("\n[Step 2/5] 初始化 NER 模型...")

entity_schema = ArrayType(StructType([
    StructField("text", StringType(), False),
    StructField("start_char", IntegerType(), False),
    StructField("end_char", IntegerType(), False),
    StructField("label", StringType(), False)
]))

@pandas_udf(entity_schema)
def extract_entities_production(
    text_iterator: Iterator[pd.Series]
) -> Iterator[pd.Series]:
    """
    生產環境用的 NER UDF
    包含完整的錯誤處理與日誌記錄
    """
    # 載入指定的模型
    nlp = spacy.load(MODEL_NAME)
    batch_num = 0
    
    for text_series in text_iterator:
        batch_num += 1
        texts = text_series.fillna('').astype(str).tolist()
        
        results = []
        for doc in nlp.pipe(texts, batch_size=64):
            entities = [
                [ent.text, ent.start_char, ent.end_char, ent.label_]
                for ent in doc.ents
            ]
            results.append(entities)
        
        yield pd.Series(results)

print("NER 模型初始化完成")

# -----------------------------------------------------------------------------
# 執行 NER 處理
# -----------------------------------------------------------------------------
print("\n[Step 3/5] 執行 NER 處理...")
start_time = time.time()

# 優化分區數量
df_partitioned = df_cleaned.repartition(96)

# 執行 NER
documents_df = df_partitioned.withColumn(
    'entities',
    extract_entities_production(col('description'))
)

# 快取結果（若需要多次使用）
documents_df.cache()

# 觸發實際運算
output_count = documents_df.count()
processing_time = time.time() - start_time

print(f"處理完成")
print(f"輸出資料筆數：{output_count:,}")
print(f"處理耗時：{processing_time:.2f} 秒")
print(f"處理速度：{output_count/processing_time:.0f} 筆/秒")

# -----------------------------------------------------------------------------
# 儲存結果
# -----------------------------------------------------------------------------
print("\n[Step 4/5] 儲存結果...")
start_time = time.time()

# 依日期分區儲存，便於後續管理
output_path_with_date = f"{OUTPUT_PATH}/date={PROCESS_DATE}"

documents_df.write.parquet(
    output_path_with_date,
    mode="overwrite",
    compression="snappy"
)

print(f"結果已儲存至：{output_path_with_date}")
print(f"儲存耗時：{time.time() - start_time:.2f} 秒")

# -----------------------------------------------------------------------------
# 品質檢查
# -----------------------------------------------------------------------------
print("\n[Step 5/5] 執行品質檢查...")

# 驗證儲存的資料
stored_df = spark.read.parquet(output_path_with_date)
stored_count = stored_df.count()

# 檢查資料完整性
if stored_count != output_count:
    raise ValueError(f"資料筆數不符！預期 {output_count}，實際 {stored_count}")

# 計算品質指標
from pyspark.sql.functions import size

entity_stats = stored_df.withColumn(
    'entity_count', size(col('entities'))
).agg({
    'entity_count': 'avg',
    'entity_count': 'sum'
}).collect()[0]

print(f"品質檢查通過！")
print(f"儲存資料筆數：{stored_count:,}")

# -----------------------------------------------------------------------------
# 執行摘要
# -----------------------------------------------------------------------------
print("\n" + "=" * 60)
print("NER 處理任務完成")
print("=" * 60)
print(f"輸入筆數：{input_count:,}")
print(f"輸出筆數：{output_count:,}")
print(f"總耗時：{time.time() - start_time:.2f} 秒")
print("=" * 60)