wehub-resource-sync
328 lines
9.8 KiB
Go
328 lines
9.8 KiB
Go
// Copyright 2024 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 functions
|
|
|
|
import (
|
|
"encoding/binary"
|
|
"fmt"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/cockroachdb/apd/v3"
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
|
|
"github.com/dolthub/doltgresql/server/functions/framework"
|
|
pgtypes "github.com/dolthub/doltgresql/server/types"
|
|
"github.com/dolthub/doltgresql/utils"
|
|
)
|
|
|
|
// initNumeric registers the functions to the catalog.
|
|
func initNumeric() {
|
|
framework.RegisterFunction(numeric_in)
|
|
framework.RegisterFunction(numeric_out)
|
|
framework.RegisterFunction(numeric_recv)
|
|
framework.RegisterFunction(numeric_send)
|
|
framework.RegisterFunction(numerictypmodin)
|
|
framework.RegisterFunction(numerictypmodout)
|
|
framework.RegisterFunction(numeric_cmp)
|
|
}
|
|
|
|
// numeric_in represents the PostgreSQL function of numeric type IO input.
|
|
var numeric_in = framework.Function3{
|
|
Name: "numeric_in",
|
|
Return: pgtypes.Numeric,
|
|
Parameters: [3]*pgtypes.DoltgresType{pgtypes.Cstring, pgtypes.Oid, pgtypes.Int32},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
|
|
input := val1.(string)
|
|
typmod := val3.(int32)
|
|
dec, _, err := apd.NewFromString(strings.TrimSpace(input))
|
|
if err != nil {
|
|
return nil, pgtypes.ErrInvalidSyntaxForType.New("numeric", input)
|
|
}
|
|
return pgtypes.GetNumericValueWithTypmod(dec, typmod)
|
|
},
|
|
}
|
|
|
|
// numeric_out represents the PostgreSQL function of numeric type IO output.
|
|
var numeric_out = framework.Function1{
|
|
Name: "numeric_out",
|
|
Return: pgtypes.Cstring,
|
|
Parameters: [1]*pgtypes.DoltgresType{pgtypes.Numeric},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
|
|
typ := t[0]
|
|
dec := val.(*apd.Decimal)
|
|
tm := typ.GetAttTypMod()
|
|
res, err := pgtypes.GetNumericValueWithTypmod(dec, tm)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return res.Text('f'), nil
|
|
},
|
|
}
|
|
|
|
// numeric_recv represents the PostgreSQL function of numeric type IO receive.
|
|
var numeric_recv = framework.Function3{
|
|
Name: "numeric_recv",
|
|
Return: pgtypes.Numeric,
|
|
Parameters: [3]*pgtypes.DoltgresType{pgtypes.Internal, pgtypes.Oid, pgtypes.Int32},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, _ [4]*pgtypes.DoltgresType, val1, val2, val3 any) (any, error) {
|
|
data := val1.([]byte)
|
|
if data == nil {
|
|
return nil, nil
|
|
}
|
|
//typmod := val3.(int32)
|
|
if len(data) == 0 {
|
|
return nil, nil
|
|
}
|
|
reader := utils.NewWireReader(data)
|
|
var d *apd.Decimal
|
|
|
|
// 1. Read Header
|
|
ndigits := reader.ReadInt16()
|
|
weight := reader.ReadInt16()
|
|
sign := reader.ReadInt16()
|
|
dscale := reader.ReadInt16()
|
|
|
|
// 2. Handle Special Values (NaN, Inf)
|
|
// These usually manifest as specific bit patterns in the header
|
|
switch uint16(sign) {
|
|
case 0xC000: // pgNumericNaN
|
|
d.Form = apd.NaN
|
|
return d, nil
|
|
case 0xD000: // pgNumericPosInf
|
|
d.Form = apd.Infinite
|
|
return d, nil
|
|
case 0xF000: // pgNumericNegInf
|
|
d.Form = apd.Infinite
|
|
d.Negative = true
|
|
return d, nil
|
|
}
|
|
|
|
// 3. Handle Finite Values
|
|
if ndigits == 0 {
|
|
d.SetInt64(0)
|
|
return d, nil
|
|
}
|
|
|
|
// Read base-10000 digits
|
|
digits := make([]int16, ndigits)
|
|
for i := 0; i < int(ndigits); i++ {
|
|
digits[i] = reader.ReadInt16()
|
|
}
|
|
|
|
// 4. Convert base-10000 to string for *apd.Decimal
|
|
// Each digit is exactly 4 characters wide (except potentially the first)
|
|
var sb strings.Builder
|
|
if sign == 16384 {
|
|
sb.WriteByte('-')
|
|
}
|
|
|
|
for i, digit := range digits {
|
|
// Calculate how many 10000-base digits are before the decimal
|
|
// 'weight' is the index of the first digit, where 0 is 10^0 in base 10000
|
|
if i == int(weight)+1 {
|
|
sb.WriteByte('.')
|
|
}
|
|
|
|
sDigit := strconv.Itoa(int(digit))
|
|
// Pad with leading zeros if not the very first digit
|
|
if l := len(sDigit); l < 4 {
|
|
padding := 4 - l
|
|
for p := 0; p < padding; p++ {
|
|
sb.WriteByte('0')
|
|
}
|
|
}
|
|
sb.WriteString(sDigit)
|
|
}
|
|
|
|
// If weight is larger than digits, we need trailing zeros
|
|
if int(weight) >= len(digits) {
|
|
for i := 0; i < int(weight)-len(digits)+1; i++ {
|
|
sb.WriteString("0000")
|
|
}
|
|
}
|
|
|
|
dec, _, err := sql.DecimalHighPrecisionCtx.NewFromString(sb.String())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return sql.DecimalRound(dec, int32(dscale))
|
|
},
|
|
}
|
|
|
|
// numeric_send represents the PostgreSQL function of numeric type IO send.
|
|
var numeric_send = framework.Function1{
|
|
Name: "numeric_send",
|
|
Return: pgtypes.Bytea,
|
|
Parameters: [1]*pgtypes.DoltgresType{pgtypes.Numeric},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
|
|
num := val.(*apd.Decimal)
|
|
typmod := t[0].GetAttTypMod()
|
|
writer := utils.NewWireWriter()
|
|
if num.Form == apd.Finite {
|
|
// Short-circuit if this is the zero value
|
|
if num.IsZero() && num.Exponent == 0 {
|
|
writer.WriteBytes([]byte{0, 0, 0, 0, 0, 0, 0, 0})
|
|
return writer.BufferData(), nil
|
|
}
|
|
// There's a way to do this more efficiently, but we can do that work once this becomes a performance issue.
|
|
// This is based on the terminology used in Postgres' `numeric.c` file
|
|
decStr := num.Text('f')
|
|
isNegative := false
|
|
if strings.HasPrefix(decStr, "-") {
|
|
isNegative = true
|
|
decStr = decStr[1:]
|
|
}
|
|
// Split the integer and fractional parts
|
|
var intPart string
|
|
var fractPart string
|
|
if idx := strings.Index(decStr, "."); idx != -1 {
|
|
intPart = decStr[:idx]
|
|
fractPart = decStr[idx+1:]
|
|
} else {
|
|
intPart = decStr
|
|
}
|
|
// Find the "dscale", which is the number of digits in the fractional part
|
|
var dscale int16
|
|
if typmod != -1 {
|
|
_, dscale32 := pgtypes.GetPrecisionAndScaleFromTypmod(typmod)
|
|
dscale = int16(dscale32)
|
|
} else {
|
|
dscale = int16(len(fractPart))
|
|
}
|
|
// Pad the integer and fractional parts so that we can take groups of 4 numbers
|
|
if intPart == "0" {
|
|
intPart = ""
|
|
} else if len(intPart)%4 != 0 {
|
|
intPart = strings.Repeat("0", 4-(len(intPart)%4)) + intPart
|
|
}
|
|
if len(fractPart)%4 != 0 {
|
|
// remove trailing zeroes on right side before filling it.
|
|
fractPart = strings.TrimRightFunc(fractPart, func(r rune) bool {
|
|
return r == '0'
|
|
})
|
|
fractPart = fractPart + strings.Repeat("0", 4-(len(fractPart)%4))
|
|
}
|
|
// Write the "ndigits" first, or the number of base-10000 digits
|
|
writer.WriteInt16(int16((len(intPart) / 4) + (len(fractPart) / 4)))
|
|
// Write the "weight", which is the number of base-10000 digits in the integer part subtracted by 1
|
|
writer.WriteInt16(int16((len(intPart) / 4) - 1))
|
|
// Write the "sign"
|
|
if isNegative {
|
|
writer.WriteInt16(16384)
|
|
} else {
|
|
writer.WriteInt16(0)
|
|
}
|
|
// Write the "dscale"
|
|
writer.WriteInt16(dscale)
|
|
// Write all of the digits
|
|
fullPart := intPart + fractPart
|
|
for i := 0; i < len(fullPart); i += 4 {
|
|
part, err := strconv.Atoi(fullPart[i : i+4])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
writer.WriteInt16(int16(part))
|
|
}
|
|
} else {
|
|
var buf []byte
|
|
wp := len(buf)
|
|
buf = append(buf, 0, 0, 0, 0, 0, 0, 0, 0)
|
|
if num.Form == apd.NaN {
|
|
binary.BigEndian.PutUint64(buf[wp:], pgNumericNaN)
|
|
} else if num.Form == apd.Infinite {
|
|
if num.Negative {
|
|
binary.BigEndian.PutUint64(buf[wp:], pgNumericNegInf)
|
|
} else {
|
|
binary.BigEndian.PutUint64(buf[wp:], pgNumericPosInf)
|
|
}
|
|
}
|
|
if typmod == -1 {
|
|
binary.BigEndian.PutUint16(buf[6:], uint16(32))
|
|
}
|
|
writer.WriteBytes(buf)
|
|
}
|
|
return writer.BufferData(), nil
|
|
},
|
|
}
|
|
|
|
// numerictypmodin represents the PostgreSQL function of numeric type IO typmod input.
|
|
var numerictypmodin = framework.Function1{
|
|
Name: "numerictypmodin",
|
|
Return: pgtypes.Int32,
|
|
Parameters: [1]*pgtypes.DoltgresType{pgtypes.CstringArray},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
|
|
arr := val.([]any)
|
|
if len(arr) == 0 {
|
|
return nil, pgtypes.ErrTypmodArrayMustBe1D.New()
|
|
} else if len(arr) > 2 {
|
|
return nil, pgtypes.ErrInvalidTypMod.New("NUMERIC")
|
|
}
|
|
|
|
p, err := strconv.ParseInt(arr[0].(string), 10, 32)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
precision := int32(p)
|
|
scale := int32(0)
|
|
if len(arr) == 2 {
|
|
s, err := strconv.ParseInt(arr[1].(string), 10, 32)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
scale = int32(s)
|
|
}
|
|
return pgtypes.GetTypmodFromNumericPrecisionAndScale(precision, scale)
|
|
},
|
|
}
|
|
|
|
// numerictypmodout represents the PostgreSQL function of numeric type IO typmod output.
|
|
var numerictypmodout = framework.Function1{
|
|
Name: "numerictypmodout",
|
|
Return: pgtypes.Cstring,
|
|
Parameters: [1]*pgtypes.DoltgresType{pgtypes.Int32},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, _ [2]*pgtypes.DoltgresType, val any) (any, error) {
|
|
typmod := val.(int32)
|
|
precision, scale := pgtypes.GetPrecisionAndScaleFromTypmod(typmod)
|
|
return fmt.Sprintf("(%v,%v)", precision, scale), nil
|
|
},
|
|
}
|
|
|
|
// numeric_cmp represents the PostgreSQL function of numeric type compare.
|
|
var numeric_cmp = framework.Function2{
|
|
Name: "numeric_cmp",
|
|
Return: pgtypes.Int32,
|
|
Parameters: [2]*pgtypes.DoltgresType{pgtypes.Numeric, pgtypes.Numeric},
|
|
Strict: true,
|
|
Callable: func(ctx *sql.Context, _ [3]*pgtypes.DoltgresType, val1, val2 any) (any, error) {
|
|
ab := val1.(*apd.Decimal)
|
|
bb := val2.(*apd.Decimal)
|
|
return int32(pgtypes.NumericCompare(ab, bb)), nil
|
|
},
|
|
}
|
|
|
|
const (
|
|
pgNumericNaN = 0x00000000c0000000
|
|
pgNumericPosInf = 0x00000000d0000000
|
|
pgNumericNegInf = 0x00000000f0000000
|
|
)
|