Livebook から Ecto で TiDB のベクトル型を扱う

## はじめに

TiDBは、PingCAP社が開発したオープンソースの分散型 **NewSQL** データベースです

:::note info
NewSQL: RDB と NoSQL の両方の特長を備えたデータベース
:::

MySQL互換のプロトコルとSQL構文をサポートし、水平スケーリングや高可用性、トランザクション処理（ACID特性）を備えています

https://pingcap.co.jp/campaign-generic-m/

フルマネージドの Cloud サービスも提供されており、 TiDB Cloud Serverless は無料から利用可能です

https://pingcap.co.jp/pricing/

本記事では Livebook から Ecto で TiDB に接続し、 DB 操作する手法を紹介します

また、主に AI 関連で使用するベクトル型の扱い方についても紹介します

なお、本記事では TiDB のはじめ方、利用手順は紹介しません

TiDB のはじめ方については以下の記事が参考になりました

https://zenn.dev/churadata/articles/1531050fe66886

実装したノートブックはこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/livebooks/tidb/tidb_ecto.livemd

## TiDB の準備

TiDB Cloud Serverless でクラスターを作成しておきます

本記事の範囲内であれば無料で利用可能です

クラスターが作成されると、以下のような概要画面が表示されます

![スクリーンショット 2025-02-17 14.51.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/45dc0dfc-744b-4b2c-93bc-e48ddf17750b.png)

右上の "Connect" ボタンをクリックしてください

以下のようにデータベースの接続情報が表示されます

![スクリーンショット 2025-02-17 15.04.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/507b1fc7-c198-4e97-ab07-0cbb64154559.png)

右中央の "Generate Password" ボタンをクリックすると、ランダムなパスワードが生成されます

![スクリーンショット 2025-02-17 15.05.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/a92b5243-a104-4501-9be8-7998cc39ad7b.png)

Livebook からはこの接続情報を使って接続します

パスワードはモーダルを閉じると次回以降表示されないので安全な場所に記録しておきましょう

## Livebook からのデータベース操作

### セットアップ

Livebook で新しいノートブックを開き、セットアップセルで以下のコードを実行します

```elixir
Mix.install([
  {:ecto, "~> 3.12"},
  {:ecto_sql, "~> 3.12"},
  {:myxql, "~> 0.7"},
  {:jason, "~> 1.4"},
  {:kino, "~> 0.14"}
])
```

TiDB は MySQL 互換であるため、MySQL 操作用モジュールの MyXQL をインストールしています

https://github.com/elixir-ecto/myxql

### データベースへの接続

Repo モジュールを用意します

```elixir
defmodule Repo do
  use Ecto.Repo,
    otp_app: :my_notebook,
    adapter: Ecto.Adapters.MyXQL
end
```

データベースへの接続情報を入力するためのフォームを作成します

```elixir
hostname_input = Kino.Input.text("HOSTNAME")
username_input = Kino.Input.text("USERNAME")
password_input = Kino.Input.password("PASSWORD")
database_input = Kino.Input.text("DATABASE")

Kino.Layout.grid(
  [hostname_input, username_input, password_input, database_input],
  columns: 4
)
```

表示されたフォームに TiDB の接続情報を入力します

![スクリーンショット 2025-02-19 15.39.17.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/d611b1b8-8c21-448c-9aa1-8a546369ba64.png)

接続情報を使って TiDB に接続します

暗号化通信のため、 `ssl` で証明書のパスを指定しています（以下のコードは Ubuntu の場合のパスです）

```elixir
opts = [
  hostname: Kino.Input.read(hostname_input),
  port: 4000,
  username: Kino.Input.read(username_input),
  password: Kino.Input.read(password_input),
  database: Kino.Input.read(database_input),
  ssl: [cacertfile: "/etc/ssl/certs/ca-certificates.crt"]
]

Kino.start_child({Repo, opts})
```

### テーブル作成

マイグレーションモジュールを定義します

今回はベクトル型も使いたいので、 `embedding` 列をベクトル型で定義しました

ベクトルインデックスを使う場合はサイズ指定が必須なので、 `size: 4` で指定しています

ベクトルインデックスには TiFlash レプリカも必要なので作成します

https://docs.pingcap.com/ja/tidbcloud/tiflash-overview

