CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: GPU を使う主要なパッケージで CUDA が有効か確かめる方法について

CUDA を利用する Python のパッケージは、多くの場合それ自体のバージョンだけでなく対応している CUDA Toolkit や cuDNN のバージョンまで気にする必要がある。 なんだか環境やコンテナイメージを作るたびに確認する方法や互換性について調べている気がするので、以下のパッケージについてまとめておく。

  • PyTorch
  • TensorFlow (2.x)
  • CuPy
  • CuDF

使った環境は次のとおり。

$ lsb_release -a
No LSB modules are available.
Distributor ID: Ubuntu
Description:    Ubuntu 18.04.5 LTS
Release:    18.04
Codename:   bionic
$ uname -r
4.15.0-124-generic
$ python -V       
Python 3.8.6
$ pip list | egrep -ie "^(torch|tensorflow|cupy|cudf) "
cudf                      0.17.0             
cupy                      8.1.0              
tensorflow                2.4.0              
torch                     1.7.1

もくじ

下準備

あらかじめ各種パッケージをインストールして Python の REPL を起動しておく。

$ python

PyTorch

PyTorch では、次のようにする。

>>> import torch
>>> torch.cuda.is_available()
True

CUDA との互換性を調べる方法

最新バージョンについては以下のページを参照する。

pytorch.org

以前のバージョンを調べたいときは以下。

pytorch.org

TensorFlow (2.x)

TensorFlow では以下のようにする。

>>> import tensorflow as tf
>>> len(tf.config.list_physical_devices('GPU')) > 0
True

CUDA との互換性を調べる方法

以下のページを参照する。 日本語のページは翻訳の関係で新しいバージョンが表示されないことがあるため注意する。

www.tensorflow.org

CuPy

CuPy はインポートでエラーにならなければ使える。

>>> import cupy as cp

あえて関数にするとしたら、こんな感じ?

def cupy_is_available():
    try:
        import cupy as cp
        return True
    except ImportError:
        return False

CUDA との互換性を調べる方法

CuPy は CUDA のバージョンごとにバイナリが用意されているので、それを使えば良い。

github.com

CuDF

CuDF も、インポートでエラーにならなければ使えるはず。

>>> import cudf

少なくとも現行の安定バージョン (0.17) では、ソースコードを確認するとインポートするタイミングでチェックが走っているようなので。

github.com

CUDA との互換性を調べる方法

互換性を調べるときは以下のインストールページを確認するのが良いのかな。

rapids.ai

いじょう。

MSI GeForce RTX 3090 VENTUS 3X 24G OC グラフィックスボード VD7357

MSI GeForce RTX 3090 VENTUS 3X 24G OC グラフィックスボード VD7357

  • 発売日: 2020/10/09
  • メディア: Personal Computers

Python: 主成分分析を重み付き和への分解と解釈した場合の可視化

読んでいる本の中に、主成分分析 (Principal Component Analysis; PCA) はデータを重み付き和に分解していると解釈することもできる、という記述があった。 なるほどーと思ったので、今回はそれについて試してみた。

使った環境は次のとおり。

$ sw_vers 
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H15
$ python -V
Python 3.8.6

下準備

下準備として、あらかじめ必要なパッケージをインストールしておく。

$ pip install scikit-learn matplotlib

使うデータセットを確認する

今回は Labeled Faces in the Wild データセットを用いる。 これは、著名人の顔画像を切り抜いたデータセットになっている。

以下のサンプルコードでは、データセットの先頭 10 件の画像をプロットしている。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from sklearn.datasets import fetch_lfw_people
from matplotlib import pyplot as plt


def main():
    # Labeled Faces in the Wild データセット
    people = fetch_lfw_people(min_faces_per_person=20)

    # データセットの情報
    print(f'dataset shape: {people.images.shape}')
    print(f'number of classes: {len(people.target_names)}')

    # 先頭の 10 件をグレイスケールで表示してみる
    fig, axes = plt.subplots(2, 5,
                             figsize=(12, 6),
                             subplot_kw={'xticks': (),
                                         'yticks': ()})
    mappings = zip(people.target, people.images, axes.ravel())
    for target, image, ax in mappings:
        ax.imshow(image, cmap='gray')
        ax.set_title(people.target_names[target])

    plt.show()


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python lfw.py

