github.com/apache/beam/sdks/v2@v2.48.2/go/examples/complete/autocomplete/autocomplete.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  package main
    17  
    18  import (
    19  	"context"
    20  	"flag"
    21  	"os"
    22  	"regexp"
    23  
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/textio"
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam/log"
    27  	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
    28  	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/top"
    29  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/beamx"
    30  	"github.com/apache/beam/sdks/v2/go/pkg/beam/x/debug"
    31  )
    32  
    33  // TODO(herohde) 5/30/2017: fully implement https://github.com/apache/beam/blob/master/examples/java/src/main/java/org/apache/beam/examples/complete/AutoComplete.java
    34  
    35  var (
    36  	input = flag.String("input", os.ExpandEnv("$GOPATH/src/github.com/apache/beam/sdks/go/data/haiku/old_pond.txt"), "Files to read.")
    37  	n     = flag.Int("top", 3, "Number of completions")
    38  )
    39  
    40  var wordRE = regexp.MustCompile(`[a-zA-Z]+('[a-z])?`)
    41  
    42  func init() {
    43  	register.Function2x0(extractFn)
    44  	register.Function2x1(less)
    45  	register.Emitter1[string]()
    46  }
    47  
    48  func extractFn(line string, emit func(string)) {
    49  	for _, word := range wordRE.FindAllString(line, -1) {
    50  		emit(word)
    51  	}
    52  }
    53  
    54  func less(a, b string) bool {
    55  	return len(a) < len(b)
    56  }
    57  
    58  func main() {
    59  	flag.Parse()
    60  	beam.Init()
    61  
    62  	ctx := context.Background()
    63  
    64  	log.Info(ctx, "Running autocomplete")
    65  
    66  	p := beam.NewPipeline()
    67  	s := p.Root()
    68  	lines, err := textio.Immediate(s, *input)
    69  	if err != nil {
    70  		log.Exitf(ctx, "Failed to read %v: %v", *input, err)
    71  	}
    72  	words := beam.ParDo(s, extractFn, lines)
    73  
    74  	hits := top.Largest(s, words, *n, less)
    75  	debug.Print(s, hits)
    76  
    77  	if err := beamx.Run(ctx, p); err != nil {
    78  		log.Exitf(ctx, "Failed to execute job: %v", err)
    79  	}
    80  }