SlickのバルクインサートSQLをshapelessで自動生成

65744 ワード

はじめに

SlickはScalaでよく利用されているデータベース用のライブラリーであり、SQLのようなシンタックスをScalaの言語内DSLとして提供する。ただSlickはコレクションとなった複数の同一型のデータのインサートが低速であることが知られている。たとえば さんの記事Slick(MySQL)でBulk Upsertを実装するでは次のような方法で高速なINSERTクエリーを自作する方法が紹介されている。

SimpleDBIO { session =>
  val statement = session.connection.prepareStatement(sqlStatement)
  books.zipWithIndex.foreach {
    case (row, i) =>
      val rowSize = 3
      val offset  = i * rowSize
      statement.setInt(1 + offset, row.id)
      statement.setString(2 + offset, row.name)
      statement.setString(3 + offset, row.author)
  }
  statement.executeUpdate()
}

このコードではまずSlickの機能でプリペアードステートメントなINSERTクエリー[1]を取得し、それに対してバルクで挿入したいデータbooks: Seq[Book]をループで回してプレースホルダーに対応するインデックスにrow.idrow.nameなどのデータを挿入している。
最終的には筆者の実装もこれと同じテクニックを利用して高速化を達成するが、一方でこのsetIntsetStringのような低レベル[2]な関数を直接プログラマーが記述するため次のような問題が考えられる。

  • インデックスを間違えたりテーブル定義上intなカラムにsetStringすればランタイムエラーになる
  • テーブル定義やBookのデータ構造が変化した場合には、このクエリー生成部分も修正しなければランタイムエラーとなる

これらのことからSlick標準の++=と比べて効率はよいが一方で保守性は低下したと言わざるを得ない。
この記事ではこのような低レベルなクエリーの生成をshapelessを利用して完全に自動化し、人間が危険なインデックス管理やsetStringのメソッド選択をする余地を消去した。さらにshapeless-3というScala 3版のshapelessを利用してScala 3に対応してある[3]。また、最後に詳しく述べるがSlickの++=と比べて10倍程度の高速化が達成された。
この記事で紹介するコードの完全なものは下記のGitHubリポジトリーに置かれている。

この記事について質問やコメントなどがあれば気軽に連絡してほしい。

型クラスBulkInsertable[A]

まずは、最初に述べたようなデータベースに保存されうるBookのような、とある型Aがバルクインサートできることを示す型クラスBulkInsertableを次のように定義する。

BulkInsertable.scala
trait BulkInsertable[A] {
  def set(statement: PreparedStatement, a: A): State[Int, Unit]

  def parameterLength: Int
}

BulkInsertableには2つの機能がある。

  1. set …… プリペアードステートメントに型Aの値(= a)をセットする
  2. parameterLength …… 型Aに必要なプレースホルダーの個数

まず(1)のsetについては低レベルな関数と述べたsetStringsetIntの抽象化版ということになる。ただしこの関数はインデックスを取らず、代わりに返り値の型がステートモナド(State)になっている。このステートモナドの状態がIntであり、ここでインデックスを管理している。
そして次のparameterLengthはパラメータの数で、たとえば次のようなケースクラスBookについて考える。

case class Book(id: Int, name: String, author: String)

これをインサートするSQLは次のようになる。

INSERT INTO book (`id`, `name`, `author`) VALUES  ('1', 'MyBook1', 'Yamada')

したがってBookパラメーター数は3となる。これは(?, ?, ?)のようなプリペアードステートメントのプレースホルダーを生成するために用いる。Bookの場合はこれでいいとして、次のようなやや複雑な型についても考える。

case class BookWithOwner(ownerName: String, book: Book)

このBookWithOwnerを1つのテーブルに押し込んだと考えると、parameterLength1 + 3 = 4を返す必要がある[4]

BulkInsertableインスタンス定義

次に型クラスBulkInsertableのインスタンスを定義する。

SetParameterを用いたプリミティブ型のインスタンス

