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

## 概要

Databricks から Confluent に Protobuf 形式の 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/9d74e9a6-4edb-30d5-eaeb-f5e9e8729226.png)


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

```text
syntax = "proto3";
package ksql;

message users {
  int64 registertime = 1;
  string userid = 2;
  string regionid = 3;
  string gender = 4;
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/baa5c37f-def4-8cca-414c-250298534e2b.png)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/5828f4de-e72f-22c3-52a0-c1b533841b67.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)

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

```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)

### 8. ノートブックを同じディレクトリに`user.proto`ファイルを作成

```proto
syntax = "proto3";
package ksql;

message users {
  int64 registertime = 1;
  string userid = 2;
  string regionid = 3;
  string gender = 4;
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a55dcbc5-0294-aade-a631-fda278b27082.png)

### 9. Desc ファイルである`proto.desc`を作成

```sh
%sh sudo apt-get update
```

```sh
%sh sudo apt-get install -y protobuf-compiler
```

```sh
%sh protoc --include_imports \
--descriptor_set_out=user.desc \
user.proto
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/dde41729-e63a-10af-8202-aad4a3db40f0.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 = 100003
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/74694270-1691-40fe-afe0-63024fa19193.png)

### 3.  Databricks に Topic と Desc ファイルに関する情報を変数にセット

```python
topic_name = "protobuf_without_sr_01"
desc_file_paht = "./user.desc"
message_name = "ksql.users"
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/ae67733a-b8f4-47f2-5a87-3b3671d484cd.png)


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

```python
# 下記ドキュメントを参考に message-indexes を指定
# https://docs.confluent.io/cloud/current/sr/fundamentals/serdes-develop/index.html#wire-format
message_indexes = b"\x00"

# 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)
wire_header += message_indexes
print(wire_header)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/dedd4825-4291-f105-5032-48e0dece9500.png)

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

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

from pyspark.sql.functions import col, struct, lit
from pyspark.sql.protobuf.functions import to_protobuf


src_df = spark.readStream.table(src_table_name)

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

# magic-byte, schema-id, message-indexes, Avro のデータをもつ value 列を定義
src_df = src_df.withColumn(
    "value", struct("registertime", "userid", "regionid", "gender")
)
src_df = src_df.withColumn(
    "value",
    to_protobuf(
        col("value"),
        messageName=message_name,
        descFilePath=desc_file_paht,
    ),
)

# value 列に Confluent Protobuf Wire Header を追加
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/04a2c8ef-bd83-1a02-2d37-b3a6362f0bf1.png)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f85c15b8-c74b-d768-1d5f-e0eac1f90598.png)