すると、次のようなプロットが得られる。

f:id:momijiame:20201207001809p:plain
Labeled Faces in the Wild データセットの画像例

主成分分析したときの主軸 (Principal Axes) を可視化する

一般的に、主成分分析では得られた主成分得点 (Principal Component Score) に着目することが多い。 一方で、データを主成分得点に変換するときに用いる主軸 (Principal Axes) の情報も得られる。 試しに、主軸を主成分ごとに可視化してみよう。

以下のサンプルコードでは、先頭の 10 主成分について主軸をプロットしている。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from sklearn.datasets import fetch_lfw_people
from sklearn.decomposition import PCA
from matplotlib import pyplot as plt


def main():
    people = fetch_lfw_people(min_faces_per_person=20)

    print(f'dataset shape: {people.images.shape}')
    print(f'number of classes: {len(people.target_names)}')

    x, y = people.data, people.target
    image_shape = people.images.shape[1:]

    # PCA で先頭の 100 成分を取り出す
    pca = PCA(n_components=100,
              random_state=42)
    x_pca = pca.fit_transform(x)

    print(f'transformed shape: {x_pca.shape}')
    print(f'principal axes shape: {pca.components_.shape}')

    # 主軸 (Principal Axes) を可視化する
    fig, axes = plt.subplots(2, 5,
                             figsize=(12, 6),
                             subplot_kw={'xticks': (),
                                         'yticks': ()})
    mappings = zip(pca.components_, axes.ravel())
    for i, (component, ax) in enumerate(mappings):
        ax.imshow(component.reshape(image_shape))
        ax.set_title(f'components: {i + 1}')

    plt.show()


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python plotax.py  
dataset shape: (3023, 62, 47)
number of classes: 62
transformed shape: (3023, 100)
principal axes shape: (100, 2914)

すると、以下のようなプロットが得られる。

f:id:momijiame:20201207225939p:plain
主軸 (Principal Axes) の可視化

主成分得点と主軸から画像を再構成する

元のデータは、先ほど得られた主軸と主成分得点という二つの要素から再構成できる。 ここから、主成分分析はデータを主軸と主成分得点に分解していると解釈できる。

以下のサンプルコードでは使う主成分の数を変化させながら主成分得点と主軸を使って画像を再構成している。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from sklearn.datasets import fetch_lfw_people
from sklearn.decomposition import PCA
from matplotlib import pyplot as plt


def main():
    people = fetch_lfw_people(min_faces_per_person=20)

    print(f'dataset shape: {people.images.shape}')
    print(f'number of classes: {len(people.target_names)}')

    x, y = people.data, people.target
    image_shape = people.images.shape[1:]

    fig, axes = plt.subplots(5, 5,
                             figsize=(12, 6),
                             subplot_kw={'xticks': (),
                                         'yticks': ()})

    # 元画像をプロットする
    axes[0, 0].set_title('original image')
    for i in range(5):
        axes[i, 0].imshow(people.images[i],
                          cmap='gray')

    # 利用する主成分の次元ごとに処理する
    for i, n_components in enumerate([10, 50, 100, 500],
                                     start=1):
        # 主成分得点に変換する
        pca = PCA(n_components=n_components,
                  random_state=42)
        x_pca = pca.fit_transform(x)
        # 主成分得点を元の画像に逆変換する
        x_pca_reversed = pca.inverse_transform(x_pca)

        # 逆変換した画像をプロットする
        axes[0, i].set_title(f'{n_components} components')
        for j in range(5):
            ax = axes[j, i]
            ax.imshow(x_pca_reversed[j].reshape(image_shape),
                      cmap='gray')

    plt.show()


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python pcarev.py 
dataset shape: (3023, 62, 47)
number of classes: 62

すると、以下のようなプロットが得られる。

f:id:momijiame:20201207230507p:plain
主成分得点と主軸から再構成した画像

