如何基于T-SQL查询或SQL Server存储过程创建Spark DataFrame

从pyodbc游标提取Schema的方法 #

pyodbc游标自带的description属性可以获取结果集的列元数据，包括列名、数据类型等信息，我们可以基于这个属性生成Spark可识别的Schema结构。

步骤1：提取列元数据并映射Spark类型 #

执行查询后，先从游标中提取列信息，再将SQL Server数据类型映射为Spark对应类型：

import pyodbc
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, TimestampType, DecimalType

connection_string = "Driver={SQL Server};Server=LOCALHOST;Database=master;Trusted_Connection=yes;"
db_connection = pyodbc.connect(connection_string)

sql_query = """
SET NOCOUNT ON
DECLARE @TBL_TEST AS TABLE (
    column_1 INT NOT NULL PRIMARY KEY CLUSTERED IDENTITY(1, 1),
    column_2 VARCHAR(10) NOT NULL,
    column_3 VARCHAR(20) NULL,
    column_4 INT NOT NULL
)

INSERT INTO @TBL_TEST (column_2, column_3, column_4)
VALUES
('test1_col2', 'test1_col3', 100),
('test2_col2', 'test2_col3', 200),
('test3_col2', NULL, 300)

SET NOCOUNT OFF
SELECT t.* FROM @TBL_TEST AS t
"""

cursor = db_connection.cursor()
cursor.execute(sql_query)

# 获取列名列表
columns = [col[0] for col in cursor.description]
# 自定义SQL Server到Spark的类型映射（可根据实际需求扩展）
type_mapping = {
    int: IntegerType(),
    str: StringType(),
    pyodbc.Date: TimestampType(),
    pyodbc.Decimal: DecimalType()
}

# 生成Spark Schema
spark_schema = StructType()
sample_row = cursor.fetchone()
for idx, col_name in enumerate(columns):
    data_type = type(sample_row[idx])
    spark_type = type_mapping.get(data_type, StringType())
    spark_schema.add(StructField(col_name, spark_type, nullable=True))

# 重新执行查询获取完整数据
cursor.execute(sql_query)
rows = cursor.fetchall()

cursor.close()
db_connection.close()

步骤2：构建Spark DataFrame #

将pyodbc返回的行对象转为列表，结合生成的Schema创建DataFrame：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("TSQLToSpark").getOrCreate()

# 转换行对象为列表格式
data = [list(row) for row in rows]
df = spark.createDataFrame(data, schema=spark_schema)

df.show()
df.printSchema()

其他可行实现方案 #

方案1：直接使用Spark JDBC执行T-SQL查询/存储过程 #

Spark JDBC原生支持执行自定义T-SQL语句，包括调用存储过程，无需通过pyodbc中转，还能自动推断Schema并优化数据读取：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("SparkJDBCForTSQL").getOrCreate()

jdbc_url = "jdbc:sqlserver://LOCALHOST;databaseName=master;integratedSecurity=true;"
df = spark.read \
    .format("jdbc") \
    .option("url", jdbc_url) \
    .option("query", sql_query)  # 可替换为存储过程调用，比如"EXEC dbo.YourStoredProc @Param=123"
    .option("driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver") \
    .load()

df.show()
df.printSchema()

方案2：存储过程结果写入临时表后读取 #

如果存储过程无法直接通过JDBC的query参数执行，可先将结果写入SQL Server临时表，再让Spark读取：

1. 执行存储过程并写入临时表（可嵌入到JDBC的query参数中）：

EXEC dbo.YourTargetProc @Param1=1;
SELECT * INTO ##GlobalTempResult FROM dbo.ProcResultTable;

2. Spark读取全局临时表：

df = spark.read \
    .format("jdbc") \
    .option("url", jdbc_url) \
    .option("dbtable", "##GlobalTempResult") \
    .option("driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver") \
    .load()

注：使用全局临时表（##前缀）可跨连接会话访问，适合Spark的多线程读取场景。

方案3：借助Pandas+Apache Arrow自动转换 #

通过pyodbc将数据转为Pandas DataFrame，再利用Spark的Apache Arrow支持快速转换为Spark DataFrame，无需手动处理Schema：

import pandas as pd
from pyspark.sql import SparkSession

# pyodbc获取数据并转为Pandas DataFrame
cursor.execute(sql_query)
pandas_df = pd.DataFrame.from_records(cursor.fetchall(), columns=columns)

# 启用Arrow加速转换
spark = SparkSession.builder.appName("ArrowTSQLConversion").getOrCreate()
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")

spark_df = spark.createDataFrame(pandas_df)
spark_df.show()

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