Stream Your Cosmos DB Changes To Databricks with Spark 3 | by Christopher Tao | Towards Data Science

Cosmos DB is becoming more and more popular because of its low latency response, which makes it more suitable as a transactional database solution. However, the data in the Cosmos DB may also need to be utilised for data analysis and reporting. In that case, we need to use a better-suited tool. Spark is usually considered the best choice for Big Data Analysis.
Databricks as a modern cloud data platform is built on top of Spark and provides a lot of cool features to improve both the development experiences and the actual data processing performance. Therefore, it is used by many companies across a wide range of industries.
Microsoft used to have tremendous good resources and references for the Spark 2 connectors for Cosmos DB. However, when it comes to Databricks DBR (Databricks Runtime Version) 8.x which are based on Spark 3, we have to use the corresponding Spark 3 connector. Unfortunately, there are not too much documentation and examples to follow yet.
In this article, I’ll provide a simple example of this scenario and some frequent issues we have met and their solution.
Structured Streaming is a scalable and fault-tolerant stream processing engine built on the Spark SQL engine. You can express your streaming computation the same way you would express a batch computation on static data. The Spark SQL engine will take care of running it incrementally and continuously and updating the final result as streaming data continues to arrive.
The official document of Cosmos DB Change Feed suggests Azure Functions to receive the incremental data. With Databricks, we don’t have to add one more layer because Spark Streaming can be a perfect solution for this.
Now, follow me on these step by step tutorials. It won’t be difficult and you could establish your Spark Streaming to the Cosmos DB change feed in about 10 minutes.
We need to have both the URI Endpoint and the key from the Cosmos DB that we are going to read. It is recommended to use a Read-only key because we are not going to write anything back to the Cosmos DB. Granting minimum required permission is always a good manner.
The URL can be found on the Cosmos DB Overview page in the Azure portal.
On the same page, you also need to note what is the name of the Database.
Then, the Read-only key can be found in the “keys” tab.
It is also highly recommended to put these connection details into an Azure Key Vault so that we don’t put them in the Databricks Notebooks as plain code, which will become a severe security hole.
Before we can use the connector, we need to install the library onto the cluster. Go to the “Compute” tab in the Databricks workspace and choose the cluster you want to use. Then, navigate to the “Libraries” tab and click “Install New”.
After that, select Maven as the library source and input the coordinates:
com.azure.cosmos.spark:azure-cosmos-spark_3-1_2-12:4.0.0
Please be noticed that the version might be different depending on when you are reading my article :) It would be better to check the latest version from Azure Documentation.
Once we put the Cosmos DB Read-only key into the Azure Key Vault, we need to create a secret scope in the Databricks workspace. This is not the focus of this article so that I’ll skip it. If you don’t know how to create a Databricks secret scope, please follow the documentation as follows.
docs.microsoft.com
Before we can create the Python Dictionary as the change feed configuration object, we need to read the connection details from the Azure Key Vault.
cosmosDatabaseName = dbutils.secrets.get('YOUR-KEY-VAUL', 'DB_NAME')cosmosContainerName = dbutils.widgets.get("container_name")cosmosEndpoint = dbutils.secrets.get('YOUR-KEY-VAUL', 'URI_ENDPOINT')cosmosMasterKey = dbutils.secrets.get('YOUR-KEY-VAUL', 'READ-ONLY-KEY')
It is optional whether to use widgets for the container name. In my case, I want to parameterise the container name so that this notebook can be reused later on.
Now, we can construct the configuration object.
# Initiate Cosmos Connection Config ObjectchangeFeedCfg = {  "spark.cosmos.accountEndpoint": cosmosEndpoint,  "spark.cosmos.accountKey": cosmosMasterKey,  "spark.cosmos.database": cosmosDatabaseName,  "spark.cosmos.container": cosmosContainerName,  "spark.cosmos.read.partitioning.strategy": "Default",  "spark.cosmos.read.inferSchema.enabled" : "false",  "spark.cosmos.changeFeed.startFrom" : "Now",  "spark.cosmos.changeFeed.mode" : "Incremental"  "spark.cosmos.changeFeed.maxItemCountPerTriggerHint" : "50000"}
Some configuration objects need some explanation:
spark.cosmos.read.partitioning.strategy
This is to define the partitioning strategy used (Default, Custom, Restrictive or Aggressive). Usually, we can use the default settings so that the number of partitions will be dynamically calculated to keep it optimised.
spark.cosmos.read.inferSchema.enabled
I would suggest NOT to enable the inferring schema although it looks pretty smart. When we have a relatively large number of documents in the Cosmos DB, there will be a high chance that the inferred schema will be wrong. Also, if the schema of the Cosmos Container is not consistent, it is almost guaranteed to lead to schema error.
We can read the raw JSON string and use PySpark to process it later.
spark.cosmos.changeFeed.startFrom
This item tells the connector where to start. Usually, we can set it to “Beginning” so that the first batch of the stream will start from the beginning of the table.
However, if we’re reading from a huge container. It would more clever to read since “Now” and load the historical data in a separate job to increase the performance and robustness.
spark.cosmos.changeFeed.mode
This has to be set to Incremental because we are going to read from the change feed.
spark.cosmos.changeFeed.maxItemCountPerTriggerHint
This is optional. However, it will be useful if your updates on the Cosmos DB side is not consistent. For example, some batches are small while some others are large. This option will try to process a limited number of items per trigger so that it would not blow up your cluster straight away when there is an extraordinary large batch arrived.
To simplify this tutorial, we don’t have to have a complex write configuration object for the Spark Stream. The only thing we must have is the checkpoint location.
writeCfg = {  "checkpointLocation": "/FileStore/checkpoints/"}
Checkpoint is very important because if our Spark Stream is stopped either planned or unplanned, it will memorise the checkpoint so that we can read from there.
The Cosmos Spark Connector will create a sub-directory so that we don’t have to specify the name for the current particular container that we are reading from.
We need to create a table before we can write the incremental data that we received from the Cosmos DB Change Feed. If you’re not using the inferring schema feature like me, the table should be defined as follow.
CREATE TABLE IF NOT EXISTS <schema_name>.<table_name> (  _rawBody STRING,  id STRING,  _ts LONG,  _etag STRING)USING DELTALOCATION '<dbfs_location>';
If you are using the inferring schema, you will need to define the table with corresponding columns.
Now we can use the configuration objects to establish the stream. The following code will establish the stream and read it into a Spark Dataframe.
df = spark.readStream.format("cosmos.oltp.changeFeed").options(**changeFeedCfg).load()
You may want to do some transformation on your Dataframe. After that, we can write it to the table we have just created.
df.writeStream.format('delta').outputMode('append').options(**writeCfg).table('<schema_name>.<table_name>')
We have actually finished the notebook. If you run it, it will establish the Spark Stream and starts to receive any changes from the Cosmos DB and write it to a table.
However, it is highly recommended to put everything in a job so that we can use a job cluster rather than an interactive one that is shared with all other activities. Also, it will be more confident in terms the reliability if we run the streaming as a job.
Navigate to the “Jobs” in the Databricks workspace. Input the Task name and choose the notebook that we have just created.
Don’t forget to add the connector as the dependent library.
During our development, because the connector is relatively immature in terms of its documentation, we have found the following two problems. Hopefully, this could save you some time if you met them, too.
If you have a relatively large container in the cosmos DB, you might want to start the streaming at the current timestamp and load the historical data in a separate activity.
However, if the change feed has been ever set to Beginning, and the checkpoint was not cleared, the stream may not be able to start since “Now” because the checkpoint “tells” it: “Hey, you should start from where we up to previously!”
"spark.cosmos.changeFeed.startFrom" : "Beginning"# v.s."spark.cosmos.changeFeed.startFrom" : "Now"
The solution is to manually remove the checkpoint from the DBFS, or change the checkpoint to a new path. Then, the argument “Now” will work properly.
You may see this exception after the stream has been run for a certain amount of time.
azure_cosmos_spark.io.netty.handler.timeout.ReadTimeoutException
One of the reasons for this exception could be that the current running nodes of the cluster are too less so that the availability will be a bottleneck. Also, if there is suddenly a large batch coming in, a relatively small cluster will not be able to finish the processing on time.
Therefore, the solution has two steps.
Firstly, we can increase the “Min Workers” to at least 2. If you’re really dealing with a high velocity, give it more.
The second solution will be adding a retry policy for the job. So, when the job failed, it will re-run it after a while.
We can find the retry policy in the “Advanced options”
Then, it is your decision to set how many times it will retry and how long it will wait until the retry.
In this article, I have used a simple example to demonstrate how to use the Azure Cosmos DB Spark 3 OLTP connector to establish a Spark Structured Streaming to read from the Cosmos DB change feed. I believe this is the easiest way to read the incremental change from a Cosmos DB into a Databricks workspace for further data analytics.
medium.com
If you feel my articles are helpful, please consider joining Medium Membership to support me and thousands of other writers! (Click the link above)