再構成に使う主成分の数が増えるほど、だんだんと鮮明な画像が得られていることがわかる。 ちなみに、今回使ったデータセットでは 500 主成分を使った場合に累積寄与率が 99% を越えていた。

とはいえ、再構成するのに scikit-learn の API を使うだけだと面白くない。 なので、以下のサンプルコードでは scikit-learn の API を使わずに画像を再構成している。 具体的には主成分得点と主軸のドット積を取った上で、元データの平均を足せば良い。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from sklearn.datasets import fetch_lfw_people
from sklearn.decomposition import PCA
from matplotlib import pyplot as plt


def main():
    people = fetch_lfw_people(min_faces_per_person=20)

    print(f'dataset shape: {people.images.shape}')
    print(f'number of classes: {len(people.target_names)}')

    x, y = people.data, people.target
    image_shape = people.images.shape[1:]

    fig, axes = plt.subplots(5, 5,
                             figsize=(12, 6),
                             subplot_kw={'xticks': (),
                                         'yticks': ()})

    axes[0, 0].set_title('original image')
    for i in range(5):
        axes[i, 0].imshow(people.images[i],
                          cmap='gray')

    for i, n_components in enumerate([10, 50, 100, 500],
                                     start=1):
        pca = PCA(n_components=n_components,
                  random_state=42)
        x_pca = pca.fit_transform(x)
        # inverse_transform() を使わずに逆変換してみる
        x_pca_reversed = np.dot(x_pca, pca.components_) + pca.mean_

        # 逆変換した画像をプロットする
        axes[0, i].set_title(f'{n_components} components')
        for j in range(5):
            ax = axes[j, i]
            ax.imshow(x_pca_reversed[j].reshape(image_shape),
                      cmap='gray')

    plt.show()


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python pcarev2.py 
dataset shape: (3023, 62, 47)
number of classes: 62

先ほどと同様のプロットが得られる。

f:id:momijiame:20201207230752p:plain
主成分得点と主軸から再構成した画像 (scikit-learn API を使わない場合)

いじょう。

動作中の Docker コンテナからイメージを作る

Docker コンテナを使って技術検証をしているときに、色々と試行錯誤している場面では、ある手順から作業をやり直したくなることがある。 すべての作業がすぐに終わるなら特に問題にはならないものの、時間がかかる場合には初めからやり直したときに大きなロスが生じる。 そんなときは、動作中の Docker コンテナから Docker イメージを作ることで時間の節約ができる。 具体的には docker commit サブコマンドを使うことで実現できる。

今回、使った環境は次のとおり。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H15
$ docker version                     
Client: Docker Engine - Community
 Cloud integration: 1.0.2
 Version:           19.03.13
 API version:       1.40
 Go version:        go1.13.15
 Git commit:        4484c46d9d
 Built:             Wed Sep 16 16:58:31 2020
 OS/Arch:           darwin/amd64
 Experimental:      false

Server: Docker Engine - Community
 Engine:
  Version:          19.03.13
  API version:      1.40 (minimum version 1.12)
  Go version:       go1.13.15
  Git commit:       4484c46d9d
  Built:            Wed Sep 16 17:07:04 2020
  OS/Arch:          linux/amd64
  Experimental:     false
 containerd:
  Version:          v1.3.7
  GitCommit:        8fba4e9a7d01810a393d5d25a3621dc101981175
 runc:
  Version:          1.0.0-rc10
  GitCommit:        dc9208a3303feef5b3839f4323d9beb36df0a9dd
 docker-init:
  Version:          0.18.0
  GitCommit:        fec3683

もくじ

作業中を模した Docker コンテナを用意する

とりあえず、何でも構わないのでコンテナを立ち上げておく。

$ docker container run \
  --rm \
  -it ubuntu:18.04 \
  bash

立ち上げた以下のコンテナを、何らかの検証で作業中のコンテナに見立てる。

root@f499b0f5b72b:/#

たとえば、次のようにファイルを書き出しておこう。 このファイルがあることで、作業が途中まで完了していることの目印とする。

root@f499b0f5b72b:/# echo "Hello, World" > /var/tmp/greet.txt

動作中のコンテナからイメージを作る

