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 を作成

```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/7efc5cf8-f8e5-2e03-8ef1-ac0f92d2b5e3.png)

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

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

### 3. Databricks 上でノートブックと同じ階層のディレクトリに`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/2123802d-6a16-d6e4-9b66-7b04af2a11cb.png)

### 4. Databricks 上で`user.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/ccb31bb0-1939-d8b5-bd11-b2e8cb4d95c4.png)

## データを取得する手順

### 1. Databricks 上で 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/6ce615e5-5267-5fdb-5619-44c3dd2cda32.png)

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/e78944eb-f14a-cc73-acad-662809ae22aa.png)

### 3. 除外するバイト数を変数にセット

```python
magic_bytes = 1
schema_id_bytes = 4
message_indexes_bytes = 1
offset_bytes = str(magic_bytes + schema_id_bytes + message_indexes_bytes)
start_bytes = str(int(offset_bytes) + 1)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/7b2d6fbd-7479-c0d7-6a5f-c0321a3d13fb.png)


### 4. Condluent からデータを取得

```python
from pyspark.sql.functions import col, expr
from pyspark.sql.protobuf.functions import from_protobuf

# Confluent から読み込み
df = (
    spark.readStream.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("subscribe", topic_name)
    .option("startingOffsets", "earliest")
    .load()
)

# Key を Deserialize
key_des_df = df.withColumn(
    "key",
    col("Key").cast("string"),
)

# Value から Magic Byte と Schema ID を除去
no_magic_df = key_des_df.withColumn(
    "value",
    expr(f"substring(value, {start_bytes}, length(value) - {offset_bytes})"),
)

# Magic Byte と Schema ID を取り除いた Value を Deserialize
decoded_df = no_magic_df.withColumn(
    "decoded_value",
    from_protobuf(
        col("value"),
        messageName=message_name,
        descFilePath=desc_file_paht,
    ),
)

# Deserialize した Value のカラムを展開
tgt_cols = [
    "key",
    "decoded_value.*",
    "topic",
    "partition",
    "offset",
    "timestamp",
    "timestampType",
]
src_df = decoded_df.select(tgt_cols)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/ac137354-a6b0-c91a-c9c8-617f8c25abb3.png)


### 5. データを表示

```python
src_df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c23fe267-c628-cf66-fedc-872dbfa54c1a.png)

## 事後処理

### Databricks ノートブックの実行を停止

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/7544912e-1a82-a6ad-2dc2-c07d5723b85c.png)