こちらの続編です。

以前の例では、生成したSQLを実行するために、OpenAI APIから返却される複数のSQLを実行する際、SQLセルにSQLを貼り付ける必要がありました。これは、PySparkのspark.sqlがマルチステートメントのSQLに対応していないためです。

これですと、自動化しにくいのでOpenAI APIから返却される複数SQLをそのまま実行できるようにします。こちらの記事で説明されているテクニックを使います。区切り文字でSQLを分割し、それぞれでspark.sqlを実行するというアプローチです。

何も指定しないと、区切り文字は;ですが、SQL自身にセミコロンが含まれていると不都合が生じるので、他とかぶらなそうな区切り文字を含めるようにします。

区切り文字を定義します。

Python

sql_line_delimiter = "----databricks----"

区切り文字でSQLを分割して実行する関数を定義します。

Python

def execute_queries(sql):
  queries = sql.split(sql_line_delimiter)

  for query in queries:
    if query.strip() != "":
      spark.sql(query)

SQLに区切り文字を含めてもらいます。これもプロンプトエンジニアリング。

Python

sql = openai_completion(f"""
Databricksのhive_metastoreのデータベースtakaakiyayoi_dbにopenai_testというテーブルを作成するSQLを生成してください。
カラムは4つで文字列と日付型を含めてください。ダミーのレコードを10行追加するSQLを生成してください。
文字列には日本語を使用してください。SQLの区切り文字は{sql_line_delimiter}にしてください
""")
print(sql)

ちゃんと区切り文字が挿入されています。素敵。

----databricks----
CREATE TABLE takaakiyayoi_db.openai_test (
  id string,
  name string,
  date_target date,
  comment string
);

----databricks----
INSERT INTO takaakiyayoi_db.openai_test VALUES
('1163', '山田太郎', '2019-10-31', '本日は晴天なり'),
('1164', '田中花子', '2019-07-12', '雨の中お出かけします'),
('1165', '佐藤二郎', '2020-01-13', 'おしゃべり大好きです'),
('1166', '鈴木一郎', '2019-12-25', 'クリスマスは楽しいね'),
('1167', '斉藤三郎', '2020-02-03', '旅行に行きたいな'),
('1168', '中村平太', '2020-06-15', '今日で卒業式です'),
('1169', '加藤孝子', '2020-04-25', 'BBQ大会です'),
('1170', '山口雄太', '2020-08-14', '入社後初仕事です'),
('1171', '吉田春香', '2020-09-19', '友人に会いに行きます'),
('1172', '後藤翔太', '2021-01-01', '新年を迎えます');

これを実行すると、テーブルが作成され行がインサートされます。

Python

execute_queries(sql)

テーブルが作成されました。

Python

sql = openai_completion("テーブルtakaakiyayoi_db.openai_testから5行取得するSQLを生成してください。")
print(sql)
display(spark.sql(sql))

さらに一歩前進。

Databricksクイックスタートガイド

Databricksクイックスタートガイド

Databricks無料トライアル

Databricks無料トライアル