Dataflow を試してみる

python --version

$ tar zxvf google-cloud-sdk-192.0.0-darwin-x86_64.tar.gz
$ ./google-cloud-sdk/install.sh

$ gcloud init

$ /usr/bin/ruby -e "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install)"

$ brew install maven

$ vi ~/.bash_profile

export GOOGLE_APPLICATION_CREDENTIALS=”ダウンロードした秘密鍵のファイルパス”

$ source ~/.bash_profile

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
   xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <groupId>ca.test.dataflow</groupId>
   <artifactId>ca-test-dataflow</artifactId>
   <version>0.0.1-SNAPSHOT</version>
   <properties>
       <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
   </properties>
   <build>
       <plugins>
           <plugin>
               <groupId>org.apache.maven.plugins</groupId>
               <artifactId>maven-compiler-plugin</artifactId>
               <version>3.7.0</version>
               <configuration>
                   <source>1.8</source>
                   <target>1.8</target>
               </configuration>
           </plugin>
           <plugin>
               <groupId>org.apache.maven.plugins</groupId>
               <artifactId>maven-eclipse-plugin</artifactId>
               <configuration>
                   <downloadSources>true</downloadSources>
                   <downloadJavadocs>true</downloadJavadocs>
               </configuration>
           </plugin>
       </plugins>


       <pluginManagement>
           <plugins>
               <plugin>
                   <groupId>org.codehaus.mojo</groupId>
                   <artifactId>exec-maven-plugin</artifactId>
                   <version>1.4.0</version>
                   <configuration>
                       <cleanupDaemonThreads>false</cleanupDaemonThreads>
                   </configuration>
               </plugin>
           </plugins>
       </pluginManagement>
   </build>


   <dependencies>
       <dependency>
           <groupId>org.apache.beam</groupId>
           <artifactId>beam-runners-google-cloud-dataflow-java</artifactId>
           <version>2.3.0</version>
       </dependency>
       <dependency>
           <groupId>org.apache.beam</groupId>
           <artifactId>beam-sdks-java-core</artifactId>
           <version>2.3.0</version>
       </dependency>
       <dependency>
           <groupId>org.apache.beam</groupId>
           <artifactId>beam-runners-direct-java</artifactId>
           <version>2.3.0</version>
           <scope>runtime</scope>
       </dependency>
       <dependency>
           <groupId>org.apache.beam</groupId>
           <artifactId>beam-sdks-java-io-google-cloud-platform</artifactId>
           <version>2.3.0</version>
       </dependency>
       <dependency>
           <groupId>org.slf4j</groupId>
           <artifactId>slf4j-jdk14</artifactId>
           <version>1.7.25</version>
       </dependency>
   </dependencies>
</project>

package ca.test.dataflow;


import static com.google.datastore.v1.client.DatastoreHelper.makeKey;
import static com.google.datastore.v1.client.DatastoreHelper.makeValue;


import java.time.LocalDateTime;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;


import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.io.gcp.datastore.DatastoreIO;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubIO;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.PCollection;


import com.google.datastore.v1.Entity;
import com.google.datastore.v1.Key;


/**
* PubsubToDatastore
*/
@SuppressWarnings("serial")
public class PubsubToDatastore {
  
   private static String TIME_FORMAT = "yyyy-MM-dd HH:mm:ss.SSS";
   private static String TIME_ZONE = "Asia/Tokyo";
   private static DateTimeFormatter dtFormatter = DateTimeFormatter.ofPattern(TIME_FORMAT);


   static class AddDatetimeFn extends DoFn {
      
       private String toNamespace;
       private String toKind;
       private String toNamePrefix;
      
       public AddDatetimeFn(String toNamespace, String toKind, String toNamePrefix) {
          
           this.toNamespace = toNamespace;
           this.toKind = toKind;
           this.toNamePrefix = toNamePrefix;
       }


