Python如何从DataFrame的SQL片段列中解析列名，生成ID与列的唯一组合

实现思路 #

通过SQL语法解析工具提取列名，相比正则匹配可以完美规避SQL字符串常量、关键字的误识别问题，可靠性更高，我们用Python的sqlglot库实现，步骤如下：

• 安装sqlglot依赖
• 包装SQL片段为可解析的完整SQL语句，提取语法树中的列引用
• 按ID分组去重得到最终结果

代码实现 #

import pandas as pd
import sqlglot
from sqlglot import exp

# 构造样例DataFrame，可替换为你自己的数据源
df = pd.DataFrame({
    "ID": [1, 1, 1, 2],
    "CONDITION": [
        'STATUS = "ABCD"',
        'STATUS = "DEFG"',
        'CASE WHEN STATUS = "XYZ" AND MOD = 1 THEN "CORRECT" ELSE "INCORRECT" END',
        'CODEX = 100'
    ]
})

def extract_sql_columns(sql_snippet: str) -> set:
    column_set = set()
    try:
        # 将SQL片段包装为SELECT语句的WHERE条件，满足语法解析要求
        parsed_sql = sqlglot.parse_one(f"SELECT * FROM dummy_table WHERE {sql_snippet}")
        # 遍历语法树提取所有列引用
        for column in parsed_sql.find_all(exp.Column):
            column_set.add(column.name)
    except Exception:
        # 解析异常可根据业务需求自定义处理逻辑，这里直接返回空集合
        pass
    return column_set

# 提取列名并展开为长表
df = df.assign(COLUMN=df["CONDITION"].apply(extract_sql_columns)).explode("COLUMN")
# 去重得到最终结果
result_df = df[["ID", "COLUMN"]].drop_duplicates().reset_index(drop=True)

输出结果 #

简易正则方案（仅适用于简单场景） #

如果不想引入第三方依赖，可以用正则实现，但要注意无法完全规避误识别问题：

import re

SQL_KEYWORDS = {"CASE", "WHEN", "THEN", "ELSE", "END", "AND", "OR", "NOT", "NULL", "TRUE", "FALSE"}

def extract_columns_regex(sql_snippet: str) -> set:
    # 先移除所有引号包裹的字符串，避免匹配到字符串内的内容
    cleaned_sql = re.sub(r'["\'].*?["\']', '', sql_snippet)
    # 匹配符合SQL列名规范的标识符
    matched = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', cleaned_sql)
    # 过滤SQL关键字
    return {item for item in matched if item.upper() not in SQL_KEYWORDS}

内容的提问来源于stack exchange，提问作者user2653353