importjava.io.IOException;importorg.apache.beam.examples.common.WriteOneFilePerWindow;importorg.apache.beam.sdk.Pipeline;importorg.apache.beam.sdk.io.gcp.pubsub.PubsubIO;importorg.apache.beam.sdk.options.Default;importorg.apache.beam.sdk.options.Description;importorg.apache.beam.sdk.options.PipelineOptionsFactory;importorg.apache.beam.sdk.options.StreamingOptions;importorg.apache.beam.sdk.options.Validation.Required;importorg.apache.beam.sdk.transforms.windowing.FixedWindows;importorg.apache.beam.sdk.transforms.windowing.Window;importorg.joda.time.Duration;publicclassPubSubToGcs{/*   * Define your own configuration options. Add your own arguments to be processed   * by the command-line parser, and specify default values for them.   */publicinterfacePubSubToGcsOptionsextendsStreamingOptions{@Description("The Cloud Pub/Sub topic to read from.")@RequiredStringgetInputTopic();voidsetInputTopic(Stringvalue);@Description("Output file's window size in number of minutes.")@Default.Integer(1)IntegergetWindowSize();voidsetWindowSize(Integervalue);@Description("Path of the output file including its filename prefix.")@RequiredStringgetOutput();voidsetOutput(Stringvalue);}publicstaticvoidmain(String[]args)throwsIOException{// The maximum number of shards when writing output.intnumShards=1;PubSubToGcsOptionsoptions=PipelineOptionsFactory.fromArgs(args).withValidation().as(PubSubToGcsOptions.class);options.setStreaming(true);Pipelinepipeline=Pipeline.create(options);pipeline// 1) Read string messages from a Pub/Sub topic..apply("Read PubSub Messages",PubsubIO.readStrings().fromTopic(options.getInputTopic()))// 2) Group the messages into fixed-sized minute intervals..apply(Window.into(FixedWindows.of(Duration.standardMinutes(options.getWindowSize()))))// 3) Write one file to GCS for every window of messages..apply("Write Files to GCS",newWriteOneFilePerWindow(options.getOutput(),numShards));// Execute the pipeline and wait until it finishes running.pipeline.run().waitUntilFinish();}}

importargparsefromdatetimeimportdatetimeimportloggingimportrandomfromapache_beamimport(DoFn,GroupByKey,io,ParDo,Pipeline,PTransform,WindowInto,WithKeys,)fromapache_beam.options.pipeline_optionsimportPipelineOptionsfromapache_beam.transforms.windowimportFixedWindowsclassGroupMessagesByFixedWindows(PTransform):"""A composite transform that groups Pub/Sub messages based on publish time    and outputs a list of tuples, each containing a message and its publish time.    """def__init__(self,window_size,num_shards=5):# Set window size to 60 seconds.self.window_size=int(window_size*60)self.num_shards=num_shardsdefexpand(self,pcoll):return(pcoll# Bind window info to each element using element timestamp (or publish time).|"Window into fixed intervals"            >>WindowInto(FixedWindows(self.window_size))|"Add timestamp to windowed elements" >>ParDo(AddTimestamp())# Assign a random key to each windowed element based on the number of shards.|"Add key" >>WithKeys(lambda_:random.randint(0,self.num_shards-1))# Group windowed elements by key. All the elements in the same window must fit# memory for this. If not, you need to use `beam.util.BatchElements`.|"Group by key" >>GroupByKey())classAddTimestamp(DoFn):defprocess(self,element,publish_time=DoFn.TimestampParam):"""Processes each windowed element by extracting the message body and its        publish time into a tuple.        """yield(element.decode("utf-8"),datetime.utcfromtimestamp(float(publish_time)).strftime("%Y-%m-%d %H:%M:%S.%f"),)classWriteToGCS(DoFn):def__init__(self,output_path):self.output_path=output_pathdefprocess(self,key_value,window=DoFn.WindowParam):"""Write messages in a batch to Google Cloud Storage."""ts_format="%H:%M"window_start=window.start.to_utc_datetime().strftime(ts_format)window_end=window.end.to_utc_datetime().strftime(ts_format)shard_id,batch=key_valuefilename="-".join([self.output_path,window_start,window_end,str(shard_id)])withio.gcsio.GcsIO().open(filename=filename,mode="w")asf:formessage_body,publish_timeinbatch:f.write(f"{message_body},{publish_time}\n".encode())defrun(input_topic,output_path,window_size=1.0,num_shards=5,pipeline_args=None):# Set `save_main_session` to True so DoFns can access globally imported modules.pipeline_options=PipelineOptions(pipeline_args,streaming=True,save_main_session=True)withPipeline(options=pipeline_options)aspipeline:(pipeline# Because `timestamp_attribute` is unspecified in `ReadFromPubSub`, Beam# binds the publish time returned by the Pub/Sub server for each message# to the element's timestamp parameter, accessible via `DoFn.TimestampParam`.# https://beam.apache.org/releases/pydoc/current/apache_beam.io.gcp.pubsub.html#apache_beam.io.gcp.pubsub.ReadFromPubSub|"Read from Pub/Sub" >>io.ReadFromPubSub(topic=input_topic)|"Window into" >>GroupMessagesByFixedWindows(window_size,num_shards)|"Write to GCS" >>ParDo(WriteToGCS(output_path)))if__name__=="__main__":logging.getLogger().setLevel(logging.INFO)parser=argparse.ArgumentParser()parser.add_argument("--input_topic",help="The Cloud Pub/Sub topic to read from."'"projects/<PROJECT_ID>/topics/<TOPIC_ID>".',)parser.add_argument("--window_size",type=float,default=1.0,help="Output file's window size in minutes.",)parser.add_argument("--output_path",help="Path of the output GCS file including the prefix.",)parser.add_argument("--num_shards",type=int,default=5,help="Number of shards to use when writing windowed elements to GCS.",)known_args,pipeline_args=parser.parse_known_args()run(known_args.input_topic,known_args.output_path,known_args.window_size,known_args.num_shards,pipeline_args,)