       @ProcessElement
       public void processElement(ProcessContext c) {
          
           // Pub/Sub のメッセージ String をカンマで分割して各カラムデータを取得します
           String[] columns = c.element().split(",");
          
           Entity.Builder entityBuilder = Entity.newBuilder();
          
           Integer jobId = Integer.parseInt(columns[0].trim());
           Integer rowId = Integer.parseInt(columns[1].trim());
           Key.Builder keyBuilder = makeKey(toKind, String.format("%s-%04d-%04d", toNamePrefix, jobId, rowId));
           if (toNamespace != null) {
               keyBuilder.getPartitionIdBuilder().setNamespaceId(toNamespace);
           }
           Integer wordNum = columns[2].trim().split(" ").length;
           String importDtStr = columns[3].trim();
           ZonedDateTime importDt = LocalDateTime.parse(importDtStr, dtFormatter).atZone(ZoneId.of(TIME_ZONE));
          
           // 現在の日時を取得して、追加するカラムデータを作成します
           ZonedDateTime updateDt = ZonedDateTime.now(ZoneId.of(TIME_ZONE));


           // 各カラムデータを Entity クラスの property に put します
           entityBuilder.setKey(keyBuilder.build());
           entityBuilder.putProperties("rowId", makeValue(rowId).build());
           entityBuilder.putProperties("sentence", makeValue(columns[2].trim()).build());
           entityBuilder.putProperties("wordNum", makeValue(wordNum).build());
           entityBuilder.putProperties("importTimestamp", makeValue(dtFormatter.format(importDt)).build());
           entityBuilder.putProperties("updateTimestamp", makeValue(dtFormatter.format(updateDt)).build());
          
           // 生成した Entity クラスを output に格納して返します
           c.output(entityBuilder.build());
       }
   }


   public class TransformData extends PTransform, PCollection> {
      
       private PubsubToDatastoreOptions options;
      
       public TransformData(PubsubToDatastoreOptions options) {
           this.options = options;
       }


       @Override
       public PCollection expand(PCollection line) {
           // Functionクラスを使用して各Pub/Subメッセージを Entity へ変換します
           PCollection rows = line.apply(ParDo.of(
                   new AddDatetimeFn(options.getToNamespace(), options.getToKind(), options.getToNamePrefix())));


           return rows;
       }
   }
  
   public interface PubsubToDatastoreOptions extends DataflowPipelineOptions {


       // DataflowPipelineOptions を継承してクラス特有の引数を追加します
       String getFromTopic();
       void setFromTopic(String fromTopic);
      
       String getToNamespace();
       void setToNamespace(String toNamespace);
      
       String getToKind();
       void setToKind(String toKind);
      
       String getToNamePrefix();
       void setToNamePrefix(String toNamePrefix);
   }


   public void execute(String[] args) {
      
       // 引数を渡して pipeline 実行時の設定クラスを作成します
       PubsubToDatastoreOptions options = PipelineOptionsFactory.fromArgs(args)
               .withValidation().as(PubsubToDatastoreOptions.class);
       options.setStreaming(true);


       Pipeline pipeline = Pipeline.create(options);


       pipeline
           // Pub/Sub topic からメッセージを取得します
           .apply("read from Pubsub", PubsubIO.readStrings().fromTopic(options.getFromTopic()))
           // Pub/Sub メッセージの Stringデータを Datastore の Entity に変換します
           .apply("transform Pubsub data", new TransformData(options))
           // 変換した Entity を Datastore へ書き込みます
           .apply("write to Datastore", DatastoreIO.v1().write().withProjectId(options.getProject()));


       // 構築した pipeline を実行します
       pipeline.run().waitUntilFinish();
   }
  
   public static void main(String[] args) {
      
       PubsubToDatastore pubsubToDatastore = new PubsubToDatastore();
       pubsubToDatastore.execute(args);
   }
}

package ca.test.dataflow;


import ca.test.dataflow.PortableConfiguration;


import com.google.api.services.pubsub.Pubsub;
import com.google.api.services.pubsub.model.PublishRequest;
import com.google.api.services.pubsub.model.PublishResponse;
import com.google.api.services.pubsub.model.PubsubMessage;
import com.google.common.collect.ImmutableList;


