chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,327 @@
|
||||
// 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
|
||||
)
|
||||
Reference in New Issue
Block a user