Redshift で UDF を使って YAML をパースする

2016-01-04

Redshift で YAML の配列データをパースして、Hive の lateral view、Presto の unset 的な感じで配列を展開したかったんですが、結論から言うと、実用的なレベルでは実現できませんでした。
YAML の配列を Redshift でよしなに扱いたい場合、Redshift にデータを入れる時点で配列を複数の行に展開するのが良いんじゃないかと思います。

以下、備忘録的な内容です。

テーブルの作成

次のように、各ユーザの友達情報を YAML で持つ friends テーブルを作成します。比較のため JSON でも同様のデータを保持することにします。

CREATE TABLE users (
  id integer not null
);

CREATE TABLE friends (
  user_id integer not null,
  yaml_friend_ids varchar(max),
  json_friend_ids varchar(max)
);

サンプルデータの作成

以下のスクリプトで適当に 1 万ユーザのデータを作成してみました。
Redshift で改行のあるデータを取り込むには手前にバックスラッシュが必要のようなのでそうしています。
cf. http://docs.aws.amazon.com/ja_jp/redshift/latest/dg/r_COPY_command_examples.html#r_COPY_preparing_data

# generate_redshift_sample_data.rb
require 'csv'
require 'json'
require 'yaml'
require 'randomext'

USER_COUNT = 10000
MEAN_FRIEND_COUNT = 100

random = Random.new
user_ids = (1..USER_COUNT).to_a

users_csv = CSV.open('users.csv', 'w')
friends_csv = CSV.open('friends.csv', 'w')

user_ids.each do |id|
  friend_count = random.poisson(MEAN_FRIEND_COUNT)
  users_csv << [id]
  friend_ids = user_ids.sample(friend_count).sort
  friends_csv << [id, friend_ids.to_yaml.gsub("\n", "\\\n"), friend_ids.to_json]
end

users_csv.close
friends_csv.close

% ruby generate_redshift_sample_data.rb

データのインポート

まず、先ほど作成したデータを S3 にアップします

% brew install s3cmd
% s3cmd --configure
% s3cmd put users.csv s3://arabiki-test/redshift-sample/
% s3cmd put friends.csv s3://arabiki-test/redshift-sample/

次に、Redshift でデータを取り込みます。

COPY users FROM 's3://arabiki-test/redshift-sample/users.csv'
CREDENTIALS 'aws_access_key_id=<access-key-id>;aws_secret_access_key=<secret-access-key>';

COPY friends FROM 's3://arabiki-test/redshift-sample/friends.csv'
CREDENTIALS 'aws_access_key_id=<access-key-id>;aws_secret_access_key=<secret-access-key>'
ESCAPE
REMOVEQUOTES
DELIMITER AS ',';

PyYAML のインポート

Redshift では Python で UDF を定義することができるのですが、Python は組み込みで YAML をサポートしていないので、PyYAML の zip を作成して取り込みます。
cf. http://docs.aws.amazon.com/ja_jp/redshift/latest/dg/udf-python-language-support.html

% curl -LO https://github.com/yaml/pyyaml/archive/3.10.tar.gz
% tar xvf 3.10.tar.gz
% cd pyyaml-3.10/lib
% zip -r yaml.zip yaml
% s3cmd put yaml.zip s3://arabiki-test/redshift-sample/

CREATE LIBRARY yaml LANGUAGE plpythonu FROM 's3://arabiki-test/redshift-sample/yaml.zip'
CREDENTIALS 'aws_access_key_id=<access-key-id>;aws_secret_access_key=<secret-access-key>';

UDF の定義

JSON 関数と同様の関数を定義します。
UDF には f_ prefix を付けることが推奨されているようです。
cf. http://docs.aws.amazon.com/ja_jp/redshift/latest/dg/udf-naming-udfs.html

CREATE OR REPLACE FUNCTION f_yaml_extract_array_element_text (yaml_string varchar(max), pos integer) RETURNS varchar(max) IMMUTABLE AS $$
  import yaml
  return str(yaml.load(yaml_string)[pos])
$$ LANGUAGE plpythonu;

CREATE OR REPLACE FUNCTION f_yaml_array_length (yaml_string varchar(max)) RETURNS integer IMMUTABLE AS $$
  import yaml
  return len(yaml.load(yaml_string))
$$ LANGUAGE plpythonu;

なお、IMMUTABLE などの説明は PostgreSQL のドキュメントが参考になります。
cf. http://www.postgresql.org/docs/9.5/static/sql-createfunction.html

僕は次のように解釈しました。

IMMUTABLE: 引数が同じであれば必ず同じ値を返す
STABLE: 1 つの SQL 内では引数が同じであれば同じ値を返すが、2 回実行すると異なる値を返す可能性がある（テーブルの状態に依存する関数など）
VOLATILE: 1 つの SQL 内でも異なる値を返す可能性がある

簡単なクエリを実行してみる

次のようにシンプルな SQL を実行してみるとわかるんですが、めちゃくちゃ遅いです…

SELECT
  f_yaml_array_length(yaml_friend_ids) friend_count
FROM
  friends
;

JSON だと速いんですが…

SELECT
  json_array_length(json_friend_ids) friend_count
FROM
  friends
;

PyYAML は libyaml を使うと 10 倍ぐらい速くなりそうですが、Redshift の UDF ではピュア Python なものしか受け付けてないと思うので厳しそうですね…

配列を展開する

Redshift には Hive の lateral view のような機能は存在しないようです。
なので、WITH statement を使って頑張っている方もいるようです。

今回の例だと次のような感じでしょうか。JSON 関数を使ってもかなりの時間がかかるので、YAML 用の UDF を使うと気の遠くなる時間がかかりそうです…

WITH exploded_array AS (
  SELECT
    friends.user_id,
    json_extract_array_element_text(json_friend_ids, users.id - 1) AS friend_id
  FROM
    friends, users
  WHERE
    users.id <= json_array_length(json_friend_ids)
)
SELECT
  *
FROM
  exploded_array
;

以上、備忘録でした！