Docker ホスト側で、お目当てのコンテナが動作していることを確認する。 別にコンテナが終了していても構わないんだけど、今回は --rm オプションをつけているのでコンテナは終了すると消えてしまう。

$ docker container ps
CONTAINER ID        IMAGE               COMMAND             CREATED             STATUS              PORTS               NAMES
f499b0f5b72b        ubuntu:18.04        "bash"              31 seconds ago      Up 29 seconds                           eloquent_keller

それでは、docker commit コマンドを使って、コンテナからイメージを作ろう。 最初の引数にお目当てとなるコンテナの ID ないし名前を指定して、次の引数にコンテナを元ネタにしたイメージの名前を指定する。

$ docker commit f499b0f5b72b example
sha256:0bbb09b0e9a30bfd0ba7a342eddb7bbcc6380b0542a1379c7c8e54587e24392a

イメージの一覧を確認すると、ちゃんと指定した名前でイメージができている。

$ docker image list | grep example
example                              latest                                           0bbb09b0e9a3        14 seconds ago      64.2MB

作ったイメージからコンテナを起動する

上記のイメージを使ってコンテナを立ち上げてみよう。

$ docker container run \
  --rm \
  -it example \
  bash

イメージの元ネタになったコンテナにあったはずのファイルを確認すると、次のようにちゃんとある。 つまり、作業の途中からやり直すことができている。

root@a05ba64a868b:/# cat /var/tmp/greet.txt 
Hello, World

めでたしめでたし。

参考

docs.docker.com

シェルスクリプトで数値をゼロパディングする

シェルスクリプトを書いていると、数値をゼロパディングする必要に迫られることがある。 たとえば、ファイル名や日付を処理するときに多い。 結論から先に述べると、数値のゼロパディングは printf(1) を使うことで実現できる。

使った環境は次のとおり。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H15

数値をゼロパディングする

たとえば、以下のようにシェル変数を用意する。

$ N=5

このままでは、当然のことながらパディングされていない。

$ echo $N   
5

printf(1) を使って二桁にゼロパディングするには、次のようにする。 なお、使えるフォーマットは man 1 printf を参照のこと。

$ printf "%02d\n" $N
05

シェルスクリプトの中で使う場合には、インラインで実行してやれば良いかな。

$ PADDING_N=$(printf "%02d" $N)
$ echo $PADDING_N         
05

いじょう。

GNU date で月末の日付を得る

今回は、GNU date を使って月末の日付を得る方法について。 シェルスクリプトで一ヶ月単位の処理を書こうとすると、よく調べることになるのでメモしておく。

検証に使った環境は次のとおり。

$ sw_vers                                        
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H15
$ date --version
date (GNU coreutils) 8.32
Copyright (C) 2020 Free Software Foundation, Inc.
License GPLv3+: GNU GPL version 3 or later <https://gnu.org/licenses/gpl.html>.
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.

Written by David MacKenzie.

もくじ

下準備

検証に使ったのは macOS で、デフォルトでは date(1) が BSD date になっている。 GNU date の方がオプションが多彩で扱いやすいので入れ替えよう。 Homebrew で Coreutils を入れると gdate というコマンドで GNU date が使えるようになる。

$ brew install coreutils
$ alias date=gdate

GNU date の使い方

GNU date では + 以降にフォーマット文字列を指定することで日時を整形できる。 使えるフォーマット文字列は man 1 gdate を参照のこと。

たとえば、次のようにすれば年月日をハイフンでつないだフォーマットで今日の日付を表示できる。

$ date '+%Y-%m-%d'                 
2020-11-13

先ほどの実行結果を見てわかるとおり、デフォルトでは表示される内容として現在の日時が使われる。 これは --date オプションに now を指定していることに等しい。

$ date '+%Y-%m-%d' --date 'now'    
2020-11-13

--date オプションは Human readable なフォーマットを受け付けるようになっている。 たとえば 1 month を指定すると、現在の 1 ヶ月後になる。

$ date '+%Y-%m-%d' --date '1 month'
2020-12-13

他にも、1 day agoyesterday とすれば昨日が得られる。

