[Go SDK] Convert logical type arguments to and from schema protos

sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go

@@ -0,0 +1,156 @@
+// 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 schema
+
+import (
+	"math"
+	"reflect"
+
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+)
+
+// atomicFieldType returns the schema FieldType of an atomic Go type.
+// The Go type must be one of the built in types that map directly to an
+// AtomicType: bool, uint8, int16, int32, int64, float32, float64, string
+// and []byte.
+func atomicFieldType(t reflect.Type) (*pipepb.FieldType, error) {
+	var at pipepb.AtomicType
+	if t == reflectx.ByteSlice {
+		at = pipepb.AtomicType_BYTES
+	} else if enum, ok := reflectTypeToAtomicTypeMap[t.Kind()]; ok && reflectKindToTypeMap[t.Kind()] == t {
+		at = enum
+	} else {
+		return nil, errors.Errorf("%v is not an atomic schema type", t)
+	}
+	return &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_AtomicType{
+			AtomicType: at,
+		},
+	}, nil
+}
+
+// atomicValueToProto converts a Go value of an atomic type to its schema
+// FieldType and FieldValue. The BYTE schema type is signed, so a uint8 is
+// converted through int8.
+func atomicValueToProto(rv reflect.Value) (*pipepb.FieldType, *pipepb.FieldValue, error) {
+	if !rv.IsValid() {
+		return nil, nil, errors.New("invalid value")
+	}
+	ft, err := atomicFieldType(rv.Type())
+	if err != nil {
+		return nil, nil, err
+	}
+	var av *pipepb.AtomicTypeValue
+	switch ft.GetAtomicType() {
+	case pipepb.AtomicType_BYTE:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Byte{Byte: int32(int8(rv.Uint()))}}
+	case pipepb.AtomicType_INT16:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int16{Int16: int32(rv.Int())}}
+	case pipepb.AtomicType_INT32:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int32{Int32: int32(rv.Int())}}
+	case pipepb.AtomicType_INT64:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int64{Int64: rv.Int()}}
+	case pipepb.AtomicType_FLOAT:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Float{Float: float32(rv.Float())}}
+	case pipepb.AtomicType_DOUBLE:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Double{Double: rv.Float()}}
+	case pipepb.AtomicType_STRING:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_String_{String_: rv.String()}}
+	case pipepb.AtomicType_BOOLEAN:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Boolean{Boolean: rv.Bool()}}
+	case pipepb.AtomicType_BYTES:
+		av = &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Bytes{Bytes: rv.Bytes()}}
+	default:
+		return nil, nil, errors.Errorf("unhandled atomic type %v", ft.GetAtomicType())
+	}
+	return ft, &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{AtomicValue: av}}, nil
+}
+
+// atomicValueFromProto converts a schema FieldValue of the atomic FieldType ft
+// to a Go value of the matching built in type.
+func atomicValueFromProto(ft *pipepb.FieldType, fv *pipepb.FieldValue) (reflect.Value, error) {
+	at, ok := ft.GetTypeInfo().(*pipepb.FieldType_AtomicType)
+	if !ok {
+		return reflect.Value{}, errors.Errorf("field type %v is not atomic", ft)
+	}
+	t, ok := atomicTypeToReflectType[at.AtomicType]
+	if !ok {
+		return reflect.Value{}, errors.Errorf("unknown atomic type: %v", at.AtomicType)
+	}
+	av := fv.GetAtomicValue()
+	if av == nil {
+		return reflect.Value{}, errors.Errorf("field value %v is not atomic", fv)
+	}
+	if got := atomicValueType(av); got != at.AtomicType {
+		return reflect.Value{}, errors.Errorf("atomic value %v does not match field type %v", av, at.AtomicType)
+	}
+	rv := reflect.New(t).Elem()
+	switch v := av.GetValue().(type) {
+	case *pipepb.AtomicTypeValue_Byte:
+		if v.Byte < math.MinInt8 || v.Byte > math.MaxInt8 {
+			return reflect.Value{}, errors.Errorf("BYTE value %v is out of range", v.Byte)
+		}
+		rv.SetUint(uint64(uint8(v.Byte)))
+	case *pipepb.AtomicTypeValue_Int16:
+		if v.Int16 < math.MinInt16 || v.Int16 > math.MaxInt16 {
+			return reflect.Value{}, errors.Errorf("INT16 value %v is out of range", v.Int16)
+		}
+		rv.SetInt(int64(v.Int16))
+	case *pipepb.AtomicTypeValue_Int32:
+		rv.SetInt(int64(v.Int32))
+	case *pipepb.AtomicTypeValue_Int64:
+		rv.SetInt(v.Int64)
+	case *pipepb.AtomicTypeValue_Float:
+		rv.SetFloat(float64(v.Float))
+	case *pipepb.AtomicTypeValue_Double:
+		rv.SetFloat(v.Double)
+	case *pipepb.AtomicTypeValue_String_:
+		rv.SetString(v.String_)
+	case *pipepb.AtomicTypeValue_Boolean:
+		rv.SetBool(v.Boolean)
+	case *pipepb.AtomicTypeValue_Bytes:
+		rv.SetBytes(v.Bytes)
+	}
+	return rv, nil
+}
+
+// atomicValueType returns the atomic type of the value set in av, or
+// UNSPECIFIED when no value is set.
+func atomicValueType(av *pipepb.AtomicTypeValue) pipepb.AtomicType {
+	switch av.GetValue().(type) {
+	case *pipepb.AtomicTypeValue_Byte:
+		return pipepb.AtomicType_BYTE
+	case *pipepb.AtomicTypeValue_Int16:
+		return pipepb.AtomicType_INT16
+	case *pipepb.AtomicTypeValue_Int32:
+		return pipepb.AtomicType_INT32
+	case *pipepb.AtomicTypeValue_Int64:
+		return pipepb.AtomicType_INT64
+	case *pipepb.AtomicTypeValue_Float:
+		return pipepb.AtomicType_FLOAT
+	case *pipepb.AtomicTypeValue_Double:
+		return pipepb.AtomicType_DOUBLE
+	case *pipepb.AtomicTypeValue_String_:
+		return pipepb.AtomicType_STRING
+	case *pipepb.AtomicTypeValue_Boolean:
+		return pipepb.AtomicType_BOOLEAN
+	case *pipepb.AtomicTypeValue_Bytes:
+		return pipepb.AtomicType_BYTES
+	}
+	return pipepb.AtomicType_UNSPECIFIED
+}

sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go

@@ -0,0 +1,163 @@
+// 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 schema
+
+import (
+	"math"
+	"reflect"
+	"testing"
+
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+	"github.com/google/go-cmp/cmp"
+)
+
+func TestAtomicValueProtoRoundTrip(t *testing.T) {
+	tests := []struct {
+		value any
+		want  pipepb.AtomicType
+	}{
+		{true, pipepb.AtomicType_BOOLEAN},
+		{uint8(200), pipepb.AtomicType_BYTE},
+		{int16(-300), pipepb.AtomicType_INT16},
+		{int32(70000), pipepb.AtomicType_INT32},
+		{int64(-1 << 40), pipepb.AtomicType_INT64},
+		{float32(1.5), pipepb.AtomicType_FLOAT},
+		{float64(2.25), pipepb.AtomicType_DOUBLE},
+		{"text", pipepb.AtomicType_STRING},
+		{[]byte("bytes"), pipepb.AtomicType_BYTES},
+	}
+	for _, test := range tests {
+		t.Run(reflect.TypeOf(test.value).String(), func(t *testing.T) {
+			ft, fv, err := atomicValueToProto(reflect.ValueOf(test.value))
+			if err != nil {
+				t.Fatalf("atomicValueToProto(%v) = %v, want nil error", test.value, err)
+			}
+			if got := ft.GetAtomicType(); got != test.want {
+				t.Errorf("atomicValueToProto(%v) field type = %v, want %v", test.value, got, test.want)
+			}
+			got, err := atomicValueFromProto(ft, fv)
+			if err != nil {
+				t.Fatalf("atomicValueFromProto(%v, %v) = %v, want nil error", ft, fv, err)
+			}
+			if d := cmp.Diff(test.value, got.Interface()); d != "" {
+				t.Errorf("round trip of %v: diff (-want, +got): %v", test.value, d)
+			}
+		})
+	}
+}
+
+func TestAtomicValueToProto_Unsupported(t *testing.T) {
+	type named int32
+	for _, value := range []any{int(1), uint32(1), named(1), struct{ A int32 }{}, []int32{1}} {
+		if _, _, err := atomicValueToProto(reflect.ValueOf(value)); err == nil {
+			t.Errorf("atomicValueToProto(%T) = nil error, want error", value)
+		}
+	}
+}
+
+func TestAtomicValueFromProto_Mismatch(t *testing.T) {
+	stringType := &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_STRING},
+	}
+	int32Value := &pipepb.FieldValue{
+		FieldValue: &pipepb.FieldValue_AtomicValue{
+			AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int32{Int32: 3}},
+		},
+	}
+	if _, err := atomicValueFromProto(stringType, int32Value); err == nil {
+		t.Errorf("atomicValueFromProto(STRING, int32 value) = nil error, want error")
+	}
+	rowType := &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{}}},
+	}
+	if _, err := atomicValueFromProto(rowType, int32Value); err == nil {
+		t.Errorf("atomicValueFromProto(ROW, int32 value) = nil error, want error")
+	}
+	if _, err := atomicValueFromProto(stringType, &pipepb.FieldValue{}); err == nil {
+		t.Errorf("atomicValueFromProto(STRING, null value) = nil error, want error")
+	}
+}
+
+func TestAtomicValueToProto_ByteIsSigned(t *testing.T) {
+	_, fv, err := atomicValueToProto(reflect.ValueOf(uint8(200)))
+	if err != nil {
+		t.Fatalf("atomicValueToProto(uint8(200)) = %v, want nil error", err)
+	}
+	if got := fv.GetAtomicValue().GetByte(); got != -56 {
+		t.Errorf("atomicValueToProto(uint8(200)) byte = %v, want -56", got)
+	}
+	byteType := &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_BYTE},
+	}
+	got, err := atomicValueFromProto(byteType, fv)
+	if err != nil {
+		t.Fatalf("atomicValueFromProto(BYTE, -56) = %v, want nil error", err)
+	}
+	if got.Interface() != uint8(200) {
+		t.Errorf("atomicValueFromProto(BYTE, -56) = %v, want uint8(200)", got)
+	}
+}
+
+func TestAtomicValueFromProto_Range(t *testing.T) {
+	byteType := &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_BYTE},
+	}
+	int16Type := &pipepb.FieldType{
+		TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_INT16},
+	}
+	byteValue := func(v int32) *pipepb.FieldValue {
+		return &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{
+			AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Byte{Byte: v}},
+		}}
+	}
+	int16Value := func(v int32) *pipepb.FieldValue {
+		return &pipepb.FieldValue{FieldValue: &pipepb.FieldValue_AtomicValue{
+			AtomicValue: &pipepb.AtomicTypeValue{Value: &pipepb.AtomicTypeValue_Int16{Int16: v}},
+		}}
+	}
+	tests := []struct {
+		name string
+		ft   *pipepb.FieldType
+		fv   *pipepb.FieldValue
+		want any // nil when the value is out of range.
+	}{
+		{"BYTE min", byteType, byteValue(math.MinInt8), uint8(0x80)},
+		{"BYTE max", byteType, byteValue(math.MaxInt8), uint8(0x7f)},
+		{"BYTE below min", byteType, byteValue(math.MinInt8 - 1), nil},
+		{"BYTE above max", byteType, byteValue(256), nil},
+		{"INT16 min", int16Type, int16Value(math.MinInt16), int16(math.MinInt16)},
+		{"INT16 max", int16Type, int16Value(math.MaxInt16), int16(math.MaxInt16)},
+		{"INT16 below min", int16Type, int16Value(math.MinInt16 - 1), nil},
+		{"INT16 above max", int16Type, int16Value(65537), nil},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			got, err := atomicValueFromProto(test.ft, test.fv)
+			if test.want == nil {
+				if err == nil {
+					t.Fatalf("atomicValueFromProto(%v) = %v, want error", test.fv, got)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("atomicValueFromProto(%v) = %v, want nil error", test.fv, err)
+			}
+			if got.Interface() != test.want {
+				t.Errorf("atomicValueFromProto(%v) = %v, want %v", test.fv, got, test.want)
+			}
+		})
+	}
+}