evalengine: virtual machine

go/slices2/slices.go

@@ -0,0 +1,39 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+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 slices2 contains generic Slice helpers;
+// Some of this code is sourced from https://github.com/luraim/fun (Apache v2)
+package slices2
+
+// All returns true if all elements return true for given predicate
+func All[T any](s []T, fn func(T) bool) bool {
+	for _, e := range s {
+		if !fn(e) {
+			return false
+		}
+	}
+	return true
+}
+
+// Any returns true if at least one element returns true for given predicate
+func Any[T any](s []T, fn func(T) bool) bool {
+	for _, e := range s {
+		if fn(e) {
+			return true
+		}
+	}
+	return false
+}

go/vt/dbconnpool/connection_pool.go

@@ -22,15 +22,14 @@ object to pool these DBConnections.
 package dbconnpool
 
 import (
+	"context"
 	"errors"
 	"net"
 	"sync"
 	"time"
 
 	"vitess.io/vitess/go/netutil"
 
-	"context"
-
 	"vitess.io/vitess/go/pools"
 	"vitess.io/vitess/go/stats"
 	"vitess.io/vitess/go/vt/dbconfigs"

go/vt/vtgate/evalengine/api_arithmetic_test.go

@@ -350,7 +350,7 @@ func TestArithmetics(t *testing.T) {
 			// testing for overflow of float64
 			v1:  NewFloat64(math.MaxFloat64),
 			v2:  NewFloat64(0.5),
-			err: dataOutOfRangeError(math.MaxFloat64, 0.5, "BIGINT", "/").Error(),
+			err: dataOutOfRangeError(math.MaxFloat64, 0.5, "DOUBLE", "/").Error(),
 		}},
 	}, {
 		operator: "*",

go/vt/vtgate/evalengine/api_types.go

@@ -17,10 +17,6 @@ limitations under the License.
 package evalengine
 
 import (
-	"fmt"
-	"strconv"
-
-	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/sqltypes"
 	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -175,28 +171,3 @@ func ToNative(v sqltypes.Value) (any, error) {
 	}
 	return out, err
 }
-
-func NormalizeValue(v sqltypes.Value, coll collations.ID) string {
-	typ := v.Type()
-	if typ == sqltypes.Null {
-		return "NULL"
-	}
-	if typ == sqltypes.VarChar && coll == collations.CollationBinaryID {
-		return fmt.Sprintf("VARBINARY(%q)", v.Raw())
-	}
-	if v.IsQuoted() || typ == sqltypes.Bit {
-		return fmt.Sprintf("%v(%q)", typ, v.Raw())
-	}
-	if typ == sqltypes.Float32 || typ == sqltypes.Float64 {
-		var bitsize = 64
-		if typ == sqltypes.Float32 {
-			bitsize = 32
-		}
-		f, err := strconv.ParseFloat(v.RawStr(), bitsize)
-		if err != nil {
-			panic(err)
-		}
-		return fmt.Sprintf("%v(%s)", typ, FormatFloat(typ, f))
-	}
-	return fmt.Sprintf("%v(%s)", typ, v.Raw())
-}

go/vt/vtgate/evalengine/arena.go

@@ -0,0 +1,120 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+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 evalengine
+
+import (
+	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/vt/vtgate/evalengine/internal/decimal"
+)
+
+// Arena is an arena memory allocator for eval types.
+// It allocates the types from reusable slices to prevent heap allocations.
+// After each evaluation execution, (*Arena).reset() should be called to reset the arena.
+type Arena struct {
+	aInt64   []evalInt64
+	aUint64  []evalUint64
+	aFloat64 []evalFloat
+	aDecimal []evalDecimal
+	aBytes   []evalBytes
+}
+
+func (a *Arena) reset() {
+	a.aInt64 = a.aInt64[:0]
+	a.aUint64 = a.aUint64[:0]
+	a.aFloat64 = a.aFloat64[:0]
+	a.aDecimal = a.aDecimal[:0]
+	a.aBytes = a.aBytes[:0]
+}
+
+func (a *Arena) newEvalDecimalWithPrec(dec decimal.Decimal, prec int32) *evalDecimal {
+	if cap(a.aDecimal) > len(a.aDecimal) {
+		a.aDecimal = a.aDecimal[:len(a.aDecimal)+1]
+	} else {
+		a.aDecimal = append(a.aDecimal, evalDecimal{})
+	}
+	val := &a.aDecimal[len(a.aDecimal)-1]
+	val.dec = dec
+	val.length = prec
+	return val
+}
+
+func (a *Arena) newEvalDecimal(dec decimal.Decimal, m, d int32) *evalDecimal {
+	if m == 0 && d == 0 {
+		return a.newEvalDecimalWithPrec(dec, -dec.Exponent())
+	}
+	return a.newEvalDecimalWithPrec(dec.Clamp(m-d, d), d)
+}
+
+func (a *Arena) newEvalInt64(i int64) *evalInt64 {
+	if cap(a.aInt64) > len(a.aInt64) {
+		a.aInt64 = a.aInt64[:len(a.aInt64)+1]
+	} else {
+		a.aInt64 = append(a.aInt64, evalInt64{})
+	}
+	val := &a.aInt64[len(a.aInt64)-1]
+	val.i = i
+	return val
+}
+
+func (a *Arena) newEvalUint64(u uint64) *evalUint64 {
+	if cap(a.aUint64) > len(a.aUint64) {
+		a.aUint64 = a.aUint64[:len(a.aUint64)+1]
+	} else {
+		a.aUint64 = append(a.aUint64, evalUint64{})
+	}
+	val := &a.aUint64[len(a.aUint64)-1]
+	val.u = u
+	val.hexLiteral = false
+	return val
+}
+
+func (a *Arena) newEvalFloat(f float64) *evalFloat {
+	if cap(a.aFloat64) > len(a.aFloat64) {
+		a.aFloat64 = a.aFloat64[:len(a.aFloat64)+1]
+	} else {
+		a.aFloat64 = append(a.aFloat64, evalFloat{})
+	}
+	val := &a.aFloat64[len(a.aFloat64)-1]
+	val.f = f
+	return val
+}
+
+func (a *Arena) newEvalBytesEmpty() *evalBytes {
+	if cap(a.aBytes) > len(a.aBytes) {
+		a.aBytes = a.aBytes[:len(a.aBytes)+1]
+	} else {
+		a.aBytes = append(a.aBytes, evalBytes{})
+	}
+	return &a.aBytes[len(a.aBytes)-1]
+}
+
+func (a *Arena) newEvalBinary(raw []byte) *evalBytes {
+	b := a.newEvalBytesEmpty()
+	b.tt = int16(sqltypes.VarBinary)
+	b.col = collationBinary
+	b.bytes = raw
+	return b
+}
+
+func (a *Arena) newEvalText(raw []byte, tc collations.TypedCollation) *evalBytes {
+	b := a.newEvalBytesEmpty()
+	b.tt = int16(sqltypes.VarChar)
+	b.col = tc
+	b.bytes = raw
+	return b
+}

go/vt/vtgate/evalengine/arithmetic.go

@@ -17,6 +17,7 @@ limitations under the License.
 package evalengine
 
 import (
+	"math"
 	"strings"
 
 	"golang.org/x/exp/constraints"
@@ -193,7 +194,7 @@ func mathSub_iu(v1 int64, v2 uint64) (*evalUint64, error) {
 }
 
 func mathSub_iu0(v1 int64, v2 uint64) (uint64, error) {
-	if v1 < 0 || v1 < int64(v2) {
+	if v1 < 0 {
 		return 0, dataOutOfRangeError(v1, v2, "BIGINT UNSIGNED", "-")
 	}
 	return mathSub_uu0(uint64(v1), v2)
@@ -218,12 +219,12 @@ func mathSub_ui(v1 uint64, v2 int64) (*evalUint64, error) {
 }
 
 func mathSub_ui0(v1 uint64, v2 int64) (uint64, error) {
-	if int64(v1) < v2 && v2 > 0 {
+	if v2 > 0 && v1 < uint64(v2) {
 		return 0, dataOutOfRangeError(v1, v2, "BIGINT UNSIGNED", "-")
 	}
 	// uint - (- int) = uint + int
 	if v2 < 0 {
-		return mathAdd_ui0(v1, -v2)
+		return mathAdd_uu0(v1, uint64(-v2))
 	}
 	return mathSub_uu0(v1, uint64(v2))
 }
@@ -237,7 +238,7 @@ func mathMul_ui0(v1 uint64, v2 int64) (uint64, error) {
 	if v1 == 0 || v2 == 0 {
 		return 0, nil
 	}
-	if v2 < 0 || int64(v1) < 0 {
+	if v2 < 0 {
 		return 0, dataOutOfRangeError(v1, v2, "BIGINT UNSIGNED", "*")
 	}
 	return mathMul_uu0(v1, uint64(v2))
@@ -405,11 +406,9 @@ func mathDiv_ff(v1, v2 float64) (eval, error) {
 
 func mathDiv_ff0(v1, v2 float64) (float64, error) {
 	result := v1 / v2
-	divisorLessThanOne := v2 < 1
-	resultMismatch := v2*result != v1
 
-	if divisorLessThanOne && resultMismatch {
-		return 0, dataOutOfRangeError(v1, v2, "BIGINT", "/")
+	if math.IsInf(result, 1) || math.IsInf(result, -1) {
+		return 0, dataOutOfRangeError(v1, v2, "DOUBLE", "/")
 	}
 	return result, nil
 }