optimize matrix enumeration

Mocha.sln

@@ -65,6 +65,8 @@ Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Mocha.Antlr4.Generated", "s
 EndProject
 Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Mocha.Storage.Benchmarks", "tests\Mocha.Storage.Benchmarks\Mocha.Storage.Benchmarks.csproj", "{4B63E7B0-C8C8-4206-A113-A3D4EFB7533E}"
 EndProject
+Project("{FAE04EC0-301F-11D3-BF4B-00C04F79EFBC}") = "Mocha.Query.Benchmarks", "tests\Mocha.Query.Benchmarks\Mocha.Query.Benchmarks.csproj", "{89EF2B9E-24CF-4A80-8199-284E34C11277}"
+EndProject
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "grafana", "grafana", "{5B280BB9-A7AE-46E6-B735-CF484965A736}"
 EndProject
 Project("{2150E333-8FDC-42A3-9474-1A3956D46DE8}") = "provisioning", "provisioning", "{C6E71F9B-BB5B-4946-9732-96CBA0B366B3}"
@@ -95,6 +97,7 @@ Global
 		{D56CA47A-A948-4FB5-9E16-C61E12535521} = {24F9E34A-D92A-4C0A-851F-1E864181BF97}
 		{904CC523-A2D4-4982-8A7B-A6A0F5A5EB19} = {6983D239-07DA-4DFA-9AAA-F6876029FF8D}
 		{4B63E7B0-C8C8-4206-A113-A3D4EFB7533E} = {24F9E34A-D92A-4C0A-851F-1E864181BF97}
+		{89EF2B9E-24CF-4A80-8199-284E34C11277} = {24F9E34A-D92A-4C0A-851F-1E864181BF97}
 		{5B280BB9-A7AE-46E6-B735-CF484965A736} = {D598862A-999C-40FD-A190-EBD00376D077}
 		{C6E71F9B-BB5B-4946-9732-96CBA0B366B3} = {5B280BB9-A7AE-46E6-B735-CF484965A736}
 		{8FAF9AF6-7399-41FC-96C9-43756AAD16CB} = {C6E71F9B-BB5B-4946-9732-96CBA0B366B3}
@@ -148,5 +151,9 @@ Global
 		{4B63E7B0-C8C8-4206-A113-A3D4EFB7533E}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{4B63E7B0-C8C8-4206-A113-A3D4EFB7533E}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{4B63E7B0-C8C8-4206-A113-A3D4EFB7533E}.Release|Any CPU.Build.0 = Release|Any CPU
+		{89EF2B9E-24CF-4A80-8199-284E34C11277}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{89EF2B9E-24CF-4A80-8199-284E34C11277}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{89EF2B9E-24CF-4A80-8199-284E34C11277}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{89EF2B9E-24CF-4A80-8199-284E34C11277}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 EndGlobal

src/Mocha.Query/Prometheus/PromQL/Engine/Evaluator.cs

@@ -11,7 +11,6 @@
 
 namespace Mocha.Query.Prometheus.PromQL.Engine;
 
