Skip to content

refactor(query engine): parallel inputs reading in aggregator #527

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
forshev wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(query engine): parallel inputs reading in aggregator #527

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 61 additions & 50 deletions query/exec/aggregator.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"cmp"
"fmt"
"slices"
"sync"

"github.com/ozontech/seq-db/query"
"github.com/ozontech/seq-db/query/encoding"
Expand Down Expand Up @@ -31,6 +32,8 @@ type DistributedAggregator struct {
state ExecutorState
inputs []query.RecordProducer

mu sync.Mutex

aggFunc seq.AggFunc
quantiles []float64

Expand Down Expand Up @@ -59,58 +62,13 @@ func NewDistributedAggregator(

func (a *DistributedAggregator) Next() *query.Record {
if a.state == ExecutorStateReadingInput {
// TODO: read from all inputs simultaneously (???)
var wg sync.WaitGroup

@dkharms dkharms Sep 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do I understand correctly that we will start at most #shards goroutines? Should we limit how many concurrent goroutines we can start?

for _, input := range a.inputs {
for {
r := input.Next()
if r == nil {
break
}

key := aggKey{
token: r.Vals[0].Decoded().(string),
ts: r.Vals[6].Decoded().(uint64),
}

s, exists := a.buckets[key]
if !exists {
s = seq.NewSamplesContainers()
}

if !exists {
s.Min = r.Vals[1].Decoded().(float64)
s.Max = r.Vals[2].Decoded().(float64)
} else {
s.Min = min(s.Min, r.Vals[1].Decoded().(float64))
s.Max = max(s.Max, r.Vals[2].Decoded().(float64))
}

s.Sum += r.Vals[3].Decoded().(float64)
s.Total += int64(r.Vals[4].Decoded().(uint64))

if a.aggFunc == seq.AggFuncQuantile {
for _, v := range r.Vals[7].Decoded().([]float64) {
s.InsertSample(v)
}
}

if a.aggFunc == seq.AggFuncUniqueCount {
if a.values == nil {
a.values = make(map[aggKey]map[string]struct{})
}
m, ok := a.values[key]
if !ok {
m = make(map[string]struct{})
a.values[key] = m
}
for _, v := range r.Vals[8].Decoded().([]string) {
m[v] = struct{}{}
}
}

a.buckets[key] = s
}
wg.Go(func() {
a.drainInput(input)
})
}
wg.Wait()

a.state = ExecutorStateProcessingData
}
Expand Down Expand Up @@ -171,6 +129,59 @@ func (a *DistributedAggregator) Next() *query.Record {
return r
}

func (a *DistributedAggregator) drainInput(input query.RecordProducer) {
for {
r := input.Next()
if r == nil {
break
}

key := aggKey{
token: r.Vals[0].Decoded().(string),
ts: r.Vals[6].Decoded().(uint64),
Comment on lines +140 to +141

@dkharms dkharms Sep 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need to introduce some schema for these records -- now these constants are hard-coded all over the codebase.

Maybe we can do something like:

type Col[T any] struct {
      Idx   int
      Type  DataType
}

func (r *Record) Get[T any](c Col[T]) T {
      return r.Vals[c.Idx].Decoded().(T)
}

I am fine with leaving TODO for that and creating an issue.

@forshev forshev Sep 23, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that's correct. we will have to introduce some column and schema objects

}

a.mu.Lock()

s, exists := a.buckets[key]
if !exists {
s = seq.NewSamplesContainers()
s.Min = r.Vals[1].Decoded().(float64)
s.Max = r.Vals[2].Decoded().(float64)
} else {
s.Min = min(s.Min, r.Vals[1].Decoded().(float64))
s.Max = max(s.Max, r.Vals[2].Decoded().(float64))
}

s.Sum += r.Vals[3].Decoded().(float64)
s.Total += int64(r.Vals[4].Decoded().(uint64))

if a.aggFunc == seq.AggFuncQuantile {
for _, v := range r.Vals[7].Decoded().([]float64) {
s.InsertSample(v)
}
}

if a.aggFunc == seq.AggFuncUniqueCount {
if a.values == nil {
a.values = make(map[aggKey]map[string]struct{})
}
m, ok := a.values[key]
if !ok {
m = make(map[string]struct{})
a.values[key] = m
}
for _, v := range r.Vals[8].Decoded().([]string) {
m[v] = struct{}{}
}
}

a.buckets[key] = s

a.mu.Unlock()
}
}

func (a *DistributedAggregator) Finalize() *query.Summary {
var total uint64
var firstErr error
Expand Down
Loading