Mimir Query Engine: Add instant-vector functions

pkg/streamingpromql/functions.go

@@ -11,10 +11,19 @@ import (
 	"github.com/grafana/mimir/pkg/streamingpromql/types"
 )
 
-type InstantVectorFunctionOperatorFactory func(args []types.Operator, pool *pooling.LimitingPool) (types.InstantVectorOperator, error)
+type instantVectorFunctionOperatorFactories func(args []types.Operator, pool *pooling.LimitingPool) (types.InstantVectorOperator, error)
 
-func SimpleFunctionFactory(name string, seriesDataFunc functions.InstantVectorFunction) InstantVectorFunctionOperatorFactory {
-	// Simple Function Factory is for functions that have exactly 1 argument and drop the series name.
+// transformationFunctionOperatorFactory creates an InstantVectorFunctionOperatorFactory for functions
+// that have exactly 1 argument and drop the series name.
+//
+// Parameters:
+//   - name: The name of the function.
+//   - seriesDataFunc: The function to be wrapped.
+//
+// Returns:
+//
+//	An InstantVectorFunctionOperatorFactory.
+func transformationFunctionOperatorFactory(name string, seriesDataFunc functions.InstantVectorFunction) instantVectorFunctionOperatorFactories {
 	return func(args []types.Operator, pool *pooling.LimitingPool) (types.InstantVectorOperator, error) {
 		if len(args) != 1 {
 			// Should be caught by the PromQL parser, but we check here for safety.
@@ -37,7 +46,7 @@ func SimpleFunctionFactory(name string, seriesDataFunc functions.InstantVectorFu
 	}
 }
 
-func rateFunction(args []types.Operator, pool *pooling.LimitingPool) (types.InstantVectorOperator, error) {
+func rateFunctionOperator(args []types.Operator, pool *pooling.LimitingPool) (types.InstantVectorOperator, error) {
 	if len(args) != 1 {
 		// Should be caught by the PromQL parser, but we check here for safety.
 		return nil, fmt.Errorf("expected exactly 1 argument for rate, got %v", len(args))
@@ -55,12 +64,12 @@ func rateFunction(args []types.Operator, pool *pooling.LimitingPool) (types.Inst
 	}, nil
 }
 
-var _ InstantVectorFunctionOperatorFactory = rateFunction
+var _ instantVectorFunctionOperatorFactories = rateFunctionOperator
 
 // These functions return an instant-vector.
-var instantVectorFunctions = map[string]InstantVectorFunctionOperatorFactory{
-	"acos":            SimpleFunctionFactory("acos", functions.Acos),
-	"histogram_count": SimpleFunctionFactory("histogram_count", functions.HistogramCount),
-	"histogram_sum":   SimpleFunctionFactory("histogram_sum", functions.HistogramSum),
-	"rate":            rateFunction,
+var instantVectorFunctions = map[string]instantVectorFunctionOperatorFactories{
+	"acos":            transformationFunctionOperatorFactory("acos", functions.Acos),
+	"histogram_count": transformationFunctionOperatorFactory("histogram_count", functions.HistogramCount),
+	"histogram_sum":   transformationFunctionOperatorFactory("histogram_sum", functions.HistogramSum),
+	"rate":            rateFunctionOperator,
 }

pkg/streamingpromql/functions/common.go

@@ -19,7 +19,8 @@ func DropSeriesName(seriesMetadata []types.SeriesMetadata, _ *pooling.LimitingPo
 
 type InstantVectorFunction func(seriesData types.InstantVectorSeriesData, pool *pooling.LimitingPool) (types.InstantVectorSeriesData, error)
 
-func FloatTransformationFunc(transform func(f float64) float64) InstantVectorFunction {
+// floatTransformationFunc is not needed elsewhere, so it is not exported yet
+func floatTransformationFunc(transform func(f float64) float64) InstantVectorFunction {
 	return func(seriesData types.InstantVectorSeriesData, pool *pooling.LimitingPool) (types.InstantVectorSeriesData, error) {
 		for i := range seriesData.Floats {
 			seriesData.Floats[i].F = transform(seriesData.Floats[i].F)
@@ -34,6 +35,6 @@ func FloatTransformationDropHistogramsFunc(transform func(f float64) float64) In
 		// https://prometheus.io/docs/prometheus/latest/querying/functions
 		pool.PutHPointSlice(seriesData.Histograms)
 		seriesData.Histograms = nil
-		return FloatTransformationFunc(transform)(seriesData, pool)
+		return floatTransformationFunc(transform)(seriesData, pool)
 	}
 }

pkg/streamingpromql/functions/common_test.go

@@ -32,7 +32,7 @@ func TestDropSeriesName(t *testing.T) {
 
 func TestFloatTransformationFunc(t *testing.T) {
 	transform := func(f float64) float64 { return f * 2 }
-	transformFunc := FloatTransformationFunc(transform)
+	transformFunc := floatTransformationFunc(transform)
 
 	seriesData := types.InstantVectorSeriesData{
 		Floats: []promql.FPoint{

pkg/streamingpromql/testdata/ours/native_histograms.test

@@ -30,26 +30,24 @@ load 1m
 eval range from 0 to 5m step 1m mixed_metric
 	{__name__="mixed_metric"} 1 2 3 {{count:4 sum:5 buckets:[1 2 1]}} {{count:6 sum:8 buckets:[1 4 1]}} {{count:6 sum:8 buckets:[1 4 1]}}
 
-eval range from 0 to 4m step 1m histogram_count(mixed_metric)
-{} _ _ _ 4 6
+eval instant at 3m histogram_count(mixed_metric)
+	{} 4
 
-eval range from 0 to 4m step 1m histogram_sum(mixed_metric)
-{} _ _ _ 5 8
+eval instant at 4m histogram_sum(mixed_metric)
+	{} 8
 
-# histogram_count ignores any float values
-eval instant at 2m histogram_count(mixed_metric)
+eval range from 0 to 5m step 1m histogram_count(mixed_metric)
+{} _ _ _ 4 6 6
+
+eval range from 0 to 5m step 1m histogram_sum(mixed_metric)
+{} _ _ _ 5 8 8
 
 # histogram_count ignores any float values
-eval instant at 3m histogram_count(mixed_metric)
-	{} 4
+eval instant at 2m histogram_count(mixed_metric)
 
 # histogram_sum ignores any float values
 eval instant at 2m histogram_sum(mixed_metric)
 
-# histogram_sum ignores any float values
-eval instant at 4m histogram_sum(mixed_metric)
-	{} 8
-
 clear
 
 # Test multiple histograms