从目标方言中的 SQL 查询创建 (Py)Spark dataframe

Question

很快，我的需要：从 T-SQL (SQL Server) 中或多或少复杂的查询和/或从 SQL 服务器存储过程的 output 创建一个 Spark dataframe。

据我了解，Spark 不允许以底层数据源的方言执行查询。 是的，有一种方法可以获取低级别 object 并执行存储过程，但通过这种方式我在 output 中没有 Spark DF。

因此，我想以经典的 pyodbc 方式执行查询，获取结果，然后使用 function SparkSession.createDataFrame(data, schema=None, samplingRatio=None, verifySchema=True)构建 Spark dataframe 提供数据和模式. 我可以获取数据，但无法从 output cursor 构建模式（成对列表（列名、数据类型） ）。按照一个工作示例从 SQL 的本地实例中（生成和）提取示例数据服务器：

import pyodbc

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()
rows = cursor.execute(sql_query).fetchall()
cursor.close()
db_connection.close()

print(rows)

如何从返回的 cursor 中提取模式并获得模式object 以提供给 createDataFrame() function？

请记住，我的目标是在主题上，所以也欢迎其他方式！

先感谢您！

Answer 1

如果您使用 pyodbc，则催化剂优化器生成的结果 java 字节代码仅作为一个节点（执行程序）运行，而不是整个集群。 对于更大的数据集，这会阻止集群的充分利用和性能问题。

最好用JDBC的spark驱动，微软有。

https://learn.microsoft.com/en-us/sql/connect/spark/connector?view=sql-server-ver16

将复杂的 T-SQL 创建为视图并读取它们。 这就是 spark 的用途——读取文件。 使用 JDBC 驱动程序 (spark)，如果需要，您可以通过更改分区方法并行读取。

为正确版本的 spark 安装 Marven 库。

我使用的是 Spark 版本 > 3.1。

我有一个名为 v 的视图的冒险作品数据库。

#
#  Set connection properties
#

server_name = "jdbc:sqlserver://svr4tips2030.database.windows.net"
database_name = "dbs4advwrks"
url = server_name + ";" + "databaseName=" + database_name + ";"
table_name = "dbo.vDMPrep"
user_name = "enter your user here"
password = "enter your password here"

使用 JDBC 驱动程序进行典型的 spark.read() 调用。

df = spark.read \
        .format("com.microsoft.sqlserver.jdbc.spark") \
        .option("url", url) \
        .option("dbtable", table_name) \
        .option("user", user_name) \
        .option("password", password).load()

display(df)

这是显示 dataframe 的结果。

数据框是否严格类型化？ 答案是肯定的，因为它从 SQL 服务器获取字段信息。

最后但同样重要的是，视图是否复杂？ 下图显示了 8 个表的连接和聚合以获得视图的最终结果。

总之，使用数据库中的视图为 Spark 预编译数据集。 使用微软的JDBC驱动，使用dataframe从SQL服务器读写。

至于存储过程，有一种方法是使用驱动程序来执行非查询。 我将不得不寻找代码。 请继续关注更新或第 2 部分。

Answer 2

这是答案的第二部分。 没有好的方法可以将存储过程调用的结果作为 dataframe 返回。

这是此驱动程序的 MSFT github 站点上的链接，说明不支持存储过程。

https://github.com/microsoft/sql-spark-connector/issues/21

这是一个 hack - 解决方法。

就我而言，我的 SP 将做一些工作并将其保存到临时表中。 使用上述技术来读取表格。

下面的代码删除表（如果存在）然后重新加载。

-- 
-- Sample Call
-- 
CREATE PROCEDURE dbo.StackOverFlowTest
AS
BEGIN
    DROP TABLE IF EXISTS stage.DimSalesTerritory;
    SELECT * INTO stage.DimSalesTerritory FROM dbo.DimSalesTerritory
END

这是获取低级别 JAVA 驱动程序管理器的代码。 它具有调用 SP 的属性。

#
#  Grab the low level driver manager, exec sp
#