$ date '+%Y-%m-%d' --date '1 day ago'
2020-11-12
$ date '+%Y-%m-%d' --date 'yesterday'
2020-11-13

それ以外のテクニックとして、日時をずらした上でフォーマットを駆使して意図した内容を得ることもできる。 たとえば、月初の日付がほしいなら、フォーマットの日付を 1 日にしてやれば良い。 以下では、来月の 1 日の日付を取得している。

$ date '+%Y-%m-01' --date '1 month'
2020-12-01

月末の日付を得る

さて、それでは今回の本題に入る。 たとえば、ここまでの要領で「今月末の日付」を得るには「来月の 1 日の 1 日前」を計算してやれば良いことになる。 これを一回の date(1) の実行で得るのは難しいので、インラインコマンドを駆使する。 つまり、まずは「来月の 1 日」を得た上で、それを次の date(1) の --date オプションに埋め込んでやる。

$ date '+%Y-%m-%d' --date "1 day ago $(date '+%Y-%m-01' --date '1 month')"
2020-11-30

ちゃんと今月末の日付が得られた。

次に、任意の月の月末の日付がほしいなら、インラインで実行する date(1) の --date オプションにその月の日付を指定すれば良い。 たとえば、以下では 2020 年 2 月の末日を取得している。 今年はうるう日があった。

$ date '+%Y-%m-%d' --date "1 day ago $(date '+%Y-%m-01' --date '1 month 2020-02-01')"
2020-02-29

めでたしめでたし。

シェルスクリプトの中でスクリプトのあるディレクトリを取得する

シェルスクリプトの中から、実行したスクリプトのあるディレクトリを必要とする場面はちょいちょいある。 たとえば、スクリプトの中で相対パスを使って別のファイルを読み込むような処理が典型的だと思う。 その場合、スクリプトを実行したときのカレントディレクトリを起点として相対パスを処理してしまうと上手くいかない。 代わりに、スクリプトのあるディレクトリを把握した上で、そこを起点に相対パスを処理する必要がある。 今回は、それを実現する方法についてメモしておく。

使った環境は次のとおり。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H15
$ bash --version
GNU bash, version 3.2.57(1)-release (x86_64-apple-darwin19)
Copyright (C) 2007 Free Software Foundation, Inc.
$ dirname --version
dirname (GNU coreutils) 8.32
Copyright (C) 2020 Free Software Foundation, Inc.
License GPLv3+: GNU GPL version 3 or later <https://gnu.org/licenses/gpl.html>.
This is free software: you are free to change and redistribute it.
There is NO WARRANTY, to the extent permitted by law.

Written by David MacKenzie and Jim Meyering.

もくじ

下準備

今回のサンプルコードは、macOS に標準で入っている BSD dirname でも動作する。 ただし、検証に使っている環境としては GNU Coreutils の実装を利用している。

$ brew install coreutils
$ alias dirname="gdirname"

スクリプトのあるディレクトリを取得する

早速だけど、以下にサンプルコードを示す。 以下のサンプルコードでは SCRIPT_DIR というシェル変数に、スクリプトのあるディレクトリを絶対パスで代入している。

#!/usr/bin/env bash

set -CEuo pipefail

# dirname $0 でスクリプトの存在するディレクトリが得られる
# インラインコマンドで cd して pwd すれば絶対パスが得られる
SCRIPT_DIR=$(cd $(dirname $0) && pwd)
echo ${SCRIPT_DIR}

上記を適当な名前と場所に保存しよう。 ここでは example.sh という名前で /tmp ディレクトリに保存した。

まずは、試しに Bash に絶対パスでシェルスクリプトを指定して実行してみる。

$ bash /tmp/example.sh 
/tmp

ちゃんと保存されているディレクトリの絶対パスが表示された。

続いてはルートディレクトリに移動してから相対パスでシェルスクリプトを指定してみよう。

$ cd /
$ bash ./tmp/example.sh 
/tmp

これでも、ちゃんと保存されているディレクトリの絶対パスが表示された。

ついでに、実行権限をつけて、Shebang 経由でスクリプトを実行してみよう。

$ chmod +x example.sh
$ ./tmp/example.sh    
/tmp

これも、ちゃんと表示される。

