sdks/go/pkg/beam/core/typex/fulltype.go - beam - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one or more
 // contributor license agreements.  See the NOTICE file distributed with
 // this work for additional information regarding copyright ownership.
 // The ASF licenses this file to You under the Apache License, Version 2.0
 // (the "License"); you may not use this file except in compliance with
 // the License.  You may obtain a copy of the License at
 //
 //    http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // Package typex contains full type representation for PCollections and DoFns, and
 // utilities for type checking.
 package typex

 import (
 	"fmt"
 	"reflect"
 	"strings"

 	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
 )

 // FullType represents the tree structure of data types processed by the graph.
 // It allows representation of composite types, such as KV<int, string> or
 // W<CoGBK<int, int>>, as well as "generic" such types, KV<int,T> or CoGBK<X,Y>,
 // where the free "type variables" are the fixed universal types: T, X, etc.
 type FullType interface {
 	// Class returns the class of the FullType. It is never Illegal.
 	Class() Class
 	// Type returns the Go type at the root of the tree, such as KV, X
 	// or int.
 	Type() reflect.Type
 	// Components returns the real components of the root type, if Composite.
 	Components() []FullType
 }

 type tree struct {
 	class      Class
 	t          reflect.Type
 	components []FullType
 }

 func (t *tree) Class() Class {
 	return t.class
 }

 func (t *tree) Type() reflect.Type {
 	return t.t
 }

 func (t *tree) Components() []FullType {
 	return t.components
 }

 func (t *tree) String() string {
 	switch t.class {
 	case Concrete:
 		return t.t.String()
 	case Universal:
 		return t.t.Name()
 	case Container:
 		if IsList(t.t) {
 			return fmt.Sprintf("[]%v", t.components[0])
 		}
 		return fmt.Sprintf("<invalid: %v>", t.t)
 	case Composite:
 		var args []string
 		for _, c := range t.components {
 			args = append(args, fmt.Sprintf("%v", c))
 		}
 		return fmt.Sprintf("%v<%v>", printShortComposite(t.t), strings.Join(args, ","))
 	default:
 		return fmt.Sprintf("<invalid: %v>", t.t)
 	}
 }

 func printShortComposite(t reflect.Type) string {
 	switch t {
 	case WindowedValueType:
 		return "W"
 	case CoGBKType:
 		return "CoGBK"
 	case KVType:
 		return "KV"
 	default:
 		return fmt.Sprintf("invalid(%v)", t)
 	}
 }

 // TODO(herohde) 4/20/2017: internalize fulltypes?

 // New constructs a new full type with the given elements. It panics
 // if not valid.
 func New(t reflect.Type, components ...FullType) FullType {
 	checkTypesNotNil(components)

 	class := ClassOf(t)
 	switch class {
 	case Concrete, Universal:
 		return &tree{class, t, nil}
 	case Container:
 		switch t.Kind() {
 		case reflect.Slice:
 			// We include the child type as a component for convenience.
 			return &tree{class, t, []FullType{New(t.Elem())}}
 		default:
 			panic(fmt.Sprintf("Unexpected aggregate type: %v", t))
 		}
 	case Composite:
 		switch t {
 		case KVType:
 			if len(components) != 2 {
 				panic("Invalid number of components for KV")
 			}
 			if isAnyNonKVComposite(components) {
 				panic("Invalid to nest composites inside KV")
 			}
 			return &tree{class, t, components}
 		case WindowedValueType:
 			if len(components) != 1 {
 				panic("Invalid number of components for WindowedValue")
 			}
 			if components[0].Type() == WindowedValueType {
 				panic("Invalid to nest WindowedValue")
 			}
 			return &tree{class, t, components}
 		case CoGBKType:
 			if len(components) < 2 {
 				panic("Invalid number of components for CoGBK")
 			}
 			if isAnyNonKVComposite(components) {
 				panic("Invalid to nest composites inside CoGBK")
 			}
 			return &tree{class, t, components}
 		default:
 			panic(fmt.Sprintf("Unexpected composite type: %v", t))
 		}
 	default:
 		panic(fmt.Sprintf("Invalid underlying type: %v", t))
 	}
 }

 // NOTE(herohde) 1/26/2018: we allow nested KV types and coders to support the
 // CoGBK translation (using KV<K,KV<int,[]byte>> to encode keyed raw union values)
 // and potentially other uses. We do not have a reasonable way to emit nested KV
 // values, so user functions are still limited to non-nested KVs. Universally-typed
 // KV-values might be simple to allow, for example.

 func isAnyNonKVComposite(list []FullType) bool {
 	for _, t := range list {
 		if t.Class() == Composite && t.Type() != KVType {
 			return true
 		}
 	}
 	return false
 }

 // Convenience functions.

 // IsW returns true iff the type is a WindowedValue.
 func IsW(t FullType) bool {
 	return t.Type() == WindowedValueType
 }

 // NewW constructs a new WindowedValue of the given type.
 func NewW(t FullType) FullType {
 	return New(WindowedValueType, t)
 }

 // SkipW skips a WindowedValue layer, if present. If no, returns the input.
 func SkipW(t FullType) FullType {
 	if t.Type() == WindowedValueType {
 		return t.Components()[0]
 	}
 	return t
 }

 // SkipK skips the key in a KV layer, if present. If no, returns the input.
 func SkipK(t FullType) FullType {
 	if t.Type() == KVType {
 		return t.Components()[1]
 	}
 	return t
 }

 // IsKV returns true iff the type is a KV.
 func IsKV(t FullType) bool {
 	return t.Type() == KVType
 }

 // NewKV constructs a new KV of the given key and value types.
 func NewKV(components ...FullType) FullType {
 	return New(KVType, components...)
 }

 // IsCoGBK returns true iff the type is a CoGBK.
 func IsCoGBK(t FullType) bool {
 	return t.Type() == CoGBKType
 }

 // NewCoGBK constructs a new CoGBK of the given component types.
 func NewCoGBK(components ...FullType) FullType {
 	return New(CoGBKType, components...)
 }

 // IsStructurallyAssignable returns true iff a from value is structurally
 // assignable to the to value of the given types. Types that are
 // "structurally assignable" (SA) are assignable if type variables are
 // disregarded. In other words, depending on the bindings of universal
 // type variables, types may or may not be assignable. However, types that
 // are not SA are not assignable under any bindings.
 //
 // For example:
 //
 //   SA:  KV<int,int>    := KV<int,int>
 //   SA:  KV<int,X>      := KV<int,string>  // X bound to string by assignment
 //   SA:  KV<int,string> := KV<int,X>       // Assignable only if X is already bound to string
 //   SA:  KV<int,string> := KV<X,X>         // Not assignable under any binding
 //
 //   Not SA:  KV<int,string> := KV<string,X>
 //   Not SA:  X              := KV<int,string>
 //   Not SA:  GBK(X,Y)       := KV<int,string>
 //
 func IsStructurallyAssignable(from, to FullType) bool {
 	switch from.Class() {
 	case Concrete:
 		if to.Class() == Concrete {
 			return from.Type().AssignableTo(to.Type())
 		}
 		return to.Class() == Universal
 	case Universal:
 		// Universals are not structurally assignable to Composites, such as KV or GBK.
 		return to.Class() == Universal || to.Class() == Concrete || to.Class() == Container
 	case Container:
 		if to.Class() == Container {
 			return IsList(from.Type()) && IsList(to.Type()) && IsStructurallyAssignable(from.Components()[0], to.Components()[0])
 		}
 		return to.Class() == Universal
 	case Composite:
 		if from.Type() != to.Type() {
 			return false
 		}
 		if len(from.Components()) != len(to.Components()) {
 			return false
 		}
 		for i, elm := range from.Components() {
 			if !IsStructurallyAssignable(elm, to.Components()[i]) {
 				return false
 			}
 		}
 		return true
 	default:
 		return false
 	}
 }

 // IsEqual returns true iff the types are equal.
 func IsEqual(from, to FullType) bool {
 	if from.Type() != to.Type() {
 		return false
 	}
 	if len(from.Components()) != len(to.Components()) {
 		return false
 	}
 	for i, elm := range from.Components() {
 		if !IsEqual(elm, to.Components()[i]) {
 			return false
 		}
 	}
 	return true
 }

 // IsEqualList returns true iff the lists of types are equal.
 func IsEqualList(from, to []FullType) bool {
 	if len(from) != len(to) {
 		return false
 	}
 	for i, t := range from {
 		if !IsEqual(t, to[i]) {
 			return false
 		}
 	}
 	return true
 }

 // IsBound returns true iff the type has no universal type components.
 // Nodes and coders need bound types.
 func IsBound(t FullType) bool {
 	if t.Class() == Universal {
 		return false
 	}
 	for _, elm := range t.Components() {
 		if !IsBound(elm) {
 			return false
 		}
 	}
 	return true
 }

 // Bind returns a substitution from universals to types in the given models,
 // such as {"T" -> X, "X" -> int}. Each model must be assignable to the
 // corresponding type. For example, Bind(KV<T,int>, KV<string, int>) would
 // produce {"T" -> string}.
 func Bind(types, models []FullType) (map[string]reflect.Type, error) {
 	if len(types) != len(models) {
 		return nil, errors.Errorf("typex.Bind: invalid number of models: %v, want %v", len(models), len(types))
 	}

 	m := make(map[string]reflect.Type)
 	for i := 0; i < len(types); i++ {
 		t := types[i]
 		model := models[i]

 		if !IsStructurallyAssignable(model, t) {
 			return nil, errors.Errorf("typex.Bind: %v is not assignable to %v", model, t)
 		}
 		if err := walk(t, model, m); err != nil {
 			return nil, err
 		}
 	}
 	return m, nil
 }

 func walk(t, model FullType, m map[string]reflect.Type) error {
 	switch t.Class() {
 	case Universal:
 		// By checking that the model is assignable to t, we know that they are
 		// structurally compatible. We rely on the exact reflect.Type in the
 		// Aggregate case to pick the correct binding, i.e., we do not need to
 		// construct such a type.

 		name := t.Type().Name()
 		if current, ok := m[name]; ok && current != model.Type() {
 			return errors.Errorf("bind conflict for %v: %v != %v", name, current, model.Type())
 		}
 		m[name] = model.Type()
 		return nil
 	case Composite, Container:
 		for i, elm := range t.Components() {
 			if err := walk(elm, model.Components()[i], m); err != nil {
 				return err
 			}
 		}
 		return nil
 	default:
 		return nil
 	}
 }

 // Substitute returns types identical to the given types, but with all
 // universals substituted. All free type variables must be present in the
 // substitution.
 func Substitute(list []FullType, m map[string]reflect.Type) ([]FullType, error) {
 	var ret []FullType
 	for _, t := range list {
 		repl, err := substitute(t, m)
 		if err != nil {
 			return nil, err
 		}
 		ret = append(ret, repl)
 	}
 	return ret, nil
 }

 func substitute(t FullType, m map[string]reflect.Type) (FullType, error) {
 	switch t.Class() {
 	case Universal:
 		name := t.Type().Name()
 		repl, ok := m[name]
 		if !ok {
 			return nil, errors.Errorf("substituting type %v: type not bound", name)
 		}
 		return New(repl), nil
 	case Container:
 		comp, err := substituteList(t.Components(), m)
 		if err != nil {
 			return nil, err
 		}
 		if IsList(t.Type()) {
 			return New(reflect.SliceOf(comp[0].Type()), comp...), nil
 		}
 		return nil, errors.Errorf("unexpected aggregate %v, only slices allowed", t)
 	case Composite:
 		comp, err := substituteList(t.Components(), m)
 		if err != nil {
 			return nil, err
 		}
 		return New(t.Type(), comp...), nil

 	default:
 		return t, nil
 	}
 }

 func substituteList(list []FullType, m map[string]reflect.Type) ([]FullType, error) {
 	var ret []FullType
 	for _, elm := range list {
 		repl, err := substitute(elm, m)
 		if err != nil {
 			return nil, err
 		}
 		ret = append(ret, repl)
 	}
 	return ret, nil
 }

 func checkTypesNotNil(list []FullType) {
 	for i, t := range list {
 		if t == nil {
 			panic(fmt.Sprintf("nil type at index: %v", i))
 		}
 	}
 }
	// Licensed to the Apache Software Foundation (ASF) under one or more
	// contributor license agreements. See the NOTICE file distributed with
	// this work for additional information regarding copyright ownership.
	// The ASF licenses this file to You under the Apache License, Version 2.0
	// (the "License"); you may not use this file except in compliance with
	// the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// Package typex contains full type representation for PCollections and DoFns, and
	// utilities for type checking.
	package typex

	import (
	"fmt"
	"reflect"
	"strings"

	"github.com/apache/beam/sdks/go/pkg/beam/internal/errors"
	)

	// FullType represents the tree structure of data types processed by the graph.
	// It allows representation of composite types, such as KV<int, string> or
	// W<CoGBK<int, int>>, as well as "generic" such types, KV<int,T> or CoGBK<X,Y>,
	// where the free "type variables" are the fixed universal types: T, X, etc.
	type FullType interface {
	// Class returns the class of the FullType. It is never Illegal.
	Class() Class
	// Type returns the Go type at the root of the tree, such as KV, X
	// or int.
	Type() reflect.Type
	// Components returns the real components of the root type, if Composite.
	Components() []FullType
	}

	type tree struct {
	class Class
	t reflect.Type
	components []FullType
	}

	func (t *tree) Class() Class {
	return t.class
	}

	func (t *tree) Type() reflect.Type {
	return t.t
	}

	func (t *tree) Components() []FullType {
	return t.components
	}

	func (t *tree) String() string {
	switch t.class {
	case Concrete:
	return t.t.String()
	case Universal:
	return t.t.Name()
	case Container:
	if IsList(t.t) {
	return fmt.Sprintf("[]%v", t.components[0])
	}
	return fmt.Sprintf("<invalid: %v>", t.t)
	case Composite:
	var args []string
	for _, c := range t.components {
	args = append(args, fmt.Sprintf("%v", c))
	}
	return fmt.Sprintf("%v<%v>", printShortComposite(t.t), strings.Join(args, ","))
	default:
	return fmt.Sprintf("<invalid: %v>", t.t)
	}
	}

	func printShortComposite(t reflect.Type) string {
	switch t {
	case WindowedValueType:
	return "W"
	case CoGBKType:
	return "CoGBK"
	case KVType:
	return "KV"
	default:
	return fmt.Sprintf("invalid(%v)", t)
	}
	}

	// TODO(herohde) 4/20/2017: internalize fulltypes?

	// New constructs a new full type with the given elements. It panics
	// if not valid.
	func New(t reflect.Type, components ...FullType) FullType {
	checkTypesNotNil(components)

	class := ClassOf(t)
	switch class {
	case Concrete, Universal:
	return &tree{class, t, nil}
	case Container:
	switch t.Kind() {
	case reflect.Slice:
	// We include the child type as a component for convenience.
	return &tree{class, t, []FullType{New(t.Elem())}}
	default:
	panic(fmt.Sprintf("Unexpected aggregate type: %v", t))
	}
	case Composite:
	switch t {
	case KVType:
	if len(components) != 2 {
	panic("Invalid number of components for KV")
	}
	if isAnyNonKVComposite(components) {
	panic("Invalid to nest composites inside KV")
	}
	return &tree{class, t, components}
	case WindowedValueType:
	if len(components) != 1 {
	panic("Invalid number of components for WindowedValue")
	}
	if components[0].Type() == WindowedValueType {
	panic("Invalid to nest WindowedValue")
	}
	return &tree{class, t, components}
	case CoGBKType:
	if len(components) < 2 {
	panic("Invalid number of components for CoGBK")
	}
	if isAnyNonKVComposite(components) {
	panic("Invalid to nest composites inside CoGBK")
	}
	return &tree{class, t, components}
	default:
	panic(fmt.Sprintf("Unexpected composite type: %v", t))
	}
	default:
	panic(fmt.Sprintf("Invalid underlying type: %v", t))
	}
	}

	// NOTE(herohde) 1/26/2018: we allow nested KV types and coders to support the
	// CoGBK translation (using KV<K,KV<int,[]byte>> to encode keyed raw union values)
	// and potentially other uses. We do not have a reasonable way to emit nested KV
	// values, so user functions are still limited to non-nested KVs. Universally-typed
	// KV-values might be simple to allow, for example.

	func isAnyNonKVComposite(list []FullType) bool {
	for _, t := range list {
	if t.Class() == Composite && t.Type() != KVType {
	return true
	}
	}
	return false
	}

	// Convenience functions.

	// IsW returns true iff the type is a WindowedValue.
	func IsW(t FullType) bool {
	return t.Type() == WindowedValueType
	}

	// NewW constructs a new WindowedValue of the given type.
	func NewW(t FullType) FullType {
	return New(WindowedValueType, t)
	}

	// SkipW skips a WindowedValue layer, if present. If no, returns the input.
	func SkipW(t FullType) FullType {
	if t.Type() == WindowedValueType {
	return t.Components()[0]
	}
	return t
	}

	// SkipK skips the key in a KV layer, if present. If no, returns the input.
	func SkipK(t FullType) FullType {
	if t.Type() == KVType {
	return t.Components()[1]
	}
	return t
	}

	// IsKV returns true iff the type is a KV.
	func IsKV(t FullType) bool {
	return t.Type() == KVType
	}

	// NewKV constructs a new KV of the given key and value types.
	func NewKV(components ...FullType) FullType {
	return New(KVType, components...)
	}

	// IsCoGBK returns true iff the type is a CoGBK.
	func IsCoGBK(t FullType) bool {
	return t.Type() == CoGBKType
	}

	// NewCoGBK constructs a new CoGBK of the given component types.
	func NewCoGBK(components ...FullType) FullType {
	return New(CoGBKType, components...)
	}

	// IsStructurallyAssignable returns true iff a from value is structurally
	// assignable to the to value of the given types. Types that are
	// "structurally assignable" (SA) are assignable if type variables are
	// disregarded. In other words, depending on the bindings of universal
	// type variables, types may or may not be assignable. However, types that
	// are not SA are not assignable under any bindings.
	//
	// For example:
	//
	// SA: KV<int,int> := KV<int,int>
	// SA: KV<int,X> := KV<int,string> // X bound to string by assignment
	// SA: KV<int,string> := KV<int,X> // Assignable only if X is already bound to string
	// SA: KV<int,string> := KV<X,X> // Not assignable under any binding
	//
	// Not SA: KV<int,string> := KV<string,X>
	// Not SA: X := KV<int,string>
	// Not SA: GBK(X,Y) := KV<int,string>
	//
	func IsStructurallyAssignable(from, to FullType) bool {
	switch from.Class() {
	case Concrete:
	if to.Class() == Concrete {
	return from.Type().AssignableTo(to.Type())
	}
	return to.Class() == Universal
	case Universal:
	// Universals are not structurally assignable to Composites, such as KV or GBK.
	return to.Class() == Universal \|\| to.Class() == Concrete \|\| to.Class() == Container
	case Container:
	if to.Class() == Container {
	return IsList(from.Type()) && IsList(to.Type()) && IsStructurallyAssignable(from.Components()[0], to.Components()[0])
	}
	return to.Class() == Universal
	case Composite:
	if from.Type() != to.Type() {
	return false
	}
	if len(from.Components()) != len(to.Components()) {
	return false
	}
	for i, elm := range from.Components() {
	if !IsStructurallyAssignable(elm, to.Components()[i]) {
	return false
	}
	}
	return true
	default:
	return false
	}
	}

	// IsEqual returns true iff the types are equal.
	func IsEqual(from, to FullType) bool {
	if from.Type() != to.Type() {
	return false
	}
	if len(from.Components()) != len(to.Components()) {
	return false
	}
	for i, elm := range from.Components() {
	if !IsEqual(elm, to.Components()[i]) {
	return false
	}
	}
	return true
	}

	// IsEqualList returns true iff the lists of types are equal.
	func IsEqualList(from, to []FullType) bool {
	if len(from) != len(to) {
	return false
	}
	for i, t := range from {
	if !IsEqual(t, to[i]) {
	return false
	}
	}
	return true
	}

	// IsBound returns true iff the type has no universal type components.
	// Nodes and coders need bound types.
	func IsBound(t FullType) bool {
	if t.Class() == Universal {
	return false
	}
	for _, elm := range t.Components() {
	if !IsBound(elm) {
	return false
	}
	}
	return true
	}

	// Bind returns a substitution from universals to types in the given models,
	// such as {"T" -> X, "X" -> int}. Each model must be assignable to the
	// corresponding type. For example, Bind(KV<T,int>, KV<string, int>) would
	// produce {"T" -> string}.
	func Bind(types, models []FullType) (map[string]reflect.Type, error) {
	if len(types) != len(models) {
	return nil, errors.Errorf("typex.Bind: invalid number of models: %v, want %v", len(models), len(types))
	}

	m := make(map[string]reflect.Type)
	for i := 0; i < len(types); i++ {
	t := types[i]
	model := models[i]

	if !IsStructurallyAssignable(model, t) {
	return nil, errors.Errorf("typex.Bind: %v is not assignable to %v", model, t)
	}
	if err := walk(t, model, m); err != nil {
	return nil, err
	}
	}
	return m, nil
	}

	func walk(t, model FullType, m map[string]reflect.Type) error {
	switch t.Class() {
	case Universal:
	// By checking that the model is assignable to t, we know that they are
	// structurally compatible. We rely on the exact reflect.Type in the
	// Aggregate case to pick the correct binding, i.e., we do not need to
	// construct such a type.

	name := t.Type().Name()
	if current, ok := m[name]; ok && current != model.Type() {
	return errors.Errorf("bind conflict for %v: %v != %v", name, current, model.Type())
	}
	m[name] = model.Type()
	return nil
	case Composite, Container:
	for i, elm := range t.Components() {
	if err := walk(elm, model.Components()[i], m); err != nil {
	return err
	}
	}
	return nil
	default:
	return nil
	}
	}

	// Substitute returns types identical to the given types, but with all
	// universals substituted. All free type variables must be present in the
	// substitution.
	func Substitute(list []FullType, m map[string]reflect.Type) ([]FullType, error) {
	var ret []FullType
	for _, t := range list {
	repl, err := substitute(t, m)
	if err != nil {
	return nil, err
	}
	ret = append(ret, repl)
	}
	return ret, nil
	}

	func substitute(t FullType, m map[string]reflect.Type) (FullType, error) {
	switch t.Class() {
	case Universal:
	name := t.Type().Name()
	repl, ok := m[name]
	if !ok {
	return nil, errors.Errorf("substituting type %v: type not bound", name)
	}
	return New(repl), nil
	case Container:
	comp, err := substituteList(t.Components(), m)
	if err != nil {
	return nil, err
	}
	if IsList(t.Type()) {
	return New(reflect.SliceOf(comp[0].Type()), comp...), nil
	}
	return nil, errors.Errorf("unexpected aggregate %v, only slices allowed", t)
	case Composite:
	comp, err := substituteList(t.Components(), m)
	if err != nil {
	return nil, err
	}
	return New(t.Type(), comp...), nil

	default:
	return t, nil
	}
	}

	func substituteList(list []FullType, m map[string]reflect.Type) ([]FullType, error) {
	var ret []FullType
	for _, elm := range list {
	repl, err := substitute(elm, m)
	if err != nil {
	return nil, err
	}
	ret = append(ret, repl)
	}
	return ret, nil
	}

	func checkTypesNotNil(list []FullType) {
	for i, t := range list {
	if t == nil {
	panic(fmt.Sprintf("nil type at index: %v", i))
	}
	}
	}