From 331daac288c1202a6a7c8cf7f74eea4588248cd2 Mon Sep 17 00:00:00 2001 From: Jason Fulghum Date: Fri, 28 Aug 2026 17:26:09 -0700 Subject: [PATCH] Plan specialized result wire encoders --- server/doltgres_handler.go | 61 +++++--- server/functions/float4.go | 11 +- server/functions/float8.go | 11 +- server/types/float_output.go | 86 +++++++++++ server/types/float_output_test.go | 65 ++++++++ server/wire_encoder.go | 173 +++++++++++++++++++++ server/wire_encoder_test.go | 245 ++++++++++++++++++++++++++++++ 7 files changed, 614 insertions(+), 38 deletions(-) create mode 100644 server/types/float_output.go create mode 100644 server/types/float_output_test.go create mode 100644 server/wire_encoder.go create mode 100644 server/wire_encoder_test.go diff --git a/server/doltgres_handler.go b/server/doltgres_handler.go index 2b24743ab2..6e59e7492c 100644 --- a/server/doltgres_handler.go +++ b/server/doltgres_handler.go @@ -359,7 +359,7 @@ func (h *DoltgresHandler) convertBindParameters(ctx *sql.Context, types []uint32 return nil, err } if values[i] != nil { - if formatCode == 0 { + if formatCode == pgtype.TextFormatCode { v, err := dgType.IoInput(ctx, string(values[i])) if err != nil { return nil, err @@ -650,7 +650,11 @@ func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter, return nil, err } - outputRow, err := rowToBytes(ctx, schema, row, formatCodes) + encoder, err := newWireRowEncoder(ctx, schema, formatCodes) + if err != nil { + return nil, err + } + outputRow, err := encoder.encode(ctx, row) if err != nil { return nil, err } @@ -664,6 +668,10 @@ func resultForMax1RowIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter, // and writes results into the callback function. func (h *DoltgresHandler) resultForDefaultIter(ctx *sql.Context, schema sql.Schema, iter sql.RowIter, callback func(*sql.Context, *Result) error, resultFields []pgproto3.FieldDescription, formatCodes []int16) (*Result, bool, error) { defer trace.StartRegion(ctx, "DoltgresHandler.resultForDefaultIter").End() + encoder, err := newWireRowEncoder(ctx, schema, formatCodes) + if err != nil { + return nil, false, err + } // TODO: use errguard.Go instead? pan2err := func(err *error) { @@ -764,7 +772,7 @@ func (h *DoltgresHandler) resultForDefaultIter(ctx *sql.Context, schema sql.Sche } outputRow := res.nextRowValues(len(schema)) - rErr := rowToBytesInto(ctx, schema, row, formatCodes, outputRow) + rErr := encoder.encodeInto(ctx, row, outputRow) if rErr != nil { return rErr } @@ -816,7 +824,7 @@ func (h *DoltgresHandler) resultForDefaultIter(ctx *sql.Context, schema sql.Sche return iter.Close(ctx) }) - err := eg.Wait() + err = eg.Wait() if err != nil { if printErrorStackTraces { fmt.Printf("error running query: %+v\n", err) @@ -865,31 +873,36 @@ func rowToBytesInto(ctx *sql.Context, s sql.Schema, row sql.Row, formatCodes []i for i, v := range row { if v == nil { o[i] = nil - } else if formatCodes[i] == 1 { - switch d := s[i].Type.(type) { - case *pgtypes.DoltgresType: - o[i], err = d.CallSend(ctx, v) - if err != nil { - return err - } - default: - cast := pgexprs.NewGMSCast(expression.NewLiteral(v, d)) - v, err = cast.Eval(ctx, nil) - if err != nil { - return err - } - o[i], err = cast.DoltgresType(ctx).CallSend(ctx, v) - if err != nil { - return err - } - } } else { - val, err := s[i].Type.SQL(ctx, []byte{}, v) // We use []byte{} as there's a distinction between nil and empty + o[i], err = valueToBytes(ctx, s[i].Type, formatCodes[i], v) if err != nil { return err } - o[i] = val.ToBytes() } } return nil } + +// valueToBytes applies the generic text or binary conversion for one result value. +func valueToBytes(ctx *sql.Context, typ sql.Type, formatCode int16, v any) ([]byte, error) { + var err error + if formatCode == pgtype.BinaryFormatCode { + switch d := typ.(type) { + case *pgtypes.DoltgresType: + return d.CallSend(ctx, v) + default: + cast := pgexprs.NewGMSCast(expression.NewLiteral(v, d)) + v, err = cast.Eval(ctx, nil) + if err != nil { + return nil, err + } + return cast.DoltgresType(ctx).CallSend(ctx, v) + } + } else { + val, err := typ.SQL(ctx, []byte{}, v) // We use []byte{} as there's a distinction between nil and empty + if err != nil { + return nil, err + } + return val.ToBytes(), nil + } +} diff --git a/server/functions/float4.go b/server/functions/float4.go index 2ee9b418e5..497acbabde 100644 --- a/server/functions/float4.go +++ b/server/functions/float4.go @@ -15,7 +15,6 @@ package functions import ( - "math" "strconv" "strings" @@ -62,13 +61,11 @@ var float4out = framework.Function1{ Parameters: [1]*pgtypes.DoltgresType{pgtypes.Float32}, Strict: true, Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) { - fVal := float64(val.(float32)) - if math.IsInf(fVal, 1) { - return "Infinity", nil - } else if math.IsInf(fVal, -1) { - return "-Infinity", nil + extraFloatDigits, err := pgtypes.ExtraFloatDigits(ctx) + if err != nil { + return nil, err } - return strconv.FormatFloat(fVal, 'f', -1, 32), nil + return string(pgtypes.AppendFloat32Text(nil, val.(float32), extraFloatDigits)), nil }, } diff --git a/server/functions/float8.go b/server/functions/float8.go index ad73b787f3..95294ac951 100644 --- a/server/functions/float8.go +++ b/server/functions/float8.go @@ -15,7 +15,6 @@ package functions import ( - "math" "strconv" "strings" @@ -62,13 +61,11 @@ var float8out = framework.Function1{ Parameters: [1]*pgtypes.DoltgresType{pgtypes.Float64}, Strict: true, Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) { - fVal := val.(float64) - if math.IsInf(fVal, 1) { - return "Infinity", nil - } else if math.IsInf(fVal, -1) { - return "-Infinity", nil + extraFloatDigits, err := pgtypes.ExtraFloatDigits(ctx) + if err != nil { + return nil, err } - return strconv.FormatFloat(fVal, 'f', -1, 64), nil + return string(pgtypes.AppendFloat64Text(nil, val.(float64), extraFloatDigits)), nil }, } diff --git a/server/types/float_output.go b/server/types/float_output.go new file mode 100644 index 0000000000..a67015d1fe --- /dev/null +++ b/server/types/float_output.go @@ -0,0 +1,86 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed 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 types + +import ( + "fmt" + "math" + "strconv" + + "github.com/dolthub/go-mysql-server/sql" +) + +// ExtraFloatDigits returns the PostgreSQL float-output precision setting. +func ExtraFloatDigits(ctx *sql.Context) (int, error) { + if ctx == nil { + return 1, nil + } + value, err := ctx.GetSessionVariable(ctx, "extra_float_digits") + if err != nil { + return 0, err + } + switch value := value.(type) { + case int: + return value, nil + case int64: + return int(value), nil + default: + return 0, fmt.Errorf("extra_float_digits has unexpected type %T", value) + } +} + +// AppendFloat32Text appends PostgreSQL-compatible float4 text to dst. +func AppendFloat32Text(dst []byte, value float32, extraFloatDigits int) []byte { + return appendFloatText(dst, float64(value), 32, extraFloatDigits) +} + +// AppendFloat64Text appends PostgreSQL-compatible float8 text to dst. +func AppendFloat64Text(dst []byte, value float64, extraFloatDigits int) []byte { + return appendFloatText(dst, value, 64, extraFloatDigits) +} + +// appendFloatText implements PostgreSQL's special values, shortest mode, and legacy precision mode. +func appendFloatText(dst []byte, value float64, bitSize int, extraFloatDigits int) []byte { + if math.IsInf(value, 1) { + return append(dst, "Infinity"...) + } + if math.IsInf(value, -1) { + return append(dst, "-Infinity"...) + } + if math.IsNaN(value) { + return append(dst, "NaN"...) + } + if extraFloatDigits <= 0 { + precision := 15 + extraFloatDigits + if bitSize == 32 { + precision = 6 + extraFloatDigits + } + if precision < 1 { + precision = 1 + } + return strconv.AppendFloat(dst, value, 'g', precision, bitSize) + } + + abs := math.Abs(value) + format := byte('f') + lowerFixed := 1e-4 + if bitSize == 32 { + lowerFixed = float64(float32(1e-4)) + } + if abs != 0 && (abs < lowerFixed || bitSize == 32 && abs >= 1e6 || bitSize == 64 && abs >= 1e15) { + format = 'e' + } + return strconv.AppendFloat(dst, value, format, -1, bitSize) +} diff --git a/server/types/float_output_test.go b/server/types/float_output_test.go new file mode 100644 index 0000000000..8d11e4950b --- /dev/null +++ b/server/types/float_output_test.go @@ -0,0 +1,65 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed 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 types + +import ( + "math" + "testing" +) + +// TestExtraFloatDigitsWithoutContextUsesDefault verifies context-free formatting uses PostgreSQL's default precision. +func TestExtraFloatDigitsWithoutContextUsesDefault(t *testing.T) { + got, err := ExtraFloatDigits(nil) + if err != nil { + t.Fatal(err) + } + if got != 1 { + t.Fatalf("got %d, want 1", got) + } +} + +// TestAppendFloatTextMatchesPostgres verifies default and legacy output modes against PostgreSQL 15. +func TestAppendFloatTextMatchesPostgres(t *testing.T) { + tests := []struct { + name string + got []byte + extraFloatDigits int + want string + }{ + {"float4 smallest subnormal", AppendFloat32Text(nil, math.SmallestNonzeroFloat32, 1), 1, "1e-45"}, + {"float8 smallest subnormal", AppendFloat64Text(nil, math.SmallestNonzeroFloat64, 1), 1, "5e-324"}, + {"float4 negative zero", AppendFloat32Text(nil, float32(math.Copysign(0, -1)), 1), 1, "-0"}, + {"float8 negative zero", AppendFloat64Text(nil, math.Copysign(0, -1), 1), 1, "-0"}, + {"float4 lower fixed", AppendFloat32Text(nil, 1e-4, 1), 1, "0.0001"}, + {"float4 lower exponent", AppendFloat32Text(nil, 1e-5, 1), 1, "1e-05"}, + {"float4 upper fixed", AppendFloat32Text(nil, 1e5, 1), 1, "100000"}, + {"float4 upper exponent", AppendFloat32Text(nil, 1e6, 1), 1, "1e+06"}, + {"float8 lower fixed", AppendFloat64Text(nil, 1e-4, 1), 1, "0.0001"}, + {"float8 lower exponent", AppendFloat64Text(nil, 1e-5, 1), 1, "1e-05"}, + {"float8 upper fixed", AppendFloat64Text(nil, 1e14, 1), 1, "100000000000000"}, + {"float8 upper exponent", AppendFloat64Text(nil, 1e15, 1), 1, "1e+15"}, + {"float4 legacy precision", AppendFloat32Text(nil, float32(1.17549435e-38), 0), 0, "1.17549e-38"}, + {"float8 legacy precision", AppendFloat64Text(nil, 1.234567890123456, 0), 0, "1.23456789012346"}, + {"float4 minimum precision", AppendFloat32Text(nil, float32(1.17549435e-38), -15), -15, "1e-38"}, + {"float8 minimum precision", AppendFloat64Text(nil, 1.234567890123456, -15), -15, "1"}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + if string(test.got) != test.want { + t.Fatalf("extra_float_digits=%d: got %q want %q", test.extraFloatDigits, test.got, test.want) + } + }) + } +} diff --git a/server/wire_encoder.go b/server/wire_encoder.go new file mode 100644 index 0000000000..31795a8fc4 --- /dev/null +++ b/server/wire_encoder.go @@ -0,0 +1,173 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed 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 server + +import ( + "fmt" + "strconv" + + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/encodings" + "github.com/jackc/pgx/v5/pgtype" + + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +// wireRowEncoder resolves the text/binary conversion for each result column +// once per result iterator. The fallback deliberately remains rowToBytes' +// general conversion path so domains, arrays, session-sensitive types and all +// binary formats retain their existing semantics. +type wireRowEncoder struct { + columns []wireColumnEncoder +} + +// wireColumnEncoder converts one non-NULL result value to its wire payload. +type wireColumnEncoder func(*sql.Context, any) ([]byte, error) + +// newWireRowEncoder plans one encoder for each schema column and canonical format code. +func newWireRowEncoder(ctx *sql.Context, schema sql.Schema, formatCodes []int16) (*wireRowEncoder, error) { + if len(formatCodes) != len(schema) { + return nil, fmt.Errorf("wire schema has %d columns and %d format codes", len(schema), len(formatCodes)) + } + extraFloatDigits := 1 + floatSettingLoaded := false + var err error + p := &wireRowEncoder{columns: make([]wireColumnEncoder, len(schema))} + for i, col := range schema { + if !floatSettingLoaded && formatCodes[i] == pgtype.TextFormatCode { + if pgType, ok := col.Type.(*pgtypes.DoltgresType); ok && pgType.TypType == pgtypes.TypeType_Base && + (pgType.ID == pgtypes.Float32.ID || pgType.ID == pgtypes.Float64.ID) { + extraFloatDigits, err = pgtypes.ExtraFloatDigits(ctx) + if err != nil { + return nil, err + } + floatSettingLoaded = true + } + } + p.columns[i] = planWireColumnEncoder(col.Type, formatCodes[i], extraFloatDigits) + } + return p, nil +} + +// planWireColumnEncoder selects a proven specialized text encoder or the generic path. +func planWireColumnEncoder(typ sql.Type, format int16, extraFloatDigits int) wireColumnEncoder { + pgType, ok := typ.(*pgtypes.DoltgresType) + if format != pgtype.TextFormatCode || !ok || pgType.TypType != pgtypes.TypeType_Base { + return genericWireColumnEncoder(typ, format) + } + fallback := genericWireColumnEncoder(typ, format) + switch pgType.ID { + case pgtypes.Bool.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + v, ok := value.(bool) + if !ok { + return fallback(ctx, value) + } + if v { + return []byte{'t'}, nil + } + return []byte{'f'}, nil + } + case pgtypes.Int16.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(int16); ok { + return strconv.AppendInt(nil, int64(v), 10), nil + } + return fallback(ctx, value) + } + case pgtypes.Int32.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(int32); ok { + return strconv.AppendInt(nil, int64(v), 10), nil + } + return fallback(ctx, value) + } + case pgtypes.Int64.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(int64); ok { + return strconv.AppendInt(nil, v, 10), nil + } + return fallback(ctx, value) + } + case pgtypes.Float32.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(float32); ok { + return pgtypes.AppendFloat32Text(nil, v, extraFloatDigits), nil + } + return fallback(ctx, value) + } + case pgtypes.Float64.ID: + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(float64); ok { + return pgtypes.AppendFloat64Text(nil, v, extraFloatDigits), nil + } + return fallback(ctx, value) + } + case pgtypes.Text.ID: + return stringWireColumnEncoder(fallback) + case pgtypes.VarChar.ID: + if pgType.GetAttTypMod() == -1 { + return stringWireColumnEncoder(fallback) + } + return fallback + default: + return fallback + } +} + +// stringWireColumnEncoder returns string bytes without copying when the runtime type matches. +func stringWireColumnEncoder(fallback wireColumnEncoder) wireColumnEncoder { + return func(ctx *sql.Context, value any) ([]byte, error) { + if v, ok := value.(string); ok { + return encodings.StringToBytes(v), nil + } + return fallback(ctx, value) + } +} + +// genericWireColumnEncoder preserves the existing conversion for unsupported columns. +func genericWireColumnEncoder(typ sql.Type, format int16) wireColumnEncoder { + return func(ctx *sql.Context, v any) ([]byte, error) { + return valueToBytes(ctx, typ, format, v) + } +} + +// encode allocates column metadata and encodes one row into it. +func (p *wireRowEncoder) encode(ctx *sql.Context, row sql.Row) ([][]byte, error) { + out := make([][]byte, len(row)) + if err := p.encodeInto(ctx, row, out); err != nil { + return nil, err + } + return out, nil +} + +// encodeInto encodes one row into caller-owned column metadata. +func (p *wireRowEncoder) encodeInto(ctx *sql.Context, row sql.Row, out [][]byte) error { + if len(row) != len(p.columns) || len(out) != len(row) { + return fmt.Errorf("wire row has %d values, %d encoders, and %d output slots", len(row), len(p.columns), len(out)) + } + for i, value := range row { + if value == nil { + out[i] = nil + continue + } + encoded, err := p.columns[i](ctx, value) + if err != nil { + return err + } + out[i] = encoded + } + return nil +} diff --git a/server/wire_encoder_test.go b/server/wire_encoder_test.go new file mode 100644 index 0000000000..84f2da2da7 --- /dev/null +++ b/server/wire_encoder_test.go @@ -0,0 +1,245 @@ +// Copyright 2026 Dolthub, Inc. +// +// Licensed 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 server + +import ( + "bytes" + "math" + "sync" + "testing" + + "github.com/dolthub/go-mysql-server/sql" + "github.com/jackc/pgx/v5/pgtype" + + "github.com/dolthub/doltgresql/server/config" + "github.com/dolthub/doltgresql/server/functions" + "github.com/dolthub/doltgresql/server/functions/framework" + pgtypes "github.com/dolthub/doltgresql/server/types" +) + +func TestWireRowEncoderMatchesGenericPath(t *testing.T) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{ + &sql.Column{Name: "b", Type: pgtypes.Bool}, + &sql.Column{Name: "i2", Type: pgtypes.Int16}, + &sql.Column{Name: "i4", Type: pgtypes.Int32}, + &sql.Column{Name: "i8", Type: pgtypes.Int64}, + &sql.Column{Name: "f4", Type: pgtypes.Float32}, + &sql.Column{Name: "f8", Type: pgtypes.Float64}, + &sql.Column{Name: "s", Type: pgtypes.Text}, + &sql.Column{Name: "n", Type: pgtypes.Text}, + } + rows := []sql.Row{ + {true, int16(-12), int32(123456), int64(-9876543210), float32(1.25), float64(-3.5), "hello", nil}, + {false, int16(0), int32(-1), int64(0), float32(0), float64(1.0 / 3.0), "", "世界"}, + } + plan, err := newWireRowEncoder(ctx, schema, make([]int16, len(schema))) + if err != nil { + t.Fatal(err) + } + for _, row := range rows { + want, err := rowToBytes(ctx, schema, row, nil) + if err != nil { + t.Fatal(err) + } + got, err := plan.encode(ctx, row) + if err != nil { + t.Fatal(err) + } + if len(got) != len(want) { + t.Fatalf("length: got %d want %d", len(got), len(want)) + } + for i := range got { + if !bytes.Equal(got[i], want[i]) { + t.Errorf("column %d: got %q want %q", i, got[i], want[i]) + } + } + } +} + +func TestWireRowEncoderFallsBackForNonSimpleTypes(t *testing.T) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{ + &sql.Column{Name: "limited", Type: pgtypes.VarChar.WithAttTypMod(7)}, + &sql.Column{Name: "array", Type: pgtypes.Int32Array}, + } + row := sql.Row{"abcdef", []any{int32(1), nil, int32(3)}} + plan, err := newWireRowEncoder(ctx, schema, make([]int16, len(schema))) + if err != nil { + t.Fatal(err) + } + want, err := rowToBytes(ctx, schema, row, nil) + if err != nil { + t.Fatal(err) + } + got, err := plan.encode(ctx, row) + if err != nil { + t.Fatal(err) + } + for i := range want { + if !bytes.Equal(got[i], want[i]) { + t.Fatalf("column %d: got %q want %q", i, got[i], want[i]) + } + } +} + +func TestWireRowEncoderPreservesBinaryFormat(t *testing.T) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{&sql.Column{Name: "i4", Type: pgtypes.Int32}} + row := sql.Row{int32(-123456)} + formats := []int16{pgtype.BinaryFormatCode} + plan, err := newWireRowEncoder(ctx, schema, formats) + if err != nil { + t.Fatal(err) + } + want, err := rowToBytes(ctx, schema, row, formats) + if err != nil { + t.Fatal(err) + } + got, err := plan.encode(ctx, row) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got[0], want[0]) { + t.Fatalf("got %v want %v", got[0], want[0]) + } +} + +func TestWireRowEncoderPreservesSpecialFloats(t *testing.T) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{ + &sql.Column{Name: "f4", Type: pgtypes.Float32}, + &sql.Column{Name: "f8", Type: pgtypes.Float64}, + } + rows := []sql.Row{ + {float32(math.NaN()), math.Inf(1)}, + {float32(math.Inf(-1)), math.Copysign(0, -1)}, + } + plan, err := newWireRowEncoder(ctx, schema, make([]int16, len(schema))) + if err != nil { + t.Fatal(err) + } + for _, row := range rows { + want, err := rowToBytes(ctx, schema, row, nil) + if err != nil { + t.Fatal(err) + } + got, err := plan.encode(ctx, row) + if err != nil { + t.Fatal(err) + } + for i := range want { + if !bytes.Equal(got[i], want[i]) { + t.Fatalf("column %d: got %q want %q", i, got[i], want[i]) + } + } + } +} + +func TestWireRowEncoderFormatsFiniteFloatBoundariesLikePostgres(t *testing.T) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{ + &sql.Column{Name: "f4_min", Type: pgtypes.Float32}, + &sql.Column{Name: "f4_max", Type: pgtypes.Float32}, + &sql.Column{Name: "f8_min", Type: pgtypes.Float64}, + &sql.Column{Name: "f8_max", Type: pgtypes.Float64}, + } + row := sql.Row{ + float32(1.17549435e-38), + float32(3.4028235e38), + 2.2250738585072014e-308, + 1.7976931348623157e308, + } + want := [][]byte{ + []byte("1.1754944e-38"), + []byte("3.4028235e+38"), + []byte("2.2250738585072014e-308"), + []byte("1.7976931348623157e+308"), + } + plan, err := newWireRowEncoder(ctx, schema, make([]int16, len(schema))) + if err != nil { + t.Fatal(err) + } + got, err := plan.encode(ctx, row) + if err != nil { + t.Fatal(err) + } + for i := range want { + if !bytes.Equal(got[i], want[i]) { + t.Fatalf("column %d: got %q want PostgreSQL output %q", i, got[i], want[i]) + } + } +} + +func TestWireRowEncoderRejectsMismatchedShapes(t *testing.T) { + if _, err := newWireRowEncoder(sql.NewEmptyContext(), sql.Schema{&sql.Column{Type: pgtypes.Text}}, nil); err == nil { + t.Fatal("expected mismatched schema and format code counts to return an error") + } + plan := &wireRowEncoder{columns: []wireColumnEncoder{genericWireColumnEncoder(pgtypes.Text, pgtype.TextFormatCode)}} + if err := plan.encodeInto(sql.NewEmptyContext(), sql.Row{"a"}, nil); err == nil { + t.Fatal("expected mismatched output shape to return an error") + } + if err := plan.encodeInto(sql.NewEmptyContext(), sql.Row{"a", "b"}, make([][]byte, 2)); err == nil { + t.Fatal("expected mismatched row shape to return an error") + } +} + +func BenchmarkWireRowEncoding(b *testing.B) { + initWireEncoderTestFunctions() + ctx := sql.NewEmptyContext() + schema := sql.Schema{ + &sql.Column{Name: "id", Type: pgtypes.Int32}, + &sql.Column{Name: "flag", Type: pgtypes.Bool}, + &sql.Column{Name: "value", Type: pgtypes.Float64}, + &sql.Column{Name: "name", Type: pgtypes.Text}, + } + row := sql.Row{int32(123456), true, 1234.5678, "representative table scan value"} + plan, err := newWireRowEncoder(ctx, schema, make([]int16, len(schema))) + if err != nil { + b.Fatal(err) + } + b.Run("generic", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if _, err := rowToBytes(ctx, schema, row, nil); err != nil { + b.Fatal(err) + } + } + }) + b.Run("planned", func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if _, err := plan.encode(ctx, row); err != nil { + b.Fatal(err) + } + } + }) +} + +func initWireEncoderTestFunctions() { + wireEncoderInitOnce.Do(func() { + pgtypes.Init() + config.Init() + functions.Init() + framework.Initialize(nil) + }) +} + +var wireEncoderInitOnce sync.Once