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

はじめに

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

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

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

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

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

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

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

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

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

TiDB の準備

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

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

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

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

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

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

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

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

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

セットアップ

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

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

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

データベースへの接続

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

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

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

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 の接続情報を入力します

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

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

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 レプリカも必要なので作成します

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

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

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 でもテーブルを確認できます

スキーマの定義

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

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

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

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 として扱うため、配列を文字列として渡します

%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()

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 操作と同じようにデータを取得できます

team_members = Repo.all(TeamMember)

実行結果

[
  %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 関数を使用します

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

実行結果

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

データの更新

データの更新も同様です

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

データの削除

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

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

まとめ

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

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

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