github.com/apache/beam/sdks/v2@v2.48.2/go/examples/snippets/06schemas.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 snippets
    17  
    18  import (
    19  	"fmt"
    20  	"io"
    21  	"reflect"
    22  	"time"
    23  
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    25  	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder"
    26  )
    27  
    28  // [START schema_define]
    29  
    30  type Purchase struct {
    31  	// ID of the user who made the purchase.
    32  	UserID string `beam:"userId"`
    33  	// Identifier of the item that was purchased.
    34  	ItemID int64 `beam:"itemId"`
    35  	// The shipping address, a nested type.
    36  	ShippingAddress ShippingAddress `beam:"shippingAddress"`
    37  	// The cost of the item in cents.
    38  	Cost int64 `beam:"cost"`
    39  	// The transactions that paid for this purchase.
    40  	// A slice since the purchase might be spread out over multiple
    41  	// credit cards.
    42  	Transactions []Transaction `beam:"transactions"`
    43  }
    44  
    45  type ShippingAddress struct {
    46  	StreetAddress string  `beam:"streetAddress"`
    47  	City          string  `beam:"city"`
    48  	State         *string `beam:"state"`
    49  	Country       string  `beam:"country"`
    50  	PostCode      string  `beam:"postCode"`
    51  }
    52  
    53  type Transaction struct {
    54  	Bank           string  `beam:"bank"`
    55  	PurchaseAmount float64 `beam:"purchaseAmount"`
    56  }
    57  
    58  // [END schema_define]
    59  
    60  // Validate that the interface is being implemented.
    61  var _ beam.SchemaProvider = &TimestampNanosProvider{}
    62  
    63  // [START schema_logical_provider]
    64  
    65  // TimestampNanos is a logical type using time.Time, but
    66  // encodes as a schema type.
    67  type TimestampNanos time.Time
    68  
    69  func (tn TimestampNanos) Seconds() int64 {
    70  	return time.Time(tn).Unix()
    71  }
    72  func (tn TimestampNanos) Nanos() int32 {
    73  	return int32(time.Time(tn).UnixNano() % 1000000000)
    74  }
    75  
    76  // tnStorage is the storage schema for TimestampNanos.
    77  type tnStorage struct {
    78  	Seconds int64 `beam:"seconds"`
    79  	Nanos   int32 `beam:"nanos"`
    80  }
    81  
    82  var (
    83  	// reflect.Type of the Value type of TimestampNanos
    84  	tnType        = reflect.TypeOf((*TimestampNanos)(nil)).Elem()
    85  	tnStorageType = reflect.TypeOf((*tnStorage)(nil)).Elem()
    86  )
    87  
    88  // TimestampNanosProvider implements the beam.SchemaProvider interface.
    89  type TimestampNanosProvider struct{}
    90  
    91  // FromLogicalType converts checks if the given type is TimestampNanos, and if so
    92  // returns the storage type.
    93  func (p *TimestampNanosProvider) FromLogicalType(rt reflect.Type) (reflect.Type, error) {
    94  	if rt != tnType {
    95  		return nil, fmt.Errorf("unable to provide schema.LogicalType for type %v, want %v", rt, tnType)
    96  	}
    97  	return tnStorageType, nil
    98  }
    99  
   100  // BuildEncoder builds a Beam schema encoder for the TimestampNanos type.
   101  func (p *TimestampNanosProvider) BuildEncoder(rt reflect.Type) (func(any, io.Writer) error, error) {
   102  	if _, err := p.FromLogicalType(rt); err != nil {
   103  		return nil, err
   104  	}
   105  	enc, err := coder.RowEncoderForStruct(tnStorageType)
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  	return func(iface any, w io.Writer) error {
   110  		v := iface.(TimestampNanos)
   111  		return enc(tnStorage{
   112  			Seconds: v.Seconds(),
   113  			Nanos:   v.Nanos(),
   114  		}, w)
   115  	}, nil
   116  }
   117  
   118  // BuildDecoder builds a Beam schema decoder for the TimestampNanos type.
   119  func (p *TimestampNanosProvider) BuildDecoder(rt reflect.Type) (func(io.Reader) (any, error), error) {
   120  	if _, err := p.FromLogicalType(rt); err != nil {
   121  		return nil, err
   122  	}
   123  	dec, err := coder.RowDecoderForStruct(tnStorageType)
   124  	if err != nil {
   125  		return nil, err
   126  	}
   127  	return func(r io.Reader) (any, error) {
   128  		s, err := dec(r)
   129  		if err != nil {
   130  			return nil, err
   131  		}
   132  		tn := s.(tnStorage)
   133  		return TimestampNanos(time.Unix(tn.Seconds, int64(tn.Nanos))), nil
   134  	}, nil
   135  }
   136  
   137  // [END schema_logical_provider]
   138  
   139  func LogicalTypeExample() {
   140  	// [START schema_logical_register]
   141  	beam.RegisterSchemaProvider(tnType, &TimestampNanosProvider{})
   142  	// [END schema_logical_register]
   143  }