動作原理について

あとはもう蛇足だけど、どうして先ほどのサンプルコードが動作するかの説明をしていく。

変数 $0 について

そもそも、シェルスクリプトでは $0 という変数にスクリプトが実行されたときのパスが入る。 たとえば、次のようなサンプルコードを用意してみよう。

#!/usr/bin/env bash

set -CEuo pipefail

# 最初の引数は実行するときに使ったパスになる
echo $0

上記を、先ほどと同じように色々な呼び方で実行してみよう。

$ bash /tmp/example.sh 
/tmp/example.sh
$ bash ./tmp/example.sh
./tmp/example.sh
$ ./tmp/example.sh    
./tmp/example.sh

呼び方によって絶対パスだったり相対パスだったりが入っている。

dirname(1) について

そして、dirname(1) を使うとパスの中でディレクトリの部分が得られる。 たとえば、シェルから直接呼んでみると、次のようになる。

$ dirname /tmp/example.sh 
/tmp
$ dirname ./tmp/example.sh 
./tmp

これで、ひとまずスクリプトのあるディレクトリが相対パスで得られた。 あとは、扱いやすさを考えて必要に応じて絶対パスにする。 絶対パスにするときはインラインで cd して pwd すれば良い。

$ (cd $(dirname ./tmp/example.sh) && pwd)
/tmp
$ (cd $(dirname /tmp/example.sh) && pwd) 
/tmp

いじょう。

補足

シンボリックリンクが使われている場合には、必要に応じて pwd(1) に -P オプションを指定して解決した方が良いかも。

Python: SQLAlchemy のテーブルに後からインデックスを追加する

今回は、Python の O/R マッパーである SQLAlchemy について。 テーブルを定義した時点のモデルには無かったインデックスを、後から追加する方法についてメモしておく。

なお、実務における RDBMS のスキーマ変更に関しては、Alembic のようなフレームワークを使ってバージョン管理することを強くおすすめしたい。

blog.amedama.jp

使った環境は次のとおり。

$ sw_vers                     
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H2
$ python -V
Python 3.8.6
$ python -c "import sqlite3; print(sqlite3.version)"    
2.6.0
$ sqlite3 --version           
3.28.0 2019-04-15 14:49:49 378230ae7f4b721c8b8d83c8ceb891449685cd23b1702a57841f1be40b5daapl

もくじ

下準備

あらかじめ、SQLAlchemy をインストールしておく。

$ pip install sqlalchemy

最初からテーブルの定義にインデックスがある場合

はじめに、最初からテーブルを表すモデルにインデックスの指定がある場合について確認しておく。 以下のサンプルコードには User というクラスが users というテーブルを表している。 そして、age というカラムに対応するアトリビュートに index=True が指定されている。 こうなっていると、テーブルを初期化した時点で、age というカラムにインデックスが有効となる。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.schema import Column
from sqlalchemy.types import Integer
from sqlalchemy.types import Text

ModelBase = declarative_base()


LOGGER = logging.getLogger(__name__)


class User(ModelBase):
    """ユーザのデータを模したサンプル用のモデル"""
    __tablename__ = 'users'

    # ユーザ名
    name = Column(Text, primary_key=True)
    # NOTE: 最初からインデックスを張るときは index=True オプションを使う
    age = Column(Integer, nullable=False, index=True)


def init_db(engine: Engine, drop: bool = False):
    """データベースを初期化する"""
    if drop:
        ModelBase.metadata.drop_all(engine)

    ModelBase.metadata.create_all(engine)


def get_engine(db_uri: str) -> Engine:
    engine = create_engine(
        db_uri,
        pool_recycle=3600,
        encoding='utf-8',
    )
    return engine


def main():
    # SQLAlchemy のログを DEBUG レベルで出力する
    logger = logging.getLogger('sqlalchemy')
    logger.setLevel(logging.DEBUG)
    # DEBUG レベル以上のログを出力する
    logging.basicConfig(level=logging.DEBUG)

    # データベースを初期化する
    db_uri = 'sqlite:///example.db?cache=shared'
    engine = get_engine(db_uri)
    init_db(engine, drop=True)


