Databricks にて monotonically_increasing_id 利用時の結合後の書き込み結果が想定外となる事象への暫定対応方法

## 概要

特定の条件下で、`display` 結果とテーブルへの書き込み結果に差異が生じる事象が確認されたため、暫定的な対応方法を共有します。

今回、`monotonically_increasing_id` 関数から生成される値を Pandas の Index に相当する列として扱い、2 つに分割したデータフレームを最終的に結合してテーブルへ書き込む処理を実施しました。

```mermaid
graph TD

    A[ソーステーブルから読み込み]
    A --> D[monotonically_increasing_id を使用して ROW_ID 列を追加]
    D --> E[original_dfとして定義]

    D --> G[dq_checked_dfとして定義]
    G --> K[データ品質チェック]

    E --> H[original_df と dq_checked_df を ROW_ID で LEFT JOIN]
    K --> H

    H --> J[テーブルへ書き込み]


```

しかし、テーブルへの書き込み時に想定通りの結果になりませんでした。書き込みを行うデータフレームと、実際に書き込まれたテーブルの値を照合したところ、データフレーム上では `delete_flg` 列に値が入っているのに、テーブル上では `NULL` になっていました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/3581fec2-860c-49e2-a778-13db09154ac4.png)

検証を進めた結果、以下の条件をすべて満たす場合に想定外の動作が起きる可能性があるようです。推測ですが、`monotonically_increasing_id` 列の値が、データフレームを `display` したときと、テーブルへ書き込むときで異なることが原因と考えられます。

1. `monotonically_increasing_id` 関数の出力を基に 2 つのデータフレームを結合している  
2. ソースからのフィルタリング条件が前回の処理と同じ  
3. テーブルへの書き込みを実施している  

## 対応方針

`monotonically_increasing_id` が生成する値を固定化させるため、一時的な領域にデータフレームを書き出して再読み込みする方法をとっています。

```python
# monotonically_increasing_id の値を実体化
tgt_df.write.format("delta").mode("overwrite").option("mergeSchema", True).save(volume_dir)
tgt_df = spark.read.format("delta").load(volume_dir)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/955d88eb-dc06-4ce7-ae2c-4880dbc7a74b.png)

データフレームに対してキャッシュを利用する方法もあるようですが、Serverless ではキャッシュ機能を利用できないためその方法を選択しませんでした。また、標準クラスター上でキャッシュを試した場合も、私のプログラムでは想定通り動作しなかったため、採用する際には事前の検証が必要だと考えられます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/6324374d-0dde-401d-b323-7617fca10a71.png)

出所：[サーバレスコンピュートの制限 | Databricks Documentation](https://docs.databricks.com/aws/ja/compute/serverless/limitations#キャッシングの制限)

Databricks Runtime 14.1 以降を利用できる場合には、行追跡（row tracking）機能の`_metadata.row_id`列を利用することも選択できます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b98dc253-c9c9-4a2b-885d-7a16670ce926.png)

出所：[Delta テーブルの行追跡を使用する | Databricks Documentation](https://docs.databricks.com/aws/ja/delta/row-tracking#行追跡メタデータ-フィールドのスキーマは何ですか)


## 想定外の動作の再現方法

### 事前準備

```python
%sql
CREATE CATALOG IF NOT EXISTS meget_test;
```

```python
%sql
CREATE SCHEMA IF NOT EXISTS manabian_split_test;
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a8ecd507-b85e-4c1c-84f1-443613ef5204.png)



### エラーの再現方法

#### テーブルの作成

```python
src_table = "meget_test.manabian_split_test.src_table_01"
tgt_table = "meget_test.manabian_split_test.tgt_table_01"
```

```python
spark.sql(f"DROP TABLE IF EXISTS {src_table};")
spark.sql(f"""CREATE TABLE {src_table} (
    COL_1 string,
    ingest_date date,
    delete_flg integer,
    rescued_data string
);""")

spark.sql(f"DROP TABLE IF EXISTS {tgt_table};")
spark.sql(f"""CREATE TABLE {tgt_table} (
    COL_1 string,
    ingest_date date,
    delete_flg integer,
    rescued_data string
);""")
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/09224e68-1a13-4402-897a-8a9a685bf893.png)

#### `write_to_table` ノートブックの作成と書き込み用コードの定義

```python
src_table = "meget_test.manabian_split_test.src_table_01"
tgt_table = "meget_test.manabian_split_test.tgt_table_01"

