Databricks から Confluent に JSON 形式の Topic データを書き込む際に手動でスキーマを指定して実施する方法

## 概要

Databricks から Confluent に JSON 形式の Topic データを書き込む際に手動でスキーマを指定して実施する方法を紹介します。

本記事は、以下のシリーズの一部です。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/27593d44-23cc-bd0a-98d1-bd538793cad4.png)

引用元：[Databricks と Confluent をつなぐ：データ連携の全体像を徹底解説 #Spark - Qiita](https://qiita.com/manabian/items/b79d2ad8a624065c9f57)

## 事前準備

### 1. Confluent にて Topic を作成

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b11eea6c-f379-dec9-bd87-d853886af430.png)


### 2. Confluent にて Topic にスキーマを設定

```text
{
  "$schema": "http://json-schema.org/draft-07/schema#",
  "properties": {
    "gender": {
      "connect.index": 3,
      "type": "string"
    },
    "regionid": {
      "connect.index": 2,
      "type": "string"
    },
    "registertime": {
      "connect.index": 0,
      "connect.type": "int64",
      "type": "integer"
    },
    "userid": {
      "connect.index": 1,
      "type": "string"
    }
  },
  "title": "ksql.users",
  "type": "object"
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/8ec91961-9dfe-c0c7-f56d-8f05f05fcabe.png)

### 3. Confluent にて作成したスキーマ（Subject）の Schema ID の値を取得

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/14968943-6636-4036-d9db-7800535b6a73.png)

### 4. Confluent にて API キーを主塔

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/134d657f-daeb-f6c6-67aa-99c6da0a2524.png)

### 5. Confluent にて Schema Registry API キーを取得

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/933e0ddd-5c41-7a83-031f-073f70737ab9.png)

### 6. Databricks にてデータベースのオブジェクトを作成

```sql
%sql
CREATE CATALOG IF NOT EXISTS confluent_test;
CREATE SCHEMA IF NOT EXISTS confluent_test.test;
```

```pyhon
src_table_name = "confluent_test.test.src_table"
checkpoint_volume_name = "confluent_test.test.checkpoint_01"
checkpoint_volume_dir = f"/Volumes/confluent_test/test/checkpoint_01"
```

```python
spark.sql(f"CREATE TABLE IF NOT EXISTS {src_table_name}")
spark.sql(f"CREATE VOLUME IF NOT EXISTS {checkpoint_volume_name}")
dbutils.fs.rm(checkpoint_volume_dir, True)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/070d5c9a-d30d-79d6-ac62-7f4269e2d7a8.png)

### 6. テーブルにデータを書きこみ

```python
data = [
    ("User_1", 1507214855421, "Region_1", "OTHER"),
    ("User_2", 1516353515688, "Region_2", "OTHER"),
    ("User_3", 1511384261793, "Region_3", "FEMALE")
]

columns = ["userid", "registertime", "regionid", "gender"]

df = spark.createDataFrame(data, columns)

df.write.format("delta").mode("overwrite").saveAsTable(src_table_name)

spark.table(src_table_name).display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/d1875a0f-22ed-7b9a-795a-c256efddd98c.png)

## データを書き込む手順

### 1. Databricks 上で API キーと Schema Registry API キーを変数にセット

```python
# Confluent に対する認証情報をセット
bootstrap_servers = "pkc-921aa.us-east-2.aws.confluent.cloud:9092"
sasl_username = "U4BML4OPRCaaaa"
sasl_password = "zzEjUzyUwYYJneNE240L7jgPVrHAtRBANo6wnnzCgKlGIJ9OxJqaJZWRvaaaaa"
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/6558440a-030e-95ce-fd5b-d3ec6424a31b.png)

### 2. Databricks にて Schema ID を変数にセット

```python
schema_id_num = 100004
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/94fdbc28-d12d-f397-77e0-9a68e6243b4f.png)


### 3.  Databricks を変数にセット

```python
topic_name = "json_without_sr_01"
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/ae39bdf4-4317-2210-05d4-979ee560cbdd.png)

### 4. Databricks にてWire Header を変数にセット

```python
# Confluent Wire Header の定義
magic_byte_num = 0
wire_header = bytes([magic_byte_num])
wire_header += schema_id_num.to_bytes(4, byteorder='big', signed=False)
print(wire_header)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/0cb39127-6c48-1bc5-c5a3-7f761ecc0b98.png)

### 5. Databricks にて Conflunet へ書き込みを実施

```python
# 書き込み時に利用するチェックポイントのディレクトリを指定
checkpoint_dir = checkpoint_volume_dir + "/" + topic_name

from pyspark.sql.functions import col, struct, lit ,concat, to_json

src_df = spark.readStream.table(src_table_name)

# key 列を定義
src_df = src_df.withColumn("key", col("userid"))

# magic-byte, schema-id, Avro のデータをもつ value 列を定義
src_df = src_df.withColumn("value", struct("registertime", "userid", "regionid", "gender"))
src_df = src_df.withColumn("value", to_json(col("value")))
src_df = src_df.withColumn("value", col("value").cast("binary"))
src_df = src_df.withColumn("add_bytes", lit(wire_header))
src_df = src_df.withColumn("value", concat(col("add_bytes"), col("value")))

# key 列と value 列のみを抽出
src_df = src_df.select("key", "value")

(
    src_df.writeStream.format("kafka")
    .option("kafka.bootstrap.servers", bootstrap_servers)
    .option("kafka.security.protocol", "SASL_SSL")
    .option(
        "kafka.sasl.jaas.config",
        f'kafkashaded.org.apache.kafka.common.security.plain.PlainLoginModule required username="{sasl_username}" password="{sasl_password}";',
    )
    .option("kafka.sasl.mechanism", "PLAIN")
    .option("topic", topic_name)
    .option("checkpointLocation", checkpoint_dir)
    .trigger(availableNow=True)
    .start()
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/60cab752-9655-e246-09cd-8fc3b335343c.png)

### 6. Confluent 上でデータが書き込まれたことを確認

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f70fe809-4f4f-da95-bd9b-14c4cfee6343.png)