if __name__ == '__main__':
    main()

上記を実行してみよう。 SQLAlchemy のデバッグログを有効にしているため、バックエンドに発行したクエリなどが出力される。 ここでは、それらの出力は省略している。

$ python withindex.py
...

作成されたデータベースのスキーマを確認してみよう。 すると、次のとおり age カラムにインデックスが確認できる。

$ sqlite3 example.db ".schema"            
CREATE TABLE sqlite_sequence(name,seq);
CREATE TABLE users (
    name TEXT NOT NULL, 
    age INTEGER NOT NULL, 
    PRIMARY KEY (name)
);
CREATE INDEX ix_users_age ON users (age);

最初からテーブルの定義にインデックスがない場合

続いてが本題の、最初の定義にはインデックスの指定がない場合について。 以下のサンプルコードでは、先ほど作成したテーブルを DROP して、新たにインデックスの指定がないものを作り直している。

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import logging

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.schema import Column
from sqlalchemy.types import Integer
from sqlalchemy.types import Text

ModelBase = declarative_base()


LOGGER = logging.getLogger(__name__)


class User(ModelBase):
    __tablename__ = 'users'

    name = Column(Text, primary_key=True)
    # この時点ではインデックスがない
    age = Column(Integer, nullable=False)


def init_db(engine: Engine, drop: bool = False):
    if drop:
        ModelBase.metadata.drop_all(engine)

    ModelBase.metadata.create_all(engine)


def get_engine(db_uri: str) -> Engine:
    engine = create_engine(
        db_uri,
        pool_recycle=3600,
        encoding='utf-8',
    )
    return engine


def main():
    logger = logging.getLogger('sqlalchemy')
    logger.setLevel(logging.DEBUG)
    logging.basicConfig(level=logging.DEBUG)

    db_uri = 'sqlite:///example.db?cache=shared'
    engine = get_engine(db_uri)
    init_db(engine, drop=True)


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python withoutindex.py
...

すると、次のようにインデックスのないテーブルが作り直される。

$ sqlite3 example.db ".schema"
CREATE TABLE users (
    name TEXT NOT NULL, 
    age INTEGER NOT NULL, 
    PRIMARY KEY (name)
);

この状況から、ただモデルのアトリビュートに index=True を追加して ModelBase.metadata.create_all() しても上手くいかない。

SQLAlchemy でインデックスを後から追加する

SQLAlchemy で後からインデックスを追加するには Index というクラスのインスタンスを作る。 そして、create() というメソッドを呼ぶことでインデックスが作成できるようだ。

文章で説明するよりも実際に試す方が分かりやすいと思うので REPL を使って作業していく。 なお、REPL を起動する場所は、先ほどのサンプルコードを作ったディレクトリにしよう。 あるいは、サンプルコードの場所をシェル変数の PYTHONPATH に指定してもいい。

$ python

以下がインデックスを作ることができるクラス。

>>> from sqlalchemy import Index

既存モデルのカラムに対応したインデックスを作りたいなら、次のようにする。 先ほどサンプルコードとして示したモジュールをインポートして使っている。

>>> from withoutindex import User
>>> UserAgeIndex = Index('ix_users_age', User.age)

動かすのに sqlalchemy.engine.Engine のインスタンスが必要なので用意する。

>>> from withoutindex import get_engine
>>> db_uri = 'sqlite:///example.db?cache=shared'
>>> engine = get_engine(db_uri)

あとは、先ほど作った Index クラスのインスタンスに対して create() メソッドの呼ぶだけ。

>>> UserAgeIndex.create(engine)
Index('ix_users_age', Column('age', Integer(), table=<users>, nullable=False))

テーブルの定義を確認すると、次のようにインデックスが追加されている。

$ sqlite3 example.db ".schema"
CREATE TABLE users (
    name TEXT NOT NULL, 
    age INTEGER NOT NULL, 
    PRIMARY KEY (name)
);
CREATE INDEX ix_users_age ON users (age);

いじょう。

繰り返しになるけど、実際には上記のような手動オペレーションをするのではなく、Alembic などを使ってスキーマをバージョン管理しよう。