Export quantile functions (#15190)

Export quantile functions

For use in Mimir's query engine, it would be helpful if these
functions were exported.

Co-authored-by: Björn Rabenstein <github@rabenste.in>
Signed-off-by: Joshua Hesketh <josh@hesketh.net.au>

---------

Signed-off-by: Joshua Hesketh <josh@nitrotech.org>
Signed-off-by: Joshua Hesketh <josh@hesketh.net.au>
Co-authored-by: Björn Rabenstein <github@rabenste.in>

promql/functions.go

@@ -1300,7 +1300,7 @@ func funcHistogramFraction(vals []parser.Value, args parser.Expressions, enh *Ev
 		}
 		enh.Out = append(enh.Out, Sample{
 			Metric:   sample.Metric,
-			F:        histogramFraction(lower, upper, sample.H),
+			F:        HistogramFraction(lower, upper, sample.H),
 			DropName: true,
 		})
 	}
@@ -1352,7 +1352,7 @@ func funcHistogramQuantile(vals []parser.Value, args parser.Expressions, enh *Ev
 			mb = &metricWithBuckets{sample.Metric, nil}
 			enh.signatureToMetricWithBuckets[string(enh.lblBuf)] = mb
 		}
-		mb.buckets = append(mb.buckets, bucket{upperBound, sample.F})
+		mb.buckets = append(mb.buckets, Bucket{upperBound, sample.F})
 	}
 
 	// Now deal with the histograms.