driver_manager = spark._sc._gateway.jvm.java.sql.DriverManager
connection = driver_manager.getConnection(url, user_name, password)
connection.prepareCall("EXEC dbo.StackOverFlowTest").execute()
connection.close()

使用 spark.read() 从 SP 填充的新表中检索数据。

Answer 3

我希望这适合您的用例。 同样，这不会扩展，因为它在控制节点（执行程序）上运行。 如果您有一个 5 节点集群，这将仅在 1 个节点上运行。

话虽如此，我们可以让 spark 推断数据类型。 如果查看pyodbc文档，则必须安装本地 ODBC 驱动程序。 这对工作站上的 anaconda 有利，但对 Spark 集群不利。 相反，用户pymssql模块是自包含的本机代码。 使用集群库的 PyPi 部分进行安装。

https://learn.microsoft.com/en-us/sql/connect/python/pymssql/python-sql-driver-pymssql?view=sql-server-ver16

现在我们有了驱动程序，让我们编写一个模块，该模块将从 SELECT 语句或存储过程调用 EXEC 返回数据帧。

#
#  Create function to call tsql + return df
#

# use module
import pymssql  

# define function
def call_tsql(info):
  
  # make connection
  conn = pymssql.connect(server=info[0], user=info[1], password=info[2], database=info[3])  
  
  # open cursor
  cursor = conn.cursor()  
  cursor.execute(info[4])
  
  # grab column data (name, type, ...)
  desc = cursor.description
  
  # grab data as list of tuples
  dat1 = cursor.fetchall()
  
  # close cursor
  conn.commit()
  conn.close()
  
  # extract column names
  col1 = list(map(lambda x: x[0], desc))
  
  # let spark infer data types
  df1 = spark.createDataFrame(data=dat1, schema=col1)
  
  # return dataframe
  return df1
  

此代码将仅支持一个结果集。 如果需要 MARS，多个活动结果集，请修改。

有很多评论。 简而言之，连接信息的元组和 TSQL 作为参数传入，并返回 dataframe。

是的，cursor.description 中有数据类型，但它们是经过编码的。 我没有找到好的映射。 由于您不处理大数据，因此推断架构。 否则，传递架构的 DDL 语句而不是列标题。

#
# Make call using SELECT statement
#

# tuple of info (server, user, pwd, database, query)
info = ('svr4tips2030.database.windows.net', '<your user>', '<your pwd>', 'dbs4advwrks', 'select * from dbo.vDMPrep')

# get data frame
df2 = call_tsql(info)

上图显示了推断类型，下图显示了数据。

我创建了一个非常简单的存储过程，它只是从视图中选择。

CREATE PROCEDURE [dbo].[StackOverFlowTest]
AS
BEGIN
    SELECT * FROM [dbo].[vDMPrep]
END
GO

如果我们更改执行此存储过程的调用，我们会得到相同的答案。

#
# Make call using EXEC statement
#

# tuple of info (server, user, pwd, database, query)
info = ('svr4tips2030.database.windows.net', '<your user>', '<your pwd>', 'dbs4advwrks', 'exec [dbo].[StackOverFlowTest]')

# get data frame
df2 = call_tsql(info)

随附的是示例 SPARK DDL 语句。 请参阅作为参数传递给笔记本的 JSON 文档的 file_schema 部分。

#
# Table 1 - dim.account
#

# Set parameters for notebook
parms = {
"datalake_path": "/mnt/datalake/bronze/",
"datafile_path": "/dim/account/dim-account-20200905T101530.csv",
"debug_flag": "false",
"partition_count": "2",
"file_schema": "AccountKey INT, ParentAccountKey INT, AccountCodeAlternateKey INT, ParentAccountCodeAlternateKey INT, AccountDescription STRING, AccountType STRING, Operator STRING, CustomMembers STRING, ValueType STRING, CustomMemberOptions STRING"
}

# Run notebook with selections
ret = dbutils.notebook.run("./nb-full-load-delta-table", 60, parms)

# Show return value if any
print(ret)

回顾一下，如果您遇到性能问题，请创建一个单节点集群并在该节点上扩展计算。 没有因为有多个节点，因为它们不会被pymssql模块使用。