Cloud Dataflow for Javaを実行時パラメータ付きでAppEngine for Goから実行する

Cloud Dataflow Template

Cloud Dataflow for JavaをAppEngine for Goから実行する の続きの記事です。
前回は単純にDataflow Templateを作成し、実行するのみでした。
今回は実行する時にパラメータを渡す方法を説明します。

ValueProviderの作成

実行時パラメータを受け取る場合、ロジック側で受け入れることができるようにValueProviderを定義しておきます。
以下の例では、 InputTable , OutputProjectID , OutputKind の3つの値をパラメータとして受け入れるように書いています。

public interface BigQueryToDatastoreOptions extends PipelineOptions {

		@Description("Path of the bigquery table to read from")
		@Default.String("cpb101demo1:samples.table")
		ValueProvider<String> getInputTable();

		@Description("Output destination Datastore ProjectID")
		@Default.String("cpb101demo1")
		ValueProvider<String> getOutputProjectID();

		@Description("Output destination Datastore Kind")
		@Default.String("hogeKind")
		ValueProvider<String> getOutputKind();

		void setInputTable(ValueProvider<String> value);

		void setOutputProjectID(ValueProvider<String> value);

		void setOutputKind(ValueProvider<String> value);
	}

	public static void main(String[] args) {
		BigQueryToDatastoreOptions options = PipelineOptionsFactory.fromArgs(args).withValidation()
				.as(BigQueryToDatastoreOptions.class);

		String inputTable = options.getInputTable().get();
		String projectID = options.getOutputProjectID().get();
		String kind = options.getOutputKind().get();

		Pipeline p = Pipeline.create(options);

		PCollection<KV<Integer, Iterable<TableRow>>> keywordGroups = p
				.apply(BigQueryIO.Read.named("ReadUtterance").from(inputTable)).apply(new GroupKeywords());
		
		CreateEntities createEntities = new CreateEntities();
		createEntities.setKind(kind);
		
		PCollection<Entity> entities = keywordGroups.apply(createEntities);
		entities.apply(DatastoreIO.v1().write().withProjectId(projectID));

		p.run();
	}

Dataflow Templateの作成

Dataflow Templateは以下のように作成します。

mvn compile exec:java -Dexec.mainClass=org.sinmetal.flow.BigQueryToDatastore \
              -Dexec.args="--project=cpb101demo1 \
              --stagingLocation=gs://cpb101demo1/staging \
              --inputTable=cpb101demo1:samples.table \
              --outputProjectID=cpb101demo1 \
              --outputKind=Utterance \
              --dataflowJobFile=gs://cpb101demo1-df-template/BigQueryToDatastore201702021500 \
              --runner=TemplatingDataflowPipelineRunner"

Dataflow Templateの実行

実行方法は前回の記事のApp Engineの時と同じです。
実行時パラメータをRequest Bodyの中に含めてあげれば、Dataflow Template実行時にその値が受け渡されます。

{
	"jobName": "samplejob",
	"gcsPath": "gs://cpb101demo1-df-template/BigQueryToDatastore201703081355",
	"parameters": {
		"inputTable": "topgate-ai-dev:dialog.utterance_1000",
		"outputProjectID": "cpb101demo1",
        "outputKind": "sample1356"
	},
	"environment": {
		"tempLocation": "gs://cpb101demo1/staging",
		"zone": "us-central1-f"
	}
}

パラメータをパイプラインで持ち回る方法

今回、なんだかんだで苦労したのが、実行時に渡されたパラメータをパイプラインで持ち回る方法です。
パラメータはパイプライン実行前に取得できますが、利用するのはパイプラインの途中になります。
Javaのコード上は1つのマシンで動いているように見えるDataflowですが、実際には複数のインスタンスに分かれて実行されるので、値を共有するのは地味に大変です。
ベストなのかは分からないけど、とりあえず成功した方法は、 パイプライン作成後にインスタンス変数に値をセットする です。

こんな感じで、セッターを用意しておいて...

static class CreateEntityFn extends DoFn<KV<Integer, Iterable<TableRow>>, Entity> {
	
    String kind = "";

    public void setKind(String kind) {
        this.kind = kind;
    }
}

パイプライン作成時に、セットします。

    Pipeline p = Pipeline.create(options);

    PCollection<KV<Integer, Iterable<TableRow>>> keywordGroups = p
				.apply(BigQueryIO.Read.named("ReadUtterance").from(inputTable)).apply(new GroupKeywords());
		
    CreateEntities createEntities = new CreateEntities();
    createEntities.setKind(kind);

分かってしまえば単純ですが、ここにたどり着くまでに、うろうろしてたので、これからやる人は同じ穴にはまらないように注意してください！

#おわり

実行時に読み込み元のCSVファイルを変更したいとか、出力先のDatastoreのKindを変えたいとかは、よくあるパターンだと思うので、ValueProvider使って、受け渡していきましょう！

今回使ったサンプルコードは github に置いてあります。