```elixir
defmodule Migrations.CreateTeamMemberTable do
  use Ecto.Migration

  def change do
    create table(:team_member) do
      add(:name, :string)
      add(:age, :integer)
      add(:weight, :float)
      add(:has_license, :boolean)
      add(:hash, :binary)
      add(:embedding, :vector, size: 4)  # ベクトル型
      add(:languages, :json)
      add(:skil_level, :json)
      add(:salary, :decimal)
      add(:date_of_birth, :date)
      add(:starting_time_of_work, :time)

      timestamps()
    end

    flush()

    # TiFlash レプリカの作成（ベクトルインデックス作成に必要）
    execute """
    ALTER TABLE team_member SET TIFLASH REPLICA 1
    """

    flush()

    # ベクトルインデックスの作成
    execute """
    ALTER TABLE team_member
    ADD VECTOR INDEX idx_team_member_embedding
    ((VEC_COSINE_DISTANCE(embedding)));
    """
  end
end
```

定義したマイグレーションを適用します

```elixir
Ecto.Migrator.up(Repo, 1, Migrations.CreateTeamMemberTable)
```

実行時の標準出力

```
01:53:26.083 [info] == Running 1 Migrations.CreateTeamMemberTable.change/0 forward

01:53:26.084 [info] create table team_member

01:53:26.544 [info] execute "ALTER TABLE team_member SET TIFLASH REPLICA 1\n"

01:53:27.000 [info] execute "ALTER TABLE team_member\nADD VECTOR INDEX idx_team_member_embedding\n((VEC_COSINE_DISTANCE(embedding)));\n"

01:53:28.024 [info] == Migrated 1 in 1.9s
```

マイグレーションが完了すると、 TiDB の SQL Editor でもテーブルを確認できます

![スクリーンショット 2025-02-19 17.05.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1485835/28f019d5-91a7-4fcf-8054-8ec0b4f62a92.png)

### スキーマの定義

Ecto でスキーマを定義します

ベクトル型は実態としては `float` の配列ですが、 `string` として扱います

:::note warn
`{:array, :float}` としてスキーマ定義した場合、 INSERT 時は動作しますが、 SELECT 時にパースできずエラーになります
:::

```elixir
defmodule TeamMember do
  use Ecto.Schema
  import Ecto.Changeset

  schema "team_member" do
    field(:name, :string)
    field(:age, :integer)
    field(:weight, :float)
    field(:has_license, :boolean)
    field(:hash, :binary)
    field(:embedding, :string)  # ベクトル型は string として扱う
    field(:languages, {:array, :string})
    field(:skil_level, {:map, :integer})
    field(:salary, :decimal)
    field(:date_of_birth, :date)
    field(:starting_time_of_work, :time)

    timestamps()
  end

  def changeset(team_member, attrs) do
    team_member
    |> cast(attrs, [
      :name,
      :age,
      :weight,
      :has_license,
      :hash,
      :embedding,
      :languages,
      :skil_level,
      :salary,
      :date_of_birth,
      :starting_time_of_work
    ])
    |> validate_required([:name])
  end
end
```

### データの追加

通常の Ecto 操作と同じようにデータを追加できます

ベクトル型は `string` として扱うため、配列を文字列として渡します

```elixir
%TeamMember{}
|> TeamMember.changeset(%{
  name: "Alice",
  age: 20,
  weight: 60.0,
  has_license: true,
  hash: <<0b11111111>>,
  embedding: "[0.0, 0.1, 0.2, 0.3]",
  languages: ["Japanese", "English"],
  skil_level: %{frontend: 5, backend: 3},
  salary: 1000_000,
  date_of_birth: ~D[2000-01-01],
  starting_time_of_work: ~T[08:30:00.0010]
})
|> Repo.insert()
```

```elixir
now =
  NaiveDateTime.utc_now()
  |> NaiveDateTime.truncate(:second)

entries =
  [
    %{name: "Bob", age: 20, embedding: "[0.0, 0.1, 0.2, 0.4]"},
    %{name: "John", age: 30, embedding: "[0.1, 0.1, 0.3, 0.4]"},
    %{name: "Ryo", age: 39, embedding: "[0.3, 0.4, 0.2, 0.1]"}
  ]
  |> Enum.map(fn attr ->
    Map.merge(attr, %{
      inserted_at: now,
      updated_at: now
    })
  end)

Repo.insert_all(TeamMember, entries)
```

