Databricks にて Confluent から Avro 形式の Topic のデータ取得をスキーマレジストリを利用して実施する方法

## 概要

Databricks にて Confluent から Avro 形式の 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/234f99c1-f7cb-713c-f48a-dc1d40e7ca40.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. Confluent にて Schema Registry API キーを取得

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/933e0ddd-5c41-7a83-031f-073f70737ab9.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"


# Schema Registry の接続情報
sr_url = "https://psrc-lAAAA.us-east-2.aws.confluent.cloud"
sr_api_key = "CMFFLEP4BCTMAAAA"
sr_api_secret = "9j/Bf/if0l5v7HocEJKd2xOt/z5Gr0I5M7lRpPQNslV0QOzMkLEQzgRgVBsAAAA"
schema_registry_options = {
    "confluent.schema.registry.basic.auth.credentials.source": "USER_INFO",
    "confluent.schema.registry.basic.auth.user.info": f"{sr_api_key}:{sr_api_secret}",
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/0e2e3895-cddf-f655-cf1d-7b92c98a6f25.png)

### 2. Topic 名とサブジェクト名を変数にセット

```python
topic_name = "avro_user_1"
subject_name = "avro_user_1-value"
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/90af8630-2eb8-d7b1-c875-7ab2bb7b6491.png)

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

```python
from pyspark.sql.functions import col, expr
from pyspark.sql.avro.functions import from_avro

# 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"),
)

# Magic Byte と Schema ID を取り除いた Value を Deserialize
decoded_df = key_des_df.withColumn(
    "value",
    from_avro(
        col("value"),
        options=schema_registry_options,
        subject=subject_name,
        schemaRegistryAddress=sr_url,
    ),
)

# Deserialize した Value のカラムを展開
tgt_cols = [
    "key",
    "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/7e89adf4-2f12-bccd-a4e1-51f56cd7800b.png)

### 4. データを表示

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9b4efad5-406f-fbc6-53b1-3344ae741664.png)

## 事後処理

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

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