-// TODO: Use ArrayPool
 internal class Evaluator
 {
     public long StartTimestampUnixSec { get; init; }
@@ -26,7 +25,8 @@ internal class Evaluator
 
     public MatrixResult Eval(Expression expr)
     {
-        var numSteps = (int)((EndTimestampUnixSec - StartTimestampUnixSec) / Interval.TotalSeconds) + 1;
+        var intervalSeconds = (long)Interval.TotalSeconds;
+        var numSteps = (int)((EndTimestampUnixSec - StartTimestampUnixSec) / intervalSeconds) + 1;
         MatrixResult result;
         switch (expr)
         {
@@ -132,18 +132,18 @@ public MatrixResult Eval(Expression expr)
                     var matrixSelector = (MatrixSelector)call.Args[matrixArgIndex];
                     result = new MatrixResult(matrixSelector.Series.Count());
                     var selectorOffsetSeconds = (long)matrixSelector.Offset.TotalSeconds;
-                    var selectorRangeSeconds = matrixSelector.Range.TotalSeconds;
-                    var stepRangeSeconds = (long)Math.Min(selectorRangeSeconds, Interval.TotalSeconds);
+                    var selectorRangeSeconds = (long)matrixSelector.Range.TotalSeconds;
 
                     // Reuse objects across steps to save memory allocations.
-                    // TODO: use ArrayPool
+                    var points = new List<DoublePoint>();
                     var inMatrix = new MatrixResult(1) { new Series { Metric = Labels.Empty, Points = [] } };
                     inArgs[matrixArgIndex] = inMatrix;
                     var enh = new EvalNodeHelper { Output = new VectorResult(1) };
 
                     // Process all the calls for one time series at a time.
                     foreach (var timeSeries in matrixSelector.Series)
                     {
+                        points.Clear();
                         var series = new Series
                         {
                             Metric = timeSeries.Labels.DropMetricName(),
@@ -155,7 +155,8 @@ public MatrixResult Eval(Expression expr)
                         var step = -1;
                         var refTimeStart = StartTimestampUnixSec - selectorOffsetSeconds;
                         var refTimeEnd = EndTimestampUnixSec - selectorOffsetSeconds;
-                        for (var ts = refTimeStart; ts <= refTimeEnd; ts += stepRangeSeconds)
+                        using var matrixEnumerator = new MatrixEnumerator(timeSeries.Samples);
+                        for (var ts = refTimeStart; ts <= refTimeEnd; ts += intervalSeconds)
                         {
                             step++;
                             // Set the non-matrix arguments.
@@ -172,18 +173,19 @@ public MatrixResult Eval(Expression expr)
                             var maxTs = ts;
                             var minTs = maxTs - selectorRangeSeconds;
                             // Evaluate the matrix selector for this series for this step.
-                            // TODO: optimize enumeration
-                            var points = timeSeries.Samples
-                                .Where(s => s.TimestampUnixSec >= minTs && s.TimestampUnixSec <= maxTs)
-                                .Select(s =>
-                                    new DoublePoint { TimestampUnixSec = s.TimestampUnixSec, Value = s.Value })
-                                .ToList();
+                            points = matrixEnumerator.Enumerate(minTs, maxTs, points);
 
                             if (points.Count <= 0)
                             {
                                 continue;
                             }
 
+                            _currentSamples += points.Count;
+                            if (_currentSamples > MaxSamples)
+                            {
+                                throw new TooManySamplesException();
+                            }
+
                             inMatrix[0].Points = points;
                             enh.TimestampUnixSec = ts;
                             enh.Output.Clear();
@@ -350,7 +352,6 @@ public MatrixResult Eval(Expression expr)
                             // TODO: use ArrayPool
                             Points = new List<DoublePoint>(numSteps)
                         };
-                        var intervalSeconds = (long)Interval.TotalSeconds;
                         var refTimeStart = StartTimestampUnixSec - offsetSeconds;
                         var refTimeEnd = EndTimestampUnixSec - offsetSeconds;
                         using var enumerator = timeSeries.Samples.Reverse().GetEnumerator();
@@ -427,7 +428,6 @@ public MatrixResult Eval(Expression expr)
         }
     }
 
-
     /// <summary>
     /// Evaluates the given expressions, and then for each step calls
     /// the given function with the values computed for each expression at that step.

src/Mocha.Query/Prometheus/PromQL/Engine/MatrixEnumerator.cs

@@ -0,0 +1,82 @@
+// Licensed to the .NET Core Community under one or more agreements.
+// The .NET Core Community licenses this file to you under the MIT license.
+
+using Mocha.Core.Storage.Prometheus.Metrics;
+using Mocha.Query.Prometheus.PromQL.Values;
+
+namespace Mocha.Query.Prometheus.PromQL.Engine;
+
+public class MatrixEnumerator(IEnumerable<TimeSeriesSample> samples) : IDisposable
+{
+    private readonly IEnumerator<TimeSeriesSample> _enumerator = samples.GetEnumerator();
+
+    public List<DoublePoint> Enumerate(
+        long minTs,
+        long maxTs,
+        List<DoublePoint> reusedPoints)
+    {
+        ArgumentNullException.ThrowIfNull(reusedPoints);
+
+        if (minTs >= maxTs)
+        {
+            throw new ArgumentException("minTs must be less than maxTs");
+        }
+
+        var keepFrom = 0;
+        while (keepFrom < reusedPoints.Count && reusedPoints[keepFrom].TimestampUnixSec < minTs)
+        {
+            keepFrom++;
+        }
+
+        // If there is an overlap between previous and current ranges, keep the overlapping part.
+        // If keepFrom is 0, all points are within the range, so keep them all.
+        if (keepFrom > 0)
+        {
+            reusedPoints.RemoveRange(0, keepFrom);
+        }
+        else if (keepFrom == reusedPoints.Count)
+        {
+            // No overlap, clear all points.
+            reusedPoints.Clear();
+        }
+
+        while (true)
+        {
+            // Current is uninitialized or has been fully consumed
+            if (_enumerator.Current == null)
+            {
+                if (!_enumerator.MoveNext())
+                {
+                    break;
+                }
+            }
+
+            var sample = _enumerator.Current;
+
+            // Future data, leave it for the next step
+            if (sample!.TimestampUnixSec > maxTs)
+            {
+                break;
+            }
+
+            // If the sample is within the range, add it to the points
+            if (sample.TimestampUnixSec >= minTs)
+            {
+                reusedPoints.Add(new DoublePoint { TimestampUnixSec = sample.TimestampUnixSec, Value = sample.Value });
+            }
+
+            // Move to the next sample
+            if (!_enumerator.MoveNext())
+            {
+                break;
+            }
+        }
+
+        return reusedPoints;
+    }
+
+    public void Dispose()
+    {
+        _enumerator.Dispose();
+    }
+}

tests/Mocha.Query.Benchmarks/MatrixEnumerationBenchmark.cs

@@ -0,0 +1,69 @@
+// Licensed to the .NET Core Community under one or more agreements.
+// The .NET Core Community licenses this file to you under the MIT license.
+
+using BenchmarkDotNet.Attributes;
+using Mocha.Core.Storage.Prometheus.Metrics;
+using Mocha.Query.Prometheus.PromQL.Engine;
+using Mocha.Query.Prometheus.PromQL.Values;
+
+[MemoryDiagnoser]
+public class MatrixEnumerationBenchmark
+{
+    [Params(1_000, 10_000)] public int SampleCount { get; set; }
+
+    [Params(15)] public int SampleIntervalSeconds { get; set; }
+
+    [Params(30, 60)] public int SelectorRangeSeconds { get; set; }
+
+    private List<TimeSeriesSample> _samples = null!;
+    private long _startTs;
+    private long _endTs;
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        _samples = new List<TimeSeriesSample>(SampleCount);
+
+        long ts = 0;
+        for (int i = 0; i < SampleCount; i++)
+        {
+            _samples.Add(new TimeSeriesSample { TimestampUnixSec = ts, Value = i });
+
+            ts += SampleIntervalSeconds;
+        }
+
+        _startTs = _samples[0].TimestampUnixSec;
+        _endTs = _samples[^1].TimestampUnixSec;
+    }
+
+    [Benchmark(Baseline = true)]
+    public void LinqEveryStep()
+    {
+        for (var ts = _startTs; ts <= _endTs; ts += SelectorRangeSeconds)
+        {
+            var maxTs = ts;
+            var minTs = maxTs - SelectorRangeSeconds;
+
+            var points = _samples
+                .Where(s => s.TimestampUnixSec >= minTs &&
+                            s.TimestampUnixSec <= maxTs)
+                .Select(s => new DoublePoint { TimestampUnixSec = s.TimestampUnixSec, Value = s.Value })
+                .ToList();
+        }
+    }
+
+    [Benchmark]
+    public void EnumeratorSlidingWindow()
+    {
+        var reusePoints = new List<DoublePoint>();
+        using var enumerator = new MatrixEnumerator(_samples);
+
+        for (var ts = _startTs; ts <= _endTs; ts += SelectorRangeSeconds)
+        {
+            var maxTs = ts;
+            var minTs = maxTs - SelectorRangeSeconds;
+
+            enumerator.Enumerate(minTs, maxTs, reusePoints);
+        }
+    }
+}

tests/Mocha.Query.Benchmarks/Mocha.Query.Benchmarks.csproj

@@ -0,0 +1,20 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+    <PropertyGroup>
+        <OutputType>Exe</OutputType>
+        <TargetFramework>net8.0</TargetFramework>
+        <ImplicitUsings>enable</ImplicitUsings>
+        <Nullable>enable</Nullable>
+        <EnforceCodeStyleInBuild>true</EnforceCodeStyleInBuild>
+    </PropertyGroup>
+
+    <ItemGroup>
+        <PackageReference Include="BenchmarkDotNet" Version="0.15.8" />
+    </ItemGroup>
+
+    <ItemGroup>
+      <ProjectReference Include="..\..\src\Mocha.Core\Mocha.Core.csproj" />
+      <ProjectReference Include="..\..\src\Mocha.Query\Mocha.Query.csproj" />
+    </ItemGroup>
+
+</Project>

tests/Mocha.Query.Benchmarks/Program.cs

@@ -0,0 +1,12 @@
+// Licensed to the .NET Core Community under one or more agreements.
+// The .NET Core Community licenses this file to you under the MIT license.
+
+using BenchmarkDotNet.Running;
+
+var allBenchmarks = new[]
+{
+    typeof(MatrixEnumerationBenchmark)
+};
+
+new BenchmarkSwitcher(allBenchmarks).Run(args);
+