### データの取得

通常の Ecto 操作と同じようにデータを取得できます

```elixir
team_members = Repo.all(TeamMember)
```

実行結果

```elixir
[
  %TeamMember{
    __meta__: #Ecto.Schema.Metadata<:loaded, "team_member">,
    id: 5,
    name: "Alice",
    age: 20,
    weight: 60.0,
    has_license: true,
    hash: <<255>>,
    embedding: "[0,0.1,0.2,0.3]",
    languages: ["Japanese", "English"],
    skil_level: %{"backend" => 3, "frontend" => 5},
    salary: Decimal.new("1000000"),
    date_of_birth: ~D[2000-01-01],
    starting_time_of_work: ~T[08:30:00],
    inserted_at: ~N[2025-02-19 01:55:15],
    updated_at: ~N[2025-02-19 01:55:15]
  },
  %TeamMember{
    __meta__: #Ecto.Schema.Metadata<:loaded, "team_member">,
    id: 6,
    name: "Bob",
    age: 20,
    weight: nil,
    has_license: nil,
    hash: nil,
    embedding: "[0,0.1,0.2,0.4]",
    languages: nil,
    skil_level: nil,
    salary: nil,
    date_of_birth: nil,
    starting_time_of_work: nil,
    inserted_at: ~N[2025-02-19 01:55:18],
    updated_at: ~N[2025-02-19 01:55:18]
  },
  %TeamMember{
    __meta__: #Ecto.Schema.Metadata<:loaded, "team_member">,
    id: 7,
    name: "John",
    age: 30,
    weight: nil,
    has_license: nil,
    hash: nil,
    embedding: "[0.1,0.1,0.3,0.4]",
    languages: nil,
    skil_level: nil,
    salary: nil,
    date_of_birth: nil,
    starting_time_of_work: nil,
    inserted_at: ~N[2025-02-19 01:55:18],
    updated_at: ~N[2025-02-19 01:55:18]
  },
  %TeamMember{
    __meta__: #Ecto.Schema.Metadata<:loaded, "team_member">,
    id: 8,
    name: "Ryo",
    age: 39,
    weight: nil,
    has_license: nil,
    hash: nil,
    embedding: "[0.3,0.4,0.2,0.1]",
    languages: nil,
    skil_level: nil,
    salary: nil,
    date_of_birth: nil,
    starting_time_of_work: nil,
    inserted_at: ~N[2025-02-19 01:55:18],
    updated_at: ~N[2025-02-19 01:55:18]
  }
]
```

ベクトル型は配列の文字列として取得されます

ベクトルインデックスを使用してコサイン距離を計算する場合、以下のように `fragment` 内で `VEC_COSINE_DISTANCE` 関数を使用します

```elixir
import Ecto.Query

Repo.all(
  from(tm in TeamMember,
    select: %{
      member_name: tm.name,
      member_age: tm.age,
      distance: fragment(
        "VEC_COSINE_DISTANCE(?, ?) as distance", tm.embedding, "[0.0, 0.1, 0.2, 0.3]"
      )
    },
    where: tm.age < 25,
    order_by: fragment("distance")
  )
)
```

実行結果

```elixir
[
  %{id: 5, member_name: "Alice", member_age: 20, distance: 0.0},
  %{id: 6, member_name: "Bob", member_age: 20, distance: 0.008539877396793072}
]
```

### データの更新

データの更新も同様です

```elixir
TeamMember
|> Repo.get_by!(name: "Bob")
|> TeamMember.changeset(%{
  age: 21,
  weight: 62.0,
  embedding: "[0.4, 0.3, 0.2, 0.1]"
})
|> Repo.update()
```

### データの削除

データの削除も特に問題なく実行できます

```elixir
TeamMember
|> Repo.get_by!(name: "Alice")
|> Repo.delete()
```

## まとめ

Ecto を使って TiDB のデータベースを操作することができました

また、ベクトル型のデータも問題なく扱えました

PostgreSQL での pgvector のようなものが存在しないので文字列扱いしないといけない点は注意です

https://github.com/pgvector/pgvector-elixir