From e02c3a5a051aff9856de11657f7447378f3cb04d Mon Sep 17 00:00:00 2001 From: Daniil Forshev Date: Wed, 30 Sep 2026 16:17:37 +0500 Subject: [PATCH] refactor(query engine): introduce Column type --- proxy/search/stream_search.go | 2 +- query/column.go | 119 ++++++++++++ query/exec/aggregator.go | 51 +++-- query/exec/filter.go | 18 +- query/exec/filter_test.go | 17 +- query/exec/merger.go | 121 +++--------- query/exec/merger_test.go | 328 +++++++++------------------------ query/exec/projector.go | 10 +- query/exec/projector_test.go | 4 +- storeapi/grpc_stream_search.go | 4 +- 10 files changed, 285 insertions(+), 389 deletions(-) create mode 100644 query/column.go diff --git a/proxy/search/stream_search.go b/proxy/search/stream_search.go index e8dd88101..fc21486f0 100644 --- a/proxy/search/stream_search.go +++ b/proxy/search/stream_search.go @@ -93,7 +93,7 @@ func (si *Ingestor) StreamSearch( mergedStream = exec.NewDistributedAggregator(producers, sr.Agg.Func, sr.Agg.Quantiles) } else { const seqIdColIdx = 0 - mergedDocsStream := exec.NewNMergedProducers(producers, seqIdColIdx, "", query.DataTypeSeqID, sr.Order) + mergedDocsStream := exec.NewNMergedProducers(producers, query.SeqIDColumn(seqIdColIdx), sr.Order) mergedStream = exec.NewLimiter(mergedDocsStream, uint32(sr.Size), uint32(offset)) } diff --git a/query/column.go b/query/column.go new file mode 100644 index 000000000..7443d541f --- /dev/null +++ b/query/column.go @@ -0,0 +1,119 @@ +package query + +import ( + insaneJSON "github.com/ozontech/insane-json" + + "github.com/ozontech/seq-db/seq" +) + +type Column[T any] struct { + idx int + dataType DataType + get func(*RecordVals) T +} + +func (c Column[T]) Idx() int { + return c.idx +} + +func (c Column[T]) DataType() DataType { + return c.dataType +} + +// Val returns value of r's column at c.idx index +func (c Column[T]) Val(r *Record) T { + return c.get(r.Vals[c.idx]) +} + +// RawData returns raw data of r's column at c.idx index +func (c Column[T]) RawData(r *Record) []byte { + return r.Vals[c.idx].RawData() +} + +func SeqIDColumn(idx int) Column[seq.ID] { + return Column[seq.ID]{ + idx: idx, + dataType: DataTypeSeqID, + get: (*RecordVals).AsSeqID, + } +} + +func BytesColumn(idx int) Column[[]byte] { + return Column[[]byte]{ + idx: idx, + dataType: DataTypeBytes, + get: (*RecordVals).AsBytes, + } +} + +func StringColumn(idx int) Column[string] { + return Column[string]{ + idx: idx, + dataType: DataTypeString, + get: (*RecordVals).AsString, + } +} + +func DocColumn(idx int) Column[*insaneJSON.Root] { + return Column[*insaneJSON.Root]{ + idx: idx, + dataType: DataTypeDocument, + get: (*RecordVals).AsDoc, + } +} + +func Float64Column(idx int) Column[float64] { + return Column[float64]{ + idx: idx, + dataType: DataTypeFloat64, + get: (*RecordVals).AsFloat64, + } +} + +func Uint64Column(idx int) Column[uint64] { + return Column[uint64]{ + idx: idx, + dataType: DataTypeUint64, + get: (*RecordVals).AsUint64, + } +} + +func Int64Column(idx int) Column[int64] { + return Column[int64]{ + idx: idx, + dataType: DataTypeInt64, + get: (*RecordVals).AsInt64, + } +} + +func Uint32Column(idx int) Column[uint32] { + return Column[uint32]{ + idx: idx, + dataType: DataTypeUint32, + get: (*RecordVals).AsUint32, + } +} + +func Int32Column(idx int) Column[int32] { + return Column[int32]{ + idx: idx, + dataType: DataTypeInt32, + get: (*RecordVals).AsInt32, + } +} + +func Float64ArrayColumn(idx int) Column[[]float64] { + return Column[[]float64]{ + idx: idx, + dataType: DataTypeFloat64Array, + get: (*RecordVals).AsFloat64Array, + } +} + +func StringArrayColumn(idx int) Column[[]string] { + return Column[[]string]{ + idx: idx, + dataType: DataTypeStringArray, + get: (*RecordVals).AsStringArray, + } +} diff --git a/query/exec/aggregator.go b/query/exec/aggregator.go index 95b232f0a..238b6e7aa 100644 --- a/query/exec/aggregator.go +++ b/query/exec/aggregator.go @@ -20,6 +20,23 @@ const ( ExecutorStateDone ) +var ( + aggColTokenIn = query.StringColumn(0) + aggColMinIn = query.Float64Column(1) + aggColMaxIn = query.Float64Column(2) + aggColSumIn = query.Float64Column(3) + aggColTotalIn = query.Uint64Column(4) + aggColTsIn = query.Uint64Column(6) + aggColSamplesIn = query.Float64ArrayColumn(7) + aggColValuesIn = query.StringArrayColumn(8) +) + +var ( + aggColTokenOut = query.StringColumn(0) + aggColValueOut = query.Float64Column(1) + aggColTsOut = query.Uint64Column(2) +) + // aggKey identifies a single timeseries bin: the grouping token plus the // (floored) timestamp. ts is 0 (DummyMID) for non-timeseries aggregations, so // all samples for the same token collapse into one bucket. @@ -137,8 +154,8 @@ func (a *DistributedAggregator) drainInput(input query.RecordProducer) { } key := aggKey{ - token: r.Vals[0].AsString(), - ts: r.Vals[6].AsUint64(), + token: aggColTokenIn.Val(r), + ts: aggColTsIn.Val(r), } a.mu.Lock() @@ -146,18 +163,18 @@ func (a *DistributedAggregator) drainInput(input query.RecordProducer) { s, exists := a.buckets[key] if !exists { s = seq.NewSamplesContainers() - s.Min = r.Vals[1].AsFloat64() - s.Max = r.Vals[2].AsFloat64() + s.Min = aggColMinIn.Val(r) + s.Max = aggColMaxIn.Val(r) } else { - s.Min = min(s.Min, r.Vals[1].AsFloat64()) - s.Max = max(s.Max, r.Vals[2].AsFloat64()) + s.Min = min(s.Min, aggColMinIn.Val(r)) + s.Max = max(s.Max, aggColMaxIn.Val(r)) } - s.Sum += r.Vals[3].AsFloat64() - s.Total += int64(r.Vals[4].AsUint64()) + s.Sum += aggColSumIn.Val(r) + s.Total += int64(aggColTotalIn.Val(r)) if a.aggFunc == seq.AggFuncQuantile { - for _, v := range r.Vals[7].AsFloat64Array() { + for _, v := range aggColSamplesIn.Val(r) { s.InsertSample(v) } } @@ -171,7 +188,7 @@ func (a *DistributedAggregator) drainInput(input query.RecordProducer) { m = make(map[string]struct{}) a.values[key] = m } - for _, v := range r.Vals[8].AsStringArray() { + for _, v := range aggColValuesIn.Val(r) { m[v] = struct{}{} } } @@ -199,21 +216,21 @@ func (a *DistributedAggregator) Finalize() *query.Summary { } func sortBuckets(aggFunc seq.AggFunc, buckets []*query.Record) { - // ts (Vals[2]) is the primary key (ASC), matching seq/qpr.go sortBuckets + // ts is the primary key (ASC), matching seq/qpr.go sortBuckets // where MID comes first. Within the same ts buckets are ordered by value. sortByTsValueDescNameAsc := func(left, right *query.Record) int { return cmp.Or( - cmp.Compare(left.Vals[2].AsUint64(), right.Vals[2].AsUint64()), - cmp.Compare(right.Vals[1].AsFloat64(), left.Vals[1].AsFloat64()), - cmp.Compare(left.Vals[0].AsString(), right.Vals[0].AsString()), + cmp.Compare(aggColTsOut.Val(left), aggColTsOut.Val(right)), + cmp.Compare(aggColValueOut.Val(right), aggColValueOut.Val(left)), + cmp.Compare(aggColTokenOut.Val(left), aggColTokenOut.Val(right)), ) } sortByTsValueNameAsc := func(left, right *query.Record) int { return cmp.Or( - cmp.Compare(left.Vals[2].AsUint64(), right.Vals[2].AsUint64()), - cmp.Compare(left.Vals[1].AsFloat64(), right.Vals[1].AsFloat64()), - cmp.Compare(left.Vals[0].AsString(), right.Vals[0].AsString()), + cmp.Compare(aggColTsOut.Val(left), aggColTsOut.Val(right)), + cmp.Compare(aggColValueOut.Val(left), aggColValueOut.Val(right)), + cmp.Compare(aggColTokenOut.Val(left), aggColTokenOut.Val(right)), ) } diff --git a/query/exec/filter.go b/query/exec/filter.go index 3375a101c..764fa9185 100644 --- a/query/exec/filter.go +++ b/query/exec/filter.go @@ -12,14 +12,10 @@ type FilterExpr[T any] interface { Eval(T) bool } -type ValGetter[T any] func(*query.RecordVals) T - type Filter[T any] struct { input query.RecordProducer - - colIdx int - getVal ValGetter[T] - expr FilterExpr[T] + col query.Column[T] + expr FilterExpr[T] // withTotal requests the accurate total of records that pass the filter. // When true, Finalize drains the (possibly partially consumed) input to the @@ -30,7 +26,7 @@ type Filter[T any] struct { // draining the remaining input. passed uint64 - // roots holds every record whose colIdx val has been decoded (and thus + // roots holds every record whose val has been decoded (and thus // Spawn'd an insaneJSON root). They are released back to the library pool in // Finalize. Record.Release is idempotent, so records that were forwarded // downstream (and released there too) are safe to release here as well. @@ -39,15 +35,13 @@ type Filter[T any] struct { func NewFilter[T any]( input query.RecordProducer, - colIdx int, - get ValGetter[T], + col query.Column[T], expr FilterExpr[T], withTotal bool, ) *Filter[T] { return &Filter[T]{ input: input, - colIdx: colIdx, - getVal: get, + col: col, expr: expr, withTotal: withTotal, } @@ -60,7 +54,7 @@ func (f *Filter[T]) Next() *query.Record { return nil } - passes := f.expr.Eval(f.getVal(r.Vals[f.colIdx])) + passes := f.expr.Eval(f.col.Val(r)) // The decoded root is now cached; keep a reference so Finalize can release it. f.roots = append(f.roots, r) if passes { diff --git a/query/exec/filter_test.go b/query/exec/filter_test.go index 06e1c22e4..1349ad828 100644 --- a/query/exec/filter_test.go +++ b/query/exec/filter_test.go @@ -16,7 +16,7 @@ func TestFilterEq(t *testing.T) { filterExpr := NewEq[uint32](cond) - testFilter(t, 0, (*query.RecordVals).AsUint32, filterExpr, func(r *query.Record) bool { + testFilter(t, query.Uint32Column(0), filterExpr, func(r *query.Record) bool { return r.Vals[0].AsUint32() == uint32(cond) }) } @@ -26,7 +26,7 @@ func TestFilterGt(t *testing.T) { filterExpr := NewGt[uint32](cond) - testFilter(t, 0, (*query.RecordVals).AsUint32, filterExpr, func(r *query.Record) bool { + testFilter(t, query.Uint32Column(0), filterExpr, func(r *query.Record) bool { return r.Vals[0].AsUint32() > uint32(cond) }) } @@ -36,7 +36,7 @@ func TestFilterLt(t *testing.T) { filterExpr := NewLt[uint32](cond) - testFilter(t, 0, (*query.RecordVals).AsUint32, filterExpr, func(r *query.Record) bool { + testFilter(t, query.Uint32Column(0), filterExpr, func(r *query.Record) bool { return r.Vals[0].AsUint32() < uint32(cond) }) } @@ -49,7 +49,7 @@ func TestDocumentFilter(t *testing.T) { filterExpr := NewDocFilter(field, NewEq(cond)) - testFilter(t, 1, (*query.RecordVals).AsDoc, filterExpr, func(r *query.Record) bool { + testFilter(t, query.DocColumn(1), filterExpr, func(r *query.Record) bool { field := r.Vals[1].AsDoc().Dig(field) return field.AsString() == cond }) @@ -57,8 +57,7 @@ func TestDocumentFilter(t *testing.T) { func testFilter[T any]( t *testing.T, - colIdx int, - get ValGetter[T], + col query.Column[T], filterExpr FilterExpr[T], wantFilterFunc func(*query.Record) bool, ) { @@ -74,7 +73,7 @@ func testFilter[T any]( } } - filter := NewFilter(&input, colIdx, get, filterExpr, false) + filter := NewFilter(&input, col, filterExpr, false) outputData := make([]*query.Record, 0) for r := filter.Next(); r != nil; r = filter.Next() { @@ -101,7 +100,7 @@ func TestFilterTotalDrainsInput(t *testing.T) { } } - filter := NewFilter(&input, 0, (*query.RecordVals).AsUint32, filterExpr, true) + filter := NewFilter(&input, query.Uint32Column(0), filterExpr, true) outputData := make([]*query.Record, 0) for i := 0; i < len(wantData); i++ { r := filter.Next() @@ -126,7 +125,7 @@ func TestFilterTotalErrorPropagated(t *testing.T) { err: assertErr, } - filter := NewFilter(&input, 0, (*query.RecordVals).AsUint32, filterExpr, true) + filter := NewFilter(&input, query.Uint32Column(0), filterExpr, true) for r := filter.Next(); r != nil; r = filter.Next() { } diff --git a/query/exec/merger.go b/query/exec/merger.go index e596ea174..7700dcfc0 100644 --- a/query/exec/merger.go +++ b/query/exec/merger.go @@ -7,16 +7,14 @@ import ( "github.com/ozontech/seq-db/seq" ) -type Merger struct { +type Merger[T any] struct { left, right query.RecordProducer curLeft, curRight *query.Record - colIdx int - field string - dataType query.DataType - order seq.DocsOrder - cmp func(any, any) int + col query.Column[T] + order seq.DocsOrder + cmp func(any, any) int // dedup drops records whose sort key repeats the previously emitted one. // It is enabled only for the seq.ID merge: shards may match the same @@ -26,41 +24,29 @@ type Merger struct { // dups counts records dropped by dedup so Finalize can subtract it from the merged total. dups uint64 - // roots holds records whose colIdx val has been decoded (Spawn'd an - // insaneJSON root) while comparing during the merge. They leave the merger - // once chosen, so this is the last owner and Finalize releases them. On the - // documents path the merger compares by SeqID and the vals stay undecoded, - // so roots stays empty; the DataTypeDocument path is handled for - // correctness. - roots []*query.Record - done bool } -func NewMerger( +func NewMerger[T any]( left query.RecordProducer, right query.RecordProducer, - colIdx int, - field string, - dataType query.DataType, + col query.Column[T], order seq.DocsOrder, -) *Merger { - return &Merger{ +) *Merger[T] { + return &Merger[T]{ left: left, right: right, - colIdx: colIdx, - field: field, - dataType: dataType, + col: col, order: order, - cmp: createCmpFunc(dataType), - dedup: dataType == query.DataTypeSeqID, + cmp: createCmpFunc(col.DataType()), + dedup: col.DataType() == query.DataTypeSeqID, curLeft: nil, curRight: nil, done: false, } } -func (m *Merger) Next() *query.Record { +func (m *Merger[T]) Next() *query.Record { if m.done { return nil } @@ -73,7 +59,7 @@ func (m *Merger) Next() *query.Record { if !m.dedup { return r } - val := m.extractValue(r) + val := m.col.Val(r) if m.lastVal != nil && m.cmp(val, m.lastVal) == 0 { // Skip duplicate. m.dups++ @@ -84,7 +70,7 @@ func (m *Merger) Next() *query.Record { } } -func (m *Merger) mergeNext() *query.Record { +func (m *Merger[T]) mergeNext() *query.Record { if m.curLeft == nil { m.curLeft = m.left.Next() } @@ -109,8 +95,8 @@ func (m *Merger) mergeNext() *query.Record { return r } - leftVal := m.extractValue(m.curLeft) - rightVal := m.extractValue(m.curRight) + leftVal := m.col.Val(m.curLeft) + rightVal := m.col.Val(m.curRight) compared := m.cmp(leftVal, rightVal) chooseLeft := compared <= 0 @@ -121,31 +107,15 @@ func (m *Merger) mergeNext() *query.Record { if chooseLeft { r := m.curLeft m.curLeft = m.left.Next() - m.trackRoot(r) return r } r := m.curRight m.curRight = m.right.Next() - m.trackRoot(r) return r } -func (m *Merger) Finalize() *query.Summary { - for _, r := range m.roots { - r.Release() - } - // The lookahead cursors may still hold partially consumed records whose - // colIdx val extractValue has decoded. - if m.dataType == query.DataTypeDocument { - if m.curLeft != nil { - m.curLeft.Release() - } - if m.curRight != nil { - m.curRight.Release() - } - } - +func (m *Merger[T]) Finalize() *query.Summary { left := m.left.Finalize() right := m.right.Finalize() summary := combineSummaries(left, right) @@ -155,16 +125,6 @@ func (m *Merger) Finalize() *query.Summary { return summary } -// trackRoot records a record leaving the merger if its colIdx val may have been -// decoded by extractValue, so Finalize can release the spawned insaneJSON root. -// Non-document types never decode an insaneJSON root, so tracking them is -// unnecessary (but harmless — Record.Release is a no-op for them). -func (m *Merger) trackRoot(r *query.Record) { - if m.dataType == query.DataTypeDocument { - m.roots = append(m.roots, r) - } -} - // combineSummaries merges the final summaries of two merged branches. The // totals are summed; an error from either side (if any) takes precedence. func combineSummaries(left, right *query.Summary) *query.Summary { @@ -184,34 +144,6 @@ func combineSummaries(left, right *query.Summary) *query.Summary { return summary } -func (m *Merger) extractValue(r *query.Record) any { - val := r.Vals[m.colIdx] - - switch m.dataType { - case query.DataTypeSeqID: - return val.AsSeqID() - case query.DataTypeDocument: - if m.field == "" { - return val.AsDoc() - } - return val.AsDoc().Dig(m.field).AsString() - case query.DataTypeString: - return val.AsString() - case query.DataTypeUint32: - return val.AsUint32() - case query.DataTypeUint64: - return val.AsUint64() - case query.DataTypeInt32: - return val.AsInt32() - case query.DataTypeInt64: - return val.AsInt64() - case query.DataTypeFloat64: - return val.AsFloat64() - default: - return "" - } -} - func createCmpFunc(dataType query.DataType) func(any, any) int { switch dataType { case query.DataTypeSeqID: @@ -226,9 +158,6 @@ func createCmpFunc(dataType query.DataType) func(any, any) int { return 0 } } - case query.DataTypeDocument: - // document field's values are extracted as strings - return func(a, b any) int { return cmp.Compare(a.(string), b.(string)) } case query.DataTypeUint32: return func(a, b any) int { return cmp.Compare(a.(uint32), b.(uint32)) } case query.DataTypeUint64: @@ -246,11 +175,9 @@ func createCmpFunc(dataType query.DataType) func(any, any) int { } } -func NewNMergedProducers( +func NewNMergedProducers[T any]( producers []query.RecordProducer, - colIdx int, - field string, - dataType query.DataType, + col query.Column[T], order seq.DocsOrder, ) query.RecordProducer { l := len(producers) @@ -258,17 +185,17 @@ func NewNMergedProducers( return &emptyRecordProducer{} } if l == 1 { - return NewMerger(producers[0], &emptyRecordProducer{}, colIdx, field, dataType, order) + return NewMerger(producers[0], &emptyRecordProducer{}, col, order) } if l == 2 { - return NewMerger(producers[0], producers[1], colIdx, field, dataType, order) + return NewMerger(producers[0], producers[1], col, order) } half := l / 2 - a := NewNMergedProducers(producers[:half], colIdx, field, dataType, order) - b := NewNMergedProducers(producers[half:], colIdx, field, dataType, order) + a := NewNMergedProducers(producers[:half], col, order) + b := NewNMergedProducers(producers[half:], col, order) - return NewMerger(a, b, colIdx, field, dataType, order) + return NewMerger(a, b, col, order) } type emptyRecordProducer struct{} diff --git a/query/exec/merger_test.go b/query/exec/merger_test.go index a695a117b..25c7a559d 100644 --- a/query/exec/merger_test.go +++ b/query/exec/merger_test.go @@ -1,7 +1,6 @@ package exec import ( - "fmt" "testing" "github.com/stretchr/testify/assert" @@ -11,159 +10,6 @@ import ( "github.com/ozontech/seq-db/seq" ) -func TestMergerAsc(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{ - "service-01", - "service-03", - "service-05", - }) - rightInput := makeMergerTestRecords([]string{ - "service-02", - "service-04", - "service-06", - }) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderAsc, - ) - - outputData := collectRecords(merger) - assert.Equal(t, []string{ - "service-01", "service-02", "service-03", - "service-04", "service-05", "service-06", - }, extractFieldValues(outputData, field)) -} - -func TestMergerDesc(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{ - "service-06", - "service-04", - "service-02", - }) - rightInput := makeMergerTestRecords([]string{ - "service-05", - "service-03", - "service-01", - }) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderDesc, - ) - - outputData := collectRecords(merger) - assert.Equal(t, []string{ - "service-06", "service-05", "service-04", - "service-03", "service-02", "service-01", - }, extractFieldValues(outputData, field)) -} - -func TestMergerLeftEmpty(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{}) - rightInput := makeMergerTestRecords([]string{ - "service-01", - "service-02", - }) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderAsc, - ) - - outputData := collectRecords(merger) - assert.Equal(t, []string{"service-01", "service-02"}, extractFieldValues(outputData, field)) -} - -func TestMergerRightEmpty(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{ - "service-01", - "service-02", - }) - rightInput := makeMergerTestRecords([]string{}) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderAsc, - ) - - outputData := collectRecords(merger) - assert.Equal(t, []string{"service-01", "service-02"}, extractFieldValues(outputData, field)) -} - -func TestMergerBothEmpty(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{}) - rightInput := makeMergerTestRecords([]string{}) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderAsc, - ) - - outputData := collectRecords(merger) - assert.Empty(t, outputData) -} - -func TestMergerDuplicates(t *testing.T) { - const field = "service" - - leftInput := makeMergerTestRecords([]string{ - "service-01", - "service-01", - "service-03", - }) - rightInput := makeMergerTestRecords([]string{ - "service-01", - "service-02", - "service-03", - }) - - merger := NewMerger( - &testProducer{data: leftInput}, - &testProducer{data: rightInput}, - 1, - field, - query.DataTypeDocument, - seq.DocsOrderAsc, - ) - - outputData := collectRecords(merger) - assert.Equal(t, []string{ - "service-01", "service-01", "service-01", - "service-02", "service-03", "service-03", - }, extractFieldValues(outputData, field)) -} - func TestMergerUint32(t *testing.T) { leftInput := makeMergerUint32Records([]uint32{1, 3, 5}) rightInput := makeMergerUint32Records([]uint32{2, 4, 6}) @@ -171,9 +17,7 @@ func TestMergerUint32(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeUint32, + query.Uint32Column(0), seq.DocsOrderAsc, ) @@ -188,9 +32,7 @@ func TestMergerUint64(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeUint64, + query.Uint64Column(0), seq.DocsOrderAsc, ) @@ -205,9 +47,7 @@ func TestMergerInt32(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeInt32, + query.Int32Column(0), seq.DocsOrderAsc, ) @@ -222,9 +62,7 @@ func TestMergerInt64(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeInt64, + query.Int64Column(0), seq.DocsOrderDesc, ) @@ -239,9 +77,7 @@ func TestMergerFloat64(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeFloat64, + query.Float64Column(0), seq.DocsOrderAsc, ) @@ -256,9 +92,7 @@ func TestMergerString(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeString, + query.StringColumn(0), seq.DocsOrderAsc, ) @@ -266,6 +100,21 @@ func TestMergerString(t *testing.T) { assert.Equal(t, []string{"apple", "apricot", "banana", "cherry", "date", "elder"}, extractStringValues(outputData)) } +func TestMergerStringDuplicates(t *testing.T) { + leftInput := makeMergerStringRecords([]string{"apple", "apple", "cherry"}) + rightInput := makeMergerStringRecords([]string{"apple", "banana", "elder"}) + + merger := NewMerger( + &testProducer{data: leftInput}, + &testProducer{data: rightInput}, + query.StringColumn(0), + seq.DocsOrderAsc, + ) + + outputData := collectRecords(merger) + assert.Equal(t, []string{"apple", "apple", "apple", "banana", "cherry", "elder"}, extractStringValues(outputData)) +} + func TestMergerSeqIDAsc(t *testing.T) { leftInput := makeMergerSeqIDRecords([]seq.ID{ {MID: 100, RID: 1}, @@ -281,9 +130,7 @@ func TestMergerSeqIDAsc(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), seq.DocsOrderAsc, ) @@ -309,9 +156,7 @@ func TestMergerSeqIDDesc(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), seq.DocsOrderDesc, ) @@ -335,9 +180,7 @@ func TestMergerSeqIDSameMIDDifferentRID(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), seq.DocsOrderAsc, ) @@ -358,9 +201,27 @@ func TestMergerSeqIDLeftEmpty(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput}, &testProducer{data: rightInput}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), + seq.DocsOrderAsc, + ) + + outputData := collectRecords(merger) + assert.Equal(t, []seq.ID{ + {MID: 100, RID: 1}, {MID: 200, RID: 2}, + }, extractSeqIDValues(outputData)) +} + +func TestMergerSeqIDRightEmpty(t *testing.T) { + leftInput := makeMergerSeqIDRecords([]seq.ID{ + {MID: 100, RID: 1}, + {MID: 200, RID: 2}, + }) + rightInput := makeMergerSeqIDRecords([]seq.ID{}) + + merger := NewMerger( + &testProducer{data: leftInput}, + &testProducer{data: rightInput}, + query.SeqIDColumn(0), seq.DocsOrderAsc, ) @@ -370,6 +231,21 @@ func TestMergerSeqIDLeftEmpty(t *testing.T) { }, extractSeqIDValues(outputData)) } +func TestMergerSeqIDBothEmpty(t *testing.T) { + leftInput := makeMergerSeqIDRecords([]seq.ID{}) + rightInput := makeMergerSeqIDRecords([]seq.ID{}) + + merger := NewMerger( + &testProducer{data: leftInput}, + &testProducer{data: rightInput}, + query.SeqIDColumn(0), + seq.DocsOrderAsc, + ) + + outputData := collectRecords(merger) + assert.Empty(t, outputData) +} + func TestMergerSeqIDDuplicates(t *testing.T) { leftInput := makeMergerSeqIDRecords([]seq.ID{ {MID: 100, RID: 1}, @@ -385,9 +261,7 @@ func TestMergerSeqIDDuplicates(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput, total: uint64(len(leftInput))}, &testProducer{data: rightInput, total: uint64(len(rightInput))}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), seq.DocsOrderAsc, ) @@ -412,9 +286,7 @@ func TestMergerSeqIDDuplicatesDesc(t *testing.T) { merger := NewMerger( &testProducer{data: leftInput, total: uint64(len(leftInput))}, &testProducer{data: rightInput, total: uint64(len(rightInput))}, - 0, - "", - query.DataTypeSeqID, + query.SeqIDColumn(0), seq.DocsOrderDesc, ) @@ -427,22 +299,6 @@ func TestMergerSeqIDDuplicatesDesc(t *testing.T) { assert.Equal(t, uint64(3), summary.Total) } -func makeMergerTestRecords(values []string) []*query.Record { - out := make([]*query.Record, 0, len(values)) - - for _, v := range values { - doc := fmt.Sprintf(`{"service":%q,"level":3}`, v) - out = append(out, &query.Record{ - Vals: []*query.RecordVals{ - query.NewRecordVals(query.DataTypeUint32, encoding.Uint32ToBytes(1)), - query.NewRecordVals(query.DataTypeDocument, []byte(doc)), - }, - }) - } - - return out -} - func makeMergerUint32Records(values []uint32) []*query.Record { out := make([]*query.Record, 0, len(values)) @@ -597,69 +453,62 @@ func extractSeqIDValues(records []*query.Record) []seq.ID { } func TestNewNMergedProducersEmpty(t *testing.T) { - const field = "service" - producers := []query.RecordProducer{} - merger := NewNMergedProducers(producers, 1, field, query.DataTypeDocument, seq.DocsOrderAsc) + merger := NewNMergedProducers(producers, query.SeqIDColumn(0), seq.DocsOrderAsc) outputData := collectRecords(merger) assert.Empty(t, outputData) } func TestNewNMergedProducersSingle(t *testing.T) { - const field = "service" - - input := makeMergerTestRecords([]string{ - "service-01", - "service-02", + input := makeMergerSeqIDRecords([]seq.ID{ + {MID: 200, RID: 2}, + {MID: 100, RID: 1}, }) producers := []query.RecordProducer{ &testProducer{data: input}, } - merger := NewNMergedProducers(producers, 1, field, query.DataTypeDocument, seq.DocsOrderAsc) + merger := NewNMergedProducers(producers, query.SeqIDColumn(0), seq.DocsOrderAsc) outputData := collectRecords(merger) - assert.Equal(t, []string{"service-01", "service-02"}, extractFieldValues(outputData, field)) + assert.Equal(t, []seq.ID{{MID: 200, RID: 2}, {MID: 100, RID: 1}}, extractSeqIDValues(outputData)) } func TestNewNMergedProducersThree(t *testing.T) { - const field = "service" - - producer1 := makeMergerTestRecords([]string{"service-01", "service-04"}) - producer2 := makeMergerTestRecords([]string{"service-02", "service-05"}) - producer3 := makeMergerTestRecords([]string{"service-03", "service-06"}) + producer1 := makeMergerSeqIDRecords([]seq.ID{{MID: 100, RID: 1}, {MID: 400, RID: 4}}) + producer2 := makeMergerSeqIDRecords([]seq.ID{{MID: 200, RID: 2}, {MID: 500, RID: 5}}) + producer3 := makeMergerSeqIDRecords([]seq.ID{{MID: 300, RID: 3}, {MID: 600, RID: 6}}) producers := []query.RecordProducer{ &testProducer{data: producer1}, &testProducer{data: producer2}, &testProducer{data: producer3}, } - merger := NewNMergedProducers(producers, 1, field, query.DataTypeDocument, seq.DocsOrderAsc) + merger := NewNMergedProducers(producers, query.SeqIDColumn(0), seq.DocsOrderAsc) outputData := collectRecords(merger) - assert.Equal(t, []string{ - "service-01", "service-02", "service-03", - "service-04", "service-05", "service-06", - }, extractFieldValues(outputData, field)) + assert.Equal(t, []seq.ID{ + {MID: 100, RID: 1}, {MID: 200, RID: 2}, + {MID: 300, RID: 3}, {MID: 400, RID: 4}, + {MID: 500, RID: 5}, {MID: 600, RID: 6}, + }, extractSeqIDValues(outputData)) } func TestNewNMergedProducersWithEmpty(t *testing.T) { - const field = "service" - - producer1 := makeMergerTestRecords([]string{"service-03", "service-01"}) - producer2 := makeMergerTestRecords([]string{}) - producer3 := makeMergerTestRecords([]string{"service-02"}) + producer1 := makeMergerSeqIDRecords([]seq.ID{{MID: 300, RID: 3}, {MID: 100, RID: 1}}) + producer2 := makeMergerSeqIDRecords([]seq.ID{}) + producer3 := makeMergerSeqIDRecords([]seq.ID{{MID: 200, RID: 2}}) producers := []query.RecordProducer{ &testProducer{data: producer1}, &testProducer{data: producer2}, &testProducer{data: producer3}, } - merger := NewNMergedProducers(producers, 1, field, query.DataTypeDocument, seq.DocsOrderDesc) + merger := NewNMergedProducers(producers, query.SeqIDColumn(0), seq.DocsOrderDesc) outputData := collectRecords(merger) - assert.Equal(t, []string{ - "service-03", "service-02", "service-01", - }, extractFieldValues(outputData, field)) + assert.Equal(t, []seq.ID{ + {MID: 300, RID: 3}, {MID: 200, RID: 2}, {MID: 100, RID: 1}, + }, extractSeqIDValues(outputData)) } func collectRecords(p query.RecordProducer) []*query.Record { @@ -669,12 +518,3 @@ func collectRecords(p query.RecordProducer) []*query.Record { } return out } - -func extractFieldValues(records []*query.Record, field string) []string { - out := make([]string, 0, len(records)) - for _, r := range records { - val := r.Vals[1].AsDoc().Dig(field).AsString() - out = append(out, val) - } - return out -} diff --git a/query/exec/projector.go b/query/exec/projector.go index e35bd36d1..3460459f3 100644 --- a/query/exec/projector.go +++ b/query/exec/projector.go @@ -15,18 +15,18 @@ type FieldsFilter struct { type DocProjector struct { input query.RecordProducer - colIdx int + col query.Column[*insaneJSON.Root] filter *FieldsFilter } func NewDocProjector( input query.RecordProducer, - colIdx int, + col query.Column[*insaneJSON.Root], filter *FieldsFilter, ) *DocProjector { return &DocProjector{ input: input, - colIdx: colIdx, + col: col, filter: filter, } } @@ -37,7 +37,7 @@ func (p *DocProjector) Next() *query.Record { return nil } - decoder := r.Vals[p.colIdx].AsDoc() + decoder := p.col.Val(r) var newRecord *query.Record if !p.filter.AllowList { @@ -77,7 +77,7 @@ func (p *DocProjector) makeRecordWithNewVals(old *query.Record, newRawData []byt newRecordVals := make([]*query.RecordVals, len(old.Vals)) for i := range len(old.Vals) { rawData := old.Vals[i].RawData() - if i == p.colIdx { + if i == p.col.Idx() { rawData = newRawData } newRecordVals[i] = query.NewRecordVals(old.Vals[i].Type, rawData) diff --git a/query/exec/projector_test.go b/query/exec/projector_test.go index 35812fbd1..4a3e1b213 100644 --- a/query/exec/projector_test.go +++ b/query/exec/projector_test.go @@ -19,7 +19,7 @@ func TestDocProjectorFields(t *testing.T) { ) } -func TestProjectorFieldsExcepr(t *testing.T) { +func TestProjectorFieldsExcept(t *testing.T) { testDocProjector( t, &FieldsFilter{Fields: []string{"level"}, AllowList: false}, @@ -36,7 +36,7 @@ func testDocProjector(t *testing.T, fieldsFilter *FieldsFilter, wantDocs []strin inputData := makeTestInputRecords(2) input := testProducer{data: inputData} - projector := NewDocProjector(&input, 1, fieldsFilter) + projector := NewDocProjector(&input, query.DocColumn(1), fieldsFilter) outputData := make([]*query.Record, 0) for r := projector.Next(); r != nil; r = projector.Next() { diff --git a/storeapi/grpc_stream_search.go b/storeapi/grpc_stream_search.go index 49f67400f..1992ae8d3 100644 --- a/storeapi/grpc_stream_search.go +++ b/storeapi/grpc_stream_search.go @@ -388,10 +388,10 @@ func (g *GrpcV1) buildProducer( return producer, typing, nil } if docFilter != nil { - producer = exec.NewFilter(producer, docDataColIdx, (*query.RecordVals).AsDoc, docFilter, req.WithTotal) + producer = exec.NewFilter(producer, query.DocColumn(docDataColIdx), docFilter, req.WithTotal) } if fieldsFilter != nil { - producer = exec.NewDocProjector(producer, docDataColIdx, fieldsFilter) + producer = exec.NewDocProjector(producer, query.DocColumn(docDataColIdx), fieldsFilter) } if searchParams.Limit > 0 { // set limit=limit+offset and offset=0 to merge stores' results correctly on proxy