diff --git a/CHANGES.md b/CHANGES.md index aecd194e1f22..63e3dc66e4cb 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -69,10 +69,12 @@ ## New Features / Improvements * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). +* (Go) Added `schema.RegisterLogicalTypeConversion` and `schema.RegisterLogicalTypeCoder` to register portable logical types in one call, and Go types for the standard `beam:logical_type:date:v1`, `beam:logical_type:micros_instant:v1`, `beam:logical_type:millis_instant:v1`, `beam:logical_type:decimal:v1`, `beam:logical_type:fixed_char:v1`, `beam:logical_type:var_char:v1`, `beam:logical_type:fixed_bytes:v1`, `beam:logical_type:var_bytes:v1` and `beam:logical_type:timestamp:v1` (precisions 3, 6 and 9) logical types, so schemas from other SDKs with these types can be used in cross-language pipelines ([#39684](https://github.com/apache/beam/issues/39684)). ## Breaking Changes * X behavior was changed ([#X](https://github.com/apache/beam/issues/X)). +* (Go) The row coder now encodes `int16` and `uint16` struct fields as 2 byte big endian INT16 values, matching the Java and Python SDKs. Rows with such fields written by earlier Go SDK versions decode differently with this version ([#40151](https://github.com/apache/beam/issues/40151)). ## Deprecations diff --git a/sdks/go/pkg/beam/core/graph/coder/int.go b/sdks/go/pkg/beam/core/graph/coder/int.go index 6eda98ccc38f..8ef56b015ac5 100644 --- a/sdks/go/pkg/beam/core/graph/coder/int.go +++ b/sdks/go/pkg/beam/core/graph/coder/int.go @@ -69,3 +69,34 @@ func DecodeInt32(r io.Reader) (int32, error) { } return int32(ret), nil } + +// EncodeUint16 encodes an uint16 in big endian format. +func EncodeUint16(value uint16, w io.Writer) error { + var data [2]byte + binary.BigEndian.PutUint16(data[:], value) + _, err := ioutilx.WriteUnsafe(w, data[:]) + return err +} + +// DecodeUint16 decodes an uint16 in big endian format. +func DecodeUint16(r io.Reader) (uint16, error) { + var data [2]byte + if err := ioutilx.ReadNBufUnsafe(r, data[:]); err != nil { + return 0, err + } + return binary.BigEndian.Uint16(data[:]), nil +} + +// EncodeInt16 encodes an int16 in big endian format. +func EncodeInt16(value int16, w io.Writer) error { + return EncodeUint16(uint16(value), w) +} + +// DecodeInt16 decodes an int16 in big endian format. +func DecodeInt16(r io.Reader) (int16, error) { + ret, err := DecodeUint16(r) + if err != nil { + return 0, err + } + return int16(ret), nil +} diff --git a/sdks/go/pkg/beam/core/graph/coder/int_test.go b/sdks/go/pkg/beam/core/graph/coder/int_test.go index 25dc681710d0..737340f16632 100644 --- a/sdks/go/pkg/beam/core/graph/coder/int_test.go +++ b/sdks/go/pkg/beam/core/graph/coder/int_test.go @@ -86,3 +86,35 @@ func TestEncodeDecodeInt32(t *testing.T) { } } } + +func TestEncodeDecodeInt16(t *testing.T) { + tests := []struct { + value int16 + want []byte + }{ + {-32768, []byte{0x80, 0x00}}, + {-2, []byte{0xff, 0xfe}}, + {-1, []byte{0xff, 0xff}}, + {0, []byte{0x00, 0x00}}, + {1, []byte{0x00, 0x01}}, + {999, []byte{0x03, 0xe7}}, + {32767, []byte{0x7f, 0xff}}, + } + + for _, test := range tests { + var buf bytes.Buffer + if err := EncodeInt16(test.value, &buf); err != nil { + t.Fatalf("EncodeInt16(%v) failed: %v", test.value, err) + } + if got := buf.Bytes(); !bytes.Equal(got, test.want) { + t.Errorf("EncodeInt16(%v) = %v, want %v", test.value, got, test.want) + } + actual, err := DecodeInt16(&buf) + if err != nil { + t.Fatalf("DecodeInt16(<%v>) failed: %v", test.value, err) + } + if actual != test.value { + t.Errorf("DecodeInt16(<%v>) = %v, want %v", test.value, actual, test.value) + } + } +} diff --git a/sdks/go/pkg/beam/core/graph/coder/row.go b/sdks/go/pkg/beam/core/graph/coder/row.go index f958d84154c1..0c2a6872dd69 100644 --- a/sdks/go/pkg/beam/core/graph/coder/row.go +++ b/sdks/go/pkg/beam/core/graph/coder/row.go @@ -226,3 +226,19 @@ func ReadSimpleRowHeader(fields int, r io.Reader) error { } return nil } + +// RowFieldEncoderForType returns an encoding function for values of the given +// type, using the encoding of that type as a field of a beam schema row. +// +// Returns an error if the given type cannot be a field of a beam schema row. +func RowFieldEncoderForType(rt reflect.Type) (func(any, io.Writer) error, error) { + return defaultEnc.fieldEncoderForType(rt) +} + +// RowFieldDecoderForType returns a decoding function for values of the given +// type, using the encoding of that type as a field of a beam schema row. +// +// Returns an error if the given type cannot be a field of a beam schema row. +func RowFieldDecoderForType(rt reflect.Type) (func(io.Reader) (any, error), error) { + return defaultDec.fieldDecoderForType(rt) +} diff --git a/sdks/go/pkg/beam/core/graph/coder/row_decoder.go b/sdks/go/pkg/beam/core/graph/coder/row_decoder.go index 8dabc2be9ff1..92fcc9a07b8a 100644 --- a/sdks/go/pkg/beam/core/graph/coder/row_decoder.go +++ b/sdks/go/pkg/beam/core/graph/coder/row_decoder.go @@ -241,6 +241,24 @@ func reflectDecodeUint(rv reflect.Value, r io.Reader) error { return nil } +func reflectDecodeInt16(rv reflect.Value, r io.Reader) error { + v, err := DecodeInt16(r) + if err != nil { + return errors.Wrap(err, "error decoding int16 field") + } + rv.SetInt(int64(v)) + return nil +} + +func reflectDecodeUint16(rv reflect.Value, r io.Reader) error { + v, err := DecodeUint16(r) + if err != nil { + return errors.Wrap(err, "error decoding uint16 field") + } + rv.SetUint(uint64(v)) + return nil +} + func reflectDecodeSinglePrecisionFloat(rv reflect.Value, r io.Reader) error { v, err := DecodeSinglePrecisionFloat(r) if err != nil { @@ -341,10 +359,14 @@ func (b *RowDecoderBuilder) decoderForSingleTypeReflect(t reflect.Type) (typeDec return typeDecoderFieldReflect{decode: reflectDecodeByte}, nil case reflect.String: return typeDecoderFieldReflect{decode: reflectDecodeString}, nil - case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64: + case reflect.Int, reflect.Int8, reflect.Int32, reflect.Int64: return typeDecoderFieldReflect{decode: reflectDecodeInt}, nil - case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16: + case reflect.Int16: + return typeDecoderFieldReflect{decode: reflectDecodeInt16}, nil + case reflect.Uint, reflect.Uint64, reflect.Uint32: return typeDecoderFieldReflect{decode: reflectDecodeUint}, nil + case reflect.Uint16: + return typeDecoderFieldReflect{decode: reflectDecodeUint16}, nil case reflect.Float32: return typeDecoderFieldReflect{decode: reflectDecodeSinglePrecisionFloat}, nil case reflect.Float64: @@ -413,3 +435,23 @@ type typeDecoderFieldReflect struct { // of the field value (i.e. &foo.bar) and not the field value (i.e. foo.bar). addr bool } + +// fieldDecoderForType returns a decoder for values of t, using the encoding +// of t as a row field. +func (b *RowDecoderBuilder) fieldDecoderForType(t reflect.Type) (func(io.Reader) (any, error), error) { + decf, err := b.decoderForSingleTypeReflect(t) + if err != nil { + return nil, err + } + return func(r io.Reader) (any, error) { + rv := reflect.New(t) + v := rv.Elem() + if decf.addr { + v = rv + } + if err := decf.decode(v, r); err != nil { + return nil, err + } + return rv.Elem().Interface(), nil + }, nil +} diff --git a/sdks/go/pkg/beam/core/graph/coder/row_encoder.go b/sdks/go/pkg/beam/core/graph/coder/row_encoder.go index 2756c9b8e706..515041309a0c 100644 --- a/sdks/go/pkg/beam/core/graph/coder/row_encoder.go +++ b/sdks/go/pkg/beam/core/graph/coder/row_encoder.go @@ -200,14 +200,24 @@ func (b *RowEncoderBuilder) encoderForSingleTypeReflect(t reflect.Type) (typeEnc return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { return EncodeStringUTF8(rv.String(), w) }}, nil - case reflect.Int, reflect.Int64, reflect.Int16, reflect.Int32, reflect.Int8: + case reflect.Int, reflect.Int64, reflect.Int32, reflect.Int8: return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { return EncodeVarInt(rv.Int(), w) }}, nil - case reflect.Uint, reflect.Uint64, reflect.Uint32, reflect.Uint16: + case reflect.Int16: + // INT16 schema fields are big endian, matching the Java and Python row coders. + return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { + return EncodeInt16(int16(rv.Int()), w) + }}, nil + case reflect.Uint, reflect.Uint64, reflect.Uint32: return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { return EncodeVarUint64(rv.Uint(), w) }}, nil + case reflect.Uint16: + // uint16 is stored as an INT16 schema field. + return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { + return EncodeUint16(uint16(rv.Uint()), w) + }}, nil case reflect.Float32: return typeEncoderFieldReflect{encode: func(rv reflect.Value, w io.Writer) error { return EncodeSinglePrecisionFloat(float32(rv.Float()), w) @@ -353,3 +363,21 @@ func (b *RowEncoderBuilder) encoderForStructReflect(t reflect.Type) (func(reflec return nil }, nil } + +// fieldEncoderForType returns an encoder for values of t, using the encoding +// of t as a row field. +func (b *RowEncoderBuilder) fieldEncoderForType(t reflect.Type) (func(any, io.Writer) error, error) { + encf, err := b.encoderForSingleTypeReflect(t) + if err != nil { + return nil, err + } + return func(v any, w io.Writer) error { + rv := reflect.ValueOf(v) + if encf.addr { + p := reflect.New(t) + p.Elem().Set(rv) + rv = p + } + return encf.encode(rv, w) + }, nil +} diff --git a/sdks/go/pkg/beam/core/graph/coder/row_test.go b/sdks/go/pkg/beam/core/graph/coder/row_test.go index 1549123558bb..4ab53a6a06c1 100644 --- a/sdks/go/pkg/beam/core/graph/coder/row_test.go +++ b/sdks/go/pkg/beam/core/graph/coder/row_test.go @@ -905,3 +905,96 @@ func TestRowHeader_TrailingZeroBytes(t *testing.T) { } } + +// TestReflectionRowCoderGeneration_Int16BigEndian checks that int16 and uint16 +// fields use the 2 byte big endian encoding of the INT16 schema type, which is +// what the Java and Python row coders produce. +func TestReflectionRowCoderGeneration_Int16BigEndian(t *testing.T) { + type row struct { + A int16 + B uint16 + C int16 + } + want := row{A: 999, B: 65535, C: -2} + wantBytes := []byte{ + 0x03, 0x00, // 3 fields, no nils. + 0x03, 0xe7, // A + 0xff, 0xff, // B + 0xff, 0xfe, // C + } + rt := reflect.TypeOf(want) + enc, err := RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(want, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", want, err) + } + if got := buf.Bytes(); !bytes.Equal(got, wantBytes) { + t.Fatalf("enc(%v) = %v, want %v", want, got, wantBytes) + } + dec, err := RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBuffer(wantBytes)) + if err != nil { + t.Fatalf("dec(%v) = %v, want nil error", wantBytes, err) + } + if d := cmp.Diff(want, got); d != "" { + t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", want, got, d) + } +} + +func TestRowFieldCoderForType(t *testing.T) { + type pair struct { + A int64 + B string + } + tests := []struct { + value any + wantBytes []byte + }{ + {value: int64(999), wantBytes: []byte{0xe7, 0x07}}, + {value: int16(999), wantBytes: []byte{0x03, 0xe7}}, + {value: "abc", wantBytes: []byte{0x03, 'a', 'b', 'c'}}, + {value: []byte{1, 2}, wantBytes: []byte{0x02, 1, 2}}, + {value: true, wantBytes: []byte{0x01}}, + {value: []int32{1, 2}, wantBytes: []byte{0x00, 0x00, 0x00, 0x02, 0x01, 0x02}}, + {value: pair{A: 7, B: "x"}, wantBytes: []byte{0x02, 0x00, 0x07, 0x01, 'x'}}, + } + for _, test := range tests { + rt := reflect.TypeOf(test.value) + t.Run(rt.String(), func(t *testing.T) { + enc, err := RowFieldEncoderForType(rt) + if err != nil { + t.Fatalf("RowFieldEncoderForType(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(test.value, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", test.value, err) + } + if got := buf.Bytes(); !bytes.Equal(got, test.wantBytes) { + t.Fatalf("enc(%v) = %v, want %v", test.value, got, test.wantBytes) + } + dec, err := RowFieldDecoderForType(rt) + if err != nil { + t.Fatalf("RowFieldDecoderForType(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBuffer(test.wantBytes)) + if err != nil { + t.Fatalf("dec(%v) = %v, want nil error", test.wantBytes, err) + } + if d := cmp.Diff(test.value, got); d != "" { + t.Fatalf("dec(enc(%v)) = %v\ndiff (-want, +got): %v", test.value, got, d) + } + }) + } + if _, err := RowFieldEncoderForType(reflect.TypeOf(func() {})); err == nil { + t.Errorf("RowFieldEncoderForType(func) = nil error, want error") + } + if _, err := RowFieldDecoderForType(reflect.TypeOf(func() {})); err == nil { + t.Errorf("RowFieldDecoderForType(func) = nil error, want error") + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal.go new file mode 100644 index 000000000000..2ea08a337bc3 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal.go @@ -0,0 +1,134 @@ +// 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 ( + "io" + "math" + "math/big" + "strings" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" +) + +// URNDecimal identifies the Decimal logical type. +const URNDecimal = "beam:logical_type:decimal:v1" + +// Decimal is an arbitrary precision decimal number, the +// beam:logical_type:decimal:v1 logical type. The value is Unscaled times ten +// to the power of minus Scale, so Decimal{Unscaled: 31415, Scale: 4} is +// 3.1415. A nil Unscaled is zero. +// +// The representation type is BYTES, encoded like the Java SDK's +// BigDecimalCoder: the scale as a varint, followed by the length prefixed two's +// complement big endian bytes of the unscaled value. Decimal can be a field of +// a row, but has no schema of its own and cannot be a top level element type. +type Decimal struct { + Unscaled *big.Int + Scale int32 +} + +// ParseDecimal parses a decimal in plain notation with an optional sign, such +// as -100.123. Exponents are not supported. +func ParseDecimal(s string) (Decimal, error) { + sign, body := "", s + if strings.HasPrefix(body, "+") || strings.HasPrefix(body, "-") { + sign, body = body[:1], body[1:] + } + digits, frac, _ := strings.Cut(body, ".") + all := digits + frac + if all == "" || strings.Trim(all, "0123456789") != "" { + return Decimal{}, errors.Errorf("invalid decimal %q", s) + } + unscaled, _ := new(big.Int).SetString(sign+all, 10) + return Decimal{Unscaled: unscaled, Scale: int32(len(frac))}, nil +} + +// String returns the decimal in plain notation, such as -100.123. +func (d Decimal) String() string { + unscaled := d.Unscaled + if unscaled == nil { + unscaled = new(big.Int) + } + digits := new(big.Int).Abs(unscaled).String() + if d.Scale <= 0 { + digits += strings.Repeat("0", int(-d.Scale)) + } else { + scale := int(d.Scale) + if len(digits) <= scale { + digits = strings.Repeat("0", scale-len(digits)+1) + digits + } + digits = digits[:len(digits)-scale] + "." + digits[len(digits)-scale:] + } + if unscaled.Sign() < 0 { + return "-" + digits + } + return digits +} + +func encodeDecimal(d Decimal, w io.Writer) error { + if err := coder.EncodeVarInt(int64(d.Scale), w); err != nil { + return err + } + return coder.EncodeBytes(twosComplementBytes(d.Unscaled), w) +} + +func decodeDecimal(r io.Reader) (Decimal, error) { + scale, err := coder.DecodeVarInt(r) + if err != nil { + return Decimal{}, errors.Wrap(err, "decoding decimal scale") + } + if scale < math.MinInt32 || scale > math.MaxInt32 { + return Decimal{}, errors.Errorf("decimal scale %v out of range", scale) + } + b, err := coder.DecodeBytes(r) + if err != nil { + return Decimal{}, errors.Wrap(err, "decoding decimal unscaled value") + } + return Decimal{Unscaled: bigIntFromTwosComplement(b), Scale: int32(scale)}, nil +} + +// twosComplementBytes returns the minimal big endian two's complement bytes +// of x, as java.math.BigInteger.toByteArray does. A nil x is zero. +func twosComplementBytes(x *big.Int) []byte { + if x == nil || x.Sign() == 0 { + return []byte{0} + } + if x.Sign() > 0 { + b := x.Bytes() + if b[0]&0x80 != 0 { + b = append([]byte{0}, b...) + } + return b + } + // The length is the bits of |x| - 1 plus a sign bit, rounded up to bytes. + // The bytes hold 2^(8n) + x. + abs := new(big.Int).Neg(x) + n := new(big.Int).Sub(abs, big.NewInt(1)).BitLen()/8 + 1 + v := new(big.Int).Lsh(big.NewInt(1), uint(8*n)) + return v.Add(v, x).Bytes() +} + +// bigIntFromTwosComplement returns the value of the big endian two's +// complement bytes b. +func bigIntFromTwosComplement(b []byte) *big.Int { + x := new(big.Int).SetBytes(b) + if len(b) > 0 && b[0]&0x80 != 0 { + x.Sub(x, new(big.Int).Lsh(big.NewInt(1), uint(8*len(b)))) + } + return x +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal_test.go new file mode 100644 index 000000000000..5320263b22f5 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/decimal_test.go @@ -0,0 +1,174 @@ +// 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 ( + "bytes" + "math/big" + "reflect" + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + "github.com/google/go-cmp/cmp" +) + +func decimalOf(unscaled string, scale int32) Decimal { + u, ok := new(big.Int).SetString(unscaled, 10) + if !ok { + panic("invalid unscaled value " + unscaled) + } + return Decimal{Unscaled: u, Scale: scale} +} + +func decimalEqual(a, b Decimal) bool { + zero := new(big.Int) + ua, ub := a.Unscaled, b.Unscaled + if ua == nil { + ua = zero + } + if ub == nil { + ub = zero + } + return a.Scale == b.Scale && ua.Cmp(ub) == 0 +} + +func TestDecimal_StringParse(t *testing.T) { + tests := []struct { + decimal Decimal + text string + }{ + {decimalOf("0", 1), "0.0"}, + {decimalOf("10", 1), "1.0"}, + {decimalOf("31415", 4), "3.1415"}, + {decimalOf("-100123", 3), "-100.123"}, + {decimalOf("5", 3), "0.005"}, + {decimalOf("-5", 3), "-0.005"}, + {decimalOf("42", 0), "42"}, + {decimalOf("1180591620717411303424", 0), "1180591620717411303424"}, + } + for _, test := range tests { + if got := test.decimal.String(); got != test.text { + t.Errorf("%v.String() = %q, want %q", test.decimal, got, test.text) + } + got, err := ParseDecimal(test.text) + if err != nil { + t.Fatalf("ParseDecimal(%q) = %v, want nil error", test.text, err) + } + if !decimalEqual(got, test.decimal) { + t.Errorf("ParseDecimal(%q) = %v, want %v", test.text, got, test.decimal) + } + } + if got, want := (Decimal{Unscaled: big.NewInt(42), Scale: -3}).String(), "42000"; got != want { + t.Errorf("negative scale String() = %q, want %q", got, want) + } + if got, want := (Decimal{}).String(), "0"; got != want { + t.Errorf("zero value String() = %q, want %q", got, want) + } + if got, err := ParseDecimal("+7.5"); err != nil || !decimalEqual(got, decimalOf("75", 1)) { + t.Errorf("ParseDecimal(\"+7.5\") = %v, %v, want 7.5", got, err) + } + for _, text := range []string{"", "-", ".", "1e3", "abc", "1.2.3", "1_000"} { + if got, err := ParseDecimal(text); err == nil { + t.Errorf("ParseDecimal(%q) = %v, want error", text, got) + } + } +} + +func TestTwosComplementBytes(t *testing.T) { + tests := []struct { + value string + want []byte + }{ + {"0", []byte{0x00}}, + {"1", []byte{0x01}}, + {"127", []byte{0x7f}}, + {"128", []byte{0x00, 0x80}}, + {"255", []byte{0x00, 0xff}}, + {"256", []byte{0x01, 0x00}}, + {"-1", []byte{0xff}}, + {"-128", []byte{0x80}}, + {"-129", []byte{0xff, 0x7f}}, + {"-256", []byte{0xff, 0x00}}, + {"-100123", []byte{0xfe, 0x78, 0xe5}}, + {"31415", []byte{0x7a, 0xb7}}, + } + for _, test := range tests { + x, _ := new(big.Int).SetString(test.value, 10) + got := twosComplementBytes(x) + if !bytes.Equal(got, test.want) { + t.Errorf("twosComplementBytes(%v) = %x, want %x", test.value, got, test.want) + } + if back := bigIntFromTwosComplement(test.want); back.Cmp(x) != 0 { + t.Errorf("bigIntFromTwosComplement(%x) = %v, want %v", test.want, back, test.value) + } + } + if got := twosComplementBytes(nil); !bytes.Equal(got, []byte{0x00}) { + t.Errorf("twosComplementBytes(nil) = %x, want 00", got) + } +} + +// decimalRow has the fields of the decimal case of standard_coders.yaml. +type decimalRow struct { + F_float float32 `beam:"f_float"` + F_decimal Decimal `beam:"f_decimal"` +} + +func TestDecimal_RowEncoding(t *testing.T) { + // The bytes are the examples of the decimal case of standard_coders.yaml. + tests := []struct { + row decimalRow + wantBytes []byte + }{ + {decimalRow{0, decimalOf("0", 1)}, []byte("\x02\x00\x00\x00\x00\x00\x01\x01\x00")}, + {decimalRow{1, decimalOf("10", 1)}, []byte("\x02\x00?\x80\x00\x00\x01\x01\n")}, + {decimalRow{3.1415, decimalOf("31415", 4)}, []byte("\x02\x00@I\x0eV\x04\x02z\xb7")}, + {decimalRow{-100.123, decimalOf("-100123", 3)}, []byte("\x02\x00\xc2\xc8>\xfa\x03\x03\xfex\xe5")}, + } + rt := reflect.TypeOf(decimalRow{}) + enc, err := coder.RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + for _, test := range tests { + var buf bytes.Buffer + if err := enc(test.row, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", test.row, err) + } + if got := buf.Bytes(); !bytes.Equal(got, test.wantBytes) { + t.Errorf("enc(%v) = %q, want %q", test.row, got, test.wantBytes) + } + got, err := dec(bytes.NewBuffer(test.wantBytes)) + if err != nil { + t.Fatalf("dec(%q) = %v, want nil error", test.wantBytes, err) + } + if d := cmp.Diff(test.row, got, cmp.Comparer(decimalEqual)); d != "" { + t.Errorf("dec(%q) diff (-want, +got): %v", test.wantBytes, d) + } + } + // The schema of a Decimal field is the decimal logical type with a BYTES representation. + schm, err := FromType(rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", rt, err) + } + lt := schm.GetFields()[1].GetType().GetLogicalType() + if lt.GetUrn() != URNDecimal || lt.GetRepresentation().GetAtomicType().String() != "BYTES" { + t.Errorf("FromType(%v) decimal field = %v, want %v with BYTES representation", rt, lt, URNDecimal) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go new file mode 100644 index 000000000000..fc0b7c7e6fd6 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue.go @@ -0,0 +1,149 @@ +// 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 ( + "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: + rv.SetUint(uint64(uint8(v.Byte))) + case *pipepb.AtomicTypeValue_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 +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go new file mode 100644 index 000000000000..e668751bbc19 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/fieldvalue_test.go @@ -0,0 +1,111 @@ +// 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 ( + "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) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go index 413b1f379a9d..bb4ab8e85689 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/logicaltypes.go @@ -17,9 +17,11 @@ package schema import ( "fmt" + "io" "reflect" "sync" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" ) @@ -61,12 +63,12 @@ type Registry struct { logicalTypeProviders map[reflect.Type]LogicalTypeProvider logicalTypeInterfaces []reflect.Type - // Maps logical type identifiers to their reflect.Type and the schema representation. - // the type identifier is the reflect.Type name, and included in the proto as well. + // logicalTypes maps registered logical types by identifier and argument. + // The identifier is included in the proto as the logical type URN. // We don't treat all types as "logical" types. - // ... why don't we treat all types as Logical types? - logicalTypes map[string]LogicalType - logicalTypeIdentifiers map[reflect.Type]string + logicalTypes map[logicalTypeKey]LogicalType + // logicalTypeByGoType maps Go types to their registered logical type. + logicalTypeByGoType map[reflect.Type]LogicalType // toReconcile contains a list of types that have been registered // but not yet processed. Registration actually happens on first @@ -74,6 +76,13 @@ type Registry struct { toReconcile []reflect.Type } +// logicalTypeKey identifies a registered logical type. argument is nil for +// logical types without an argument, and the argument value otherwise. +type logicalTypeKey struct { + identifier string + argument any +} + // NewRegistry creates an initialized LogicalTypeRegistry. func NewRegistry() *Registry { return &Registry{ @@ -81,9 +90,9 @@ func NewRegistry() *Registry { idToType: map[string]reflect.Type{}, syntheticToUser: map[reflect.Type]reflect.Type{}, - logicalTypes: map[string]LogicalType{}, - logicalTypeIdentifiers: map[reflect.Type]string{}, - logicalTypeProviders: map[reflect.Type]LogicalTypeProvider{}, + logicalTypes: map[logicalTypeKey]LogicalType{}, + logicalTypeByGoType: map[reflect.Type]LogicalType{}, + logicalTypeProviders: map[reflect.Type]LogicalTypeProvider{}, } } @@ -98,9 +107,30 @@ func (r *Registry) RegisterLogicalType(lt LogicalType) { if len(lt.ID()) == 0 { panic(fmt.Sprintf("invalid logical type, bad id: %v -> %v", lt.GoType(), lt.StorageType())) } + if lt.argT != nil { + if _, _, err := atomicValueToProto(lt.argV); err != nil { + panic(fmt.Sprintf("LogicalType[%v] has an invalid argument %v: %v", lt.ID(), lt.argV, err)) + } + if !lt.argV.Comparable() || !lt.argV.Equal(lt.argV) { + panic(fmt.Sprintf("LogicalType[%v] has the argument %v which does not compare equal to itself", lt.ID(), lt.argV)) + } + } // TODO add duplication checks. - r.logicalTypeIdentifiers[lt.GoType()] = lt.ID() - r.logicalTypes[lt.ID()] = lt + r.logicalTypeByGoType[lt.GoType()] = lt + r.logicalTypes[lt.key()] = lt +} + +// lookupLogicalType returns the logical type registered for identifier and, +// when argument is valid, that argument. A logical type registered without an +// argument matches any argument. +func (r *Registry) lookupLogicalType(identifier string, argument reflect.Value) (LogicalType, bool) { + if argument.IsValid() && argument.Comparable() { + if lt, ok := r.logicalTypes[logicalTypeKey{identifier: identifier, argument: argument.Interface()}]; ok { + return lt, true + } + } + lt, ok := r.logicalTypes[logicalTypeKey{identifier: identifier}] + return lt, ok } // RegisterLogicalTypeProvider allows registration of providers for interface types. @@ -162,6 +192,29 @@ func ToLogicalType(identifier string, goType, storageType reflect.Type) LogicalT return LogicalType{identifier: identifier, goT: goType, storageT: storageType} } +// ToLogicalTypeWithArgument creates a LogicalType with an argument, for +// parameterized logical types such as a fixed length string and its length. +// The argument is included in the schema representation of the logical type, +// and must be a value of a built in Go type with an atomic schema type other +// than []byte. It is matched by equality, so a floating point NaN is not a +// valid argument. +func ToLogicalTypeWithArgument(identifier string, goType, storageType reflect.Type, argument any) LogicalType { + lt := ToLogicalType(identifier, goType, storageType) + if argument != nil { + lt.argT = reflect.TypeOf(argument) + lt.argV = reflect.ValueOf(argument) + } + return lt +} + +// key returns the registry key of the logical type. +func (l *LogicalType) key() logicalTypeKey { + if l.argT == nil { + return logicalTypeKey{identifier: l.identifier} + } + return logicalTypeKey{identifier: l.identifier, argument: l.argV.Interface()} +} + func preRegLogicalTypes(r *Registry) { r.RegisterLogicalType(ToLogicalType("int", reflectx.Int, reflectx.Int64)) r.RegisterLogicalType(ToLogicalType("int8", reflectx.Int8, reflectx.Int64)) @@ -173,4 +226,124 @@ func preRegLogicalTypes(r *Registry) { func init() { preRegLogicalTypes(defaultRegistry) + registerStandardLogicalTypes(defaultRegistry) +} + +// RegisterLogicalTypeConversion registers a portable logical type identified +// by urn. Values of GoT are stored in schema rows as values of StorageT, +// converted with toStorage and fromStorage, and encoded with the row field +// encoding of StorageT. StorageT must have a schema representation: an atomic +// type, a struct of such types, or a registered logical type. A struct kind +// GoT with a non row StorageT can be used as a field of a row, but not as a +// top level schema type. +// +// RegisterLogicalTypeConversion must be called before beam.Init(), and +// conventionally is called in a package init() function. +func RegisterLogicalTypeConversion[GoT, StorageT any](urn string, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) { + registerLogicalTypeConversion(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), toStorage, fromStorage) +} + +// RegisterParameterizedLogicalTypeConversion is RegisterLogicalTypeConversion +// for a logical type with an argument, such as a timestamp and its precision. +// Each argument value is registered with its own Go type, and a schema field +// with the urn and that argument maps to GoT. See ToLogicalTypeWithArgument +// for the supported argument types. +func RegisterParameterizedLogicalTypeConversion[GoT, StorageT any](urn string, argument any, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) { + registerLogicalTypeConversion(defaultRegistry, ToLogicalTypeWithArgument(urn, typeOf[GoT](), typeOf[StorageT](), argument), toStorage, fromStorage) +} + +// RegisterLogicalTypeCoder registers a portable logical type identified by +// urn whose wire encoding differs from the row field encoding of StorageT. +// StorageT is only used as the schema representation of the logical type. +// enc and dec produce and consume the encoded bytes of a value directly, and +// must match the encoding of the other SDKs for the urn. +// +// RegisterLogicalTypeCoder must be called before beam.Init(), and +// conventionally is called in a package init() function. +func RegisterLogicalTypeCoder[GoT, StorageT any](urn string, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) { + registerLogicalTypeCoder(defaultRegistry, ToLogicalType(urn, typeOf[GoT](), typeOf[StorageT]()), enc, dec) +} + +func typeOf[T any]() reflect.Type { + return reflect.TypeOf((*T)(nil)).Elem() +} + +// registerLogicalTypeConversion registers lt with the registry and registers +// row coder providers for its Go type, built from the row field coder of the +// storage type and the conversion functions. +func registerLogicalTypeConversion[GoT, StorageT any](r *Registry, lt LogicalType, toStorage func(GoT) (StorageT, error), fromStorage func(StorageT) (GoT, error)) { + st := lt.StorageType() + r.RegisterLogicalType(lt) + coder.RegisterSchemaProviders(lt.GoType(), + func(reflect.Type) (func(any, io.Writer) error, error) { + enc, err := coder.RowFieldEncoderForType(st) + if err != nil { + return nil, err + } + return func(v any, w io.Writer) error { + s, err := toStorage(v.(GoT)) + if err != nil { + return err + } + return enc(s, w) + }, nil + }, + func(reflect.Type) (func(io.Reader) (any, error), error) { + dec, err := coder.RowFieldDecoderForType(st) + if err != nil { + return nil, err + } + return func(rd io.Reader) (any, error) { + s, err := dec(rd) + if err != nil { + return nil, err + } + return fromStorage(s.(StorageT)) + }, nil + }) +} + +// registerLogicalTypeCoder registers lt with the registry and registers enc +// and dec as the row coder providers for its Go type. +func registerLogicalTypeCoder[GoT any](r *Registry, lt LogicalType, enc func(GoT, io.Writer) error, dec func(io.Reader) (GoT, error)) { + r.RegisterLogicalType(lt) + coder.RegisterSchemaProviders(lt.GoType(), + func(reflect.Type) (func(any, io.Writer) error, error) { + return func(v any, w io.Writer) error { + return enc(v.(GoT), w) + }, nil + }, + func(reflect.Type) (func(io.Reader) (any, error), error) { + return func(rd io.Reader) (any, error) { + return dec(rd) + }, nil + }) +} + +// RegisterPassThroughLogicalType registers a logical type identified by urn +// whose values are used as their representation type. A schema field with +// the urn maps to goType and uses the row encoding of goType, and any +// argument of the logical type is ignored. Go values of goType are not mapped +// back to the logical type, so a struct with a goType field produces a schema +// field of the representation type. This suits the standard fixed and +// variable length string and bytes logical types, whose argument only +// constrains the length. +// +// RegisterPassThroughLogicalType must be called before beam.Init(), and +// conventionally is called in a package init() function. +func RegisterPassThroughLogicalType(urn string, goType reflect.Type) { + defaultRegistry.RegisterPassThroughLogicalType(urn, goType) +} + +// RegisterPassThroughLogicalType registers a logical type whose values are +// used as their representation type. See the package level function. +func (r *Registry) RegisterPassThroughLogicalType(urn string, goType reflect.Type) { + if len(urn) == 0 { + panic(fmt.Sprintf("invalid pass through logical type, bad urn for %v", goType)) + } + if _, err := r.reflectTypeToFieldType(goType); err != nil { + panic(fmt.Sprintf("pass through LogicalType[%v] has an invalid type %v: %v", urn, goType, err)) + } + lt := ToLogicalType(urn, goType, goType) + r.logicalTypes[lt.key()] = lt } diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go index 0d44e68285b5..0146d182437e 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema.go @@ -218,8 +218,7 @@ func (r *Registry) registerType(ut reflect.Type, seen map[reflect.Type]struct{}) // Lets do some recursion to register fundamental type parts. t := ut - if lID, ok := r.logicalTypeIdentifiers[t]; ok { - lt := r.logicalTypes[lID] + if lt, ok := r.logicalTypeByGoType[t]; ok { r.addToMaps(lt.StorageType(), t) return nil } @@ -327,16 +326,17 @@ func (r *Registry) FromType(ot reflect.Type) (*pipepb.Schema, error) { return r.fromType(ot) } -func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, string, error) { +// logicalTypeToFieldType returns the storage field type and the logical type +// of t, or a nil field type if t is not a logical type. +func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, LogicalType, error) { // Check if a logical type was registered that matches this struct type directly // and if so, extract the schema from it for use. - if lID, ok := r.logicalTypeIdentifiers[t]; ok { - lt := r.logicalTypes[lID] + if lt, ok := r.logicalTypeByGoType[t]; ok { ftype, err := r.reflectTypeToFieldType(lt.StorageType()) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", lID, lt.StorageType(), lt.GoType()) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", lt.ID(), lt.StorageType(), lt.GoType()) } - return ftype, lID, nil + return ftype, lt, nil } for _, lti := range r.logicalTypeInterfaces { if !t.Implements(lti) { @@ -345,18 +345,42 @@ func (r *Registry) logicalTypeToFieldType(t reflect.Type) (*pipepb.FieldType, st p := r.logicalTypeProviders[lti] st, err := p(t) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v] using provider for %v schema field", t, lti) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v] using provider for %v schema field", t, lti) } if st == nil { continue } ftype, err := r.reflectTypeToFieldType(st) if err != nil { - return nil, "", errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", "interface", st, t) + return nil, LogicalType{}, errors.Wrapf(err, "unable to convert LogicalType[%v]'s storage type %v for Go type of %v to a schema", "interface", st, t) } - return ftype, t.String(), nil + return ftype, ToLogicalType(t.String(), t, st), nil } - return nil, "", nil + return nil, LogicalType{}, nil +} + +// logicalStorageSchema returns a copy of the storage row schema of a logical +// type used as a top level type. +func logicalStorageSchema(ftype *pipepb.FieldType, lt LogicalType, ot reflect.Type) (*pipepb.Schema, error) { + schm := ftype.GetRowType().GetSchema() + if schm == nil { + return nil, errors.Errorf("LogicalType[%v] for %v has the non-row storage type %v and cannot be used as a top level schema", lt.ID(), ot, lt.StorageType()) + } + return proto.Clone(schm).(*pipepb.Schema), nil +} + +// appendLogicalOptions appends the options that let ToType recover the +// logical type of a top level schema. +func appendLogicalOptions(schm *pipepb.Schema, lt LogicalType) error { + schm.Options = append(schm.Options, logicalOption(lt.ID())) + if lt.ArgumentType() != nil { + opt, err := logicalArgumentOption(lt) + if err != nil { + return errors.Wrapf(err, "unable to convert the argument of LogicalType[%v]", lt.ID()) + } + schm.Options = append(schm.Options, opt) + } + return nil } // fromType handles if the initial type is a pointer or not WRT lookups against @@ -367,18 +391,20 @@ func (r *Registry) fromType(ot reflect.Type) (*pipepb.Schema, error) { if schm, ok := r.typeToSchema[ot]; ok { return schm, nil } - ftype, lID, err := r.logicalTypeToFieldType(ot) + ftype, lt, err := r.logicalTypeToFieldType(ot) if err != nil { return nil, err } if ftype != nil { - schm := ftype.GetRowType().GetSchema() - schm = proto.Clone(schm).(*pipepb.Schema) + schm, err := logicalStorageSchema(ftype, lt, ot) + if err != nil { + return nil, err + } if ot.Kind() == reflect.Ptr { schm.Options = append(schm.Options, optGoNillable()) } - if lID != "" { - schm.Options = append(schm.Options, logicalOption(lID)) + if err := appendLogicalOptions(schm, lt); err != nil { + return nil, err } schm.Id = getUUID(ot) r.typeToSchema[ot] = schm @@ -413,6 +439,9 @@ const ( // optGoLogical indicates that this top level schema has a logical type equivalent that need to be looked up. // It has a value type of String representing the URN for the logical type to look up. optGoLogicalUrn = "beam:schema:go:logical:v1" + // optGoLogicalArgument holds the argument of the logical type of this top level schema. + // Its value type is the argument type of the logical type. + optGoLogicalArgumentUrn = "beam:schema:go:logical_argument:v1" ) func optGoNillable() *pipepb.Option { @@ -498,6 +527,43 @@ func fromLogicalOption(opts []*pipepb.Option) (string, bool) { return lID, true } +// logicalArgumentOption returns the option holding the argument of a +// parameterized logical type. +func logicalArgumentOption(lt LogicalType) (*pipepb.Option, error) { + ft, fv, err := atomicValueToProto(lt.ArgumentValue()) + if err != nil { + return nil, err + } + return &pipepb.Option{ + Name: optGoLogicalArgumentUrn, + Type: ft, + Value: fv, + }, nil +} + +// fromLogicalArgumentOption returns the logical type argument of this top +// level type, or an invalid Value when the schema has no argument option. +func fromLogicalArgumentOption(opts []*pipepb.Option) (reflect.Value, error) { + o := checkOptions(opts, optGoLogicalArgumentUrn) + if o == nil { + return reflect.Value{}, nil + } + return atomicValueFromProto(o.GetType(), o.GetValue()) +} + +// logicalTypeArgument returns the argument of a logical type proto as a Go +// value. It returns an invalid Value when the logical type has no argument, +// or when the argument is not of an atomic type. +func logicalTypeArgument(lst *pipepb.LogicalType) (reflect.Value, error) { + if lst.GetArgumentType() == nil || lst.GetArgument().GetFieldValue() == nil { + return reflect.Value{}, nil + } + if _, ok := lst.GetArgumentType().GetTypeInfo().(*pipepb.FieldType_AtomicType); !ok { + return reflect.Value{}, nil + } + return atomicValueFromProto(lst.GetArgumentType(), lst.GetArgument()) +} + func (r *Registry) structToSchema(t reflect.Type) (*pipepb.Schema, error) { if t.Kind() != reflect.Struct { return nil, errors.Errorf("non struct type received in structToSchema: %v is kind %v", t, t.Kind()) @@ -506,14 +572,18 @@ func (r *Registry) structToSchema(t reflect.Type) (*pipepb.Schema, error) { return schm, nil } - ftype, lID, err := r.logicalTypeToFieldType(t) + ftype, lt, err := r.logicalTypeToFieldType(t) if err != nil { return nil, err } if ftype != nil { - schm := ftype.GetRowType().GetSchema() - schm = proto.Clone(schm).(*pipepb.Schema) - schm.Options = append(schm.Options, logicalOption(lID)) + schm, err := logicalStorageSchema(ftype, lt, t) + if err != nil { + return nil, err + } + if err := appendLogicalOptions(schm, lt); err != nil { + return nil, err + } schm.Id = getUUID(t) r.typeToSchema[t] = schm r.idToType[schm.GetId()] = t @@ -566,18 +636,26 @@ func (r *Registry) structFieldToField(sf reflect.StructField) (*pipepb.Field, er } func (r *Registry) reflectTypeToFieldType(ot reflect.Type) (*pipepb.FieldType, error) { - ftype, lID, err := r.logicalTypeToFieldType(ot) + ftype, lt, err := r.logicalTypeToFieldType(ot) if err != nil { return nil, err } if ftype != nil { + lst := &pipepb.LogicalType{ + Urn: lt.ID(), + Representation: ftype, + } + if lt.ArgumentType() != nil { + at, av, err := atomicValueToProto(lt.ArgumentValue()) + if err != nil { + return nil, errors.Wrapf(err, "unable to convert the argument of LogicalType[%v]", lt.ID()) + } + lst.ArgumentType = at + lst.Argument = av + } return &pipepb.FieldType{ TypeInfo: &pipepb.FieldType_LogicalType{ - LogicalType: &pipepb.LogicalType{ - Urn: lID, - Representation: ftype, - // TODO(BEAM-9615): Handle type Arguments. - }, + LogicalType: lst, }, }, nil } @@ -702,7 +780,11 @@ func (r *Registry) toType(s *pipepb.Schema) (reflect.Type, error) { return t, nil } if lID, ok := fromLogicalOption(s.GetOptions()); ok { - if lt, ok := r.logicalTypes[lID]; ok { + arg, err := fromLogicalArgumentOption(s.GetOptions()) + if err != nil { + return nil, errors.Wrapf(err, "invalid argument option for logical type %v", lID) + } + if lt, ok := r.lookupLogicalType(lID, arg); ok { return lt.GoType(), nil } } @@ -794,8 +876,15 @@ func (r *Registry) fieldTypeToReflectType(sft *pipepb.FieldType, opts []*pipepb. case *pipepb.FieldType_LogicalType: lst := sft.GetLogicalType() identifier := lst.GetUrn() - lt, ok := r.logicalTypes[identifier] + arg, err := logicalTypeArgument(lst) + if err != nil { + return nil, errors.Wrapf(err, "invalid argument for logical type %v", identifier) + } + lt, ok := r.lookupLogicalType(identifier, arg) if !ok { + if arg.IsValid() { + return nil, errors.Errorf("unknown logical type: %v with argument %v", identifier, arg.Interface()) + } return nil, errors.Errorf("unknown logical type: %v", identifier) } t = lt.GoType() diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go index 367d70e81d17..8c307f404615 100644 --- a/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/schema_test.go @@ -19,10 +19,12 @@ package schema import ( "fmt" + "math" "reflect" "testing" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" "github.com/google/go-cmp/cmp" "google.golang.org/protobuf/encoding/prototext" @@ -90,7 +92,24 @@ type nonRegisteredLogical struct { k int32 } +// parameterizedString is a logical type with an argument. The tests register +// it with the argument int32(5). +type parameterizedString string + +// parameterizedRow is a logical type with an argument and a row storage type, +// so it can be used as a top level type. +type parameterizedRow struct { + S int64 +} + +// anyString is a logical type registered without an argument. +type anyString string + var ( + parameterizedStringType = reflect.TypeOf((*parameterizedString)(nil)).Elem() + parameterizedRowType = reflect.TypeOf((*parameterizedRow)(nil)).Elem() + anyStringType = reflect.TypeOf((*anyString)(nil)).Elem() + unexportedFieldsType = reflect.TypeOf((*unexportedFields)(nil)).Elem() exportedFuncType = reflect.TypeOf((*exportedFunc)(nil)) anotherType = reflect.TypeOf((*anotherStruct)(nil)).Elem() @@ -754,6 +773,53 @@ func TestSchemaConversion(t *testing.T) { rt: reflect.TypeOf(struct { U nonRegisteredLogical }{}), + }, { + st: &pipepb.Schema{ + Fields: []*pipepb.Field{ + { + Name: "P", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{ + LogicalType: &pipepb.LogicalType{ + Urn: "schema.parameterizedString", + Representation: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_STRING, + }, + }, + ArgumentType: int32FieldType(), + Argument: int32FieldValue(5), + }, + }, + }, + }, + }, + }, + rt: reflect.TypeOf(struct { + P parameterizedString + }{}), + }, { + st: &pipepb.Schema{ + Options: []*pipepb.Option{ + logicalOption("schema.parameterizedRow"), + { + Name: optGoLogicalArgumentUrn, + Type: int32FieldType(), + Value: int32FieldValue(3), + }, + }, + Fields: []*pipepb.Field{ + { + Name: "S", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_INT64, + }, + }, + }, + }, + }, + rt: parameterizedRowType, }, } @@ -764,6 +830,8 @@ func TestSchemaConversion(t *testing.T) { preRegLogicalTypes(reg) reg.RegisterLogicalType(ToLogicalType(exportedFuncType.Elem().String(), exportedFuncType.Elem(), reflect.TypeOf(struct{ V int16 }{}))) reg.RegisterLogicalType(ToLogicalType(nonRegisteredLogicalType.String(), nonRegisteredLogicalType, reflect.TypeOf(struct{ K int32 }{}))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedString", parameterizedStringType, reflectx.String, int32(5))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedRow", parameterizedRowType, reflect.TypeOf(struct{ S int64 }{}), int32(3))) reg.RegisterType(reflect.TypeOf((*sRegisteredType)(nil))) reg.RegisterLogicalTypeProvider(reflect.TypeOf((*testInterface)(nil)).Elem(), func(t reflect.Type) (reflect.Type, error) { switch t { @@ -829,3 +897,221 @@ func hasEmbeddedField(t reflect.Type) bool { } return false } + +func int32FieldType() *pipepb.FieldType { + return &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_INT32, + }, + } +} + +func int32FieldValue(v int32) *pipepb.FieldValue { + return &pipepb.FieldValue{ + FieldValue: &pipepb.FieldValue_AtomicValue{ + AtomicValue: &pipepb.AtomicTypeValue{ + Value: &pipepb.AtomicTypeValue_Int32{Int32: v}, + }, + }, + } +} + +func TestToType_LogicalTypeArguments(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.parameterizedString", parameterizedStringType, reflectx.String, int32(5))) + reg.RegisterLogicalType(ToLogicalType("schema.anyString", anyStringType, reflectx.String)) + + stringType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{ + AtomicType: pipepb.AtomicType_STRING, + }, + } + rowType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{ + RowType: &pipepb.RowType{ + Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{{Name: "precision", Type: int32FieldType()}}, + }, + }, + }, + } + rowValue := &pipepb.FieldValue{ + FieldValue: &pipepb.FieldValue_RowValue{ + RowValue: &pipepb.Row{Values: []*pipepb.FieldValue{int32FieldValue(5)}}, + }, + } + + tests := []struct { + name string + lt *pipepb.LogicalType + want reflect.Type + wantErr bool + }{{ + name: "matching argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(5)}, + want: parameterizedStringType, + }, { + name: "unknown argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(7)}, + wantErr: true, + }, { + name: "missing argument", + lt: &pipepb.LogicalType{Urn: "schema.parameterizedString", Representation: stringType}, + wantErr: true, + }, { + name: "argument ignored for logical type registered without one", + lt: &pipepb.LogicalType{Urn: "schema.anyString", Representation: stringType, ArgumentType: int32FieldType(), Argument: int32FieldValue(7)}, + want: anyStringType, + }, { + name: "non atomic argument ignored", + lt: &pipepb.LogicalType{Urn: "schema.anyString", Representation: stringType, ArgumentType: rowType, Argument: rowValue}, + want: anyStringType, + }, { + name: "unknown urn", + lt: &pipepb.LogicalType{Urn: "schema.unknown", Representation: stringType}, + wantErr: true, + }} + for i, test := range tests { + t.Run(test.name, func(t *testing.T) { + s := &pipepb.Schema{ + Id: fmt.Sprintf("arguments-%d", i), + Fields: []*pipepb.Field{{ + Name: "f", + Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{LogicalType: test.lt}, + }, + }}, + } + got, err := reg.ToType(s) + if test.wantErr { + if err == nil { + t.Fatalf("ToType(%v) = %v, want error", prototext.Format(s), got) + } + return + } + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(s), err) + } + if ft := got.Field(0).Type; ft != test.want { + t.Errorf("ToType(%v) field type = %v, want %v", prototext.Format(s), ft, test.want) + } + }) + } +} + +// The multi types share one URN with different arguments, and multiAny is +// registered for the URN without an argument. +type multiOne struct{ S int64 } +type multiTwo struct{ S int64 } +type multiAny struct{ S int64 } + +func TestLogicalTypeArgumentPrecedence(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + storage := reflect.TypeOf(struct{ S int64 }{}) + oneType := reflect.TypeOf(multiOne{}) + twoType := reflect.TypeOf(multiTwo{}) + anyType := reflect.TypeOf(multiAny{}) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.multi", oneType, storage, int32(1))) + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.multi", twoType, storage, int32(2))) + reg.RegisterLogicalType(ToLogicalType("schema.multi", anyType, storage)) + + rowType := &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{{Name: "S", Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: pipepb.AtomicType_INT64}, + }}}, + }}}, + } + fieldSchema := func(id string, argument *pipepb.FieldValue) *pipepb.Schema { + lt := &pipepb.LogicalType{Urn: "schema.multi", Representation: rowType} + if argument != nil { + lt.ArgumentType = int32FieldType() + lt.Argument = argument + } + return &pipepb.Schema{Id: id, Fields: []*pipepb.Field{{Name: "f", Type: &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{LogicalType: lt}, + }}}} + } + topLevelSchema := func(id string, argument *pipepb.FieldValue) *pipepb.Schema { + opts := []*pipepb.Option{logicalOption("schema.multi")} + if argument != nil { + opts = append(opts, &pipepb.Option{Name: optGoLogicalArgumentUrn, Type: int32FieldType(), Value: argument}) + } + return &pipepb.Schema{Id: id, Options: opts, Fields: rowType.GetRowType().GetSchema().GetFields()} + } + + tests := []struct { + name string + s *pipepb.Schema + want reflect.Type + }{ + {"field argument 1", fieldSchema("field-1", int32FieldValue(1)), oneType}, + {"field argument 2", fieldSchema("field-2", int32FieldValue(2)), twoType}, + {"field argument 3 falls back", fieldSchema("field-3", int32FieldValue(3)), anyType}, + {"field without argument", fieldSchema("field-none", nil), anyType}, + {"top level argument 2", topLevelSchema("top-2", int32FieldValue(2)), twoType}, + {"top level argument 3 falls back", topLevelSchema("top-3", int32FieldValue(3)), anyType}, + {"top level without argument", topLevelSchema("top-none", nil), anyType}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := reg.ToType(test.s) + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(test.s), err) + } + if len(test.s.GetOptions()) == 0 { + got = got.Field(0).Type + } + if got != test.want { + t.Errorf("ToType(%v) = %v, want exactly %v", prototext.Format(test.s), got, test.want) + } + }) + } + for _, rt := range []reflect.Type{oneType, twoType, anyType} { + schm, err := reg.FromType(rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", rt, err) + } + got, err := reg.ToType(schm) + if err != nil { + t.Fatalf("ToType(FromType(%v)) = %v, want nil error", rt, err) + } + if got != rt { + t.Errorf("ToType(FromType(%v)) = %v, want exactly %v", rt, got, rt) + } + } +} + +type bytesBacked struct{ B []byte } + +func TestRegisterLogicalType_Validation(t *testing.T) { + reg := NewRegistry() + preRegLogicalTypes(reg) + bytesBackedType := reflect.TypeOf(bytesBacked{}) + reg.RegisterLogicalType(ToLogicalType("schema.bytesBacked", bytesBackedType, reflectx.ByteSlice)) + if schm, err := reg.FromType(bytesBackedType); err == nil { + t.Errorf("FromType(%v) = %v, want error for a top level logical type with non row storage", bytesBackedType, prototext.Format(schm)) + } + if _, err := reg.FromType(reflect.TypeOf(struct{ F bytesBacked }{})); err != nil { + t.Errorf("FromType(struct{F bytesBacked}) = %v, want nil error", err) + } + + for name, argument := range map[string]any{ + "NaN": math.NaN(), + "float32 NaN": float32(math.NaN()), + "[]byte": []byte("x"), + "int": int(1), + "struct": struct{ A int32 }{}, + } { + t.Run(name, func(t *testing.T) { + defer func() { + if recover() == nil { + t.Errorf("RegisterLogicalType with argument %v did not panic", argument) + } + }() + reg.RegisterLogicalType(ToLogicalTypeWithArgument("schema.invalid", anyStringType, reflectx.String, argument)) + }) + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go new file mode 100644 index 000000000000..3d3975d8a61b --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types.go @@ -0,0 +1,162 @@ +// 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 ( + "io" + "math" + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" +) + +// URNs of the standard logical types defined in schema.proto. +const ( + // URNDate identifies the Date logical type. + URNDate = "beam:logical_type:date:v1" + // URNMicrosInstant identifies the MicrosInstant logical type. + URNMicrosInstant = "beam:logical_type:micros_instant:v1" +) + +// Date is a calendar date without a time zone or time of day, the +// beam:logical_type:date:v1 logical type. It is stored as the number of days +// since 1970-01-01 in an INT64 field. Out of range fields are normalized as +// by time.Date. +type Date struct { + Year int + Month time.Month + Day int +} + +// DateOf returns the Date of t in t's location. +func DateOf(t time.Time) Date { + y, m, d := t.Date() + return Date{Year: y, Month: m, Day: d} +} + +// Time returns the start of the date in UTC. +func (d Date) Time() time.Time { + return time.Date(d.Year, d.Month, d.Day, 0, 0, 0, 0, time.UTC) +} + +// String returns the date in ISO 8601 format, such as 2006-01-02. +func (d Date) String() string { + return d.Time().Format("2006-01-02") +} + +const secondsPerDay = 24 * 60 * 60 + +func dateToStorage(d Date) (int64, error) { + return d.Time().Unix() / secondsPerDay, nil +} + +func dateFromStorage(days int64) (Date, error) { + return DateOf(time.Unix(days*secondsPerDay, 0).UTC()), nil +} + +// MicrosInstant is a timestamp with microsecond precision, the +// beam:logical_type:micros_instant:v1 logical type. It is stored as a row of +// the seconds and the microseconds since the epoch. Encoding a value with +// sub microsecond precision fails, matching the Java SDK; truncate such values +// with time.Time.Truncate first. +type MicrosInstant time.Time + +// Time returns the instant as a time.Time. +func (m MicrosInstant) Time() time.Time { + return time.Time(m) +} + +// microsInstantStorage is the storage type of MicrosInstant. +type microsInstantStorage struct { + Seconds int64 `beam:"seconds"` + Micros int64 `beam:"micros"` +} + +func microsInstantToStorage(m MicrosInstant) (microsInstantStorage, error) { + t := time.Time(m) + nanos := t.Nanosecond() + if nanos%1000 != 0 { + return microsInstantStorage{}, errors.Errorf("MicrosInstant %v has sub microsecond precision", t) + } + return microsInstantStorage{Seconds: t.Unix(), Micros: int64(nanos / 1000)}, nil +} + +func microsInstantFromStorage(s microsInstantStorage) (MicrosInstant, error) { + return MicrosInstant(time.Unix(s.Seconds, s.Micros*1000).UTC()), nil +} + +// registerStandardLogicalTypes registers the standard logical types with the +// registry. +func registerStandardLogicalTypes(r *Registry) { + registerLogicalTypeConversion(r, ToLogicalType(URNDate, typeOf[Date](), reflectx.Int64), dateToStorage, dateFromStorage) + registerLogicalTypeConversion(r, ToLogicalType(URNMicrosInstant, typeOf[MicrosInstant](), typeOf[microsInstantStorage]()), microsInstantToStorage, microsInstantFromStorage) + registerLogicalTypeCoder(r, ToLogicalType(URNMillisInstant, typeOf[MillisInstant](), reflectx.Int64), encodeMillisInstant, decodeMillisInstant) + registerLogicalTypeCoder(r, ToLogicalType(URNDecimal, typeOf[Decimal](), reflectx.ByteSlice), encodeDecimal, decodeDecimal) + r.RegisterPassThroughLogicalType(URNFixedChar, reflectx.String) + r.RegisterPassThroughLogicalType(URNVarChar, reflectx.String) + r.RegisterPassThroughLogicalType(URNFixedBytes, reflectx.ByteSlice) + r.RegisterPassThroughLogicalType(URNVarBytes, reflectx.ByteSlice) + registerTimestampLogicalTypes(r) +} + +// URNs of the standard logical types that pass through as their +// representation type. Their INT32 argument is the fixed or maximum length, +// which the Go SDK does not enforce. +const ( + // URNFixedChar identifies the fixed length string logical type. + URNFixedChar = "beam:logical_type:fixed_char:v1" + // URNVarChar identifies the variable length string logical type. + URNVarChar = "beam:logical_type:var_char:v1" + // URNFixedBytes identifies the fixed length bytes logical type. + URNFixedBytes = "beam:logical_type:fixed_bytes:v1" + // URNVarBytes identifies the variable length bytes logical type. + URNVarBytes = "beam:logical_type:var_bytes:v1" +) + +// URNMillisInstant identifies the MillisInstant logical type. +const URNMillisInstant = "beam:logical_type:millis_instant:v1" + +// MillisInstant is a timestamp with millisecond precision, the +// beam:logical_type:millis_instant:v1 logical type. Its representation is the +// milliseconds since the epoch as an INT64, encoded as 8 big endian bytes +// shifted so that the byte order matches the chronological order, which is the +// encoding of the Java SDK's InstantCoder. Encoding a value with sub +// millisecond precision fails; truncate such values with time.Time.Truncate +// first. +type MillisInstant time.Time + +// Time returns the instant as a time.Time. +func (m MillisInstant) Time() time.Time { + return time.Time(m) +} + +func encodeMillisInstant(m MillisInstant, w io.Writer) error { + t := time.Time(m) + if t.Nanosecond()%int(time.Millisecond) != 0 { + return errors.Errorf("MillisInstant %v has sub millisecond precision", t) + } + return coder.EncodeUint64(uint64(t.UnixMilli()-math.MinInt64), w) +} + +func decodeMillisInstant(r io.Reader) (MillisInstant, error) { + shifted, err := coder.DecodeUint64(r) + if err != nil { + return MillisInstant{}, err + } + return MillisInstant(time.UnixMilli(int64(shifted) + math.MinInt64).UTC()), nil +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go new file mode 100644 index 000000000000..655b51b2c46c --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/standard_logical_types_test.go @@ -0,0 +1,332 @@ +// 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 ( + "bytes" + "reflect" + "testing" + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" + "github.com/google/go-cmp/cmp" + "google.golang.org/protobuf/encoding/prototext" + "google.golang.org/protobuf/testing/protocmp" +) + +func atomicType(at pipepb.AtomicType) *pipepb.FieldType { + return &pipepb.FieldType{TypeInfo: &pipepb.FieldType_AtomicType{AtomicType: at}} +} + +func logicalFieldType(urn string, representation *pipepb.FieldType) *pipepb.FieldType { + return &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_LogicalType{ + LogicalType: &pipepb.LogicalType{ + Urn: urn, + Representation: representation, + }, + }, + } +} + +// standardRow uses the standard logical types as fields, with the same field +// names as the micros_instant case of standard_coders.yaml. +type standardRow struct { + F_timestamp MicrosInstant `beam:"f_timestamp"` + F_string string `beam:"f_string"` + F_int int64 `beam:"f_int"` + F_date Date `beam:"f_date"` +} + +func TestStandardLogicalTypes_Schema(t *testing.T) { + rt := reflect.TypeOf(standardRow{}) + want := &pipepb.Schema{ + Fields: []*pipepb.Field{ + {Name: "f_timestamp", Type: logicalFieldType(URNMicrosInstant, &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{ + {Name: "seconds", Type: atomicType(pipepb.AtomicType_INT64)}, + {Name: "micros", Type: atomicType(pipepb.AtomicType_INT64)}, + }, + }}}, + })}, + {Name: "f_string", Type: atomicType(pipepb.AtomicType_STRING)}, + {Name: "f_int", Type: atomicType(pipepb.AtomicType_INT64)}, + {Name: "f_date", Type: logicalFieldType(URNDate, atomicType(pipepb.AtomicType_INT64))}, + }, + } + got, err := FromType(rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", rt, err) + } + if d := cmp.Diff(want, got, protocmp.Transform(), protocmp.IgnoreFields(&pipepb.Schema{}, "id")); d != "" { + t.Errorf("FromType(%v) diff (-want, +got): %v", rt, d) + } + // A schema from another SDK has a different id, and must still map to the + // standard Go types. + want.Id = "standard-logical-types-test" + gotType, err := ToType(want) + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(want), err) + } + for i, wantField := range []reflect.Type{typeOf[MicrosInstant](), typeOf[string](), typeOf[int64](), typeOf[Date]()} { + if got := gotType.Field(i).Type; got != wantField { + t.Errorf("ToType field %d type = %v, want %v", i, got, wantField) + } + } +} + +func TestStandardLogicalTypes_RowEncoding(t *testing.T) { + // The bytes are the micros_instant example of standard_coders.yaml, with a + // trailing date field appended: 2020-08-13 is 18487 days since the epoch. + want := standardRow{ + F_timestamp: MicrosInstant(time.Date(2020, 8, 13, 14, 14, 14, 123456000, time.UTC)), + F_string: "2020-08-13T14:14:14.123456Z", + F_int: 1597328054123456, + F_date: Date{Year: 2020, Month: time.August, Day: 13}, + } + wantBytes := []byte("\x04\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07\x1b2020-08-13T14:14:14.123456Z\xc0\xf7\x85\xda\xae\x98\xeb\x02\xb7\x90\x01") + rt := reflect.TypeOf(want) + enc, err := coder.RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(want, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", want, err) + } + if got := buf.Bytes(); !bytes.Equal(got, wantBytes) { + t.Fatalf("enc(%v) = %q, want %q", want, got, wantBytes) + } + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBuffer(wantBytes)) + if err != nil { + t.Fatalf("dec(%q) = %v, want nil error", wantBytes, err) + } + if d := cmp.Diff(want, got, cmp.Comparer(func(a, b MicrosInstant) bool { return a.Time().Equal(b.Time()) })); d != "" { + t.Errorf("dec(enc(%v)) diff (-want, +got): %v", want, d) + } +} + +func TestMicrosInstant_SubMicrosecondPrecision(t *testing.T) { + enc, err := coder.RowEncoderForStruct(reflect.TypeOf(standardRow{})) + if err != nil { + t.Fatalf("RowEncoderForStruct = %v, want nil error", err) + } + row := standardRow{F_timestamp: MicrosInstant(time.Unix(0, 1500))} + if err := enc(row, &bytes.Buffer{}); err == nil { + t.Errorf("enc(%v) = nil error, want error for sub microsecond precision", row) + } +} + +func TestDate_Storage(t *testing.T) { + tests := []struct { + date Date + days int64 + }{ + {Date{1970, time.January, 1}, 0}, + {Date{1969, time.December, 31}, -1}, + {Date{2020, time.August, 13}, 18487}, + {Date{1600, time.February, 29}, -135081}, + {Date{9999, time.December, 31}, 2932896}, + } + for _, test := range tests { + days, err := dateToStorage(test.date) + if err != nil { + t.Fatalf("dateToStorage(%v) = %v, want nil error", test.date, err) + } + if days != test.days { + t.Errorf("dateToStorage(%v) = %v, want %v", test.date, days, test.days) + } + date, err := dateFromStorage(test.days) + if err != nil { + t.Fatalf("dateFromStorage(%v) = %v, want nil error", test.days, err) + } + if date != test.date { + t.Errorf("dateFromStorage(%v) = %v, want %v", test.days, date, test.date) + } + } + if got, want := DateOf(time.Date(2020, 8, 13, 23, 59, 59, 0, time.UTC)).String(), "2020-08-13"; got != want { + t.Errorf("DateOf(...).String() = %v, want %v", got, want) + } +} + +func TestRegisterLogicalTypeConversion(t *testing.T) { + type celsius float64 + type reading struct { + Temperature celsius + } + // Registering with the default registry mirrors what users do in init(). + RegisterLogicalTypeConversion[celsius, float64]("schema.test:celsius:v1", + func(c celsius) (float64, error) { return float64(c) + 273.15, nil }, + func(k float64) (celsius, error) { return celsius(k - 273.15), nil }) + + rt := reflect.TypeOf(reading{}) + schm, err := FromType(rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", rt, err) + } + wantField := logicalFieldType("schema.test:celsius:v1", atomicType(pipepb.AtomicType_DOUBLE)) + if d := cmp.Diff(wantField, schm.GetFields()[0].GetType(), protocmp.Transform()); d != "" { + t.Errorf("FromType(%v) field diff (-want, +got): %v", rt, d) + } + + enc, err := coder.RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(reading{Temperature: 10}, &buf); err != nil { + t.Fatalf("enc = %v, want nil error", err) + } + // 283.15 as a big endian double, after the row header. + wantBytes := []byte{0x01, 0x00, 0x40, 0x71, 0xb2, 0x66, 0x66, 0x66, 0x66, 0x66} + if got := buf.Bytes(); !bytes.Equal(got, wantBytes) { + t.Fatalf("enc(reading{10}) = %v, want %v", got, wantBytes) + } + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(&buf) + if err != nil { + t.Fatalf("dec = %v, want nil error", err) + } + if d := cmp.Diff(reading{Temperature: 10}, got, cmp.Comparer(func(a, b celsius) bool { return a-b < 1e-9 && b-a < 1e-9 })); d != "" { + t.Errorf("dec(enc(reading{10})) diff (-want, +got): %v", d) + } +} + +// millisRow has the fields of the millis_instant case of standard_coders.yaml. +type millisRow struct { + F_timestamp MillisInstant `beam:"f_timestamp"` + F_string string `beam:"f_string"` + F_int int64 `beam:"f_int"` +} + +func TestMillisInstant_RowEncoding(t *testing.T) { + // The bytes are the millis_instant example of standard_coders.yaml. + want := millisRow{ + F_timestamp: MillisInstant(time.Date(2020, 8, 13, 14, 14, 14, 123000000, time.UTC)), + F_string: "2020-08-13T14:14:14.123Z", + F_int: 1597328054123, + } + wantBytes := []byte("\x03\x00\x80\x00\x01s\xe8+\xd7k\x182020-08-13T14:14:14.123Z\xeb\xae\xaf\xc1\xbe.") + rt := reflect.TypeOf(want) + enc, err := coder.RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(want, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", want, err) + } + if got := buf.Bytes(); !bytes.Equal(got, wantBytes) { + t.Fatalf("enc(%v) = %q, want %q", want, got, wantBytes) + } + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBuffer(wantBytes)) + if err != nil { + t.Fatalf("dec(%q) = %v, want nil error", wantBytes, err) + } + if d := cmp.Diff(want, got, cmp.Comparer(func(a, b MillisInstant) bool { return a.Time().Equal(b.Time()) })); d != "" { + t.Errorf("dec(enc(%v)) diff (-want, +got): %v", want, d) + } + // Timestamps before the epoch sort before the epoch in the encoding. + var before, epoch bytes.Buffer + if err := encodeMillisInstant(MillisInstant(time.UnixMilli(-1)), &before); err != nil { + t.Fatalf("encodeMillisInstant(-1ms) = %v, want nil error", err) + } + if err := encodeMillisInstant(MillisInstant(time.UnixMilli(0)), &epoch); err != nil { + t.Fatalf("encodeMillisInstant(epoch) = %v, want nil error", err) + } + if bytes.Compare(before.Bytes(), epoch.Bytes()) >= 0 { + t.Errorf("encoding of -1ms %x is not below the encoding of the epoch %x", before.Bytes(), epoch.Bytes()) + } + if err := encodeMillisInstant(MillisInstant(time.Unix(0, 1500000)), &buf); err == nil { + t.Errorf("encodeMillisInstant(1.5ms) = nil error, want error for sub millisecond precision") + } +} + +func TestPassThroughLogicalTypes(t *testing.T) { + nullable := func(urn string, representation pipepb.AtomicType, length int32) *pipepb.FieldType { + ft := logicalFieldType(urn, atomicType(representation)) + ft.Nullable = true + ft.GetLogicalType().ArgumentType = int32FieldType() + ft.GetLogicalType().Argument = int32FieldValue(length) + return ft + } + // The schema of the char, varchar, binary and varbinary case of standard_coders.yaml. + s := &pipepb.Schema{ + Id: "pass-through-logical-types-test", + Fields: []*pipepb.Field{ + {Name: "f_char", Type: nullable(URNFixedChar, pipepb.AtomicType_STRING, 5)}, + {Name: "f_varchar", Type: nullable(URNVarChar, pipepb.AtomicType_STRING, 10)}, + {Name: "f_bytes", Type: nullable(URNFixedBytes, pipepb.AtomicType_BYTES, 5)}, + {Name: "f_varbytes", Type: nullable(URNVarBytes, pipepb.AtomicType_BYTES, 10)}, + }, + } + rt, err := ToType(s) + if err != nil { + t.Fatalf("ToType(%v) = %v, want nil error", prototext.Format(s), err) + } + stringPtr := reflect.PointerTo(typeOf[string]()) + bytesPtr := reflect.PointerTo(typeOf[[]byte]()) + for i, want := range []reflect.Type{stringPtr, stringPtr, bytesPtr, bytesPtr} { + if got := rt.Field(i).Type; got != want { + t.Errorf("ToType field %d type = %v, want %v", i, got, want) + } + } + // Decoding uses the plain string and bytes encodings, and pads nothing. + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBufferString("\x04\x01\x06\x05null?\x04null")) + if err != nil { + t.Fatalf("dec = %v, want nil error", err) + } + rv := reflect.ValueOf(got) + if s := rv.Field(0).Elem().String(); s != "null?" { + t.Errorf("f_char = %q, want %q", s, "null?") + } + if !rv.Field(1).IsNil() || !rv.Field(2).IsNil() { + t.Errorf("f_varchar and f_bytes = %v, %v, want nil", rv.Field(1), rv.Field(2)) + } + if b := rv.Field(3).Elem().Bytes(); string(b) != "null" { + t.Errorf("f_varbytes = %q, want %q", b, "null") + } + // Plain string and bytes fields keep their atomic schema types. + schm, err := FromType(reflect.TypeOf(struct { + S string + B []byte + }{})) + if err != nil { + t.Fatalf("FromType = %v, want nil error", err) + } + for i, want := range []pipepb.AtomicType{pipepb.AtomicType_STRING, pipepb.AtomicType_BYTES} { + if got := schm.GetFields()[i].GetType(); got.GetLogicalType() != nil || got.GetAtomicType() != want { + t.Errorf("FromType field %d = %v, want atomic %v", i, got, want) + } + } +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp.go new file mode 100644 index 000000000000..b26eb15729dd --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp.go @@ -0,0 +1,134 @@ +// 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 ( + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" +) + +// URNTimestamp identifies the Timestamp logical type. +const URNTimestamp = "beam:logical_type:timestamp:v1" + +// The beam:logical_type:timestamp:v1 logical type is a timestamp with a +// precision argument: the number of decimal digits of the subsecond field, so +// 3 for milliseconds, 6 for microseconds and 9 for nanoseconds. It is stored +// as a row of the seconds since the epoch and the non negative subseconds in +// units of 10^-precision seconds. The subseconds field is INT16 for a +// precision below 5 and INT32 otherwise. +// +// The Go SDK registers one Go type per common precision. Encoding a value +// with more digits of precision than the type fails; truncate such values with +// time.Time.Truncate first. +type ( + // TimestampMillis is beam:logical_type:timestamp:v1 with precision 3. + TimestampMillis time.Time + // TimestampMicros is beam:logical_type:timestamp:v1 with precision 6. + TimestampMicros time.Time + // TimestampNanos is beam:logical_type:timestamp:v1 with precision 9. + TimestampNanos time.Time +) + +// Time returns the timestamp as a time.Time. +func (t TimestampMillis) Time() time.Time { return time.Time(t) } + +// Time returns the timestamp as a time.Time. +func (t TimestampMicros) Time() time.Time { return time.Time(t) } + +// Time returns the timestamp as a time.Time. +func (t TimestampNanos) Time() time.Time { return time.Time(t) } + +// timestampShortStorage is the storage type of timestamps with a precision +// below 5. +type timestampShortStorage struct { + Seconds int64 `beam:"seconds"` + Subseconds int16 `beam:"subseconds"` +} + +// timestampStorage is the storage type of timestamps with a precision of 5 +// or more. +type timestampStorage struct { + Seconds int64 `beam:"seconds"` + Subseconds int32 `beam:"subseconds"` +} + +// timestampToStorage splits t into seconds and subseconds at the precision. +func timestampToStorage(t time.Time, precision int) (int64, int64, error) { + unit := int64(1) + for i := precision; i < 9; i++ { + unit *= 10 + } + nanos := int64(t.Nanosecond()) + if nanos%unit != 0 { + return 0, 0, errors.Errorf("timestamp %v has more than %d digits of subsecond precision", t, precision) + } + return t.Unix(), nanos / unit, nil +} + +// timestampFromStorage combines seconds and subseconds at the precision. +func timestampFromStorage(seconds, subseconds int64, precision int) (time.Time, error) { + unit, limit := int64(1), int64(1) + for i := 0; i < 9; i++ { + if i < precision { + limit *= 10 + } else { + unit *= 10 + } + } + if subseconds < 0 || subseconds >= limit { + return time.Time{}, errors.Errorf("invalid subseconds %d for timestamp with precision %d", subseconds, precision) + } + return time.Unix(seconds, subseconds*unit).UTC(), nil +} + +func timestampMillisToStorage(t TimestampMillis) (timestampShortStorage, error) { + seconds, subseconds, err := timestampToStorage(time.Time(t), 3) + return timestampShortStorage{Seconds: seconds, Subseconds: int16(subseconds)}, err +} + +func timestampMillisFromStorage(s timestampShortStorage) (TimestampMillis, error) { + t, err := timestampFromStorage(s.Seconds, int64(s.Subseconds), 3) + return TimestampMillis(t), err +} + +func timestampMicrosToStorage(t TimestampMicros) (timestampStorage, error) { + seconds, subseconds, err := timestampToStorage(time.Time(t), 6) + return timestampStorage{Seconds: seconds, Subseconds: int32(subseconds)}, err +} + +func timestampMicrosFromStorage(s timestampStorage) (TimestampMicros, error) { + t, err := timestampFromStorage(s.Seconds, int64(s.Subseconds), 6) + return TimestampMicros(t), err +} + +func timestampNanosToStorage(t TimestampNanos) (timestampStorage, error) { + seconds, subseconds, err := timestampToStorage(time.Time(t), 9) + return timestampStorage{Seconds: seconds, Subseconds: int32(subseconds)}, err +} + +func timestampNanosFromStorage(s timestampStorage) (TimestampNanos, error) { + t, err := timestampFromStorage(s.Seconds, int64(s.Subseconds), 9) + return TimestampNanos(t), err +} + +// registerTimestampLogicalTypes registers the timestamp logical types with +// the registry. +func registerTimestampLogicalTypes(r *Registry) { + registerLogicalTypeConversion(r, ToLogicalTypeWithArgument(URNTimestamp, typeOf[TimestampMillis](), typeOf[timestampShortStorage](), int32(3)), timestampMillisToStorage, timestampMillisFromStorage) + registerLogicalTypeConversion(r, ToLogicalTypeWithArgument(URNTimestamp, typeOf[TimestampMicros](), typeOf[timestampStorage](), int32(6)), timestampMicrosToStorage, timestampMicrosFromStorage) + registerLogicalTypeConversion(r, ToLogicalTypeWithArgument(URNTimestamp, typeOf[TimestampNanos](), typeOf[timestampStorage](), int32(9)), timestampNanosToStorage, timestampNanosFromStorage) +} diff --git a/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp_test.go b/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp_test.go new file mode 100644 index 000000000000..0e81b4264c39 --- /dev/null +++ b/sdks/go/pkg/beam/core/runtime/graphx/schema/timestamp_test.go @@ -0,0 +1,152 @@ +// 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 ( + "bytes" + "reflect" + "testing" + "time" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/coder" + pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" + "github.com/google/go-cmp/cmp" + "google.golang.org/protobuf/testing/protocmp" +) + +type millisTimestampRow struct { + F_timestamp TimestampMillis `beam:"f_timestamp"` +} + +type microsTimestampRow struct { + F_timestamp TimestampMicros `beam:"f_timestamp"` +} + +type nanosTimestampRow struct { + F_timestamp TimestampNanos `beam:"f_timestamp"` +} + +func timestampSchema(subseconds pipepb.AtomicType, precision int32) *pipepb.Schema { + ft := logicalFieldType(URNTimestamp, &pipepb.FieldType{ + TypeInfo: &pipepb.FieldType_RowType{RowType: &pipepb.RowType{Schema: &pipepb.Schema{ + Fields: []*pipepb.Field{ + {Name: "seconds", Type: atomicType(pipepb.AtomicType_INT64)}, + {Name: "subseconds", Type: atomicType(subseconds)}, + }, + }}}, + }) + ft.GetLogicalType().ArgumentType = int32FieldType() + ft.GetLogicalType().Argument = int32FieldValue(precision) + return &pipepb.Schema{Fields: []*pipepb.Field{{Name: "f_timestamp", Type: ft}}} +} + +func TestTimestamp_Schema(t *testing.T) { + tests := []struct { + rt reflect.Type + subseconds pipepb.AtomicType + precision int32 + }{ + {reflect.TypeOf(millisTimestampRow{}), pipepb.AtomicType_INT16, 3}, + {reflect.TypeOf(microsTimestampRow{}), pipepb.AtomicType_INT32, 6}, + {reflect.TypeOf(nanosTimestampRow{}), pipepb.AtomicType_INT32, 9}, + } + for _, test := range tests { + want := timestampSchema(test.subseconds, test.precision) + got, err := FromType(test.rt) + if err != nil { + t.Fatalf("FromType(%v) = %v, want nil error", test.rt, err) + } + if d := cmp.Diff(want, got, protocmp.Transform(), protocmp.IgnoreFields(&pipepb.Schema{}, "id")); d != "" { + t.Errorf("FromType(%v) diff (-want, +got): %v", test.rt, d) + } + want.Id = test.rt.String() + gotType, err := ToType(want) + if err != nil { + t.Fatalf("ToType(precision %d) = %v, want nil error", test.precision, err) + } + if got, wantField := gotType.Field(0).Type, test.rt.Field(0).Type; got != wantField { + t.Errorf("ToType(precision %d) field type = %v, want exactly %v", test.precision, got, wantField) + } + } + unsupported := timestampSchema(pipepb.AtomicType_INT16, 4) + unsupported.Id = "timestamp-precision-4" + if got, err := ToType(unsupported); err == nil { + t.Errorf("ToType(precision 4) = %v, want error", got) + } +} + +func TestTimestamp_RowEncoding(t *testing.T) { + // The bytes are the timestamp:v1 examples of standard_coders.yaml. + tests := []struct { + row any + wantBytes []byte + }{ + {millisTimestampRow{TimestampMillis(time.Unix(1597328054, 999000000).UTC())}, []byte("\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x03\xe7")}, + {millisTimestampRow{TimestampMillis(time.Unix(-2, 500000000).UTC())}, []byte("\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\x01\xf4")}, + {millisTimestampRow{TimestampMillis(time.Unix(1597328054, 0).UTC())}, []byte("\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x00\x00")}, + {microsTimestampRow{TimestampMicros(time.Unix(1597328054, 123456000).UTC())}, []byte("\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\xc0\xc4\x07")}, + {microsTimestampRow{TimestampMicros(time.Unix(-2, 500000000).UTC())}, []byte("\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\xa0\xc2\x1e")}, + {nanosTimestampRow{TimestampNanos(time.Unix(1597328054, 123456789).UTC())}, []byte("\x01\x00\x02\x00\xb6\x95\xd5\xf9\x05\x95\x9a\xef:")}, + {nanosTimestampRow{TimestampNanos(time.Unix(-2, 999999999).UTC())}, []byte("\x01\x00\x02\x00\xfe\xff\xff\xff\xff\xff\xff\xff\xff\x01\xff\x93\xeb\xdc\x03")}, + } + timeEqual := cmp.Options{ + cmp.Comparer(func(a, b TimestampMillis) bool { return a.Time().Equal(b.Time()) }), + cmp.Comparer(func(a, b TimestampMicros) bool { return a.Time().Equal(b.Time()) }), + cmp.Comparer(func(a, b TimestampNanos) bool { return a.Time().Equal(b.Time()) }), + } + for _, test := range tests { + rt := reflect.TypeOf(test.row) + enc, err := coder.RowEncoderForStruct(rt) + if err != nil { + t.Fatalf("RowEncoderForStruct(%v) = %v, want nil error", rt, err) + } + var buf bytes.Buffer + if err := enc(test.row, &buf); err != nil { + t.Fatalf("enc(%v) = %v, want nil error", test.row, err) + } + if got := buf.Bytes(); !bytes.Equal(got, test.wantBytes) { + t.Errorf("enc(%v) = %q, want %q", test.row, got, test.wantBytes) + } + dec, err := coder.RowDecoderForStruct(rt) + if err != nil { + t.Fatalf("RowDecoderForStruct(%v) = %v, want nil error", rt, err) + } + got, err := dec(bytes.NewBuffer(test.wantBytes)) + if err != nil { + t.Fatalf("dec(%q) = %v, want nil error", test.wantBytes, err) + } + if d := cmp.Diff(test.row, got, timeEqual); d != "" { + t.Errorf("dec(%q) diff (-want, +got): %v", test.wantBytes, d) + } + } +} + +func TestTimestamp_Errors(t *testing.T) { + enc, err := coder.RowEncoderForStruct(reflect.TypeOf(millisTimestampRow{})) + if err != nil { + t.Fatalf("RowEncoderForStruct = %v, want nil error", err) + } + row := millisTimestampRow{TimestampMillis(time.Unix(0, 1500000))} + if err := enc(row, &bytes.Buffer{}); err == nil { + t.Errorf("enc(%v) = nil error, want error for sub millisecond precision", row) + } + if _, err := timestampMillisFromStorage(timestampShortStorage{Seconds: 0, Subseconds: 1000}); err == nil { + t.Errorf("timestampMillisFromStorage(subseconds 1000) = nil error, want error") + } + if _, err := timestampNanosFromStorage(timestampStorage{Seconds: 0, Subseconds: -1}); err == nil { + t.Errorf("timestampNanosFromStorage(subseconds -1) = nil error, want error") + } +} diff --git a/sdks/go/test/regression/coders/fromyaml/fromyaml.go b/sdks/go/test/regression/coders/fromyaml/fromyaml.go index 466f361b0138..793c440f9484 100644 --- a/sdks/go/test/regression/coders/fromyaml/fromyaml.go +++ b/sdks/go/test/regression/coders/fromyaml/fromyaml.go @@ -29,16 +29,19 @@ import ( "runtime/debug" "strconv" "strings" + "time" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/window" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/exec" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/util/reflectx" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" "github.com/google/go-cmp/cmp" "golang.org/x/text/encoding/charmap" + "google.golang.org/protobuf/proto" yaml "gopkg.in/yaml.v2" ) @@ -49,11 +52,7 @@ var unimplementedCoders = map[string]bool{ } var filteredCases = []struct{ filter, reason string }{ - {"logical", "BEAM-9615: Support logical types"}, {"30ea5a25-dcd8-4cdb-abeb-5332d15ab4b9", "https://github.com/apache/beam/issues/21206: Support encoding position."}, - {"80be749a-5700-4ede-89d8-dd9a4433a3f8", "https://github.com/apache/beam/issues/19817: Support millis_instant."}, - {"800c44ae-a1b7-4def-bbf6-6217cca89ec4", "https://github.com/apache/beam/issues/19817: Support decimal."}, - {"f0ffb3a4-f46f-41ca-a942-85e3e939452a", "https://github.com/apache/beam/issues/23526: Support char/varchar, binary/varbinary."}, } // Coder is a representation a serialized beam coder. @@ -164,6 +163,7 @@ func (s *Spec) testStandardCoder() (err error) { } if d := cmp.Diff(recoded, string(out.Bytes())); d != "" { log.Printf("Encoding error: diff(-want,+got): %v\n", d) + encFails++ } } if decFails+encFails > 0 { @@ -177,6 +177,50 @@ var cmpOpts = []cmp.Option{ cmp.Transformer("bytes2string", func(in []byte) (out string) { return string(in) }), + cmp.Comparer(func(a, b schema.MicrosInstant) bool { + return a.Time().Equal(b.Time()) + }), + cmp.Comparer(func(a, b schema.MillisInstant) bool { + return a.Time().Equal(b.Time()) + }), + cmp.Comparer(func(a, b schema.Decimal) bool { + return a.String() == b.String() + }), + cmp.Comparer(func(a, b schema.TimestampMillis) bool { + return a.Time().Equal(b.Time()) + }), + cmp.Comparer(func(a, b schema.TimestampMicros) bool { + return a.Time().Equal(b.Time()) + }), + cmp.Comparer(func(a, b schema.TimestampNanos) bool { + return a.Time().Equal(b.Time()) + }), +} + +// timestampFromYAML builds a timestamp:v1 value from its seconds and +// subseconds example fields, with the subseconds in units of +// 10^-precision seconds. +func timestampFromYAML(v any, precision int) time.Time { + var seconds, subseconds int64 + for _, f := range v.(yaml.MapSlice) { + switch f.Key.(string) { + case "seconds": + seconds = int64(f.Value.(int)) + case "subseconds": + subseconds = int64(f.Value.(int)) + } + } + unit := int64(1) + for i := precision; i < 9; i++ { + unit *= 10 + } + return time.Unix(seconds, subseconds*unit).UTC() +} + +// latin1Bytes returns the bytes that a yaml string denotes. The yaml strings +// escape bytes above 0x7f as code points, which yaml decodes to UTF-8. +func latin1Bytes(s string) ([]byte, error) { + return charmap.ISO8859_1.NewEncoder().Bytes([]byte(s)) } func diff(c Coder, elem *exec.FullValue, eg yaml.MapItem) bool { @@ -301,29 +345,12 @@ func diff(c Coder, elem *exec.FullValue, eg yaml.MapItem) bool { } return pass case "beam:coder:row:v1": - fs := eg.Value.(yaml.MapSlice) - var rfs []reflect.StructField - // There are only 2 pointer examples, but they reuse field names, - // so we key off the proto hash to know which example we're handling. - ptrEg := strings.Contains(c.Payload, "51ace21c7393") - for _, rf := range fs { - name := rf.Key.(string) - t := nameToType[name] - if ptrEg { - t = reflect.PtrTo(t) - } - rfs = append(rfs, reflect.StructField{ - Name: strings.ToUpper(name[:1]) + name[1:], - Type: t, - Tag: reflect.StructTag(fmt.Sprintf("beam:\"%v\"", name)), - }) - } - rv := reflect.New(reflect.StructOf(rfs)).Elem() - for i, rf := range fs { - setField(rv, i, rf.Value) + row, err := rowFromYAML(c.Payload, eg.Value.(yaml.MapSlice)) + if err != nil { + log.Printf("unable to build the expected row: %v\n", err) + return false } - - got, want = elem.Elm, rv.Interface() + got, want = elem.Elm, row case "beam:coder:timer:v1": pass := true tm := elem.Elm.(exec.TimerRecv) @@ -411,73 +438,158 @@ func diffPane(eg any, got typex.PaneInfo) bool { return pass } -// standard_coders.yaml uses the name for type indication, except for nullability. -var nameToType = map[string]reflect.Type{ - "str": reflectx.String, - "i32": reflectx.Int32, - "f64": reflectx.Float64, - "arr": reflect.SliceOf(reflectx.String), - "f_bool": reflectx.Bool, - "f_bytes": reflect.PtrTo(reflectx.ByteSlice), - "f_map": reflect.MapOf(reflectx.String, reflect.PtrTo(reflectx.Int64)), - "f_float": reflectx.Float32, +// rowFromYAML builds the expected value of a beam:coder:row:v1 example from +// the schema in the coder payload and the field values of the example. +// The Go type is the one the Go SDK derives from the schema, so logical type +// fields have their registered Go types. +func rowFromYAML(payload string, fields yaml.MapSlice) (any, error) { + b, err := latin1Bytes(payload) + if err != nil { + return nil, err + } + var s pipepb.Schema + if err := proto.Unmarshal(b, &s); err != nil { + return nil, fmt.Errorf("unmarshalling the row schema: %v", err) + } + rt, err := schema.ToType(&s) + if err != nil { + return nil, fmt.Errorf("converting the row schema to a type: %v", err) + } + rv := reflect.New(rt).Elem() + if err := setStruct(rv, fields); err != nil { + return nil, err + } + return rv.Interface(), nil +} + +// setStruct sets the fields of the struct value rv from the yaml mapping. +func setStruct(rv reflect.Value, fields yaml.MapSlice) error { + for _, f := range fields { + name := f.Key.(string) + fv, ok := fieldByName(rv, name) + if !ok { + return fmt.Errorf("no field %q in %v", name, rv.Type()) + } + if err := setValue(fv, f.Value); err != nil { + return fmt.Errorf("field %q: %v", name, err) + } + } + return nil } -func setField(rv reflect.Value, i int, v any) { +// fieldByName returns the field of rv with the given schema field name, which +// is either the Go field name or the name in the beam struct tag. +func fieldByName(rv reflect.Value, name string) (reflect.Value, bool) { + rt := rv.Type() + for i := 0; i < rt.NumField(); i++ { + sf := rt.Field(i) + if sf.Name == name { + return rv.Field(i), true + } + if tag, ok := sf.Tag.Lookup("beam"); ok && strings.Split(tag, ",")[0] == name { + return rv.Field(i), true + } + } + return reflect.Value{}, false +} + +// setValue sets rv from a yaml example value. A nil value leaves rv at its +// zero value, which is nil for nullable fields. +func setValue(rv reflect.Value, v any) error { if v == nil { - return + return nil } - rf := rv.Field(i) - if rf.Kind() == reflect.Ptr { - // Ensure it's initialized. - rf.Set(reflect.New(rf.Type().Elem())) - rf = rf.Elem() + if rv.Kind() == reflect.Ptr { + rv.Set(reflect.New(rv.Type().Elem())) + rv = rv.Elem() + } + switch rv.Type() { + case reflect.TypeOf(schema.MicrosInstant{}): + var seconds, micros int64 + for _, f := range v.(yaml.MapSlice) { + switch f.Key.(string) { + case "seconds": + seconds = int64(f.Value.(int)) + case "micros": + micros = int64(f.Value.(int)) + } + } + rv.Set(reflect.ValueOf(schema.MicrosInstant(time.Unix(seconds, micros*1000).UTC()))) + return nil + case reflect.TypeOf(schema.MillisInstant{}): + // The example value is the shifted INT64 representation. + millis := int64(v.(int)) + math.MinInt64 + rv.Set(reflect.ValueOf(schema.MillisInstant(time.UnixMilli(millis).UTC()))) + return nil + case reflect.TypeOf(schema.Decimal{}): + d, err := schema.ParseDecimal(v.(string)) + if err != nil { + return err + } + rv.Set(reflect.ValueOf(d)) + return nil + case reflect.TypeOf(schema.TimestampMillis{}): + rv.Set(reflect.ValueOf(schema.TimestampMillis(timestampFromYAML(v, 3)))) + return nil + case reflect.TypeOf(schema.TimestampMicros{}): + rv.Set(reflect.ValueOf(schema.TimestampMicros(timestampFromYAML(v, 6)))) + return nil + case reflect.TypeOf(schema.TimestampNanos{}): + rv.Set(reflect.ValueOf(schema.TimestampNanos(timestampFromYAML(v, 9)))) + return nil } - switch rf.Kind() { + switch rv.Kind() { case reflect.String: - rf.SetString(v.(string)) - case reflect.Int32: - rf.SetInt(int64(v.(int))) + rv.SetString(v.(string)) + case reflect.Int16, reflect.Int32, reflect.Int64: + rv.SetInt(int64(v.(int))) case reflect.Float32: c, err := strconv.ParseFloat(v.(string), 32) if err != nil { - panic(err) + return err } - rf.SetFloat(c) + rv.SetFloat(c) case reflect.Float64: c, err := strconv.ParseFloat(v.(string), 64) if err != nil { - panic(err) + return err } - rf.SetFloat(c) + rv.SetFloat(c) + case reflect.Bool: + rv.SetBool(v.(bool)) case reflect.Slice: - if rf.Type() == reflectx.ByteSlice { - rf.Set(reflect.ValueOf([]byte(v.(string)))) - break + if rv.Type() == reflectx.ByteSlice { + rv.SetBytes([]byte(v.(string))) + return nil } - // Value is a []any with string values. - var arr []string - for _, a := range v.([]any) { - arr = append(arr, a.(string)) + items := v.([]any) + sv := reflect.MakeSlice(rv.Type(), len(items), len(items)) + for i, item := range items { + if err := setValue(sv.Index(i), item); err != nil { + return err + } } - rf.Set(reflect.ValueOf(arr)) - case reflect.Bool: - rf.SetBool(v.(bool)) + rv.Set(sv) case reflect.Map: - // only f_map presently, which is always map[string]*int64 - rm := reflect.MakeMap(rf.Type()) - for _, a := range v.(yaml.MapSlice) { - rk := reflect.ValueOf(a.Key.(string)) - rv := reflect.Zero(rf.Type().Elem()) - if a.Value != nil { - rv = reflect.New(reflectx.Int64) - rv.Elem().SetInt(int64(a.Value.(int))) + mv := reflect.MakeMap(rv.Type()) + for _, entry := range v.(yaml.MapSlice) { + key := reflect.New(rv.Type().Key()).Elem() + if err := setValue(key, entry.Key); err != nil { + return err + } + val := reflect.New(rv.Type().Elem()).Elem() + if err := setValue(val, entry.Value); err != nil { + return err } - rm.SetMapIndex(rk, rv) + mv.SetMapIndex(key, val) } - rf.Set(rm) - + rv.Set(mv) + case reflect.Struct: + return setStruct(rv, v.(yaml.MapSlice)) + default: + return fmt.Errorf("unsupported field type %v", rv.Type()) } + return nil } func (s *Spec) parseCoder(c Coder) string { @@ -486,10 +598,14 @@ func (s *Spec) parseCoder(c Coder) string { for _, comp := range c.Components { compIDs = append(compIDs, s.parseCoder(comp)) } + payload, err := latin1Bytes(c.Payload) + if err != nil { + panic(fmt.Sprintf("invalid coder payload %q: %v", c.Payload, err)) + } s.coderPBs[id] = &pipepb.Coder{ Spec: &pipepb.FunctionSpec{ Urn: c.Urn, - Payload: []byte(c.Payload), + Payload: payload, }, ComponentCoderIds: compIDs, } diff --git a/website/www/site/content/en/documentation/programming-guide.md b/website/www/site/content/en/documentation/programming-guide.md index 3f37f45bace6..6fe2a6609c26 100644 --- a/website/www/site/content/en/documentation/programming-guide.md +++ b/website/www/site/content/en/documentation/programming-guide.md @@ -3729,7 +3729,12 @@ In Java, a logical type is specified as a subclass of the `LogicalType` class. A {{< paragraph class="language-go">}} In Go, a logical type is specified with a custom implementation of the `beam.SchemaProvider` interface. For example, the logical type provider representing nanosecond timestamps -might be implemented as follows +might be implemented as follows. +A logical type shared with other SDKs is identified by a URN and is +registered with `schema.RegisterLogicalTypeConversion` from the +`github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime/graphx/schema` package, +which takes the URN and the two functions converting between the Go type and +its storage type, and builds the row coder from the storage type. {{< /paragraph >}} {{< paragraph class="language-typescript">}} @@ -3784,8 +3789,15 @@ other than to handle `MicrosInstant`. {{< /paragraph >}} {{< paragraph class="language-go">}} -Currently the Go SDK provides minimal convenience logical types, -other than to handle additional integer primitives, and `time.Time`. +The Go SDK registers Go types for the standard logical types of the Beam +schema model in the `schema` package: `schema.Date`, `schema.MillisInstant`, +`schema.MicrosInstant`, `schema.TimestampMillis`, `schema.TimestampMicros`, +`schema.TimestampNanos` and `schema.Decimal`. The fixed and variable length +string and bytes logical types map to `string` and `[]byte`. A struct field of +one of these types produces the standard logical type in the schema, and rows +from other SDKs with these logical types decode into them. `time.Time` keeps +its Go specific encoding, so use one of the timestamp types for a field shared +with another SDK. {{< /paragraph >}} ##### **EnumerationType**