sql: add width_bucket builtin

Implements the width_bucket() builtin function

Details on the Postgres implementation can be found here:
https://www.postgresql.org/docs/11/functions-math.html

Resolves cockroachdb#38855

Release note (sql change): add the width_bucket builtin function.

docs/generated/sql/functions.md

@@ -637,6 +637,10 @@ has no relationship with the commit order of concurrent transactions.</p>
 <tr><td><code>trunc(val: <a href="decimal.html">decimal</a>) &rarr; <a href="decimal.html">decimal</a></code></td><td><span class="funcdesc"><p>Truncates the decimal values of <code>val</code>.</p>
 </span></td></tr>
 <tr><td><code>trunc(val: <a href="float.html">float</a>) &rarr; <a href="float.html">float</a></code></td><td><span class="funcdesc"><p>Truncates the decimal values of <code>val</code>.</p>
+</span></td></tr>
+<tr><td><code>width_bucket(operand: <a href="decimal.html">decimal</a>, b1: <a href="decimal.html">decimal</a>, b2: <a href="decimal.html">decimal</a>, count: <a href="int.html">int</a>) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>return the bucket number to which operand would be assigned in a histogram having count equal-width buckets spanning the range b1 to b2.</p>
+</span></td></tr>
+<tr><td><code>width_bucket(operand: anyelement, thresholds: anyelement[]) &rarr; <a href="int.html">int</a></code></td><td><span class="funcdesc"><p>return the bucket number to which operand would be assigned given an array listing the lower bounds of the buckets; returns 0 for an input less than the first lower bound; the thresholds array must be sorted, smallest first, or unexpected results will be obtained</p>
 </span></td></tr></tbody>
 </table>
 

pkg/sql/sem/builtins/builtins.go

@@ -2348,6 +2348,42 @@ may increase either contention or retry errors, or both.`,
 		}, "Truncates the decimal values of `val`."),
 	),
 
+	"width_bucket": makeBuiltin(defProps(),
+		tree.Overload{
+			Types: tree.ArgTypes{{"operand", types.Decimal}, {"b1", types.Decimal},
+				{"b2", types.Decimal}, {"count", types.Int}},
+			ReturnType: tree.FixedReturnType(types.Int),
+			Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+				operand, _ := args[0].(*tree.DDecimal).Float64()
+				b1, _ := args[1].(*tree.DDecimal).Float64()
+				b2, _ := args[2].(*tree.DDecimal).Float64()
+				count := int(tree.MustBeDInt(args[3]))
+				return tree.NewDInt(tree.DInt(widthBucket(operand, b1, b2, count))), nil
+			},
+			Info: "return the bucket number to which operand would be assigned in a histogram having count " +
+				"equal-width buckets spanning the range b1 to b2.",
+		},
+		tree.Overload{
+			Types:      tree.ArgTypes{{"operand", types.Any}, {"thresholds", types.AnyArray}},
+			ReturnType: tree.FixedReturnType(types.Int),
+			Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+				operand := args[0]
+				thresholds := tree.MustBeDArray(args[1])
+
+				for i, v := range thresholds.Array {
+					if operand.Compare(ctx, v) < 0 {
+						return tree.NewDInt(tree.DInt(i)), nil
+					}
+				}
+
+				return tree.NewDInt(tree.DInt(thresholds.Len())), nil
+			},
+			Info: "return the bucket number to which operand would be assigned given an array listing the " +
+				"lower bounds of the buckets; returns 0 for an input less than the first lower bound; the " +
+				"thresholds array must be sorted, smallest first, or unexpected results will be obtained",
+		},
+	),
+
 	// Array functions.
 
 	"string_to_array": makeBuiltin(arrayPropsNullableArgs(),
@@ -4654,6 +4690,22 @@ func rpad(s string, length int, fill string) (string, error) {
 	return buf.String(), nil
 }
 
+func widthBucket(operand float64, b1 float64, b2 float64, count int) int {
+	if operand > b2 {
+		return count + 1
+	}
+
+	if operand < b1 {
+		return 0
+	}
+
+	width := (b2 - b1) / float64(count)
+	difference := operand - b1
+	bucket := int(math.Floor(difference/width) + 1)
+
+	return bucket
+}
+
 // CleanEncodingName sanitizes the string meant to represent a
 // recognized encoding. This ignores any non-alphanumeric character.
 //

pkg/sql/sem/builtins/builtins_test.go

@@ -234,3 +234,27 @@ func TestLPadRPad(t *testing.T) {
 		}
 	}
 }
+
+func TestFloatWidthBucket(t *testing.T) {
+	testCases := []struct {
+		operand  float64
+		b1       float64
+		b2       float64
+		count    int
+		expected int
+	}{
+		{0.5, 2, 3, 5, 0},
+		{8, 2, 3, 5, 6},
+		{1.5, 1, 3, 2, 1},
+		{5.35, 0.024, 10.06, 5, 3},
+		{1, 1, 10, 2, 1},  // minimum should be inclusive
+		{10, 1, 10, 2, 3}, // maximum should be exclusive
+	}
+
+	for _, tc := range testCases {
+		got := widthBucket(tc.operand, tc.b1, tc.b2, tc.count)
+		if got != tc.expected {
+			t.Errorf("expected %d, found %d", tc.expected, got)
+		}
+	}
+}