From 1b3d72166c94b8c442281ee6459ddc7c136814e9 Mon Sep 17 00:00:00 2001
From: Gabriel Darbord <gabriel.darbord@inria.fr>
Date: Mon, 28 Oct 2024 14:52:05 +0100
Subject: [PATCH] Add `TopN` query

---
 .../OTelTopNQuery.class.st                    | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/src/Famix-OpenTelemetry-Importer/OTelTopNQuery.class.st b/src/Famix-OpenTelemetry-Importer/OTelTopNQuery.class.st
index ca77795..db464c5 100644
--- a/src/Famix-OpenTelemetry-Importer/OTelTopNQuery.class.st
+++ b/src/Famix-OpenTelemetry-Importer/OTelTopNQuery.class.st
@@ -1,9 +1,52 @@
+"
+I group spans with my `selector` block, sort them with my `comparator` (or do not sort if not provided), and keep only the top `count`.
+"
 Class {
 	#name : 'OTelTopNQuery',
 	#superclass : 'OTelQuery',
 	#instVars : [
+		'comparator',
 		'count'
 	],
 	#category : 'Famix-OpenTelemetry-Importer',
 	#package : 'Famix-OpenTelemetry-Importer'
 }
+
+{ #category : 'accessing' }
+OTelTopNQuery >> comparator [
+	"If not specified, do not sort"
+
+	^ comparator
+]
+
+{ #category : 'accessing' }
+OTelTopNQuery >> comparator: aTwoArgsBlock [
+
+	comparator := aTwoArgsBlock
+]
+
+{ #category : 'transforming' }
+OTelTopNQuery >> count [
+
+	^ count ifNil: [ count := 1 ]
+]
+
+{ #category : 'transforming' }
+OTelTopNQuery >> count: anInteger [
+
+	count := anInteger max: 1
+]
+
+{ #category : 'transforming' }
+OTelTopNQuery >> value: model [
+	"Only keep the top n of each group in the model."
+
+	(model allSpans groupedBy: self selector) do: [ :spans |
+		| sortedSpans |
+		sortedSpans := self comparator
+			               ifNil: [ spans ]
+			               ifNotNil: [ :comp | spans sorted: comp ].
+		self count + 1 to: spans size do: [ :i |
+			self removeSpan: (sortedSpans at: i) ] ].
+	^ model
+]