from pyspark.sql.functions import expr, monotonically_increasing_id

tgt_df = spark.table(src_table)

# ROW_ID 列を追加
tmp_rowid_col_name = "tmp_rowid_col_name"
tgt_df = tgt_df.withColumn(
    tmp_rowid_col_name,
    monotonically_increasing_id(),
)

# 異なるデータフレームとして定義を実施して、 orignal_df の方のカラムを DROP
orignal_df = tgt_df
orignal_df = orignal_df.drop(
    "delete_flg",
)

# ROW ID 相当の列結合後に DROP
joined_df = orignal_df.alias("orignal_df").join(
    tgt_df.alias("dq_checked_df"),
    orignal_df[tmp_rowid_col_name] == tgt_df[tmp_rowid_col_name],
    "left",
)
joined_df = joined_df.select(
    "orignal_df.*",
    f"dq_checked_df.delete_flg",
)
joined_df = joined_df.drop(tmp_rowid_col_name)

# テーブルに書き込み
output_tbl_name = joined_df
joined_df.write.format("delta").mode("append").option("mergeSchema", "true").saveAsTable(tgt_table)
```

```python
# 書き込み後のデータフレームを display
print("out_df")
joined_df.display()
```

```python
# 書き込み後のテーブルの確認
df = spark.table(tgt_table)
df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b1547954-2717-42db-b134-9555faf70c07.png)

#### 1 回目のデータ準備

```python
import datetime

schema_01 = """
COL_1 string,
ingest_date date,
delete_flg integer,
rescued_data string
"""
data_01 = [
    {
        "COL_1": "P0",
        "ingest_date": datetime.date(2020, 1, 2),
        "delete_flg": 0,
        "rescued_data": "",
    }
]
df = spark.createDataFrame(data_01, schema_01)
df.write.mode("overwrite").saveAsTable(src_table)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c94ba4e3-c1f4-4fb3-9570-466e3798253c.png)

#### 1 回目の書き込み処理

```python
dbutils.notebook.run("./write_to_table", 0)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/14e73409-356a-4699-89b0-402dfa45a157.png)

#### display の結果とテーブルの比較 (1 回目)

想定通りの結果となります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/60889587-3c47-4587-a277-d41afe00944a.png)

#### 2 回目のデータ準備

```python
import datetime

schema_01 = """
COL_1 string,
ingest_date date,
delete_flg integer,
rescued_data string
"""
data_01 = [
    {
        "COL_1": "P1",
        "ingest_date": datetime.date(2020, 1, 2),
        "delete_flg": 0,
        "rescued_data": "",
    }
]
df = spark.createDataFrame(data_01, schema_01)
df.write.mode("append").saveAsTable(src_table)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/42789c0e-1de4-4695-9143-70fd621461da.png)

#### 2 回目の書き込み処理

```python
dbutils.notebook.run("./write_to_table", 0)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/e5bd4714-fe33-4965-947e-b22ab064fa4c.png)

#### display の結果とテーブルの比較 (2 回目)

想定外の結果となりました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/d861c4fd-b283-4ae4-a3f0-89aba6de0d51.png)



## 暫定対応方法

### 事前準備

```python
%sql
CREATE CATALOG IF NOT EXISTS manabian_split_test;
```

```python
%sql
CREATE SCHEMA IF NOT EXISTS manabian_split_test;
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a8ecd507-b85e-4c1c-84f1-443613ef5204.png)



### エラーの再現方法

#### テーブルの再作成

```python
src_table = "meget_test.manabian_split_test.src_table_01"
tgt_table = "meget_test.manabian_split_test.tgt_table_01"
```

```python
spark.sql(f"DROP TABLE IF EXISTS {src_table};")
spark.sql(f"""CREATE TABLE {src_table} (
    COL_1 string,
    ingest_date date,
    delete_flg integer,
    rescued_data string
);""")

spark.sql(f"DROP TABLE IF EXISTS {tgt_table};")
spark.sql(f"""CREATE TABLE {tgt_table} (
    COL_1 string,
    ingest_date date,
    delete_flg integer,
    rescued_data string
);""")
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/09224e68-1a13-4402-897a-8a9a685bf893.png)



#### Volume の作成