import org.apache.beam.sdk.options.Default;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;


import java.io.IOException;
import java.time.ZoneId;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;


/**
* MemoryToPubsub
*/
public class MemoryToPubsub {


   private static String TIME_FORMAT = "yyyy-MM-dd HH:mm:ss.SSS";
   private static String TIME_ZONE = "Asia/Tokyo";
   private static DateTimeFormatter dtFormatter = DateTimeFormatter.ofPattern(TIME_FORMAT);
  
   public class PubsubJob {
      
       private String toTopic;
       private Integer jobId;
       private Integer rowNum;
      
       public PubsubJob(String toTopic, int jobId, int rowNum) {
          
           this.toTopic = toTopic;
           this.jobId = jobId;
           this.rowNum = rowNum;
       }
  
       public void publishData() throws IOException {
          
           Pubsub pubsub = PortableConfiguration.createPubsubClient();
           for (int i=1; i<=rowNum; i++) {
               String rowId = String.valueOf(i);
               // Pub/Sub へ送信するメッセージを作成
               // 例: "1,1,JobId: 1  rowId: 1 no test desu,2018-05-01 14:02:38.420"
               String message = jobId + "," + rowId + ",JobId: " + jobId + "  rowId: " + rowId
                       + " no test desu," + dtFormatter.format(ZonedDateTime.now(ZoneId.of(TIME_ZONE)));
              
               PubsubMessage pubsubMessage = new PubsubMessage();
               pubsubMessage.encodeData(message.getBytes("UTF-8"));
               List messages = ImmutableList.of(pubsubMessage);
               PublishRequest publishRequest = new PublishRequest().setMessages(messages);
              
               PublishResponse publishResponse = pubsub.projects().topics()
                       .publish(toTopic, publishRequest)
                       .execute();
           }
       }
      
       public void doJob() {
          
           try {
               publishData();
           } catch (IOException e) {
               e.printStackTrace();
           }
       }
   }
  
   public interface MemoryToPubsubOptions extends PipelineOptions {
      
       // 送信先のトピック("projects/<プロジェクト名>/topics/<トピック名>")
       String getToTopic();
       void setToTopic(String toTopic);
      
       // 並列処理job数
       @Default.Integer(1)
       Integer getJobNum();
       void setJobNum(Integer jobNum);
      
       // job あたりのデータ件数
       @Default.Integer(1)
       Integer getRowNum();
       void setRowNum(Integer RowNum);
   }
  
   public void execute(String[] args) {
      
       MemoryToPubsubOptions options = PipelineOptionsFactory.fromArgs(args).as(MemoryToPubsubOptions.class);
      
       String toTopic = options.getToTopic();
       int jobNum = options.getJobNum();
       int rowNum = options.getRowNum();


       List jobs = new ArrayList();
       for (int i=1; i<=jobNum; i++) {
           jobs.add(new PubsubJob(toTopic, i, rowNum));
       }
       jobs.parallelStream().forEach(x -> {
           x.doJob();
       });
   }


   public static void main(String[] args) {
      
       MemoryToPubsub memoryToPubsub = new MemoryToPubsub();
       memoryToPubsub.execute(args);
   }
}

mvn compile exec:java \
   -Dexec.mainClass=ca.test.dataflow.PubsubToDatastore \
   -Dexec.args=" \
   --project=<プロジェクト名> \
   --stagingLocation=gs://<バケット名>/staging \
   --autoscalingAlgorithm=NONE \
   --numWorkers=1 \
   --fromTopic=projects/<プロジェクト名>/topics/<トピック名> \
   --toNamespace= \
   --toKind= \
   --toNamePrefix= \
   --runner=DataflowRunner"

mvn -e compile exec:java \
   -Dexec.mainClass=ca.test.dataflow.MemoryToPubsub \
   -Dexec.args=" \
   --toTopic=projects/<プロジェクト名>/topics/<トピック名> \
   --jobNum=1 --rowNum=1"