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

## 概要

Databricks から Confluent に Avro 形式の Topic データを書き込む際にスキーマレジストリを利用して実施する方法を紹介します。

スキーマを指定せずに書き込みを実施できると想定していたのですが、 2024年 12 月 26 日時点ではスキーマを指定しないとエラーとなってしまいました。本記事では、スキーマを指定する手順となっておりますが、`jsonFormatSchema`引数の指定が不要になる可能性もあります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/1953fd3b-590f-7c2a-3f1d-664bae0172e8.png)

```text
org.apache.spark.SparkException: Job aborted due to stage failure: Task 2 in stage 436.0 failed 4 times, most recent failure: Lost task 2.3 in stage 436.0 (TID 568) (10.139.64.10 executor 0): org.spark_project.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException: Schema not found; error code: 40403
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9b27b70d-0baf-af87-b3cd-983ac8cd2a2d.png)

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

![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/e06c4dd6-898e-9c47-355c-c066f4d4b240.png)


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

```text
{
  "connect.name": "ksql.users",
  "fields": [
    {
      "name": "registertime",
      "type": "long"
    },
    {
      "name": "userid",
      "type": "string"
    },
    {
      "name": "regionid",
      "type": "string"
    },
    {
      "name": "gender",
      "type": "string"
    }
  ],
  "name": "users",
  "namespace": "ksql",
  "type": "record"
}
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/5155e32e-9aca-8794-0189-f0103e9c0d96.png)

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

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

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

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

### 5. 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"

# 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/9c32745b-272b-41cc-9eb7-5e0b254c3d87.png)

### 2.  Databricks に Topic と Subject を変数にセット

```python
topic_name = "avro_with_sr_01"
value_subject_name = topic_name + "-value"
```

### 3. スキーマを変数にセット

```python
json_format_schema = """
{
  "connect.name": "ksql.users",
  "fields": [
    {
      "name": "registertime",
      "type": "long"
    },
    {
      "name": "userid",
      "type": "string"
    },
    {
      "name": "regionid",
      "type": "string"
    },
    {
      "name": "gender",
      "type": "string"
    }
  ],
  "name": "users",
  "namespace": "ksql",
  "type": "record"
}
"""
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/887807d6-8b7f-3bce-4366-0b17d55f56df.png)

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/64690d04-ef5b-68b5-7a8b-aaa12891cfc0.png)

```python
from pyspark.sql.functions import col, struct, lit
from pyspark.sql.avro.functions import to_avro

src_df = spark.readStream.table(src_table_name)

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

src_df = src_df.withColumn(
    "value",
    struct("registertime", "userid", "regionid", "gender"),
)
src_df = src_df.withColumn(
    "value",
    to_avro(
        data=col("value"),
        options=schema_registry_options,
        schemaRegistryAddress=sr_url,
        subject=lit(value_subject_name),
        jsonFormatSchema=json_format_schema,
    ),
)
# key 列と value 列のみを抽出
src_df = src_df.select("key", "value")
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/2d0539d0-45ca-5614-25fe-0fd2c0075a3c.png)


```python
(
    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/6ca7b3f8-ffa4-44f7-320b-f94ecbd48b04.png)


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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/727742fd-5a88-834c-86f1-2bebec472821.png)

## 発生したエラー

### 1. `to_avro`関数に対して`jsonFormatSchema`引数を指定せずに`subject`引数を指定するとエラーが発生しました。

```python
from pyspark.sql.functions import col, struct, lit
from pyspark.sql.avro.functions import to_avro

src_df = spark.readStream.table(src_table_name)

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

src_df = src_df.withColumn(
    "value",
    struct("registertime", "userid", "regionid", "gender"),
)
src_df = src_df.withColumn(
    "value",
    to_avro(
        data=col("value"),
        options=schema_registry_options,
        schemaRegistryAddress=sr_url,
        subject=lit(value_subject_name),
        # jsonFormatSchema=json_format_schema,
    ),
)

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

src_df.display()
```

```text
org.apache.spark.SparkException: Job aborted due to stage failure: Task 2 in stage 436.0 failed 4 times, most recent failure: Lost task 2.3 in stage 436.0 (TID 568) (10.139.64.10 executor 0): org.spark_project.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException: Schema not found; error code: 40403
```

```text
Caused by: org.spark_project.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException: Schema not found; error code: 40403
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/9b27b70d-0baf-af87-b3cd-983ac8cd2a2d.png)