とりあえずプリミティブなInt型などについてのBulkInsertable[Int]setIntを利用して定義すればよさそうではあるが、Slickは便利なことにSetParameter[A]という型クラスを持っており、これを使えばプリミティブ型については次のように簡単に定義することができる。

BulkInsertableInstances.scala
implicit def setParameterInstance[A](implicit
  setParameter: SetParameter[A]
): BulkInsertable[A] =
  new BulkInsertable[A] {
    def set(statement: PreparedStatement, a: A): State[Int, Unit] =
      State { s =>
        val positionedParameters = new PositionedParameters(statement)
        positionedParameters.pos = s - 1
        (s + 1, setParameter(a, positionedParameters))
      }

    def parameterLength: Int = 1
  }

SlickのSetParameterPositionedParametersというPreparedStatementに加えて現在のインデックスを状態として持つクラスを利用している。しかし今回の実装では、暗黙的な状態を利用したくなかったのでステートモナドを利用することとし、PositionedParameterssetParameterが済んだら都度捨てることとした。
またparameterLengthについては常に1となっている。これはプリミティブな型でプレースホルダーがいらないということはないので1とした。

Scala 2のジェネリックインスタンス

shapelessはScalaのマクロを利用して、ケースクラスのようなユーザー(= プログラマー)が定義したデータ構造をHListと呼ばれる長さが変化するタプルのような構造に変換して処理できる[5]。たとえば先ほどのデータ構造Bookは次のような定義であった。

case class Book(id: Int, name: String, author: String)

これは次のようなHListに対応する。

Int :: String :: String :: HNil

このようにshapelessを利用すると任意のケースクラスをこのような対応するHListに変換したり、逆にこのHListから元の型を復活させたりする機能を提供する[6]。したがってまずはHListに関するインスタンスを次のように定義する。

scala-2/BulkInsertableGenericInstances.scala
implicit val hNilInstance: BulkInsertable[HNil] = new BulkInsertable[HNil] {
  def set(statement: PreparedStatement, a: HNil): State[Int, Unit] =
    State(s => (s, ()))

  def parameterLength: Int = 0
}

implicit def hConsInstance[H, T <: HList](implicit
  head: BulkInsertable[H],
  tail: BulkInsertable[T]
): BulkInsertable[H :: T] = new BulkInsertable[H :: T] {
  def set(statement: PreparedStatement, a: H :: T): State[Int, Unit] =
    for {
      _ <- head.set(statement, a.head)
      _ <- tail.set(statement, a.tail)
    } yield ()

  def parameterLength: Int = head.parameterLength + tail.parameterLength
}

まずHNilのケースでは何も生成する必要はないし、クエリーにも反映されないことからインデックスの更新も何もせずに終了する。一方で::のケースについて見ていく、まずsetメソッドではhead/tailに分解してそれぞれsetを呼び出している。setの返り値はステートモナドなので、このようにforで繋ぐことでインデックス更新を伝搬する。またparameterLengthhead/tailのそれぞれを足し算すればOKである。
最後に型Aから対応するL <: HListへ変換したり戻したりする部分を次のように定義して終了となる。

scala-2/BulkInsertableGenericInstances.scala
implicit def hListInstance[A, L <: HList](implicit
  gen: Generic.Aux[A, L],
  hList: Lazy[BulkInsertable[L]]
): BulkInsertable[A] = new BulkInsertable[A] {
  def set(statement: PreparedStatement, a: A): State[Int, Unit] =
    hList.value.set(statement, gen.to(a))

  def parameterLength: Int = hList.value.parameterLength
}

ここでgen: Generic.Aux[A, L]とは、型L <: HListが型Aに対応するHListであれば値の検索に成功するようなimplicitパラメーターになっている。そしてこのgenを利用してaHListにしたり戻したりすればよい。

Scala 3のジェネリックインスタンス

まず言っておくこととして、Scala 2版のshapelessとScala 3版のshapeless-3は全く互換性がなく、そもそもScala 3がリリースされてまだ時間が経ってないためかshapeless-3は機能が大幅に足りていない。なので同じような名前のライブラリーではあるが次のようにコードの見た目は全く違うものとなる。

