add sqrt function (#515)

chaisql · Dec 25, 2023 · 6b241e2 · 6b241e2
1 parent cdac7cf
commit 6b241e2
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 9 deletions.
diff --git a/cmd/chai/doc/functions.go b/cmd/chai/doc/functions.go
@@ -19,18 +19,19 @@ var builtinDocs = functionDocs{
 	"typeof":   "The typeof function returns the type of arg1.",
 	"len":      "The len function returns length of the arg1 expression if arg1 evals to string, array or object, either returns NULL.",
 	"coalesce": "The coalesce function returns the first non-null argument. NULL is returned if all arguments are null.",
-	"random":   "The random function returns a random number between math.MinInt64 and math.MaxInt64.",
 }
 
 var mathDocs = functionDocs{
-	"abs":   "Returns the absolute value of arg1.",
-	"acos":  "Returns the arcosine, in radiant, of arg1.",
-	"acosh": "Returns the inverse hyperbolic cosine of arg1.",
-	"asin":  "Returns the arsine, in radiant, of arg1.",
-	"asinh": "Returns the inverse hyperbolic sine of arg1.",
-	"atan":  "Returns the arctangent, in radians, of arg1.",
-	"atan2": "Returns the arctangent of arg1/arg2, using the signs of the two to determine the quadrant of the return value.",
-	"floor": "Returns the greatest integer value less than or equal to arg1.",
+	"abs":    "Returns the absolute value of arg1.",
+	"acos":   "Returns the arcosine, in radiant, of arg1.",
+	"acosh":  "Returns the inverse hyperbolic cosine of arg1.",
+	"asin":   "Returns the arsine, in radiant, of arg1.",
+	"asinh":  "Returns the inverse hyperbolic sine of arg1.",
+	"atan":   "Returns the arctangent, in radians, of arg1.",
+	"atan2":  "Returns the arctangent of arg1/arg2, using the signs of the two to determine the quadrant of the return value.",
+	"floor":  "Returns the greatest integer value less than or equal to arg1.",
+	"random": "The random function returns a random number between math.MinInt64 and math.MaxInt64.",
+	"sqrt":   "The sqrt function returns the square root of arg1.",
 }
 
 var stringsDocs = functionDocs{

diff --git a/internal/expr/functions/math.go b/internal/expr/functions/math.go
@@ -24,6 +24,7 @@ var mathFunctions = Definitions{
 	"atan":   atan,
 	"atan2":  atan2,
 	"random": random,
+	"sqrt":   sqrt,
 }
 
 var floor = &ScalarDefinition{
@@ -175,3 +176,19 @@ var random = &ScalarDefinition{
 		return types.NewIntegerValue(randomNum), nil
 	},
 }
+
+var sqrt = &ScalarDefinition{
+	name:  "sqrt",
+	arity: 1,
+	callFn: func(args ...types.Value) (types.Value, error) {
+		if args[0].Type() != types.DoubleValue && args[0].Type() != types.IntegerValue {
+			return types.NewNullValue(), nil
+		}
+		v, err := document.CastAs(args[0], types.DoubleValue)
+		if err != nil {
+			return nil, err
+		}
+		res := math.Sqrt(types.As[float64](v))
+		return types.NewDoubleValue(res), nil
+	},
+}
diff --git a/internal/expr/functions/testdata/math_functions.sql b/internal/expr/functions/testdata/math_functions.sql
@@ -121,3 +121,17 @@ NULL
 2.356194490192345
 ! math.atan2('foo', 1)
 'cannot cast "foo" as double'
+
+-- test: math.sqrt
+> math.sqrt(NULL)
+NULL
+> math.sqrt(4)
+2.0
+> math.sqrt(81)
+9.0
+> math.sqrt(15)
+3.872983346207417
+> math.sqrt(1.1)
+1.0488088481701516
+> math.sqrt('foo')
+NULL