Merge branch 'main' into refactor-lokistack-upsert-handler

go.mod

@@ -115,7 +115,7 @@ require (
 
 require (
 	github.com/Azure/go-autorest/autorest v0.11.29
-	github.com/DataDog/sketches-go v1.4.2
+	github.com/DataDog/sketches-go v1.4.4
 	github.com/DmitriyVTitov/size v1.5.0
 	github.com/IBM/go-sdk-core/v5 v5.13.1
 	github.com/IBM/ibm-cos-sdk-go v1.10.0

go.sum

@@ -231,8 +231,8 @@ github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/DataDog/datadog-go v2.2.0+incompatible/go.mod h1:LButxg5PwREeZtORoXG3tL4fMGNddJ+vMq1mwgfaqoQ=
 github.com/DataDog/datadog-go v3.2.0+incompatible/go.mod h1:LButxg5PwREeZtORoXG3tL4fMGNddJ+vMq1mwgfaqoQ=
-github.com/DataDog/sketches-go v1.4.2 h1:gppNudE9d19cQ98RYABOetxIhpTCl4m7CnbRZjvVA/o=
-github.com/DataDog/sketches-go v1.4.2/go.mod h1:xJIXldczJyyjnbDop7ZZcLxJdV3+7Kra7H1KMgpgkLk=
+github.com/DataDog/sketches-go v1.4.4 h1:dF52vzXRFSPOj2IjXSWLvXq3jubL4CI69kwYjJ1w5Z8=
+github.com/DataDog/sketches-go v1.4.4/go.mod h1:XR0ns2RtEEF09mDKXiKZiQg+nfZStrq1ZuL1eezeZe0=
 github.com/DataDog/zstd v1.3.5/go.mod h1:1jcaCB/ufaK+sKp1NBhlGmpz41jOoPQ35bpF36t7BBo=
 github.com/DmitriyVTitov/size v1.5.0 h1:/PzqxYrOyOUX1BXj6J9OuVRVGe+66VL4D9FlUaW515g=
 github.com/DmitriyVTitov/size v1.5.0/go.mod h1:le6rNI4CoLQV1b9gzp1+3d7hMAD/uu2QcJ+aYbNgiU0=

pkg/logql/quantile_over_time_sketch.go

@@ -71,6 +71,12 @@ func (q ProbabilisticQuantileVector) ToProto() *logproto.QuantileSketchVector {
 	return &logproto.QuantileSketchVector{Samples: samples}
 }
 
+func (q ProbabilisticQuantileVector) Release() {
+	for _, s := range q {
+		s.F.Release()
+	}
+}
+
 func ProbabilisticQuantileVectorFromProto(proto *logproto.QuantileSketchVector) (ProbabilisticQuantileVector, error) {
 	out := make([]ProbabilisticQuantileSample, len(proto.Samples))
 	var s ProbabilisticQuantileSample
@@ -107,6 +113,9 @@ func (m ProbabilisticQuantileMatrix) Merge(right ProbabilisticQuantileMatrix) (P
 func (ProbabilisticQuantileMatrix) Type() promql_parser.ValueType { return QuantileSketchMatrixType }
 
 func (m ProbabilisticQuantileMatrix) Release() {
+	for _, vec := range m {
+		vec.Release()
+	}
 	quantileVectorPool.Put(m)
 }
 

pkg/logql/quantile_over_time_sketch_test.go

@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql"
 	"github.com/stretchr/testify/require"
 
 	"github.com/grafana/loki/pkg/logproto"
@@ -188,3 +189,38 @@ func (ev *sliceStepEvaluator) Next() (ok bool, ts int64, r StepResult) {
 	ok = ev.cur < len(ev.slice)
 	return
 }
+
+func BenchmarkQuantileBatchRangeVectorIteratorAt(b *testing.B) {
+	for _, tc := range []struct {
+		numberSamples int64
+	}{
+		{numberSamples: 1},
+		{numberSamples: 1_000},
+		{numberSamples: 100_000},
+	} {
+		b.Run(fmt.Sprintf("%d-samples", tc.numberSamples), func(b *testing.B) {
+			r := rand.New(rand.NewSource(42))
+
+			key := "group"
+			// similar to Benchmark_RangeVectorIterator
+			it := &quantileSketchBatchRangeVectorIterator{
+				batchRangeVectorIterator: &batchRangeVectorIterator{
+					window: map[string]*promql.Series{
+						key: {Floats: make([]promql.FPoint, tc.numberSamples)},
+					},
+				},
+			}
+			for i := int64(0); i < tc.numberSamples; i++ {
+				it.window[key].Floats[i] = promql.FPoint{T: i, F: r.Float64()}
+			}
+
+			b.ResetTimer()
+			b.ReportAllocs()
+
+			for n := 0; n < b.N; n++ {
+				_, r := it.At()
+				r.QuantileSketchVec().Release()
+			}
+		})
+	}
+}

pkg/logql/sketch/quantile.go

@@ -3,8 +3,11 @@ package sketch
 import (
 	"errors"
 	"fmt"
+	"sync"
 
 	"github.com/DataDog/sketches-go/ddsketch"
+	"github.com/DataDog/sketches-go/ddsketch/mapping"
+	"github.com/DataDog/sketches-go/ddsketch/store"
 	"github.com/influxdata/tdigest"
 
 	"github.com/grafana/loki/pkg/logproto"
@@ -16,6 +19,7 @@ type QuantileSketch interface {
 	Quantile(float64) (float64, error)
 	Merge(QuantileSketch) (QuantileSketch, error)
 	ToProto() *logproto.QuantileSketch
+	Release()
 }
 
 type QuantileSketchFactory func() QuantileSketch
@@ -38,8 +42,17 @@ type DDSketchQuantile struct {
 	*ddsketch.DDSketch
 }
 
+const relativeAccuracy = 0.01
+
+var ddsketchPool = sync.Pool{
+	New: func() any {
+		m, _ := mapping.NewCubicallyInterpolatedMapping(relativeAccuracy)
+		return ddsketch.NewDDSketchFromStoreProvider(m, store.DefaultProvider)
+	},
+}
+
 func NewDDSketch() *DDSketchQuantile {
-	s, _ := ddsketch.NewDefaultDDSketch(0.01)
+	s := ddsketchPool.Get().(*ddsketch.DDSketch)
 	return &DDSketchQuantile{s}
 }
 
@@ -68,6 +81,11 @@ func (d *DDSketchQuantile) ToProto() *logproto.QuantileSketch {
 	}
 }
 
+func (d *DDSketchQuantile) Release() {
+	d.DDSketch.Clear()
+	ddsketchPool.Put(d.DDSketch)
+}
+
 func DDSketchQuantileFromProto(buf []byte) (*DDSketchQuantile, error) {
 	sketch := NewDDSketch()
 	err := sketch.DDSketch.DecodeAndMergeWith(buf)
@@ -127,6 +145,8 @@ func (d *TDigestQuantile) ToProto() *logproto.QuantileSketch {
 	}
 }
 
+func (d *TDigestQuantile) Release() {}
+
 func TDigestQuantileFromProto(proto *logproto.TDigest) *TDigestQuantile {
 	q := &TDigestQuantile{tdigest.NewWithCompression(proto.Compression)}
 

vendor/modules.txt

@@ -158,7 +158,7 @@ github.com/AzureAD/microsoft-authentication-library-for-go/apps/internal/options
 github.com/AzureAD/microsoft-authentication-library-for-go/apps/internal/shared
 github.com/AzureAD/microsoft-authentication-library-for-go/apps/internal/version
 github.com/AzureAD/microsoft-authentication-library-for-go/apps/public
-# github.com/DataDog/sketches-go v1.4.2
+# github.com/DataDog/sketches-go v1.4.4
 ## explicit; go 1.15
 github.com/DataDog/sketches-go/ddsketch
 github.com/DataDog/sketches-go/ddsketch/encoding