trait BulkInsertableGenericInstances { self: BulkInsertableInstances =>
  implicit def bulkInsertableGenInsance[A](implicit inst: K0.ProductInstances[BulkInsertable, A]): BulkInsertable[A] =
    new BulkInsertable[A] {
      def set(statement: PreparedStatement, a: A): State[Int, Unit] = {
        inst.foldLeft(a)(State(s => (s, ())): State[Int, Unit]) {
          [t] => (acc: State[Int, Unit], bk: BulkInsertable[t], x: t) =>
            acc >> bk.set(statement, x)
        }
      }

      def parameterLength: Int = inst.unfold(0) {
        [t] => (acc: Int, bk: BulkInsertable[t]) =>
          // The second value of this tuple is never used so it's safe for now.
          (acc + bk.parameterLength, Some(null.asInstanceOf[t]))
      }._1
    }
}

inst: ProductInstances[BulkInsertable, A]は、さきほどScala 2側で説明したようなAに対応するようなHListが見つかった場合にinstという値が得られる。しかしScala 2とは違ってshapeless-3ではHListのようなAに対応する具体的な構造にアクセスできず、代わりにinstfoldLeftmapのようなリストに対する操作を提供する。
setに関してはこのinstを用いてAfoldLeftで回せばOKである[7]。一方でparameterLengthではsetと違って、型Aの値を得られないため次のunfoldという入力した関数の返り値の2番目がNoneにならない限りinstをループしてくれる機能を利用する。

inline def unfold[Acc](i: Acc)(
  f: [t] => (Acc, BulkInsertable[t]) => (Acc, Option[t])
): (Acc, Option[T])

本来unfoldは最終的に型Aの値を作りだすための機能であり、Noneが出現した時点でループを終了してしまう。今回の用途では型Aの値を作るつもりはないが、一方でinstが持つ全てのBulkInsertable[t]についてそれのparameterLengthを足し算する必要があり、途中でループを止められては上手くいかない。そこで強引ではあるがSome(null.asInstanceOf[t])unfoldを止めないようにしつつ、accに各要素のparameterLengthを足し算している。当然unfoldの結果の_2へアクセスすればNullPointerExceptionなどのランタイムエラーが生じる危険があるため、直ちに_1を取得してして危険な_2を葬っておく。

未定義なCoproductインスタンス

実はここまでの紹介にあるジェネリックなインスタンスは片手落ちである。先に述べたようなケースクラスBookのような型の“積”はHListのような方法で対処できるものの、次のような型の“和”には対処できない。

sealed abstract class Color(val value: String)

case object Red extends Color("red")
case object Blue extends Color("blue")
case object Green extends Color("green")

結論から説明するとBulkInsertableはこのような型の和に対するインスタンスをあえて自動生成しない。つまり、もしプログラマーが特に手動で何もせずColorのようなデータモデルをBulkInsertableで処理しようとするとコンパイルエラーとなる。その理由を説明するために、別の型の和である次の例についても考えてみる。

sealed trait DeviceType

case class IOS(version: String, isIPad: Boolean)
case class Android(version: String, vender: String, isPixel: Boolean)

このDeviceTypeIOSAndroidの和となっているが、それぞれで持っているフィールドの数や型が異なる。このようなデータをどのようにテーブルへ詰めこむかを自動的に判断するのは難しい。したがってこのようなDeviceTypeを持つようなデータモデルのインスタンスは必要ならプログラマーが手動で定義することとして、shapelessを使った自動的な生成は行わない。

semiauto対応

バージョン0.2.0より古い場合、BulkInsertableのインスタンスはshapelessにより全自動で導出されていた。しかしScala 3(バージョン3.1.2)においては、このようなマクロを利用したインスタンス自動生成がコンパイル速度を飛躍的に低下させることが知られている[8]。このような場合、大きなケースクラスのインスタンスをたとえば次のように半手動(semiauto)で生成すると高速化する。

