chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,169 @@
|
||||
// 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 types
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
|
||||
"github.com/cockroachdb/errors"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/id"
|
||||
)
|
||||
|
||||
// CreateArrayTypeFromBaseType create array type from given type. This also sets the `Array` on the given base type if
|
||||
// it has not already been set.
|
||||
func CreateArrayTypeFromBaseType(baseType *DoltgresType) *DoltgresType {
|
||||
align := TypeAlignment_Int
|
||||
if baseType.Align == TypeAlignment_Double {
|
||||
align = TypeAlignment_Double
|
||||
}
|
||||
var arrayID id.Type
|
||||
if baseType.Array == nil || baseType.Array == internalNullType || baseType.Array.ID == id.NullType {
|
||||
arrayID = id.NewType(baseType.ID.SchemaName(), "_"+baseType.ID.TypeName())
|
||||
} else {
|
||||
arrayID = baseType.Array.ID
|
||||
}
|
||||
arrayType := &DoltgresType{
|
||||
ID: arrayID,
|
||||
TypLength: int16(-1),
|
||||
PassedByVal: false,
|
||||
TypType: TypeType_Base,
|
||||
TypCategory: TypeCategory_ArrayTypes,
|
||||
IsPreferred: false,
|
||||
IsDefined: true,
|
||||
Delimiter: ",",
|
||||
RelID: id.Null,
|
||||
SubscriptFunc: toFuncID("array_subscript_handler", toInternal("internal")),
|
||||
Elem: baseType,
|
||||
Array: internalNullType,
|
||||
InputFunc: toFuncID("array_in", toInternal("cstring"), toInternal("oid"), toInternal("int4")),
|
||||
OutputFunc: toFuncID("array_out", toInternal("anyarray")),
|
||||
ReceiveFunc: toFuncID("array_recv", toInternal("internal"), toInternal("oid"), toInternal("int4")),
|
||||
SendFunc: toFuncID("array_send", toInternal("anyarray")),
|
||||
ModInFunc: baseType.ModInFunc,
|
||||
ModOutFunc: baseType.ModOutFunc,
|
||||
AnalyzeFunc: toFuncID("array_typanalyze", toInternal("internal")),
|
||||
Align: align,
|
||||
Storage: TypeStorage_Extended,
|
||||
NotNull: false,
|
||||
BaseTypeType: internalNullType,
|
||||
TypMod: -1,
|
||||
NDims: 0,
|
||||
TypCollation: baseType.TypCollation,
|
||||
DefaulBin: "",
|
||||
Default: "",
|
||||
Acl: nil,
|
||||
Checks: nil,
|
||||
InternalName: fmt.Sprintf("%s[]", baseType.Name()), // This will be set to the proper name in ToArrayType
|
||||
attTypMod: baseType.attTypMod, // TODO: check
|
||||
CompareFunc: toFuncID("btarraycmp", toInternal("anyarray"), toInternal("anyarray")),
|
||||
SerializationFunc: serializeTypeArray,
|
||||
DeserializationFunc: deserializeTypeArray,
|
||||
}
|
||||
if baseType.Array == nil || baseType.Array == internalNullType || baseType.Array.ID == id.NullType {
|
||||
baseType.Array = arrayType
|
||||
}
|
||||
return arrayType
|
||||
}
|
||||
|
||||
// serializeTypeArray handles serialization from the standard representation to our serialized representation that is
|
||||
// written in Dolt.
|
||||
func serializeTypeArray(ctx *sql.Context, t *DoltgresType, val any) ([]byte, error) {
|
||||
return serializeArray(ctx, val.([]any), t.ArrayBaseType())
|
||||
}
|
||||
|
||||
// deserializeTypeArray handles deserialization from the Dolt serialized format to our standard representation used by
|
||||
// expressions and nodes.
|
||||
func deserializeTypeArray(ctx *sql.Context, t *DoltgresType, data []byte) (any, error) {
|
||||
return deserializeArray(ctx, data, t.ArrayBaseType())
|
||||
}
|
||||
|
||||
// deserializeArray serializes an array of given base type.
|
||||
func serializeArray(ctx *sql.Context, vals []any, baseType *DoltgresType) ([]byte, error) {
|
||||
bb := bytes.Buffer{}
|
||||
// Write the element count to a buffer. We're using an array since it's stack-allocated, so no need for pooling.
|
||||
var elementCount [4]byte
|
||||
binary.LittleEndian.PutUint32(elementCount[:], uint32(len(vals)))
|
||||
bb.Write(elementCount[:])
|
||||
// Create an array that contains the offsets for each value. Since we can't update the offset portion of the buffer
|
||||
// as we determine the offsets, we have to track them outside the buffer. We'll overwrite the buffer later with the
|
||||
// correct offsets. The last offset represents the end of the slice, which simplifies the logic for reading elements
|
||||
// using the "current offset to next offset" strategy. We use a byte slice since the buffer only works with byte
|
||||
// slices.
|
||||
offsets := make([]byte, (len(vals)+1)*4)
|
||||
bb.Write(offsets)
|
||||
// The starting offset for the first element is Count(uint32) + (NumberOfElementOffsets * sizeof(uint32))
|
||||
currentOffset := uint32(4 + (len(vals)+1)*4)
|
||||
for i := range vals {
|
||||
// Write the current offset
|
||||
binary.LittleEndian.PutUint32(offsets[i*4:], currentOffset)
|
||||
// Handle serialization of the value
|
||||
// TODO: ARRAYs may be multidimensional, such as ARRAY[[4,2],[6,3]], which isn't accounted for here
|
||||
serializedVal, err := baseType.SerializeValue(ctx, vals[i])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Handle the nil case and non-nil case
|
||||
if serializedVal == nil {
|
||||
bb.WriteByte(1)
|
||||
currentOffset += 1
|
||||
} else {
|
||||
bb.WriteByte(0)
|
||||
bb.Write(serializedVal)
|
||||
currentOffset += 1 + uint32(len(serializedVal))
|
||||
}
|
||||
}
|
||||
// Write the final offset, which will equal the length of the serialized slice
|
||||
binary.LittleEndian.PutUint32(offsets[len(offsets)-4:], currentOffset)
|
||||
// Get the final output, and write the updated offsets to it
|
||||
outputBytes := bb.Bytes()
|
||||
copy(outputBytes[4:], offsets)
|
||||
return outputBytes, nil
|
||||
}
|
||||
|
||||
// deserializeArray deserializes an array of given base type.
|
||||
func deserializeArray(ctx *sql.Context, data []byte, baseType *DoltgresType) ([]any, error) {
|
||||
// Check for the nil value, then ensure the minimum length of the slice
|
||||
if len(data) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if len(data) < 4 {
|
||||
return nil, errors.Errorf("deserializing non-nil array value has invalid length of %d", len(data))
|
||||
}
|
||||
// Grab the number of elements and construct an output slice of the appropriate size
|
||||
elementCount := binary.LittleEndian.Uint32(data)
|
||||
output := make([]any, elementCount)
|
||||
// Read all elements
|
||||
for i := uint32(0); i < elementCount; i++ {
|
||||
// We read from i+1 to account for the element count at the beginning
|
||||
offset := binary.LittleEndian.Uint32(data[(i+1)*4:])
|
||||
// If the value is null, then we can skip it, since the output slice default initializes all values to nil
|
||||
if data[offset] == 1 {
|
||||
continue
|
||||
}
|
||||
// The element data is everything from the offset to the next offset, excluding the null determinant
|
||||
nextOffset := binary.LittleEndian.Uint32(data[(i+2)*4:])
|
||||
o, err := baseType.DeserializeValue(ctx, data[offset+1:nextOffset])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
output[i] = o
|
||||
}
|
||||
// Returns all read elements
|
||||
return output, nil
|
||||
}
|
||||
Reference in New Issue
Block a user