```python
volume_name = "meget_test.manabian_split_test.volume_01"

spark.sql(f"CREATE VOLUME IF NOT EXISTS {volume_name};")

dbutils.fs.rm("/Volumes/meget_test/manabian_split_test/volume_01")
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c86c4e2f-3fbe-4fcc-8155-f1c06665be8a.png)



#### `write_to_table__treated` ノートブックの作成と書き込み用コードの定義

```python
src_table = "meget_test.manabian_split_test.src_table_01"
tgt_table = "meget_test.manabian_split_test.tgt_table_01"
volume_dir = "/Volumes/meget_test/manabian_split_test/volume_01"

from pyspark.sql.functions import expr, monotonically_increasing_id

tgt_df = (
    spark.table(src_table)
    # Hint 前回と異なる抽出条件の際には想定通りの動作となる
    # .where("COL_1 = 'P1'")
)

# ROW_ID 列を追加
tmp_rowid_col_name = "tmp_rowid_col_name"
tgt_df = tgt_df.withColumn(
    tmp_rowid_col_name,
    monotonically_increasing_id(),
)

# monotonically_increasing_id の値を実体化
tgt_df.write.format("delta").mode("overwrite").option("mergeSchema", True).save(volume_dir)
tgt_df = spark.read.format("delta").load(volume_dir)


# ROW_ID 列を追加
test_col = "test_col"
tgt_df = tgt_df.withColumn(
    test_col,
    expr("cast(tmp_rowid_col_name AS STRING)"),
)


# ROW_ID 列を追加
tmp_rowid_col_name = "tmp_rowid_col_name"
tgt_df = tgt_df.withColumn(
    tmp_rowid_col_name,
    monotonically_increasing_id(),
)

# 異なるデータフレームとして定義を実施して、 orignal_df の方のカラムを DROP
orignal_df = tgt_df
orignal_df = orignal_df.drop(
    "delete_flg",
)

# ROW ID 相当の列結合後に DROP
joined_df = orignal_df.alias("orignal_df").join(
    tgt_df.alias("dq_checked_df"),
    orignal_df[tmp_rowid_col_name] == tgt_df[tmp_rowid_col_name],
    "left",
)
joined_df = joined_df.select(
    "orignal_df.*",
    f"dq_checked_df.delete_flg",
)
joined_df = joined_df.drop(tmp_rowid_col_name)

# テーブルに書き込み
output_tbl_name = joined_df
joined_df.write.format("delta").mode("append").option("mergeSchema", "true").saveAsTable(tgt_table)
```

```python
# 書き込み後のデータフレームを display
print("out_df")
joined_df.display()
```

```python
# 書き込み後のテーブルの確認
df = spark.table(tgt_table)
df.display()
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/91d7606e-a464-405f-9d15-b5a45a9adbf2.png)



#### 1 回目のデータ準備

```python
import datetime

schema_01 = """
COL_1 string,
ingest_date date,
delete_flg integer,
rescued_data string
"""
data_01 = [
    {
        "COL_1": "P0",
        "ingest_date": datetime.date(2020, 1, 2),
        "delete_flg": 0,
        "rescued_data": "",
    }
]
df = spark.createDataFrame(data_01, schema_01)
df.write.mode("overwrite").saveAsTable(src_table)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/8676d810-14df-41c8-9e26-f47d8f0a9abf.png)

#### 1 回目の書き込み処理

```python
dbutils.notebook.run("./write_to_table__treated", 0)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/018c44e1-321b-4945-a9cc-cb45feaf989d.png)

#### display の結果とテーブルの比較 (1 回目)

想定通りの結果になりました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a3e4a877-ed85-4edf-8649-6843180ed6aa.png)



#### 2 回目のデータ準備

```python
import datetime

schema_01 = """
COL_1 string,
ingest_date date,
delete_flg integer,
rescued_data string
"""
data_01 = [
    {
        "COL_1": "P1",
        "ingest_date": datetime.date(2020, 1, 2),
        "delete_flg": 0,
        "rescued_data": "",
    }
]
df = spark.createDataFrame(data_01, schema_01)
df.write.mode("append").saveAsTable(src_table)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/142be79e-e56a-45bd-bb62-5871ff008449.png)

#### 2 回目の書き込み処理

```python
dbutils.notebook.run("./write_to_table", 0)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/22a2e0c0-c4cd-4515-ae2a-a38a6ad37bc9.png)

#### display の結果とテーブルの比較 (2 回目)

想定通りの結果が得られました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/a369a0c4-cae9-4426-8f36-91645ed8710a.png)