case class BigUser(
  // 非常に大きなフィールド……
)

object BigUser {
  implicit val instance: BulkInsertable[BigUser] = BulkInsertable.semiauto
}

Scala 2のshapelessはマクロを使っても比較的高速なインスタンス生成が実行できるが、一方で現在のScala 3との相互利用を考慮したときに、全自動がデフォルト動作であると深刻なコンパイル速度低下に繋がる恐れがあることから、バージョン0.2.0から半手動インスタンス導出をデフォルト動作とするように変更した。

AutoDerivedBulkInsertable

この対応を行うにあたり、まず次のようなAutoDerivedBulkInsertableを定義する。

AutoDerivedBulkInsertable.scala
abstract class AutoDerivedBulkInsertable[A] extends BulkInsertable[A]

このようにAutoDerivedBulkInsertableBulkInsertableのサブクラスとなっており、これはScalaのimplicitパラメーター検索が自己再帰により失敗することを防ぐために定義している。このAutoDerivedBulkInsertableのコンパニオンオブジェクトへ従来の全自動インスタンス導出のためのコードを次のように移動させる。

  • Scala 2
    scala-2/AutoDerivedBulkInsertable.scala
    object AutoDerivedBulkInsertable {
      implicit def hListInstance[A, L <: HList](implicit
        gen: Generic.Aux[A, L],
        hList: Lazy[BulkInsertable[L]]
      ): AutoDerivedBulkInsertable[A] = new AutoDerivedBulkInsertable[A] {
        ???
      }
    }
    
  • Scala 3
    scala-3/AutoDerivedBulkInsertable.scala
    object AutoDerivedBulkInsertable {
      implicit def bulkInsertableGenInstance[A](implicit inst: K0.ProductInstances[BulkInsertable, A]): AutoDerivedBulkInsertable[A] =
        new AutoDerivedBulkInsertable[A] {
          ???
        }
    }
    

そしてBulkInsertableのコンパニオンオブジェクトにAutoDerivedBulkInsertableのインスタンスを利用して半自動でインスタンスを提供するsemiautoを定義しておく。

BulkInsertable.scala
  final def semiauto[A](implicit instance: AutoDerivedBulkInsertable[A]): BulkInsertable[A] =
    instance

このsemiautoは半手動なのでimplicitが付いていないことが重要である。
また、過去のバージョンのように全自動でインスタンスを導出したい場合はimport AutoDerivedBulkInsertable.*すればインスタンスがスコープに展開されて全自動導出となる。

BulkInsertableからのINSERTクエリー作成

あとはここから実際のクエリーを組み立てればよい。下記のbulkInsertはバルクインサートしたいデータを受けとり、実際にインサートした個数をDBIO[Int]の型で返す。

BulkInsert.scala
trait BulkInsert[A] {
  protected def tableQuery: TableQuery[? <: Table[A]]

  def bulkInsert(dms: Seq[A])(implicit
    bulkInsertable: BulkInsertable[A]
  ): DBIO[Int] = dms match {
    case Nil =>
      DBIO.successful(0)

    case h +: ts =>
      SimpleDBIO { session =>
        val placeholder = (1 to bulkInsertable.parameterLength).map(_ => "?").mkString("(", ",", ")")
        val placeholders = (1 until dms.length).map(_ => placeholder)
        val sql = (tableQuery.insertStatement +: placeholders).mkString(",")
        val statement = session.connection.prepareStatement(sql)

        ts
          .foldLeft(bulkInsertable.set(statement, h)) { (acc, dm) =>
            bulkInsertable.set(statement, dm) >> acc
          }
          .run(1)
          .value

        statement.executeUpdate()
      }
  }
}

