こちらの続編です。
以前の例では、生成した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))
さらに一歩前進。