@@ -1374,14 +1374,14 @@ func funcHistogramQuantile(vals []parser.Value, args parser.Expressions, enh *Ev
 		}
 		enh.Out = append(enh.Out, Sample{
 			Metric:   sample.Metric,
-			F:        histogramQuantile(q, sample.H),
+			F:        HistogramQuantile(q, sample.H),
 			DropName: true,
 		})
 	}
 
 	for _, mb := range enh.signatureToMetricWithBuckets {
 		if len(mb.buckets) > 0 {
-			res, forcedMonotonicity, _ := bucketQuantile(q, mb.buckets)
+			res, forcedMonotonicity, _ := BucketQuantile(q, mb.buckets)
 			enh.Out = append(enh.Out, Sample{
 				Metric: mb.metric,
 				F:      res,

promql/quantile.go

@@ -51,20 +51,22 @@ var excludedLabels = []string{
 	labels.BucketLabel,
 }
 
-type bucket struct {
+// Bucket represents a bucket of a classic histogram. It is used internally by the promql
-	upperBound float64
+// package, but it is nevertheless exported for potential use in other PromQL engines.
-	count      float64
+type Bucket struct {
+	UpperBound float64
+	Count      float64
 }
 
-// buckets implements sort.Interface.
+// Buckets implements sort.Interface.
-type buckets []bucket
+type Buckets []Bucket
 
 type metricWithBuckets struct {
 	metric  labels.Labels
-	buckets buckets
+	buckets Buckets
 }
 
-// bucketQuantile calculates the quantile 'q' based on the given buckets. The
+// BucketQuantile calculates the quantile 'q' based on the given buckets. The
 // buckets will be sorted by upperBound by this function (i.e. no sorting
 // needed before calling this function). The quantile value is interpolated
 // assuming a linear distribution within a bucket. However, if the quantile
@@ -95,7 +97,14 @@ type metricWithBuckets struct {
 // and another bool to indicate if small differences between buckets (that
 // are likely artifacts of floating point precision issues) have been
 // ignored.
-func bucketQuantile(q float64, buckets buckets) (float64, bool, bool) {
+//
+// Generically speaking, BucketQuantile is for calculating the
+// histogram_quantile() of classic histograms. See also: HistogramQuantile
+// for native histograms.
+//
+// BucketQuantile is exported as a useful quantile function over a set of
+// given buckets. It may be used by other PromQL engine implementations.
+func BucketQuantile(q float64, buckets Buckets) (float64, bool, bool) {
 	if math.IsNaN(q) {
 		return math.NaN(), false, false
 	}
@@ -105,17 +114,17 @@ func bucketQuantile(q float64, buckets buckets) (float64, bool, bool) {
 	if q > 1 {
 		return math.Inf(+1), false, false
 	}
-	slices.SortFunc(buckets, func(a, b bucket) int {
+	slices.SortFunc(buckets, func(a, b Bucket) int {
 		// We don't expect the bucket boundary to be a NaN.
-		if a.upperBound < b.upperBound {
+		if a.UpperBound < b.UpperBound {
 			return -1
 		}
-		if a.upperBound > b.upperBound {
+		if a.UpperBound > b.UpperBound {
 			return +1
 		}
 		return 0
 	})
-	if !math.IsInf(buckets[len(buckets)-1].upperBound, +1) {
+	if !math.IsInf(buckets[len(buckets)-1].UpperBound, +1) {
 		return math.NaN(), false, false
 	}
 
@@ -125,33 +134,33 @@ func bucketQuantile(q float64, buckets buckets) (float64, bool, bool) {
 	if len(buckets) < 2 {
 		return math.NaN(), false, false
 	}
-	observations := buckets[len(buckets)-1].count
+	observations := buckets[len(buckets)-1].Count
 	if observations == 0 {
 		return math.NaN(), false, false
 	}
 	rank := q * observations
-	b := sort.Search(len(buckets)-1, func(i int) bool { return buckets[i].count >= rank })
+	b := sort.Search(len(buckets)-1, func(i int) bool { return buckets[i].Count >= rank })
 
 	if b == len(buckets)-1 {
-		return buckets[len(buckets)-2].upperBound, forcedMonotonic, fixedPrecision
+		return buckets[len(buckets)-2].UpperBound, forcedMonotonic, fixedPrecision
 	}
-	if b == 0 && buckets[0].upperBound <= 0 {
+	if b == 0 && buckets[0].UpperBound <= 0 {
-		return buckets[0].upperBound, forcedMonotonic, fixedPrecision
+		return buckets[0].UpperBound, forcedMonotonic, fixedPrecision
 	}
 	var (
 		bucketStart float64
-		bucketEnd   = buckets[b].upperBound
+		bucketEnd   = buckets[b].UpperBound
-		count       = buckets[b].count
+		count       = buckets[b].Count
 	)
 	if b > 0 {
-		bucketStart = buckets[b-1].upperBound
+		bucketStart = buckets[b-1].UpperBound
-		count -= buckets[b-1].count
+		count -= buckets[b-1].Count
-		rank -= buckets[b-1].count
+		rank -= buckets[b-1].Count
 	}
 	return bucketStart + (bucketEnd-bucketStart)*(rank/count), forcedMonotonic, fixedPrecision
 }
 
-// histogramQuantile calculates the quantile 'q' based on the given histogram.
+// HistogramQuantile calculates the quantile 'q' based on the given histogram.
 //
 // For custom buckets, the result is interpolated linearly, i.e. it is assumed
 // the observations are uniformly distributed within each bucket. (This is a
@@ -186,7 +195,13 @@ func bucketQuantile(q float64, buckets buckets) (float64, bool, bool) {
 // If q>1, +Inf is returned.
 //
 // If q is NaN, NaN is returned.
-func histogramQuantile(q float64, h *histogram.FloatHistogram) float64 {
+//
+// HistogramQuantile is for calculating the histogram_quantile() of native
+// histograms. See also: BucketQuantile for classic histograms.
+//
+// HistogramQuantile is exported as it may be used by other PromQL engine
+// implementations.
+func HistogramQuantile(q float64, h *histogram.FloatHistogram) float64 {
 	if q < 0 {
 		return math.Inf(-1)
 	}
@@ -297,11 +312,11 @@ func histogramQuantile(q float64, h *histogram.FloatHistogram) float64 {
 	return -math.Exp2(logUpper + (logLower-logUpper)*(1-fraction))
 }
 
-// histogramFraction calculates the fraction of observations between the
+// HistogramFraction calculates the fraction of observations between the
 // provided lower and upper bounds, based on the provided histogram.
 //
-// histogramFraction is in a certain way the inverse of histogramQuantile.  If
+// HistogramFraction is in a certain way the inverse of histogramQuantile.  If
-// histogramQuantile(0.9, h) returns 123.4, then histogramFraction(-Inf, 123.4, h)
+// HistogramQuantile(0.9, h) returns 123.4, then HistogramFraction(-Inf, 123.4, h)
 // returns 0.9.
 //
 // The same notes with regard to interpolation and assumptions about the zero
@@ -328,7 +343,10 @@ func histogramQuantile(q float64, h *histogram.FloatHistogram) float64 {
 // If lower or upper is NaN, NaN is returned.
 //
 // If lower >= upper and the histogram has at least 1 observation, zero is returned.
-func histogramFraction(lower, upper float64, h *histogram.FloatHistogram) float64 {
+//
+// HistogramFraction is exported as it may be used by other PromQL engine
+// implementations.
+func HistogramFraction(lower, upper float64, h *histogram.FloatHistogram) float64 {
 	if h.Count == 0 || math.IsNaN(lower) || math.IsNaN(upper) {
 		return math.NaN()
 	}
@@ -434,12 +452,12 @@ func histogramFraction(lower, upper float64, h *histogram.FloatHistogram) float6
 // coalesceBuckets merges buckets with the same upper bound.
 //
 // The input buckets must be sorted.
-func coalesceBuckets(buckets buckets) buckets {
+func coalesceBuckets(buckets Buckets) Buckets {
 	last := buckets[0]
 	i := 0
 	for _, b := range buckets[1:] {
-		if b.upperBound == last.upperBound {
+		if b.UpperBound == last.UpperBound {
-			last.count += b.count
+			last.Count += b.Count
 		} else {
 			buckets[i] = last
 			last = b
@@ -476,11 +494,11 @@ func coalesceBuckets(buckets buckets) buckets {
 //
 // We return a bool to indicate if this monotonicity was forced or not, and
 // another bool to indicate if small deltas were ignored or not.
-func ensureMonotonicAndIgnoreSmallDeltas(buckets buckets, tolerance float64) (bool, bool) {
+func ensureMonotonicAndIgnoreSmallDeltas(buckets Buckets, tolerance float64) (bool, bool) {
 	var forcedMonotonic, fixedPrecision bool
-	prev := buckets[0].count
+	prev := buckets[0].Count
 	for i := 1; i < len(buckets); i++ {
-		curr := buckets[i].count // Assumed always positive.
+		curr := buckets[i].Count // Assumed always positive.
 		if curr == prev {
 			// No correction needed if the counts are identical between buckets.
 			continue
@@ -489,14 +507,14 @@ func ensureMonotonicAndIgnoreSmallDeltas(buckets buckets, tolerance float64) (bo
 			// Silently correct numerically insignificant differences from floating
 			// point precision errors, regardless of direction.
 			// Do not update the 'prev' value as we are ignoring the difference.
-			buckets[i].count = prev
+			buckets[i].Count = prev
 			fixedPrecision = true
 			continue
 		}
 		if curr < prev {
 			// Force monotonicity by removing any decreases regardless of magnitude.
 			// Do not update the 'prev' value as we are ignoring the decrease.
-			buckets[i].count = prev
+			buckets[i].Count = prev
 			forcedMonotonic = true
 			continue
 		}

promql/quantile_test.go

@@ -24,29 +24,29 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 	eps := 1e-12
 
 	for name, tc := range map[string]struct {
-		getInput       func() buckets // The buckets can be modified in-place so return a new one each time.
+		getInput       func() Buckets // The buckets can be modified in-place so return a new one each time.
 		expectedForced bool
 		expectedFixed  bool
 		expectedValues map[float64]float64
 	}{
 		"simple - monotonic": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 10,
+						UpperBound: 10,
-						count:      10,
+						Count:      10,
 					}, {
-						upperBound: 15,
+						UpperBound: 15,
-						count:      15,
+						Count:      15,
 					}, {
-						upperBound: 20,
+						UpperBound: 20,
-						count:      15,
+						Count:      15,
 					}, {
-						upperBound: 30,
+						UpperBound: 30,
-						count:      15,
+						Count:      15,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      15,
+						Count:      15,
 					},
 				}
 			},
@@ -60,23 +60,23 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 			},
 		},
 		"simple - non-monotonic middle": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 10,
+						UpperBound: 10,
-						count:      10,
+						Count:      10,
 					}, {
-						upperBound: 15,
+						UpperBound: 15,
-						count:      15,
+						Count:      15,
 					}, {
-						upperBound: 20,
+						UpperBound: 20,
-						count:      15.00000000001, // Simulate the case there's a small imprecision in float64.
+						Count:      15.00000000001, // Simulate the case there's a small imprecision in float64.
 					}, {
-						upperBound: 30,
+						UpperBound: 30,
-						count:      15,
+						Count:      15,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      15,
+						Count:      15,
 					},
 				}
 			},
@@ -90,41 +90,41 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 			},
 		},
 		"real example - monotonic": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 1,
+						UpperBound: 1,
-						count:      6454661.3014166197,
+						Count:      6454661.3014166197,
 					}, {
-						upperBound: 5,
+						UpperBound: 5,
-						count:      8339611.2001912938,
+						Count:      8339611.2001912938,
 					}, {
-						upperBound: 10,
+						UpperBound: 10,
-						count:      14118319.2444762159,
+						Count:      14118319.2444762159,
 					}, {
-						upperBound: 25,
+						UpperBound: 25,
-						count:      14130031.5272856522,
+						Count:      14130031.5272856522,
 					}, {
-						upperBound: 50,
+						UpperBound: 50,
-						count:      46001270.3030008152,
+						Count:      46001270.3030008152,
 					}, {
-						upperBound: 64,
+						UpperBound: 64,
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					}, {
-						upperBound: 80,
+						UpperBound: 80,
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					}, {
-						upperBound: 100,
+						UpperBound: 100,
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					}, {
-						upperBound: 250,
+						UpperBound: 250,
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					}, {
-						upperBound: 1000,
+						UpperBound: 1000,
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      46008473.8585563600,
+						Count:      46008473.8585563600,
 					},
 				}
 			},
@@ -138,41 +138,41 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 			},
 		},
 		"real example - non-monotonic": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 1,
+						UpperBound: 1,
-						count:      6454661.3014166225,
+						Count:      6454661.3014166225,
 					}, {
-						upperBound: 5,
+						UpperBound: 5,
-						count:      8339611.2001912957,
+						Count:      8339611.2001912957,
 					}, {
-						upperBound: 10,
+						UpperBound: 10,
-						count:      14118319.2444762159,
+						Count:      14118319.2444762159,
 					}, {
-						upperBound: 25,
+						UpperBound: 25,
-						count:      14130031.5272856504,
+						Count:      14130031.5272856504,
 					}, {
-						upperBound: 50,
+						UpperBound: 50,
-						count:      46001270.3030008227,
+						Count:      46001270.3030008227,
 					}, {
-						upperBound: 64,
+						UpperBound: 64,
-						count:      46008473.8585563824,
+						Count:      46008473.8585563824,
 					}, {
-						upperBound: 80,
+						UpperBound: 80,
-						count:      46008473.8585563898,
+						Count:      46008473.8585563898,
 					}, {
-						upperBound: 100,
+						UpperBound: 100,
-						count:      46008473.8585563824,
+						Count:      46008473.8585563824,
 					}, {
-						upperBound: 250,
+						UpperBound: 250,
-						count:      46008473.8585563824,
+						Count:      46008473.8585563824,
 					}, {
-						upperBound: 1000,
+						UpperBound: 1000,
-						count:      46008473.8585563898,
+						Count:      46008473.8585563898,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      46008473.8585563824,
+						Count:      46008473.8585563824,
 					},
 				}
 			},
@@ -186,53 +186,53 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 			},
 		},
 		"real example 2 - monotonic": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 0.005,
+						UpperBound: 0.005,
-						count:      9.6,
+						Count:      9.6,
 					}, {
-						upperBound: 0.01,
+						UpperBound: 0.01,
-						count:      9.688888889,
+						Count:      9.688888889,
 					}, {
-						upperBound: 0.025,
+						UpperBound: 0.025,
-						count:      9.755555556,
+						Count:      9.755555556,
 					}, {
-						upperBound: 0.05,
+						UpperBound: 0.05,
-						count:      9.844444444,
+						Count:      9.844444444,
 					}, {
-						upperBound: 0.1,
+						UpperBound: 0.1,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 0.25,
+						UpperBound: 0.25,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 0.5,
+						UpperBound: 0.5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 1,
+						UpperBound: 1,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 2.5,
+						UpperBound: 2.5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 5,
+						UpperBound: 5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 10,
+						UpperBound: 10,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 25,
+						UpperBound: 25,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 50,
+						UpperBound: 50,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 100,
+						UpperBound: 100,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      9.888888889,
+						Count:      9.888888889,
 					},
 				}
 			},
@@ -246,53 +246,53 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 			},
 		},
 		"real example 2 - non-monotonic": {
-			getInput: func() buckets {
+			getInput: func() Buckets {
-				return buckets{
+				return Buckets{
 					{
-						upperBound: 0.005,
+						UpperBound: 0.005,
-						count:      9.6,
+						Count:      9.6,
 					}, {
-						upperBound: 0.01,
+						UpperBound: 0.01,
-						count:      9.688888889,
+						Count:      9.688888889,
 					}, {
-						upperBound: 0.025,
+						UpperBound: 0.025,
-						count:      9.755555556,
+						Count:      9.755555556,
 					}, {
-						upperBound: 0.05,
+						UpperBound: 0.05,
-						count:      9.844444444,
+						Count:      9.844444444,
 					}, {
-						upperBound: 0.1,
+						UpperBound: 0.1,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 0.25,
+						UpperBound: 0.25,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 0.5,
+						UpperBound: 0.5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 1,
+						UpperBound: 1,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 2.5,
+						UpperBound: 2.5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 5,
+						UpperBound: 5,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 10,
+						UpperBound: 10,
-						count:      9.888888889001, // Simulate the case there's a small imprecision in float64.
+						Count:      9.888888889001, // Simulate the case there's a small imprecision in float64.
 					}, {
-						upperBound: 25,
+						UpperBound: 25,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: 50,
+						UpperBound: 50,
-						count:      9.888888888999, // Simulate the case there's a small imprecision in float64.
+						Count:      9.888888888999, // Simulate the case there's a small imprecision in float64.
 					}, {
-						upperBound: 100,
+						UpperBound: 100,
-						count:      9.888888889,
+						Count:      9.888888889,
 					}, {
-						upperBound: math.Inf(1),
+						UpperBound: math.Inf(1),
-						count:      9.888888889,
+						Count:      9.888888889,
 					},
 				}
 			},
@@ -308,7 +308,7 @@ func TestBucketQuantile_ForcedMonotonicity(t *testing.T) {
 	} {
 		t.Run(name, func(t *testing.T) {
 			for q, v := range tc.expectedValues {
-				res, forced, fixed := bucketQuantile(q, tc.getInput())
+				res, forced, fixed := BucketQuantile(q, tc.getInput())
 				require.Equal(t, tc.expectedForced, forced)
 				require.Equal(t, tc.expectedFixed, fixed)
 				require.InEpsilon(t, v, res, eps)