まずはINSERT INTO table_nameのようなSQLを生成するためにSlickのtableQueryを抽象メンバーとして持っておく。
bulkInsertはSlickの++=と同様にSeq[A]を引数に取る。この関数は引数のシーケンスがNilの場合とそうでない場合に場合分けしている。この理由は次のようになる。

  • Slickが生成するtableQuery.insertStatementにはデフォルトで1つのプレースホルダーが入っているため、Nilの場合はデータベースへの問い合わせができないし、実際する必要がない
    • したがってDBIO.successful(0)で終了する
  • またh +: tsのケースでは実際にSQLを組み立てるが、そのあとfoldLeftを利用している。foldLeftは初期値が必要なので、そこにhを利用できて便利である

このときプレースホルダーの組み立てにparameterLengthを利用し、あとはfoldLeftsetしながらステートモナドを合成していけばプリペアードステートメントに適切な値を入力できる。

SlickのScala 3対応

現在リリースされているSlick 3.3.3はScala 3対応されておらず、Scala 2のマクロでコンパイルエラーとなってしまう。具体的にはTableQueryがマクロとなっており問題になるので、SlickのPR #2187にあるScala 3対応のTableQueryを持ってきてScala 3の場合にのみ上書きするようにして無理やり動作させる。#2187のTableQueryからは多少改変したが、正直これについては適当いじってはコンパイルを繰り返し、コンパイルが完全に通るまでにやっただけという感じなので、詳細は下記のコードを読んでほしい。

とにかくこのように定義したScala3CompatTableQueryをScala 3側のソースコードフォルダーに設置しておく。

scala-3/Scala3CompatTableQuery.scala
trait Scala3CompatTableQuery {
  inline def TableQuery[E <: AbstractTable[_]]: TableQuery[E] = ${ TableQueryImpl.applyExpr[E] }
}

一方でScala 2側には同名だが実装が空なトレイトを作っておけばよい。

scala-2/Scala3CompatTableQuery.scala
trait Scala3CompatTableQuery

そして使うところでこのScala3CompatTableQueryを継承すると無事にScala 3ではSlickのScala 2のマクロが上書きされて利用されなくなりコンパイルエラーを回避できる。

UserTestDAO.scala
object UserTestDAO extends BulkInsert[UserDataModel] with Scala3CompatTableQuery {
  private val databaseConfig: DatabaseConfig[JdbcProfile] =
    DatabaseConfig.forConfig("testMySQL")

  override protected val profile = databaseConfig.profile

  import profile.api.*

  // ......
}

ベンチマーク

ベンチマークは次のテーブルに今回作成したbulkInsertメソッドとSlickの++=でデータを10,000件インサートしてはその都度テーブルの内容を全て消去するという操作を10回行って平均を取るという方法にした。
Slickで利用するデータモデルと利用したテーブルは次のようになる。

UserDataModel.scala
case class UserDataModel(
  id: Int,
  name: Option[String],
  info: UserInfoDataModel,
  createdAt: Date
)

case class UserInfoDataModel(
  height: Double,
  weight: Double
)
mysql> SHOW CREATE TABLE users\G
*************************** 1. row ***************************
       Table: users
