Change ddsketch mapping to improve performance. (#11561)

**What this PR does / why we need it**:
We've found that the index mapping of our quantile over time
approximation has a big impact on the CPU. Changing the implementation
gives us around 50%. The mapping is used during the `At` calls to the
iterator.

```
› benchstat logarithmic.log cubic.log sort.log pool.log 
goos: linux
goarch: amd64
pkg: github.com/grafana/loki/pkg/logql
cpu: AMD Ryzen 7 3700X 8-Core Processor             
                                                     │ logarithmic.log │              cubic.log               │              sort.log               │              pool.log              │
                                                     │     sec/op      │    sec/op      vs base               │    sec/op     vs base               │   sec/op     vs base               │
QuantileBatchRangeVectorIteratorAt/1-samples-16            819.7n ± 2%   1052.5n ±  5%  +28.40% (p=0.002 n=6)   1055.0n ± 1%  +28.71% (p=0.002 n=6)   303.7n ± 4%  -62.96% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/1000-samples-16         60.34µ ± 8%    49.32µ ± 13%  -18.26% (p=0.002 n=6)    45.94µ ± 4%  -23.86% (p=0.002 n=6)   24.97µ ± 3%  -58.61% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/100000-samples-16       3.032m ± 3%    1.319m ±  1%  -56.50% (p=0.002 n=6)    1.316m ± 4%  -56.58% (p=0.002 n=6)   1.278m ± 3%  -57.86% (p=0.002 n=6)
geomean                                                    53.13µ         40.91µ        -23.00%                  39.96µ       -24.79%                 21.32µ       -59.87%

                                                     │ logarithmic.log │              cubic.log               │               sort.log               │             pool.log              │
                                                     │      B/op       │     B/op      vs base                │     B/op      vs base                │    B/op     vs base               │
QuantileBatchRangeVectorIteratorAt/1-samples-16            368.00 ± 0%    368.00 ± 0%       ~ (p=1.000 n=6) ¹    368.00 ± 0%       ~ (p=1.000 n=6) ¹   32.00 ± 0%  -91.30% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/1000-samples-16         4048.0 ± 0%    4048.0 ± 0%       ~ (p=1.000 n=6) ¹    3920.0 ± 0%  -3.16% (p=0.002 n=6)     104.0 ± 0%  -97.43% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/100000-samples-16       6192.0 ± 0%    6192.0 ± 0%       ~ (p=1.000 n=6) ¹    5936.0 ± 0%  -4.13% (p=0.002 n=6)     202.0 ± 5%  -96.74% (p=0.002 n=6)
geomean                                                   2.048Ki        2.048Ki       +0.00%                   1.998Ki       -2.45%                   87.60       -95.82%
¹ all samples are equal

                                                     │ logarithmic.log │              cubic.log              │               sort.log               │             pool.log              │
                                                     │    allocs/op    │  allocs/op   vs base                │  allocs/op   vs base                 │ allocs/op   vs base               │
QuantileBatchRangeVectorIteratorAt/1-samples-16             8.000 ± 0%    8.000 ± 0%       ~ (p=1.000 n=6) ¹    8.000 ± 0%        ~ (p=1.000 n=6) ¹   2.000 ± 0%  -75.00% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/1000-samples-16         27.000 ± 0%   27.000 ± 0%       ~ (p=1.000 n=6) ¹   23.000 ± 0%  -14.81% (p=0.002 n=6)     5.000 ± 0%  -81.48% (p=0.002 n=6)
QuantileBatchRangeVectorIteratorAt/100000-samples-16       42.000 ± 0%   42.000 ± 0%       ~ (p=1.000 n=6) ¹   34.000 ± 0%  -19.05% (p=0.002 n=6)     9.000 ± 0%  -78.57% (p=0.002 n=6)
geomean                                                     20.86         20.86       +0.00%                    18.43       -11.65%                   4.481       -78.51%
¹ all samples are equal

```

 

**Checklist**
- [ ] Reviewed the
[`CONTRIBUTING.md`](https://github.com/grafana/loki/blob/main/CONTRIBUTING.md)
guide (**required**)
- [ ] Documentation added
- [x] Tests updated
- [ ] `CHANGELOG.md` updated
- [ ] If the change is worth mentioning in the release notes, add
`add-to-release-notes` label
- [ ] Changes that require user attention or interaction to upgrade are
documented in `docs/sources/setup/upgrade/_index.md`
- [ ] For Helm chart changes bump the Helm chart version in
`production/helm/loki/Chart.yaml` and update
`production/helm/loki/CHANGELOG.md` and
`production/helm/loki/README.md`. [Example
PR](d10549e)
- [ ] If the change is deprecating or removing a configuration option,
update the `deprecated-config.yaml` and `deleted-config.yaml` files
respectively in the `tools/deprecated-config-checker` directory.
[Example
PR](0d4416a)

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