github.com/apache/beam/sdks/v2@v2.48.2/go/examples/forest/forest.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  // forest is an example that shows that pipeline construction is normal Go
    17  // code -- the pipeline "forest" is created recursively and uses a global
    18  // variable -- and that a pipeline may contain non-connected parts.
    19  //
    20  // The pipeline generated has the shape of a forest where the output of each
    21  // singleton leaf is flattened together over several rounds. This is most
    22  // clearly seen via a visual representation of the pipeline, such as the one
    23  // produced by the 'dot' runner.
    24  //
    25  // Running the pipeline logs "1", "2", "3", etc for each leaf in the forest.
    26  // Note that different runners may produces different or non-deterministic
    27  // orders.
    28  package main
    29  
    30  // beam-playground:
    31  //   name: Forest
    32  //   description: An example that shows that pipeline construction is normal Go
    33  //     code -- the pipeline "forest" is created recursively and uses a global
    34  //     variable -- and that a pipeline may contain non-connected parts. The pipeline
    35  //     generated has the shape of a forest where the output of each singleton leaf
    36  //     is flattened together over several rounds. This is most clearly seen via a
    37  //     visual representation of the pipeline, such as the one produced by the 'dot' runner.
    38  //   multifile: false
    39  //   context_line: 59
    40  //   categories:
    41  //     - Flatten
    42  //     - Branching
    43  //   complexity: MEDIUM
    44  //   tags:
    45  //     - pipeline
    46  //     - branch
    47  
    48  import (
    49  	"context"
    50  	"flag"
    51  
    52  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    53  	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
    54  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
    55  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/debug"
    56  )
    57  
    58  var (
    59  	n     = flag.Int("count", 2, "Number of trees")
    60  	depth = flag.Int("depth", 3, "Depth of each tree")
    61  )
    62  
    63  func tree(s beam.Scope, depth int) beam.PCollection {
    64  	if depth <= 0 {
    65  		return leaf(s)
    66  	}
    67  	a := tree(s, depth-1)
    68  	b := tree(s, depth-1)
    69  	c := tree(s, depth-2)
    70  	return beam.Flatten(s, a, b, c)
    71  }
    72  
    73  var count = 0
    74  
    75  func leaf(s beam.Scope) beam.PCollection {
    76  	count++
    77  	return beam.Create(s, count) // singleton PCollection<int>
    78  }
    79  
    80  func main() {
    81  	flag.Parse()
    82  	beam.Init()
    83  
    84  	ctx := context.Background()
    85  
    86  	log.Info(ctx, "Running forest")
    87  
    88  	// Build a forest of processing nodes with flatten "branches".
    89  	p := beam.NewPipeline()
    90  	s := p.Root()
    91  	for i := 0; i < *n; i++ {
    92  		t := tree(s, *depth)
    93  		debug.Print(s, t)
    94  	}
    95  
    96  	if err := beamx.Run(ctx, p); err != nil {
    97  		log.Exitf(ctx, "Failed to execute job: %v", err)
    98  	}
    99  }