Create Table: CREATE TABLE `users` (
  `id` int NOT NULL,
  `name` text,
  `height` double NOT NULL,
  `weight` double NOT NULL,
  `created_at` date NOT NULL,
  PRIMARY KEY (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci

そしてベンチマークコードはJMHを用いて次のようになっている。

Benchmarks.scala
@State(Scope.Thread)
@BenchmarkMode(Array(Mode.SingleShotTime))
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Warmup(iterations = 2)
@Measurement(iterations = 10)
@Fork(value = 1, warmups = 1)
class Benchmarks {
  import AutoDerivedBulkInsertable.*

  val num = 10000
  val dms: Seq[UserDataModel] = createDataModels(num)

  @Setup(Level.Trial)
  def setupTrial(): Unit = {
    UserTestDAO.dropTableIfExists()
    UserTestDAO.createTable()
  }

  @TearDown(Level.Iteration)
  def tearIteration(): Unit = {
    UserTestDAO.delete()
  }

  @Benchmark
  def benchSlickInsertAllJmh(): Unit = {
    UserTestDAO.addAll(dms)
  }

  @Benchmark
  def benchBulkInsertJmh(): Unit = {
    UserTestDAO.run(UserTestDAO.bulkInsert(dms))
  }
}

本当はデータベースサーバーとScalaアプリケーションサーバーを別々に用意して実験したほうがいいが、そういう環境を用意できなかった。多少雑にはなるがGitHub ActionsのDockerでMySQL 8を動作させ、同じインスタンスでsbtで起動したベンチマークを動作させたところ、次のような結果となった。

  • Scala 2
    [info] Benchmark                          Mode  Cnt     Score     Error  Units
    [info] Benchmarks.benchBulkInsertJmh        ss   10   254.737 ±  55.462  ms/op
    [info] Benchmarks.benchSlickInsertAllJmh    ss   10  2353.399 ± 129.191  ms/op
    
  • Scala 3
    [info] Benchmark                          Mode  Cnt     Score    Error  Units
    [info] Benchmarks.benchBulkInsertJmh        ss   10   294.616 ± 53.436  ms/op
    [info] Benchmarks.benchSlickInsertAllJmh    ss   10  2534.799 ± 88.160  ms/op
    

約8〜9倍程度の高速化となっており、場合によっては10倍程度スコアに差がつくこともある。

まとめ

以前からSlickの++=は性能が悪いことは議論されていた。どのようにすると高速になるのかは目処がたってなかったが、@taket0ra1さんの記事でやり方が明らかとなったので今回それを自動化した。今後はプロダクションのコードにこれを導入してより実際のアプリケーションに近い形で性能評価を進めたい。
この記事の内容とはやや関係ないが、そもそも今から新規のScalaアプリケーションでSlickを採用するべきかどうかであったり、あるいはSlickも今後利用するべきかというと微妙ではある。しかし現状、SlickはLightbend社がメンテナンスしているといったこともあって利用者は一定数いると考えられ、自分も含めて直ちにSlickから脱却できるのかというとそれは難しい。

謝辞

この記事を書くにあたって、 さんには次のような様々な情報を頂いたので感謝したい。

  • SlickのSetParameterについて
  • shapeless-3の情報
  • SlickのScala 3対応方法
  • Scala 3のコンパイル速度に関してsemiautoの必要性

参考文献

脚注
  1. この記事にはあまり関係ないが、プリペアードステートメント(Prepared statement)とはSQLのうち、IDなどのデータが入りうる部分をprintf%sのようなフォーマット指定子のように渡してやることで、外部の文字列によりプログラマーが意図しないSQLの意味論を変化させるような攻撃(SQLインジェクション)を困難にするために用いられる技術である。 ↩︎

  2. 念のため述べておくと、この記事では「具体的な機械の実装に近い」とか「抽象度が低い」といった意味合いで低レベルという語を利用するのであって、プログラムの善し悪しとは関係ない。 ↩︎

  3. 現在SlickはScala 3対応が完了していないが、SlickのPRからコードを拝借することによって無理やりScala 3でもSlickが動作するようにした。したがってSlickがこのあとScala 3対応した場合、直ちにこのプログラムを利用できる。 ↩︎

  4. このような場合、リレーショナルデータベースでは通常ownerテーブルとbookテーブルを別々に作るとは思う。さらにいえば、ScalaやDDDなどの慣例上もデータモデルとドメインモデルを別々にするといった理由でデータモデルにあたるSlickのデータ型でこのような複雑な記述をするとも考えにくいとは思う。 ↩︎

  5. この記事でも軽くは解説するが、もしshapelessを用いたプログラミングをより詳しく知りたい場合、拙著となるが“ダミー値”を自動で作成する型クラスで詳しく解説したのでそちらを参照してほしい。 ↩︎

  6. 記事には直接関係ないが、このようなプログラミング手法のことを datatype-generic programming やあるいは generic programming と呼ぶ。 ↩︎

  7. ちなみにacc >> bk.set(statement, x)とはacc.flatMap(_ => bk.set(statement, x))と同じである。 ↩︎

  8. たとえばScala 3のinlineによるcompile時間増大に対処する方法に詳細な情報がある。 ↩︎