在Databricks Python脚本中无法从Cosmos DB选取_ts等自动生成字段

问题描述 #

在Databricks编写Python脚本读取Cosmos DB数据时，执行自定义查询select * from c后，返回结果仅包含自定义字段，无法获取_ts、_etag、_rid等Cosmos DB自动生成的系统字段。相关代码如下：

read_config = {
    "spark.cosmos.accountEndpoint": url,
    "spark.cosmos.accountKey": key,
    "spark.cosmos.database": database,
    "spark.cosmos.container": container,
    "spark.cosmos.read.customQuery" : "select * from c "
}

rsltdf = spark.read.format("cosmos.oltp").options(**read_config).load()
display(rsltdf)

当前环境信息 #

• Databricks运行时版本：7.3 LTS（包含Apache Spark 3.0.1、Scala 2.12）
• 依赖库：com.azure.cosmos.spark:azure-cosmos-spark_3-3_2-12:4.18.1、azure-cosmos、pyDocumentDB

解决方法 #

方法1：显式指定系统字段 #

修改自定义查询语句，明确列出需要获取的系统字段，示例代码如下：

read_config = {
    "spark.cosmos.accountEndpoint": url,
    "spark.cosmos.accountKey": key,
    "spark.cosmos.database": database,
    "spark.cosmos.container": container,
    "spark.cosmos.read.customQuery" : "select c.*, c._ts, c._etag, c._rid from c "
}

rsltdf = spark.read.format("cosmos.oltp").options(**read_config).load()
display(rsltdf)

方法2：启用系统属性自动推断 #

在读取配置中添加spark.cosmos.read.inferSchema.systemProperties.enabled参数并设置为true，这样使用select *就能自动包含系统字段，示例代码如下：

read_config = {
    "spark.cosmos.accountEndpoint": url,
    "spark.cosmos.accountKey": key,
    "spark.cosmos.database": database,
    "spark.cosmos.container": container,
    "spark.cosmos.read.customQuery" : "select * from c ",
    "spark.cosmos.read.inferSchema.systemProperties.enabled": "true"
}

rsltdf = spark.read.format("cosmos.oltp").options(**read_config).load()
display(rsltdf)

说明：该配置参数适配你使用的azure-cosmos-spark 4.x版本，它控制连接器是否将Cosmos DB的系统属性纳入Schema推断范围，启用后无需修改查询语句即可获取系统字段。

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