diff --git a/docs/generated/sql/functions.md b/docs/generated/sql/functions.md
index d7edaad2aa84..73d8f710d6f5 100644
--- a/docs/generated/sql/functions.md
+++ b/docs/generated/sql/functions.md
@@ -639,6 +639,12 @@ has no relationship with the commit order of concurrent transactions.
trunc(val: decimal) → decimal | Truncates the decimal values of val.
|
trunc(val: float) → float | Truncates the decimal values of val.
+ |
+width_bucket(operand: decimal, b1: decimal, b2: decimal, count: int) → int | return the bucket number to which operand would be assigned in a histogram having count equal-width buckets spanning the range b1 to b2.
+ |
+width_bucket(operand: int, b1: int, b2: int, count: int) → int | return the bucket number to which operand would be assigned in a histogram having count equal-width buckets spanning the range b1 to b2.
+ |
+width_bucket(operand: anyelement, thresholds: anyelement[]) → int | 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
|
diff --git a/pkg/sql/logictest/testdata/logic_test/builtin_function b/pkg/sql/logictest/testdata/logic_test/builtin_function
index dd83db2477c6..48f66f3cacf9 100644
--- a/pkg/sql/logictest/testdata/logic_test/builtin_function
+++ b/pkg/sql/logictest/testdata/logic_test/builtin_function
@@ -2241,6 +2241,30 @@ SELECT count(*) = 1 FROM crdb_internal.check_consistency(true, '\xff', '')
----
true
+# Tests for width_bucket builtin
+query I
+SELECT width_bucket(8.0, 2.0, 3.0, 5)
+----
+6
+
+query I
+SELECT width_bucket(5.35, 0.024, 10.06, 5)
+----
+3
+
+query I
+SELECT width_bucket(7, 3, 11, 5)
+----
+3
+
+query I
+SELECT width_bucket(now(), array['yesterday', 'today', 'tomorrow']::timestamptz[])
+----
+2
+
+query error pq: width_bucket\(\): Operand and thresholds must be of the same type
+SELECT width_bucket(1, array['a', 'h', 'l', 'z']);
+
# Sanity check pg_type_is_visible.
query BBB
SELECT pg_type_is_visible('int'::regtype), pg_type_is_visible(NULL), pg_type_is_visible(99999)
diff --git a/pkg/sql/sem/builtins/builtins.go b/pkg/sql/sem/builtins/builtins.go
index 1a82b61f4872..f36de0c047ca 100644
--- a/pkg/sql/sem/builtins/builtins.go
+++ b/pkg/sql/sem/builtins/builtins.go
@@ -2371,6 +2371,60 @@ 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.Int}, {"b1", types.Int},
+ {"b2", types.Int}, {"count", types.Int}},
+ ReturnType: tree.FixedReturnType(types.Int),
+ Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+ operand := float64(tree.MustBeDInt(args[0]))
+ b1 := float64(tree.MustBeDInt(args[1]))
+ b2 := float64(tree.MustBeDInt(args[2]))
+ 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])
+
+ if !operand.ResolvedType().Equivalent(thresholds.Array[0].ResolvedType()) {
+ return tree.NewDInt(0), errors.New("Operand and thresholds must be of the same type")
+ }
+
+ 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(),
@@ -4687,6 +4741,25 @@ func rpad(s string, length int, fill string) (string, error) {
return buf.String(), nil
}
+// widthBucket returns the bucket number to which operand would be assigned in a histogram having count
+// equal-width buckets spanning the range b1 to b2
+func widthBucket(operand float64, b1 float64, b2 float64, count int) int {
+ bucket := 0
+ if (b1 < b2 && operand > b2) || (b1 > b2 && operand < b2) {
+ return count + 1
+ }
+
+ if (b1 < b2 && operand < b1) || (b1 > b2 && 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.
//
diff --git a/pkg/sql/sem/builtins/builtins_test.go b/pkg/sql/sem/builtins/builtins_test.go
index fc6ba638c90f..cef0032ba90f 100644
--- a/pkg/sql/sem/builtins/builtins_test.go
+++ b/pkg/sql/sem/builtins/builtins_test.go
@@ -234,3 +234,31 @@ 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},
+ {-3.0, -5, 5, 10, 3},
+ {1, 1, 10, 2, 1}, // minimum should be inclusive
+ {10, 1, 10, 2, 3}, // maximum should be exclusive
+ {4, 10, 1, 4, 3},
+ {11, 10, 1, 4, 0},
+ {0, 10, 1, 4, 5},
+ }
+
+ 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)
+ }
+ }
+}