A simple way to import TiSpark into Databricks to load TiDB data
Published: 2022-09-16 07:30:43

[**TiDB**](https://docs.pingcap.com/tidb/stable/overview) is an open-source NewSQL database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads. It is MySQL compatible and features horizontal scalability, strong consistency, and high availability. 

[**TiSpark**](https://github.com/pingcap/tispark) is a thin layer built for running Apache Spark on top of TiDB/TiKV to answer the complex OLAP queries. It takes advantage of both the Spark platform and the distributed TiKV cluster and seamlessly glues to TiDB, the distributed OLTP database, to provide a Hybrid Transactional/Analytical Processing (HTAP) solution to serve as a one-stop solution for both online transactions and analysis.

[**Databricks**](https://docs.databricks.com/getting-started/introduction/index.html) is a cloud-based collaborative data science, data engineering, and data analytics platform that combines the best of data warehouses and data lakes into a lakehouse architecture.

With the flexible and strong expanding ability of Databricks, you can install TiSpark in Databrick to gain special advantages in TiSpark, such as faster reading and writing, transactions, and so on.  This article will show how to use TiSpark in Databricks to handle TiDB data.

## Step 1: Deploy a TiDB in your own space

TiDB supports a tool named [**TiUP**](https://docs.pingcap.com/tidb/dev/tiup-documentation-guide) to [**quickly build the test cluster**](https://docs.pingcap.com/tidb/dev/quick-start-with-tidb) on a single machine.

> **Note:**
> Your machine must have a **Public IP** for Databricks to access. Besides, this article uses a single instance TiDB cluster, so you don't need to config hosts for TiDB, PD, and TiKV.

1. Install TiUP.```
curl --proto '=https' --tlsv1.2 -sSf https://tiup-mirrors.pingcap.com/install.sh | sh
```

2. Declare the global environment variable.```
source ${your_shell_profile}
```

3. Start the TiDB cluster.```
tiup playground --host 0.0.0.0
```

    After the TiDB cluster deploys completely, you will see this.

    ```
CLUSTER START SUCCESSFULLY, Enjoy it ^-^
To connect TiDB: mysql --comments --host 127.0.0.1 --port 4001 -u root -p (no password)
To connect TiDB: mysql --comments --host 127.0.0.1 --port 4000 -u root -p (no password)
To view the dashboard: http://127.0.0.1:2379/dashboard
PD client endpoints: [127.0.0.1:2379 127.0.0.1:2382 127.0.0.1:2384]
To view Prometheus: http://127.0.0.1:9090
To view Grafana: http://127.0.0.1:3000
```

4. Select **PD address** via MySQL client. Record the **PD address** under the instance column. This address, which is an intranet IP usually, is used to communicate with internal instances.

    ```
mysql -u root -P 4000 -h 127.0.0.1
mysql> select * from INFORMATION_SCHEMA.CLUSTER_INFO;
+---------+-----------------+-----------------+-------------+------------------------------------------+---------------------------+-----------------+-----------+
| TYPE    | INSTANCE        | STATUS_ADDRESS  | VERSION     | GIT_HASH                                 | START_TIME                | UPTIME          | SERVER_ID |
+---------+-----------------+-----------------+-------------+------------------------------------------+---------------------------+-----------------+-----------+
...
| pd      | 172.*.*.*:2379  | 172.*.*.*:2379  | 6.1.0       | d82f4fab6cf37cd1eca9c3574984e12a7ae27c42 | 2022-07-13T13:25:54+08:00 | 2h31m44.004814s |         0 |
...
+---------+-----------------+-----------------+-------------+------------------------------------------+---------------------------+-----------------+-----------+
```

5. Import sample data.

    ```
tiup bench tpcc --warehouses 1 prepare
```

6. Check the result of importing.

    ```
mysql -u root -P 4000 -h 127.0.0.1
mysql> use test;
mysql> show tables;
+----------------+
| Tables_in_test |
+----------------+
| customer       |
| district       |
| history        |
| item           |
| new_order      |
| order_line     |
| orders         |
| stock          |
| warehouse      |
+----------------+
9 rows in set (0.00 sec)
```

## Step 2: Install TiSpark in Databricks

> **Prerequisites:** 
> A Databricks account, used to login Databricks workspace. If you don't have one, you can click [**here**](https://databricks.com/try-databricks?_ga=2.149545268.366754959.1655811799-681974717.1650447133) to find out how to try Databricks to get a free trial.

Databricks supports a custom [**init script**](https://docs.databricks.com/clusters/init-scripts.html#example-cluster-scoped-init-scripts), which is a shell script and runs during startup of each cluster node before the Apache Spark driver or worker JVM starts. Here we use the init script to install TiSpark.

1. [**Log in**](https://accounts.cloud.databricks.com/login) to Databricks and open your workspace.

2. Create a new **Python NoteBook**.

3. Copy the following scripts into notebook.```
dbutils.fs.mkdirs("dbfs:/databricks/scripts/")
dbutils.fs.put(
"/databricks/scripts/tispark-install.sh",
"""
#!/bin/bash
wget --quiet -O /mnt/driver-daemon/jars/tispark-assembly-3.2_2.12-3.1.0-SNAPSHOT.jar https://github.com/pingcap/tispark/releases/download/v3.1.0/tispark-assembly-3.2_2.12-3.1.0.jar 
""", 
True)
```

4. Attach to a Spark cluster and **run cell**. Then the scripts used to install TiSpark will be stored in **[DBFS](https://docs.databricks.com/data/databricks-file-system.html)**.

5. Click **Compute** on the sidebar.

6. Choose a cluster that you want to run with TiSpark.

7. Click **Edit** to config cluster.

8. In **Configuration** panel, set **Databricks Runtime Version** to "10.4 LTS".

9. In **Advanced options**, add `dbfs:/databricks/scripts/tispark-install.sh` to Init Scripts.

    ![Advanced options](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/gwsf3ud07zdoaobh5d76.png)


## Step 3: Set TiSpark Configurations in Spark Config

Once setting **Init Scrpits**, you need to add some configurations for TiSpark in **Spark Config**.

1. Add the following configuration. `{pd address}` is we recorded in **Step 1**. For more information about TiSpark conf, you can check **[TiSpark Configurations list](https://github.com/pingcap/tispark/blob/master/docs/configuration.md)**.

    ```
spark.sql.extensions org.apache.spark.sql.TiExtensions
spark.tispark.pd.addresses {pd address}
spark.sql.catalog.tidb_catalog org.apache.spark.sql.catalyst.catalog.TiCatalog
spark.sql.catalog.tidb_catalog.pd.addresses {pd address}
```

2. (Optional) If the `{pd address}` of your TiDB cluster is different with the **Public IP** of the machine, you need to add a special conf to build a host mapping between `{pd address}` (which equivalents to **Intranet IP**) and **Public IP**.
    ```
spark.tispark.tikv.host_mapping {pd address}:{Public IP}
```

    ![Advanced options](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/7iej3v7ymx9ecpkx4h3f.png)

3. Click **Confirm and restart** to enable configuration.

## Step 4: Handle your TiDB data in Databricks with TiSpark

After the cluster with TiSpark has been started, you can create a new notebook, attach it to this cluster, and start operating TiDB data with TiSpark in Databricks directly.

1. Create a **Scala notebook** and attach the Spark cluster with TiSpark.

2. Use tidb_catalog to enable TiSpark in SparkSession.

    ```
spark.sql("use tidb_catalog")
```

3. Use `SELECT` SQL to read TiDB data.

    ```
spark.sql("select * from test.stock limit 10").show
```
    ![query result](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/kffcmvvgkqkya7tjszck.png)

4. Use Spark DataSource API to write TiDB data.

    a. Because TiSpark doesn't support DDL, you need to create a table in TiDB before writing in Databricks. Here use MySQL client to create a `best_sotck` table on your own space.
    ```
mysql -uroot -P4000 -h127.0.0.1
mysql> use test;
mysql> create table best_stock (s_i_id int(11), s_quantity int(11))
```

    b. Set TiDB options, such as address, password, port and so on.

    > **Descriptions of fields**
> **tidb.addr:** TiDB address. It is the same as the PD IP we recorded in this Step 1.
> **tidb.password:** TiDB password of user. 
> **tidb.port:** The port of TiDB. 
> **tidb.user:** The user used to connect with TiDB cluster. 

    ```
val **tidbOptions**: Map[String, String] = Map(
  "tidb.addr" -> "{tidb address}",
  "tidb.password" -> "",
  "tidb.port" -> "4000",
  "tidb.user" -> "root")
```

    c. Select data and write back to TiDB with the specified option.

    > **Descriptions of fields**
> **format:** Specify "TiDB" for TiSpark.
> **database:** The destination database of writing.
> **table: **The destination table of writing.
> **Mode:** The datasource writing mode.

    ```
val DF = spark.sql("select s_i_id, s_quantity from test.stock where s_quantity>99 ")
DF.write
.format("tidb")
.option("database", "test")
.option("table", "best_stock")
.options(tidbOptions)
.mode("append")
.save()
```

    ![insert](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/reqabnutft1oqvwg8ctx.png)


    d. Check the writing data by `SELECT` SQL.
spark.sql("select * from test.best_stock limit 10").show

    ![check result](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/vsismgjrfw3kasnrwjb3.png)


5. Use `DELETE` SQL to delete TiDB data and check with `SELECT` SQL.
    ```
spark.sql("delete from test.best_stock where s_quantity > 99")
spark.sql("select * from test.best_stock").show
```

    ![check result](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/y49kdp1rf5w8vsw5ingg.png)


## Conclusion
In this article, we use TiSpark in Databricks to access TiDB data. The key steps are:

1. Install TiSpark in Databricks via init scripts.

2. Set TiSpark Configurations in Databricks.

If you are interested in TiSpark, you can find more information on our [**TiSpark homepage**](https://docs.pingcap.com/tidb/stable/tispark-overview). Welcome to share any ideas or PR on the [**TiSpark GitHub repository**](https://github.com/pingcap/tispark).