はじめに

StremlitとDatabricks連携については既に@taka_yayoi さんの記事「StreamlitからDatabricksでサービングしている機械学習モデルを呼び出す」がありましたが、今回は単純にデータを取得する方法を試してみました。

接続方法

今回はDatabricks SQL Connectorを利用しました。pyodbcよりも簡単でライブラリーもpipで取得できます。

準備

Databricks環境の以下の情報が必要になります。

Workspace ホスト名
HTTP Path名
Token

SQLウェアハウスを作成し、「接続の詳細」からホスト名とHTTP Pathを取得できます。またTokenも同じ画面の右下の「パーソナルアクセストークンを取得」から取得できます。

データ準備

次にデータ取得のためのテーブルを作成します。すでに取得したいデータがある場合はスキップしてください。
Databricksのノートブックもしくは、SQLエディタから以下を実行します。

CREATE TABLE IF NOT EXISTS small_table
USING CSV
OPTIONS (path "/databricks-datasets/asa/small/small.csv", header "true", mode "FAILFAST")

データエクスプローラでテーブルを確認します。

コード作成

今回は２つのファイルを用意します。

requirements.txt

databricks-sql-connector

DATABRICKS_SERVER_HOSTNAME,DATABRICKS_HTTP_PATH,DATABRICKS_TOKENは後ほどStreamlitのSecretsに登録します。

app.py

import streamlit as st
from databricks import sql
import pandas as pd
import os

@st.experimental_singleton
def init_connection():
    return sql.connect(
        server_hostname = st.secrets["DATABRICKS_SERVER_HOSTNAME"],
        http_path       = st.secrets["DATABRICKS_HTTP_PATH"],
        access_token    = st.secrets["DATABRICKS_TOKEN"])

conn = init_connection()


@st.experimental_memo(ttl=600)
def run_query(query):
    with conn.cursor() as cur:
        cur.execute(query)
        return cur.fetchall()

# Print results.
st.title("Streamlit & Databricks Demo")
st.subheader("Databricks SQL Connectorを使って、Databricksのデータにアクセスする")

full = run_query("select * from small_table limit 5")
st.table(full)

rows = run_query("select Origin, count(*) as count from small_table group by Origin")
df = pd.DataFrame(rows, columns=["Origin","count"])
st.bar_chart(data=df, x="Origin", y="count", width=0, height=0, use_container_width=True)

Gitリポジトリの作成とコードのプッシュ

Streamlit CloudはGitと連携するため、リポジトリを作成しコードを保存しておきます。
今回作成したリポジトリーはこちらです。クローンしてご利用ください。

Streamlit Appの作成

今回はStreamlit App作成方法は割愛します。他にたくさんの記事がありますのでそちらをご覧ください。

Secretsの作成

Streamlit Appがデプロイされましたら、SettingsからSecretsを作成します。今回は最初に準備した以下の３つを記述します。

Appの表示

無事にDatabricksからデータを取得して表示できました。

Databricks SQL Connectorについて

従来のpyodbcドライバーに比べて非常にシンプルな設定で利用できます。ただし大量のデータには向いていないのでデータ量が多い場合はpyodbcドライバーを利用した方がおすすめです。

SQL Connectorの詳細はこちらをご覧ください。
https://learn.microsoft.com/ja-jp/azure/databricks/dev-tools/python-sql-connector

Streamlit から Databricks に接続してデータを取得するには？