Tips and Tricks for using Python with Databricks Connect
Tags: apachespark, databricks, pyspark, jupyter
Published: 2020-10-06 15:59:40

---
title: Tips and Tricks for using Python with Databricks Connect
published: true
description: 
tags: 
  - ApacheSpark
  - Databricks
  - PySpark
  - Jupyter
//cover_image: https://direct_url_to_image.jpg
---

[Databricks Connect](https://docs.databricks.com/dev-tools/databricks-connect.html) is awesome. It makes development for Apache Spark a joy. Without it I'd probably never be any effective and abandon Spark a long time ago.

# Why

Databricks Connect is a magical local instance of Spark. Your machine will think it's using local installation of Spark, however in reality it will use a remote Databricks instance. Why would you want it? Because often you need **access to large datasets** that don't fit on your machine, and often you need **bigger compute clusters** to work with data.

Traditionally, it is possible to achieve it by having some small subsets of data locally and crunching through them, then deploying remotely to execute on a different dataset, but that's not always possible, depending on the data you work with.

# Installing

I assume you have already installed [Miniconda](https://docs.conda.io/en/latest/miniconda.html) (or Anaconda if you don't care about your PC) and you can follow basic steps in the official [Databricks Connect docs](https://docs.databricks.com/dev-tools/databricks-connect.html).

I also assume you've installed Jupyter. If not, install from conda - `conda install -c conda-forge notebook`. I recommend installing this in your base environment.

Download and install [JDK 8](https://www.oracle.com/uk/java/technologies/javase/javase-jdk8-downloads.html) if you haven't already. Don't use higher versions as Spark is running under Java 8. Also make sure that you are installing x64 version of the SDK. Add `bin` to `PATH`, and create `JAVA_HOME` environment variable. On windows this usually looks like:

```
PATH=C:\Program Files (x86)\Java\jdk1.8.0_261\bin
JAVA_HOME=C:\Program Files (x86)\Java\jdk1.8.0_261
```

Create Conda environment with python version `3.7` and not `3.5` like in the original article (it's probably outdated):

```bash
conda create --name dbconnect python=3.7
```

activate the environment

```bash
conda activate dbconnect
```

and install tools v6.6:

```bash
pip install -U databricks-connect==6.6.*
```

Your cluster needs to have two variable configured in order for Databricks Connect to work:

- `spark.databricks.service.server.enabled` needs to be set to `true`
- `spark.databricks.service.port` needs to be set to a port (you need this later).



![](http://www.gavryli.uk/posts/2020/05/python-databricks-connect-tricks/clusterconfig.png)

You should be able to use Databricks Connect now from IDE you want.

# Jupyter Weirdness

I like to have Jupyter installed in my base conda environment once, and not duplicated across all the environments I'm creating. If this is the case, when you run Jupyter server, it won't display the newly created environment for Databricks Connect, as it doesn't pick up new environments automatically.

To solve this, install `ipykernel` (Jupyter kernel integration) into Databricks Connect environment:

```bash
conda install ipykernel
```

Instruct Jupyter that current environment needs to be added as a kernel:

```bash
python -m ipykernel install --user --name dbconnect --display-name "Databricks Connect (dbconnect)"
```

Go back to the base environment where you have installed Jupyter and start again:

```bash
conda activate base
jupyter kernel
```

The kernel will be displayed in the list.

### Jupyter Hints

You can enable automatic `DataFrame` display functionality by setting `spark.conf.set("spark.sql.repl.eagerEval.enabled", True)`.

There are [extensions available](https://github.com/ipython-contrib/jupyter_contrib_nbextensions) you can install to extend Jupyter functionality.

To customise Jupyter look&feel, use [this repository](https://github.com/dunovank/jupyter-themes).

# Using Multiple Databricks Clusters

Once databricks-connect is configured from the command line by specifying all the parameters, including cluster id, you are tied to that cluster, unless it's reconfigured again. To use a different cluster, a new conda environment can be created, and configured again. However, if all of your clusters are in the same databricks workspace, you can use the following trick to switch between clusters.

First, create a cluster map somewhere in code:

```python
clusters = {
    "dev": {
        "id": "cluster id",
        "port": "port"
    },
    "prod": {
        "id": "cluster id",
        "port": "port"
    }
}
```

Write a function that you can call from your notebook:

```python
def use_cluster(cluster_name: str):
    """
    When running via Databricks Connect, specify to which cluster to connect instead of the default cluster.
    This call is ignored when running in Databricks environment.
    :param cluster_name: Name of the cluster as defined in the beginning of this file.
    """
    real_cluster_name = spark.conf.get("spark.databricks.clusterUsageTags.clusterName", None)

    # do not configure if we are already running in Databricks
    if not real_cluster_name:
        cluster_config = clusters.get(cluster_name)
        log.info(f"attaching to cluster '{cluster_name}' (id: {cluster_config['id']}, port: {cluster_config['port']})")

        spark.conf.set("spark.driver.host", "127.0.0.1")
        spark.conf.set("spark.databricks.service.clusterId", cluster_config["id"])
        spark.conf.set("spark.databricks.service.port", cluster_config["port"])

```

Pass cluster name from the map to `use_cluster` - this will select an appropriate cluster before executing the code. The good thing about it is you can leave the call in Databricks notebook, as it will be ignored when running in their environment.



# Local vs Remote

## Checking if notebook is running locally or in Databricks

The trick here is to check if one of the databricks-specific functions (like displayHTML) is in the IPython user namespace:

```python
def _check_is_databricks() -> bool:
    user_ns = ip.get_ipython().user_ns
    return "displayHTML" in user_ns
```

## Getting Spark Session

Databricks notebooks initialise `spark` variable automatically, therefore you can decide whether to return it or create a new local session:

```python
def _get_spark() -> SparkSession:
    user_ns = ip.get_ipython().user_ns
    if "spark" in user_ns:
        return user_ns["spark"]
    else:
        spark = SparkSession.builder.getOrCreate()
        user_ns["spark"] = spark
        return spark
```

## Faking Display Function

Databricks has a nice `display()` function that renders dataframes. We don't have that locally, but we can fake it:

```python
def _get_display() -> Callable[[DataFrame], None]:
    fn = ip.get_ipython().user_ns.get("display")
    return fn or _display_with_json
```

## DBUtils

Depending where you run, you can create a local instance of `dbutils` or get pre-initialised one when running in Databricks:

```python
def _get_dbutils(spark: SparkSession):
    try:
        from pyspark.dbutils import DBUtils
        dbutils = DBUtils(spark)
    except ImportError:
        import IPython
        dbutils = IPython.get_ipython().user_ns.get("dbutils")
        if not dbutils:
            log.warning("could not initialise dbutils!")
    return dbutils
```

## Putting it all Together

You can put all of the above in a single python file, reference it from every notebook, and they will work regardless where you run:

`dbconnect.py`:

```python
from typing import Any, Tuple, Callable

from pyspark.sql import SparkSession, DataFrame
import logging
import IPython as ip
from pyspark.sql.types import StructType, ArrayType
import pyspark.sql.functions as f

clusters = {
    "dev": {
        "id": "cluster id",
        "port": "port"
    },
    "prod": {
        "id": "cluster id",
        "port": "port"
    }
}

# Logging

class SilenceFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> int:
        return False


logging.basicConfig(format="%(asctime)s|%(levelname)s|%(name)s|%(message)s", level=logging.INFO)
logging.getLogger("py4j.java_gateway").addFilter(SilenceFilter())
log = logging.getLogger("dbconnect")

def _check_is_databricks() -> bool:
    user_ns = ip.get_ipython().user_ns
    return "displayHTML" in user_ns


def _get_spark() -> SparkSession:
    user_ns = ip.get_ipython().user_ns
    if "spark" in user_ns:
        return user_ns["spark"]
    else:
        spark = SparkSession.builder.getOrCreate()
        user_ns["spark"] = spark
        return spark


def _display(df: DataFrame) -> None:
    df.show(truncate=False)


def _display_with_json(df: DataFrame) -> None:
    for column in df.schema:
        t = type(column.dataType)
        if t == StructType or t == ArrayType:
            df = df.withColumn(column.name, f.to_json(column.name))
    df.show(truncate=False)


def _get_display() -> Callable[[DataFrame], None]:
    fn = ip.get_ipython().user_ns.get("display")
    return fn or _display_with_json


def _get_dbutils(spark: SparkSession):
    try:
        from pyspark.dbutils import DBUtils
        dbutils = DBUtils(spark)
    except ImportError:
        import IPython
        dbutils = IPython.get_ipython().user_ns.get("dbutils")
        if not dbutils:
            log.warning("could not initialise dbutils!")
    return dbutils


# initialise Spark variables
is_databricks: bool = _check_is_databricks()
spark: SparkSession = _get_spark()
display = _get_display()
dbutils = _get_dbutils(spark)


def use_cluster(cluster_name: str):
    """
    When running via Databricks Connect, specify to which cluster to connect instead of the default cluster.
    This call is ignored when running in Databricks environment.
    :param cluster_name: Name of the cluster as defined in the beginning of this file.
    """
    real_cluster_name = spark.conf.get("spark.databricks.clusterUsageTags.clusterName", None)

    # do not configure if we are already running in Databricks
    if not real_cluster_name:
        cluster_config = clusters.get(cluster_name)
        log.info(f"attaching to cluster '{cluster_name}' (id: {cluster_config['id']}, port: {cluster_config['port']})")

        spark.conf.set("spark.driver.host", "127.0.0.1")
        spark.conf.set("spark.databricks.service.clusterId", cluster_config["id"])
        spark.conf.set("spark.databricks.service.port", cluster_config["port"])

```

In your notebook:

```python
from dbconnect import spark, dbutils, use_cluster, display

use_cluster("dev")

# ...

df = spark.table("....")	# use spark variable

display(df)	# display DataFrames

# etc...
```



# Fixing Out-of-Memory Issues

Often when using Databricks Connect you might encounter an error like `Java Heap Space etc. etc. etc.`. This simply means your local spark node (driver) is running out of memory, which by default is 2Gb. If you need more memory, it's easy to increase it.

First, find out where PySpark's home directory is:

```bash
❯ databricks-connect get-spark-home
c:\users\ivang\miniconda3\envs\hospark\lib\site-packages\pyspark
```

This should have a subfolder `conf` (create it if it doesn't exist). And a file `spark-defaults.conf` (again, create if doesn't exist). Full file path would be `c:\users\ivang\miniconda3\envs\hospark\lib\site-packages\pyspark\conf\spark-defaults.conf`. Add a line

```
spark.driver.memory 8g
```

Which increases driver memory to 8 gigabytes. 

![](http://www.gavryli.uk/posts/2020/05/python-databricks-connect-tricks/sparkconf.png)



# Monitoring Jobs

Unfortunately, I couldn't find a good way to monitor jobs from DBC environment. There's [Big Data Tools](https://www.jetbrains.com/help/idea/big-data-tools-support.html) plugin for IntelliJ, that in theory supports Spark job monitoring, and considering DBC runs a virtual local cluster, I though it would work. However, no luck in any configuration.

The best way to monitor jobs I found is to use Spark UI from the cluster on Databricks you're connected to.

This article was originally published [on my blog](https://www.aloneguid.uk/posts/2020/05/python-databricks-connect-tricks/).