【BigQuery】ANY TYPEの変数の扱い

背景

• 複数の機械学習パイプラインを運用していた場合、それらパイプラインで共通的に使用するプロシージャが存在していたりする。
• 機械学習パイプラインの例でいえば、学習用データの作成を担うプロシージャを共通プロシージャとして開発する。
• 各パイプラインは、引数の中身を切り替えることによって共通プロシージャから所望のデータを取得する。
• そのとき、プロシージャに引き渡す引数において、パイプラインによっては同一名の引数だったとしてもデータ型を任意に変えたい場合がある

解決したい課題

• あるテーブルから所望の条件に合致するデータを取得するようなプロシージャがあったとする。
• 以下の例では、get_conditionを引数としてプロシージャに引き渡し、get_condition変数を使用して、とあるテーブルから条件に合致するデータを取得している。
• ここでは、get_conditionはSTRING型である。

CREATE OR REPLACE PROCEDURE 'project.dataset.get_data_procedure'(
    get_condition STRING
)
BEGIN
    SELECT
        *
    FROM table
    WHERE column = get_condition
END;

CALL `project.dataset.get_data_procedure`("hoge");

• しかし、別のパイプラインは複数のAND条件でデータを取得したいとする。
• つまり、引数の型を配列（ARRAY<STRING>型）に変更し、配列に複数の条件を入れた上で引き渡したい。
• 引数に配列を指定する場合には、上記のプロシージャは以下のように書き換えなければならない。

CREATE OR REPLACE PROCEDURE 'project.dataset.get_data_procedure'(
    get_condition ARRAY<STRING>
)
BEGIN
    SELECT
        *
    FROM table
    WHERE column IN UNNEST(get_condition)
END;

CALL `project.dataset.get_data_procedure`(["hoge", "fuga"]);

• 通常であれば、引数の型が異なる場合、別々のプロシージャを用意する必要があるだろう。
• ここでは、同一のプロシージャで異なるデータ型の入力を可能にすることを目指す。つまり、Pythonなんかでいうところの、ANY TYPEのようなデータ型を引数として扱いたい。

解決策

• まず、BigQueryではそもそも、引数の型にANYを指定することは出来ない。
• そこで最初に考えられるアイデアは、STRING型のまま配列を引き渡す、ということをやる。
• 具体的には、プロシージャを呼び出す際に以下のように記述する。

CALL `project.dataset.get_data_procedure`('["hoge", "fuga"]');

• これによって、単一条件か複数条件かのいずれかの条件をプロシージャに引き渡すことが可能になる。
• しかし、これだけで全てが解決するわけではない。次に問題になるのが、WHERE句の部分です。
• 前述した2つのサンプルではWHERE句の書き方が異なるし、もし'["hoge", "fuga"]'という文字列が入力された場合、前者の場合には期待するデータは得られないだろうし、後者の場合では構文エラーとなるだろう。
• そこで、まずは配列形式の文字列かどうかを判断し、配列であればINを、そうでない場合には=を使うようにする。
• 配列かどうかを判断する簡単な方法は、文字列にカンマが含まれるかどうかで判断できるだろう（大括弧があるかどうかでもよいだろう）。
• そうすると、プロシージャは次のように書き換えられる。

CREATE OR REPLACE PROCEDURE 'project.dataset.get_data_procedure'(
    get_condition STRING
)
BEGIN
    SELECT
        *
    FROM table
    WHERE
        CASE
            WHEN CONTAINS_SUBSTR(get_condition, ',') THEN column IN UNNEST(get_condition)
            ELSE column = get_condition
        END
END;

CALL `project.dataset.get_data_procedure`('["hoge", "fuga"]');

• 次に問題になってくるのが、UNNEST(get_condition)の部分である。
• get_conditionはSTRING型であるため、これをARRAY<STRING>型に変換しなければならない。
• これは次のように書くことで実現できる

ARRAY(
    SELECT
        REPLACE(TRIM(elem, '[] '), '"', '')
    FROM
        UNNEST(SPLIT(get_condition, ',')) AS elem
)

• 各要素の説明：

  • SPLITで、文字列をカンマ区切りで分解し、一度1行に転置
  • TRIMで、各行に含まれる大括弧（[]）とスペースを削除
  • REPLACEで、ダブルクウォート（"）を削除
  • 最後にARRAYで改めて配列型に再構成

• 最終的に、所望の要件を実現するプロシージャは以下のようになる。

CREATE OR REPLACE PROCEDURE 'project.dataset.get_data_procedure'(
    get_condition STRING
)
BEGIN
    SELECT
        *
    FROM table
    WHERE
        CASE
            WHEN CONTAINS_SUBSTR(get_condition, ',') THEN (
                column IN UNNEST(
                    ARRAY(
                        SELECT
                            REPLACE(TRIM(elem, '[] '), '"', '')
                        FROM
                            UNNEST(SPLIT(get_condition, ',')) AS elem
                    )
                )
            )
            ELSE column = get_condition
        END
END;

CALL `project.dataset.get_data_procedure`('["hoge", "fuga"]');

おまけ

• 今回は配列形式で記述された文字列を引数に引き渡すときの問題を取り上げましたが、もし配列形式ではなく、より単純な書き方で複数の値を引き渡す場合には、プロシージャをもっとシンプルに書くことができます。
• 例えば、そもそも代入する文字列をhoge,fugaというように、最初から大括弧、ダブルクウォート、スペースを無くしたとします。
• すると、前述したWHERE句の書き方は以下のようになります。

CREATE OR REPLACE PROCEDURE 'project.dataset.get_data_procedure'(
    get_condition STRING
)
BEGIN
    SELECT
        *
    FROM table
    WHERE
        CASE
            WHEN CONTAINS_SUBSTR(get_condition, ',') THEN (
                column in UNNEST(SPLIT(get_condition, ','))
            )
            ELSE column = get_condition
        END
END;

CALL `project.dataset.get_data_procedure`('hoge,fuga');

• このように、STRING型であるが故に複数の値の持たせ方は様々なパターンが考えられます。
• 共通的に利用するプロシージャは開発効率を高める一方で、チーム内でのルール作りやその周知徹底が大事になってきます。
• また、自由に書けてしまうことにも問題があるため、ある程度UDF関数などを活用して入力データに制限を掛けることも必要でしょう。