Revealed: A ridiculously easy way to integrate Azure Cosmos DB with Azure Databricks | by Prashanth Xavier | Towards Data Science

Buddy our novice Data Engineer who recently discovered the ultimate cheat-sheet to read and write files in Databricks is now leveling up in the Azure world.
In this article, you will discover how to seamlessly integrate Azure Cosmos DB with Azure Databricks. Azure Cosmos DB is a key service in the Azure cloud platform that provides a NoSQL-like database for modern applications.
As a Data Engineer or a Data Scientist, you may want to use Azure Cosmos DB for serving your data that is modeled and prepared using Azure Databricks or you may want to analyze the data that already exists in Azure Cosmos DB using Databricks. Whatever your purpose simply follow this 3 step guide to get started.
For the uninitiated, Azure Cosmos DB worthy of the name is Microsoft’s multi-model database that can manage data at a planet-scale. It belongs to the “NoSQL Database as a Service” stack like its counterpart AWS DynamoDB.
Inside Cosmos DB, each piece of data called an item is stored inside schema-agnostic containers, which means that you don’t need to adhere to any particular schema for your data.
Cosmos DB supports multi-model APIs like MongoDB, Cassandra API, Gremlin API, and the default Core SQL API.
The Core SQL API provides you with JSON like NoSQL document store, which you can easily query using an SQL-like language.
Despite its fancy name and overwhelming features, Cosmos DB is basically a data store, a data store that we can read from and write to.
Through its seamless integration with a plethora of Azure services, Azure Databricks is just the right tool for the job.
In order to execute this exercise you must have an Azure subscription with Cosmos DB and Databricks services running. If you don’t have one, follow the steps below to get it and create the services for Free!
If you have an existing Azure subscription skip to the next section.
If you do not have an Azure subscription get a free trial here, it’s quite easy and takes less than 2 minutes. (you will need to give your credit card information, but don’t worry you will not be charged for anything)
Now, all we need is a Cosmos DB account and a Databricks workspace.
Microsoft makes it easier and easier to deploy services on Azure using quick starter templates.
Follow the link to the quick starter template to deploy Azure Cosmos DB, click on Deploy to Azure, this opens up the Azure portal on the browser. Review the steps and create your service. The Cosmos DB account will be ready before your next cup of coffee
Once the account is created you will need to create a database and a container in which your data will be stored. Follow the example below to create a Database called AdventureWorks and a Container named ratings.
Navigate to your deployed Cosmos DB account and click on Data Explorer →New Container → name your database AdventureWorks →your container ratings → Partition key as /rating → select Throughput manual and set it to 1000.
The final step is to copy Azure-Cosmos DB Read-Write Key from the Keys section as shown below. You will need this to authenticate your Databricks workspace to access the Cosmos DB account.
Similarly, follow the link below to deploy Azure Databricks workspace. Once the workspace is deployed all you need to do is create a cluster.
Navigate to your Databricks workspace, on the left-hand side select Clusters, and create a new cluster(Choose version Scala 2.11 this is the version that supports the cosmosdb library we will install next). Select Cluster-Mode as Standard (try Single Node if there are any hiccups with a standard cluster).
Wow, that was faster than a speeding ticket!
Install the latest azure-cosmosdb-spark library on your Databricks cluster.Write data to Cosmos DB using the Spark connector.Read data from Cosmos DB using the Spark connector.
Install the latest azure-cosmosdb-spark library on your Databricks cluster.
Write data to Cosmos DB using the Spark connector.
Read data from Cosmos DB using the Spark connector.
Step 1: Initial setup and installation.
In the left-hand navigation go to your Clusters and choose the cluster that you created. If it’s not running, start it now.
Select the Libraries tab and select Install New. Under the installation dialog box, select Maven and plug in the following:
Coordinates:com.databricks.training:databricks-cosmosdb-spark2.2.0-scala2.11:1.0.0
Repository:https://files.training.databricks.com/repo
Note: Currently Databricks Runtime Version with Scala 2.11 is supported, so ensure that your cluster is running with Scala 2.11 version.
Step 2:Write data to Azure Cosmos DB
In order to write data to Azure Cosmos DB, we need to prepare a configuration that tells Spark how to connect and authenticate to the Cosmos container.
URI = "cosmos-db-uri"PrimaryKey = "your-cosmos-db-key"CosmosDatabase = "AdventureWorks"CosmosCollection = "ratings"writeConfig = {  "Endpoint": URI,  "Masterkey": PrimaryKey,  "Database": CosmosDatabase,  "Collection": CosmosCollection,  "writingBatchSize":"1000",  "Upsert": "true"}
In a new Notebook, create the above configuration variables.URI and PrimaryKey are the values we copied from the Keys section of your Azure Cosmos DB account earlier.
Now we just need some data. Upload the ratings sample csv file(download from GitHub here)into your Databricks workspace.
ratingsDF = spark.read.format("csv")\            .option("header","true")\            .load("/FileStore/tables/ratings.csv")
The above command reads a sample ratings csv file with 100 rows into a dataframe.
(ratingsDF.write  .mode("overwrite")  .format("com.microsoft.azure.cosmosdb.spark")  .options(**writeConfig)  .save())
Finally, the above command simply pushes the data into our Cosmos DB container using the Spark DataFrame writer.
Verify that the data is actually in your Cosmos DB container by going into the Cosmos DB account →Data Explorer →AdventureWorks →ratings.
That was a breeze!
Step 3: Query data from Azure Cosmos DB from Databricks
In this step, we will read data that we just pushed into the ratings container. We will first need to create a new configuration setting called readConfig, which will contain a custom query to read 10 rows of data.
URI = "cosmos-db-uri"PrimaryKey = "your-cosmos-db-key"CosmosDatabase = "AdventureWorks"CosmosCollection = "ratings"query = "Select top 10 * from ratings"readConfig = {  "Endpoint": URI,  "Masterkey": PrimaryKey,  "Database": CosmosDatabase,  "Collection": CosmosCollection,  "query_custom": query}
The custom query is a simple SQL query that will be executed on the Cosmos DB container. You can also write complex queries involving where conditions. For this exercise, a simple Select would suffice.
And for the final revelation, once the readConfig is set we simply read the data from Cosmos DB into a dataframe using the Spark DataFrame writer.
readDF = (spark.read  .format("com.microsoft.azure.cosmosdb.spark")  .options(**readConfig)  .load())
Like reading and writing data, you can very much follow a similar pattern to update records in your Cosmos DB container. Simply read the data you want to update into a dataframe and, manipulate this dataframe with your updates. Once done simply push the updated dataframe using the writeConfig. You can find an example in this GitHub link to the full source code.
Azure Cosmos DB and Azure Databricks are flagship services from Microsoft and coupled together they become much more powerful and knowing how to do that will make you resourceful.
Key takeaways from this article:
What is Azure Cosmos DB and how to Create a Cosmos DB account?
How to Create an Azure Databricks workspace?
How to Write data to Azure Cosmos DB?
How to Read data from Azure Cosmos DB?
All of that without spending a dollar, completely Free!
What is your take on Azure Cosmos DB? Share your experience and if the article was insightful share it with your peers.
Link to Github source code:
For a handy Spark cheat sheet check out:
towardsdatascience.com
For a basic Introduction to SQL check out: