Skip to content

Commit

Permalink
sql: add width_bucket builtin
Browse files Browse the repository at this point in the history
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 #38855

Release note (sql change): add the width_bucket builtin function.
  • Loading branch information
kevinbarbour committed Aug 2, 2019
1 parent ca9184d commit 18375f8
Show file tree
Hide file tree
Showing 3 changed files with 80 additions and 0 deletions.
4 changes: 4 additions & 0 deletions docs/generated/sql/functions.md
Original file line number Diff line number Diff line change
Expand Up @@ -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>

Expand Down
52 changes: 52 additions & 0 deletions pkg/sql/sem/builtins/builtins.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down Expand Up @@ -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.
//
Expand Down
24 changes: 24 additions & 0 deletions pkg/sql/sem/builtins/builtins_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}

0 comments on commit 18375f8

Please sign in to comment.