From c5804e0d976a689901db3134120dc1424379b6ef Mon Sep 17 00:00:00 2001 From: Andrei Cheboksarov <37665782+cheb0@users.noreply.github.com> Date: Tue, 28 Jul 2026 11:40:45 +0400 Subject: [PATCH] batch execution AND/OR/AND NOT/range --- cmd/seq-db/seq-db.go | 8 +- config/config.go | 9 + frac/active_index.go | 20 +- frac/common/seal_params.go | 2 +- frac/config.go | 14 +- frac/fraction_test.go | 47 ++- frac/processor/aggregator_test.go | 4 + frac/processor/batch_eval_tree.go | 223 ++++++++++++ frac/processor/batch_eval_tree_test.go | 254 ++++++++++++++ frac/processor/eval_tree.go | 14 + frac/processor/search.go | 204 +++++------ frac/sealed/lids/block_test.go | 10 +- frac/sealed/lids/iterator_batched_asc.go | 6 +- frac/sealed/lids/iterator_batched_desc.go | 6 +- frac/sealed_index.go | 54 ++- node/batch.go | 110 ++++-- node/batch_ops.go | 197 +++++++++++ node/batch_ops_test.go | 403 ++++++++++++++++++++++ node/batch_test.go | 59 ++++ node/node.go | 4 +- node/node_and.go | 54 +++ node/node_nand.go | 57 +++ node/node_or.go | 124 +++++++ node/node_static.go | 88 +++++ node/util.go | 81 +++++ 25 files changed, 1912 insertions(+), 140 deletions(-) create mode 100644 frac/processor/batch_eval_tree.go create mode 100644 frac/processor/batch_eval_tree_test.go create mode 100644 node/batch_ops.go create mode 100644 node/batch_ops_test.go create mode 100644 node/util.go diff --git a/cmd/seq-db/seq-db.go b/cmd/seq-db/seq-db.go index 76d7cfb6a..85e6dcb80 100644 --- a/cmd/seq-db/seq-db.go +++ b/cmd/seq-db/seq-db.go @@ -275,7 +275,7 @@ func startStore( DocBlocksZstdLevel: cfg.Compression.DocBlockZstdCompressionLevel, DocBlockSize: int(cfg.DocsSorting.DocBlockSize), TokenFreqThresholdPercentage: cfg.Sealing.Tokens.FreqThresholdPercentage, - LIDsBitmapThreshold: cfg.Sealing.Lids.BitmapThreshold, + LIDsBitmapThreshold: cfg.Sealing.Lids.BitmapThreshold, }, Fraction: frac.Config{ Search: frac.SearchConfig{ @@ -285,6 +285,12 @@ func startStore( MaxGroupTokens: cfg.Limits.Aggregation.GroupTokens, MaxTIDsPerFraction: cfg.Limits.Aggregation.FractionTokens, }, + QueryOptimization: frac.QueryOptimizationConfig{ + BatchExecution: frac.BatchExecutionConfig{ + Enabled: cfg.QueryOptimization.BatchExecution.Enabled, + CostThreshold: cfg.QueryOptimization.BatchExecution.CostThreshold, + }, + }, }, SkipSortDocs: !cfg.DocsSorting.Enabled, KeepWalFile: false, diff --git a/config/config.go b/config/config.go index f55bc9e3d..507318859 100644 --- a/config/config.go +++ b/config/config.go @@ -176,6 +176,15 @@ type Config struct { } `config:"aggregation"` } `config:"limits"` + QueryOptimization struct { + BatchExecution struct { + Enabled bool `config:"enabled"` + // CostThreshold is the minimum estimated non-batched execution cost required to enable batch-at-a-time query + // evaluation. Suggestion is to use value which is greater than 3 x LID block size. + CostThreshold int `config:"cost_threshold" default:"50000"` + } `config:"batch_execution"` + } `config:"query_optimization"` + CircuitBreaker struct { Bulk struct { // Checkout [CircuitBreaker] for more information. diff --git a/frac/active_index.go b/frac/active_index.go index 9a07fc554..d5d1ebea5 100644 --- a/frac/active_index.go +++ b/frac/active_index.go @@ -110,6 +110,9 @@ func (dp *activeDataProvider) Search(params processor.SearchParams) (*seq.QPR, e params.To = min(params.To, dp.info.To) aggLimits := processor.AggLimits(dp.config.Search.AggLimits) + queryOpt := processor.QueryOptimizationConfig{ + BatchExecution: processor.BatchExecutionConfig(dp.config.Search.QueryOptimization.BatchExecution), + } sw := stopwatch.New() @@ -132,7 +135,7 @@ func (dp *activeDataProvider) Search(params processor.SearchParams) (*seq.QPR, e qprs := make([]*seq.QPR, 0, len(indexes)) for _, si := range indexes { - qpr, err := processor.IndexSearch(dp.ctx, params, &si, aggLimits, sw) + qpr, err := processor.IndexSearch(dp.ctx, params, &si, aggLimits, queryOpt, sw) if err != nil { return nil, err } @@ -248,6 +251,10 @@ func (si *activeTokenIndex) GetTIDsByTokenExpr(t parser.Token) ([]uint32, error) return si.tokenList.FindPattern(si.ctx, t) } +func (si *activeTokenIndex) GetFreqsByTIDs(tids []uint32, field string) []uint32 { + return make([]uint32, len(tids)) +} + func (si *activeTokenIndex) GetLIDsFromTIDs(tids []uint32, _ lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.Node { nodes := make([]node.Node, 0, len(tids)) for _, tid := range tids { @@ -259,6 +266,17 @@ func (si *activeTokenIndex) GetLIDsFromTIDs(tids []uint32, _ lids.Counter, minLI return nodes } +func (si *activeTokenIndex) GetBatchedLIDsFromTIDs(tids []uint32, _ lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.BatchedNode { + nodes := make([]node.BatchedNode, 0, len(tids)) + for _, tid := range tids { + tlids := si.tokenList.Provide(tid) + unmapped := tlids.GetLIDs(si.mids, si.rids) + inverse := inverseLIDs(unmapped, si.inverser, minLID, maxLID) + nodes = append(nodes, node.NewStaticBatched(inverse, order.IsReverse())) + } + return nodes +} + func inverseLIDs(unmapped []uint32, inv *inverser, minLID, maxLID uint32) []uint32 { result := make([]uint32, 0, len(unmapped)) for _, v := range unmapped { diff --git a/frac/common/seal_params.go b/frac/common/seal_params.go index f7633ae84..4105e7559 100644 --- a/frac/common/seal_params.go +++ b/frac/common/seal_params.go @@ -13,7 +13,7 @@ type SealParams struct { DocBlocksZstdLevel int // DocBlocksZstdLevel is the zstd compress level of each document block. LIDBlockSize int - LIDsBitmapThreshold int // LIDsBitmapThreshold is the minimum number of LIDs in the lid list to serialize as bitmap. + LIDsBitmapThreshold int // LIDsBitmapThreshold is the minimum number of LIDs in the lid list to serialize as bitmap. TokenBlockSize int TokenFreqThresholdPercentage float64 DocBlockSize int // DocBlockSize is decompressed payload size of document block. diff --git a/frac/config.go b/frac/config.go index e91392aaf..e64632cf7 100644 --- a/frac/config.go +++ b/frac/config.go @@ -8,7 +8,8 @@ type Config struct { } type SearchConfig struct { - AggLimits AggLimits + AggLimits AggLimits + QueryOptimization QueryOptimizationConfig } type AggLimits struct { @@ -17,3 +18,14 @@ type AggLimits struct { MaxGroupTokens int // MaxGroupTokens max AggQuery.GroupBy unique values. MaxTIDsPerFraction int // MaxTIDsPerFraction max number of tokens per fraction. } + +type QueryOptimizationConfig struct { + BatchExecution BatchExecutionConfig +} + +type BatchExecutionConfig struct { + Enabled bool + // CostThreshold is the minimum estimated non-batched iteration + // cost required to enable batch-at-a-time query evaluation. + CostThreshold int +} diff --git a/frac/fraction_test.go b/frac/fraction_test.go index c5cb906b2..d167b2b8e 100644 --- a/frac/fraction_test.go +++ b/frac/fraction_test.go @@ -72,6 +72,12 @@ func (s *FractionTestSuite) TearDownSuiteCommon() { func (s *FractionTestSuite) SetupTestCommon() { s.config = &frac.Config{} + s.config.Search.QueryOptimization = frac.QueryOptimizationConfig{ + BatchExecution: frac.BatchExecutionConfig{ + Enabled: true, + CostThreshold: 1000, + }, + } s.tokenizers = map[seq.TokenizerType]tokenizer.Tokenizer{ seq.TokenizerTypeKeyword: tokenizer.NewKeywordTokenizer(20, false, true), seq.TokenizerTypeText: tokenizer.NewTextTokenizer(20, false, true, 100), @@ -1358,7 +1364,7 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: midTime, }, - // AND operator queries + // AND operator queries (intersection) { name: "message:request AND message:failed", query: "message:request AND message:failed", @@ -1368,6 +1374,24 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: toTime, }, + { + name: "service:gateway AND level:5", + query: "service:gateway AND level:5", + filter: func(doc *testDoc) bool { + return doc.service == gateway && doc.level == 5 + }, + fromTime: fromTime, + toTime: toTime, + }, + { + name: "service:gateway AND level:5 AND message:processing (time range)", + query: "service:gateway AND level:5 AND message:processing", + filter: func(doc *testDoc) bool { + return doc.service == gateway && doc.level == 5 && strings.Contains(doc.message, "processing") + }, + fromTime: fromTime, + toTime: midTime, + }, { name: "service:gateway AND message:processing AND message:retry AND level:5", query: "service:gateway AND message:processing AND message:retry AND level:5", @@ -1417,7 +1441,7 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { toTime: toTime, }, { - name: "complex AND+OR", + name: "complex AND+OR 2", query: "(service:gateway OR service:proxy OR service:scheduler) AND " + "(message:request OR message:failed) AND (level:1 OR level:2 OR level:3)", filter: func(doc *testDoc) bool { @@ -1428,6 +1452,25 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: toTime, }, + // AND NOT + { + name: "service:gateway AND NOT message:request", + query: "service:gateway AND NOT message:request", + filter: func(doc *testDoc) bool { + return doc.service == gateway && !strings.Contains(doc.message, "request") + }, + fromTime: fromTime, + toTime: midTime, + }, + { + name: "service:gateway AND NOT message:request AND NOT level:3", + query: "service:gateway AND NOT message:request AND NOT level:3", + filter: func(doc *testDoc) bool { + return doc.service == gateway && !strings.Contains(doc.message, "request") && doc.level != 3 + }, + fromTime: fromTime, + toTime: midTime, + }, { name: "service:gateway AND NOT (message:request OR message:timed OR level:[0 to 3])", query: "service:gateway AND NOT (message:request OR message:timed OR level:[0 to 3])", diff --git a/frac/processor/aggregator_test.go b/frac/processor/aggregator_test.go index ae074afc0..644932d5b 100644 --- a/frac/processor/aggregator_test.go +++ b/frac/processor/aggregator_test.go @@ -156,6 +156,10 @@ func (m *MockTokenIndex) GetValByTID(tid uint32, _ string) []byte { return []byte(strconv.Itoa(int(tid))) } +func (m *MockTokenIndex) GetFreqsByTIDs(tids []uint32, _ string) []uint32 { + return make([]uint32, len(tids)) +} + type IDSourcePair struct { LID node.LID Source uint32 diff --git a/frac/processor/batch_eval_tree.go b/frac/processor/batch_eval_tree.go new file mode 100644 index 000000000..116bdf098 --- /dev/null +++ b/frac/processor/batch_eval_tree.go @@ -0,0 +1,223 @@ +package processor + +import ( + "errors" + "fmt" + + "github.com/ozontech/seq-db/metric/stopwatch" + "github.com/ozontech/seq-db/node" + "github.com/ozontech/seq-db/parser" + "github.com/ozontech/seq-db/seq" +) + +var errBatchingUnsupported = errors.New("batching unsupported") + +// maxBatchedTIDsPerLeaf limits number of TIDs for a single OrMulti node +const maxBatchedTIDsPerLeaf = 5 + +type leafTIDsCache map[parser.Token][]uint32 + +// tryBuildBatchEvalTree tries to build a batched eval tree if possible. +// +// Returns errBatchingUnsupported when the non-batched path should be used. +func tryBuildBatchEvalTree( + root *parser.ASTNode, + ti tokenIndex, + queryOpts QueryOptimizationConfig, + minLID, maxLID uint32, + stats *searchStats, + order seq.DocsOrder, + sw *stopwatch.Stopwatch, +) (node.BatchedNode, error) { + if !queryOpts.BatchExecution.Enabled { + return nil, errBatchingUnsupported + } + + if !astSupportsBatching(root) { + return nil, errBatchingUnsupported + } + + cache := make(leafTIDsCache) + cost, err := calculateQueryIterationCost(root, ti, cache) + if err != nil { + return nil, err + } + + threshold := queryOpts.BatchExecution.CostThreshold + if threshold <= 0 || cost <= uint64(threshold) { + return nil, errBatchingUnsupported + } + + return buildBatchEvalTree(root, minLID, maxLID, stats, order.IsDesc(), + func(token parser.Token) (node.BatchedNode, error) { + return evalBatchLeaf(ti, token, cache, sw, stats, minLID, maxLID, order) + }, + ) +} + +func astSupportsBatching(root *parser.ASTNode) bool { + if root == nil { + return false + } + + switch root.Value.(type) { + case *parser.Literal: + literal := root.Value.(*parser.Literal) + // currently batching supports only simple terms like 'field:A' + return len(literal.Terms) == 1 && literal.Terms[0].Kind == parser.TermText + case *parser.Range: + return true + case *parser.Logical: + logical := root.Value.(*parser.Logical) + // batching is not supported for NOT nodes yet + if logical.Operator == parser.LogicalNot { + return false + } + for i := range root.Children { + if !astSupportsBatching(root.Children[i]) { + return false + } + } + return true + default: + return false + } +} + +func calculateQueryIterationCost(root *parser.ASTNode, ti tokenIndex, cache leafTIDsCache) (uint64, error) { + if root == nil { + return 0, fmt.Errorf("empty AST") + } + + switch token := root.Value.(type) { + case *parser.Literal: + return leafIterationCost(ti, token.Field, token, cache) + case *parser.Range: + return leafIterationCost(ti, token.Field, token, cache) + case *parser.Logical: + if len(root.Children) == 0 { + return 0, nil + } + childCosts := make([]uint64, len(root.Children)) + for i, child := range root.Children { + c, err := calculateQueryIterationCost(child, ti, cache) + if err != nil { + return 0, err + } + childCosts[i] = c + } + if len(childCosts) != 2 { + return 0, fmt.Errorf("logical operator has unsupported count of children: %d", len(childCosts)) + } + switch token.Operator { + case parser.LogicalAnd: + return min(childCosts[0], childCosts[1]), nil + case parser.LogicalNAnd: + return childCosts[0] + childCosts[1], nil + case parser.LogicalOr: + return childCosts[0] + childCosts[1], nil + default: + return 0, fmt.Errorf("unsupported logical operator for cost estimation: %v", token.Operator) + } + default: + return 0, fmt.Errorf("unsupported token type for cost estimation") + } +} + +func leafIterationCost(ti tokenIndex, field string, token parser.Token, cache leafTIDsCache) (uint64, error) { + tids, err := ti.GetTIDsByTokenExpr(token) + if err != nil { + return 0, err + } + if len(tids) > maxBatchedTIDsPerLeaf { + return 0, errBatchingUnsupported + } + cache[token] = tids + if len(tids) == 0 { + return 0, nil + } + + freqs := ti.GetFreqsByTIDs(tids, field) + var cost uint64 + for _, freq := range freqs { + cost += uint64(freq) + } + return cost, nil +} + +// buildBatchEvalTree builds a BatchedNode eval tree using already-validated leaf TIDs. +func buildBatchEvalTree( + root *parser.ASTNode, + minLID, maxLID uint32, + stats *searchStats, + desc bool, + newBatchLeaf func(parser.Token) (node.BatchedNode, error), +) (node.BatchedNode, error) { + if root == nil { + return nil, fmt.Errorf("empty AST") + } + + children := make([]node.BatchedNode, 0, len(root.Children)) + for _, child := range root.Children { + childNode, err := buildBatchEvalTree(child, minLID, maxLID, stats, desc, newBatchLeaf) + if err != nil { + return nil, err + } + children = append(children, childNode) + } + + switch token := root.Value.(type) { + case *parser.Literal: + return newBatchLeaf(token) + case *parser.Range: + return newBatchLeaf(token) + case *parser.Logical: + stats.NodesTotal++ + switch token.Operator { + case parser.LogicalAnd: + return node.NewAndBatched(children[0], children[1], desc), nil + case parser.LogicalOr: + return node.NewOrBatched(children[0], children[1], desc), nil + case parser.LogicalNAnd: + return node.NewNAndBatched(children[0], children[1], desc), nil + default: + return nil, fmt.Errorf("unsupported logical operator for batched eval: %v", token.Operator) + } + default: + return nil, fmt.Errorf("unknown token type for batched eval") + } +} + +func evalBatchLeaf( + ti tokenIndex, + token parser.Token, + cache leafTIDsCache, + sw *stopwatch.Stopwatch, + stats *searchStats, + minLID, maxLID uint32, + order seq.DocsOrder, +) (node.BatchedNode, error) { + stats.LeavesTotal++ + + tids, ok := cache[token] + if !ok { + var err error + tids, err = ti.GetTIDsByTokenExpr(token) + if err != nil { + return nil, err + } + } + + if len(tids) == 0 { + stats.NodesTotal++ + return node.EmptyBatched(), nil + } + + m := sw.Start("get_batched_lids_from_tids") + batchedLIDs := ti.GetBatchedLIDsFromTIDs(tids, stats, minLID, maxLID, order) + m.Stop() + + stats.NodesTotal++ + + return node.NewOrBatchedMulti(batchedLIDs, order.IsDesc()), nil +} diff --git a/frac/processor/batch_eval_tree_test.go b/frac/processor/batch_eval_tree_test.go new file mode 100644 index 000000000..e522b064a --- /dev/null +++ b/frac/processor/batch_eval_tree_test.go @@ -0,0 +1,254 @@ +package processor + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ozontech/seq-db/frac/sealed/lids" + "github.com/ozontech/seq-db/metric/stopwatch" + "github.com/ozontech/seq-db/node" + "github.com/ozontech/seq-db/parser" + "github.com/ozontech/seq-db/seq" +) + +func TestASTSupportsBatching(t *testing.T) { + t.Run("single field", func(t *testing.T) { + q, err := parser.ParseSeqQL(`service:"foo"`, nil) + require.NoError(t, err) + assert.True(t, astSupportsBatching(q.Root)) + }) + + t.Run("and of fields", func(t *testing.T) { + q, err := parser.ParseSeqQL(`service:"foo" AND level:"error"`, nil) + require.NoError(t, err) + assert.True(t, astSupportsBatching(q.Root)) + }) + + t.Run("nested and", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND (b:2 AND c:3)`, nil) + require.NoError(t, err) + assert.True(t, astSupportsBatching(q.Root)) + }) + + t.Run("or", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 OR b:2`, nil) + require.NoError(t, err) + assert.True(t, astSupportsBatching(q.Root)) + }) + + t.Run("not", func(t *testing.T) { + q, err := parser.ParseSeqQL(`NOT a:1`, nil) + require.NoError(t, err) + assert.False(t, astSupportsBatching(q.Root)) + }) + + t.Run("and with or child", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND (b:2 OR c:3)`, nil) + require.NoError(t, err) + assert.True(t, astSupportsBatching(q.Root)) + }) +} + +type testTokenIndex struct { + tids map[string][]uint32 + freqs map[uint32]uint32 +} + +func (d *testTokenIndex) GetValByTID(tid uint32, _ string) []byte { + panic("not implemented") +} + +func (d *testTokenIndex) GetLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.Node { + panic("not implemented") +} + +func (d *testTokenIndex) GetBatchedLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.BatchedNode { + nodes := make([]node.BatchedNode, len(tids)) + for i := range tids { + nodes[i] = node.EmptyBatched() + } + return nodes +} + +func (d *testTokenIndex) GetTIDsByTokenExpr(token parser.Token) ([]uint32, error) { + key := parser.GetField(token) + ":" + parser.GetHint(token) + return d.tids[key], nil +} + +func (d *testTokenIndex) GetFreqsByTIDs(tids []uint32, _ string) []uint32 { + freqs := make([]uint32, len(tids)) + for i, tid := range tids { + freqs[i] = d.freqs[tid] + } + return freqs +} + +func TestQueryIterationCost(t *testing.T) { + index := &testTokenIndex{ + tids: map[string][]uint32{ + "a:1": {1}, + "b:2": {2}, + "c:3": {3}, + }, + freqs: map[uint32]uint32{ + 1: 80_000, + 2: 120_000, + 3: 40_000, + }, + } + + t.Run("leaf", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1`, nil) + require.NoError(t, err) + cost, err := calculateQueryIterationCost(q.Root, index, make(leafTIDsCache)) + require.NoError(t, err) + assert.Equal(t, uint64(80_000), cost) + }) + + t.Run("and", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND b:2`, nil) + require.NoError(t, err) + cost, err := calculateQueryIterationCost(q.Root, index, make(leafTIDsCache)) + require.NoError(t, err) + assert.Equal(t, uint64(80_000), cost) + }) + + t.Run("or", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 OR b:2`, nil) + require.NoError(t, err) + cost, err := calculateQueryIterationCost(q.Root, index, make(leafTIDsCache)) + require.NoError(t, err) + assert.Equal(t, uint64(200_000), cost) + }) + + t.Run("and not", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND NOT b:2`, nil) + require.NoError(t, err) + cost, err := calculateQueryIterationCost(q.Root, index, make(leafTIDsCache)) + require.NoError(t, err) + assert.Equal(t, uint64(200_000), cost) + }) + + t.Run("nested and-or", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND (b:2 OR c:3)`, nil) + require.NoError(t, err) + cost, err := calculateQueryIterationCost(q.Root, index, make(leafTIDsCache)) + require.NoError(t, err) + assert.Equal(t, uint64(80_000), cost) + }) +} + +func TestTryBuildBatchEvalTree(t *testing.T) { + const threshold = 50_000 + queryOpt := QueryOptimizationConfig{BatchExecution: BatchExecutionConfig{Enabled: true, CostThreshold: threshold}} + sw := stopwatch.New() + + denseIndex := &testTokenIndex{ + tids: map[string][]uint32{ + "a:1": {1}, + "b:2": {2}, + }, + freqs: map[uint32]uint32{ + 1: 120_000, + 2: 120_000, + }, + } + sparseIndex := &testTokenIndex{ + tids: map[string][]uint32{ + "a:1": {1}, + "b:2": {2}, + }, + freqs: map[uint32]uint32{ + 1: 1_000, + 2: 2_000, + }, + } + + t.Run("dense and query enables batching", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND b:2`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, denseIndex, queryOpt, 1, 100, stats, seq.DocsOrderDesc, sw) + require.NoError(t, err) + assert.NotNil(t, tree) + assert.Equal(t, 2, stats.LeavesTotal) + assert.Equal(t, 3, stats.NodesTotal) // 2 leaves + 1 AND + }) + + t.Run("disabled skips batching", func(t *testing.T) { + disabled := QueryOptimizationConfig{BatchExecution: BatchExecutionConfig{Enabled: false, CostThreshold: threshold}} + q, err := parser.ParseSeqQL(`a:1 AND b:2`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, denseIndex, disabled, 1, 100, stats, seq.DocsOrderDesc, sw) + require.ErrorIs(t, err, errBatchingUnsupported) + assert.Nil(t, tree) + assert.Equal(t, 0, stats.LeavesTotal) + assert.Equal(t, 0, stats.NodesTotal) + }) + + t.Run("sparse and query disables batching", func(t *testing.T) { + q, err := parser.ParseSeqQL(`a:1 AND b:2`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, sparseIndex, queryOpt, 1, 100, stats, seq.DocsOrderDesc, sw) + require.ErrorIs(t, err, errBatchingUnsupported) + assert.Nil(t, tree) + assert.Equal(t, 0, stats.LeavesTotal) + assert.Equal(t, 0, stats.NodesTotal) + }) + + t.Run("exactly at threshold disables batching", func(t *testing.T) { + index := &testTokenIndex{ + tids: map[string][]uint32{"a:1": {1}}, + freqs: map[uint32]uint32{1: threshold}, + } + q, err := parser.ParseSeqQL(`a:1`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, index, queryOpt, 1, 100, stats, seq.DocsOrderDesc, sw) + require.ErrorIs(t, err, errBatchingUnsupported) + assert.Nil(t, tree) + assert.Equal(t, 0, stats.LeavesTotal) + assert.Equal(t, 0, stats.NodesTotal) + }) + + t.Run("too many tids disables batching without mutating stats", func(t *testing.T) { + index := &testTokenIndex{ + tids: map[string][]uint32{ + "a:1": {1}, + "b:2": {10, 11, 12, 13, 14, 15}, + }, + freqs: map[uint32]uint32{ + 1: 120_000, + 10: 20_000, + 11: 20_000, + 12: 20_000, + 13: 20_000, + 14: 20_000, + 15: 20_000, + }, + } + q, err := parser.ParseSeqQL(`a:1 AND b:2`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, index, queryOpt, 1, 100, stats, seq.DocsOrderDesc, sw) + require.ErrorIs(t, err, errBatchingUnsupported) + assert.Nil(t, tree) + assert.Equal(t, 0, stats.LeavesTotal) + assert.Equal(t, 0, stats.NodesTotal) + }) + + t.Run("not query disables batching without mutating stats", func(t *testing.T) { + q, err := parser.ParseSeqQL(`NOT a:1`, nil) + require.NoError(t, err) + stats := &searchStats{} + tree, err := tryBuildBatchEvalTree(q.Root, denseIndex, queryOpt, 1, 100, stats, seq.DocsOrderDesc, sw) + require.ErrorIs(t, err, errBatchingUnsupported) + assert.Nil(t, tree) + assert.Equal(t, 0, stats.LeavesTotal) + assert.Equal(t, 0, stats.NodesTotal) + }) +} diff --git a/frac/processor/eval_tree.go b/frac/processor/eval_tree.go index f9b0e0905..0cc05b153 100644 --- a/frac/processor/eval_tree.go +++ b/frac/processor/eval_tree.go @@ -111,6 +111,20 @@ type AggLimits struct { MaxTIDsPerFraction int } +// QueryOptimizationConfig controls search-time query optimization decisions. +type QueryOptimizationConfig struct { + BatchExecution BatchExecutionConfig +} + +// BatchExecutionConfig controls batch-at-a-time query evaluation. +type BatchExecutionConfig struct { + // Enabled is the master switch for batch-at-a-time query evaluation. + Enabled bool + // CostThreshold is the minimum estimated non-batched iteration + // cost required to enable batch-at-a-time query evaluation. + CostThreshold int +} + type iteratorLimit struct { // limit value limit int diff --git a/frac/processor/search.go b/frac/processor/search.go index 2fd616e98..3602d04d7 100644 --- a/frac/processor/search.go +++ b/frac/processor/search.go @@ -34,7 +34,9 @@ type idsIndex interface { type tokenIndex interface { GetValByTID(tid uint32, field string) []byte GetTIDsByTokenExpr(token parser.Token) ([]uint32, error) + GetFreqsByTIDs(tids []uint32, field string) []uint32 GetLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.Node + GetBatchedLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.BatchedNode } type searchIndex interface { @@ -47,6 +49,7 @@ type searchBuffers struct { mids []seq.MID rids []seq.RID lids []node.LID + tmp []uint32 } var searchBuffersPool = sync.Pool{ @@ -55,19 +58,19 @@ var searchBuffersPool = sync.Pool{ // Currently, we drain up to 4k lids from eval tree, but with proper batching enabled // we can get as much as whole LID block can have (currently, 64k lids) lids: make([]node.LID, 0, consts.DefaultLIDBlockCap), + tmp: make([]uint32, 0, consts.DefaultLIDBlockCap), mids: make([]seq.MID, 0, consts.DefaultLIDBlockCap), rids: make([]seq.RID, 0, consts.DefaultLIDBlockCap), } }, } -const maxLidsToDrain = 4096 - func IndexSearch( ctx context.Context, params SearchParams, index searchIndex, aggLimits AggLimits, + queryOpt QueryOptimizationConfig, sw *stopwatch.Stopwatch, ) (qpr *seq.QPR, err error) { stats := &searchStats{} @@ -76,18 +79,53 @@ func IndexSearch( minLID, maxLID := getLIDsBorders(params, index) m.Stop() - m = sw.Start("eval_leaf") - evalTree, err := buildEvalTree(params.AST, minLID, maxLID, stats, params.Order.IsReverse(), - func(token parser.Token) (node.Node, error) { - return evalLeaf(index, token, sw, stats, minLID, maxLID, params.Order) - }, - ) + m = sw.Start("get_skip_lids") + skipLIDs, hasSkipLIDs, release, err := index.GetSkipLIDs(minLID, maxLID, params.Order.IsReverse()) + defer func() { + err = errors.Join(err, release()) + }() m.Stop() - if err != nil { return nil, err } + m = sw.Start("build_batch_eval_tree") + // TODO(cheb0) skipmasks block batched execution + var evalTree node.BatchedNode + if !hasSkipLIDs { + evalTree, err = tryBuildBatchEvalTree( + params.AST, index, queryOpt, minLID, maxLID, stats, params.Order, sw, + ) + } else { + err = errBatchingUnsupported + } + m.Stop() + if err != nil && !errors.Is(err, errBatchingUnsupported) { + return nil, err + } + + if errors.Is(err, errBatchingUnsupported) { + m = sw.Start("eval_leaf") + var nodeTree node.Node + nodeTree, err = buildEvalTree(params.AST, minLID, maxLID, stats, params.Order.IsReverse(), + func(token parser.Token) (node.Node, error) { + return evalLeaf(index, token, sw, stats, minLID, maxLID, params.Order) + }, + ) + if err != nil { + return nil, err + } + m.Stop() + + if hasSkipLIDs { + m = sw.Start("eval_skip_lids") + nodeTree = evalSkipLIDs(nodeTree, skipLIDs, stats) + m.Stop() + } + + evalTree = node.NewBatcherNode(nodeTree, params.Order.IsDesc()) + } + defer func(start time.Time) { stats.TreeDuration += time.Since(start) }(time.Now()) if util.IsCancelled(ctx) { @@ -116,22 +154,6 @@ func IndexSearch( } } - m = sw.Start("get_skip_lids") - skipLIDs, hasSkipLIDs, release, err := index.GetSkipLIDs(minLID, maxLID, params.Order.IsReverse()) - defer func() { - err = errors.Join(err, release()) - }() - m.Stop() - if err != nil { - return nil, err - } - - if hasSkipLIDs { - m = sw.Start("eval_skip_lids") - evalTree = evalSkipLIDs(evalTree, skipLIDs, stats) - m.Stop() - } - m = sw.Start("iterate_eval_tree") total, ids, histMap, aggs, err := iterateEvalTree(ctx, params, index, evalTree, aggSupplier, sw) m.Stop() @@ -175,35 +197,11 @@ func IndexSearch( return qpr, nil } -func batcher(evalTree node.Node, buf []node.LID, desc bool) func(need int) []node.LID { - if batchNode, ok := tryConvertToBatchedTree(evalTree); ok { - return func(need int) []node.LID { - buf = batchNode.NextBatch(need).CopyLIDs(desc, buf[:0]) - if len(buf) > need { - buf = buf[:need] - } - return buf - } - } - - return func(need int) []node.LID { - buf = buf[:0] - for range min(maxLidsToDrain, need) { - lid := evalTree.Next() - if lid.IsNull() { - break - } - buf = append(buf, lid) - } - return buf - } -} - func iterateEvalTree( ctx context.Context, params SearchParams, idsIndex idsIndex, - evalTree node.Node, + evalTree node.BatchedNode, aggSupplier func() ([]Aggregator, error), sw *stopwatch.Stopwatch, ) (int, seq.IDSources, HistMap, []Aggregator, error) { @@ -226,8 +224,8 @@ func iterateEvalTree( mids := buffers.mids rids := buffers.rids - - batchedEvalTree := batcher(evalTree, buffers.lids, params.Order.IsDesc()) + lidsBuf := buffers.lids[:cap(buffers.lids)] + tmpBuf := buffers.tmp[:cap(buffers.tmp)] timerEval := sw.Timer("eval_tree_next") timerMID := sw.Timer("get_mid") @@ -248,61 +246,86 @@ func iterateEvalTree( break } - maxBatchSize := needIDs + remaining := needIDs if needScanAllRange || params.Downsample > 1 { // if full range scan is required OR downsampling is active, // we must fetch as many LIDs as possible in one batch. - maxBatchSize = math.MaxUint32 + remaining = math.MaxInt32 } timerEval.Start() - lidsBatch := batchedEvalTree(maxBatchSize) + batch := evalTree.NextBatch() timerEval.Stop() - if len(lidsBatch) == 0 { + if batch.IsEmpty() { break } - total += len(lidsBatch) + iter := batch.ManyIter(params.Order.IsDesc()) - if lidsBatch = sample(lidsBatch); len(lidsBatch) == 0 { - continue - } + for remaining > 0 { + if util.IsCancelled(ctx) { + return total, ids, hist, aggs, ctx.Err() + } - if hasHist || needIDs > 0 { - timerMID.Start() - mids = idsIndex.GetMIDs(lidsBatch, mids[:0]) - timerMID.Stop() + if !needScanAllRange && params.Limit-len(ids) < 1 { + break + } + + timerEval.Start() + n := iter.CopyLIDs(lidsBuf[:min(remaining, len(lidsBuf))], tmpBuf[:min(remaining, len(tmpBuf))]) + timerEval.Stop() + + if n == 0 { + break + } + + lidsBatch := lidsBuf[:n] + total += n + remaining -= n - if hasHist { - timerHist.Start() - hist.Update(mids) - timerHist.Stop() + lidsBatch = sample(lidsBatch) + + if len(lidsBatch) == 0 { + continue } - if needIDs > 0 { - needLIDs := min(needIDs, len(lidsBatch)) + needIDs = params.Limit - len(ids) + if hasHist || needIDs > 0 { + timerMID.Start() + mids = idsIndex.GetMIDs(lidsBatch, mids[:0]) + timerMID.Stop() - timerRID.Start() - rids = idsIndex.GetRIDs(lidsBatch[:needLIDs], rids[:0]) - timerRID.Stop() + if hasHist { + timerHist.Start() + hist.Update(mids) + timerHist.Stop() + } - // fill IDs for search - for i := 0; i < needLIDs; i++ { - id := seq.ID{MID: mids[i], RID: rids[i]} - if i == 0 || lastID != id { // lids increase monotonically, it's enough to compare current id with the last one - ids = append(ids, seq.IDSource{ID: id}) + if needIDs > 0 { + needLIDs := min(needIDs, len(lidsBatch)) + + timerRID.Start() + rids = idsIndex.GetRIDs(lidsBatch[:needLIDs], rids[:0]) + timerRID.Stop() + + // fill IDs for search + for i := 0; i < needLIDs; i++ { + id := seq.ID{MID: mids[i], RID: rids[i]} + if i == 0 || lastID != id { // lids increase monotonically, it's enough to compare current id with the last one + ids = append(ids, seq.IDSource{ID: id}) + } + lastID = id } - lastID = id } } - } - // Update aggregators - if params.HasAgg() { - var err error - if aggs, err = updateAggs(aggs, lidsBatch, aggSupplier, timerAgg); err != nil { - return total, ids, hist, aggs, err + // Update aggregators + if params.HasAgg() { + var err error + if aggs, err = updateAggs(aggs, lidsBatch, aggSupplier, timerAgg); err != nil { + return total, ids, hist, aggs, err + } } } } @@ -352,17 +375,6 @@ func sampler(n uint32) func(in []node.LID) []node.LID { } } -func tryConvertToBatchedTree(evalTree node.Node) (node.BatchedNode, bool) { - switch it := evalTree.(type) { - case *lids.IteratorDesc: - return lids.NewBatchedIteratorDesc(it), true - case *lids.IteratorAsc: - return lids.NewBatchedIteratorAsc(it), true - default: - return nil, false - } -} - // getLIDsBorders return min and max LID borders (including) for search func getLIDsBorders(params SearchParams, idsIndex idsIndex) (uint32, uint32) { if idsIndex.Len() == 0 { diff --git a/frac/sealed/lids/block_test.go b/frac/sealed/lids/block_test.go index d2f3d9d6c..ad4f40768 100644 --- a/frac/sealed/lids/block_test.go +++ b/frac/sealed/lids/block_test.go @@ -192,10 +192,14 @@ func ToArray(b node.LIDBatch) []uint32 { return nil } out := make([]uint32, 0, b.Len()) - for _, lid := range b.CopyLIDs(true, nil) { - out = append(out, lid.Unpack()) + it := b.Iter() + for { + lid, ok := it.Next() + if !ok { + return out + } + out = append(out, lid) } - return out } func TestBlockPack_ReuseBuffer(t *testing.T) { diff --git a/frac/sealed/lids/iterator_batched_asc.go b/frac/sealed/lids/iterator_batched_asc.go index da9efed51..8d991360b 100644 --- a/frac/sealed/lids/iterator_batched_asc.go +++ b/frac/sealed/lids/iterator_batched_asc.go @@ -63,11 +63,11 @@ func (it *BatchedIteratorAsc) loadNextLIDsBlock() { it.blockIndex-- } -func (it *BatchedIteratorAsc) NextBatch(need int) node.LIDBatch { - return it.NextBatchGeq(need, node.NewAscZeroLID()) +func (it *BatchedIteratorAsc) NextBatch() node.LIDBatch { + return it.NextBatchGeq(node.NewAscZeroLID()) } -func (it *BatchedIteratorAsc) NextBatchGeq(_ int, nextID node.LID) node.LIDBatch { +func (it *BatchedIteratorAsc) NextBatchGeq(nextID node.LID) node.LIDBatch { for { if it.batch.IsEmpty() { if !it.tryNextBlock { diff --git a/frac/sealed/lids/iterator_batched_desc.go b/frac/sealed/lids/iterator_batched_desc.go index 318e312c1..f9c344fc5 100644 --- a/frac/sealed/lids/iterator_batched_desc.go +++ b/frac/sealed/lids/iterator_batched_desc.go @@ -63,11 +63,11 @@ func (it *BatchedIteratorDesc) loadNextLIDsBlock() { it.blockIndex++ } -func (it *BatchedIteratorDesc) NextBatch(need int) node.LIDBatch { - return it.NextBatchGeq(need, node.NewDescZeroLID()) +func (it *BatchedIteratorDesc) NextBatch() node.LIDBatch { + return it.NextBatchGeq(node.NewDescZeroLID()) } -func (it *BatchedIteratorDesc) NextBatchGeq(_ int, nextID node.LID) node.LIDBatch { +func (it *BatchedIteratorDesc) NextBatchGeq(nextID node.LID) node.LIDBatch { for { if it.batch.IsEmpty() { if !it.tryNextBlock { diff --git a/frac/sealed_index.go b/frac/sealed_index.go index b9cc83f70..29fda1f60 100644 --- a/frac/sealed_index.go +++ b/frac/sealed_index.go @@ -111,6 +111,9 @@ func (dp *sealedDataProvider) Fetch(ids []seq.ID, noSkipMasks bool) ([][]byte, e func (dp *sealedDataProvider) Search(params processor.SearchParams) (*seq.QPR, error) { aggLimits := processor.AggLimits(dp.config.Search.AggLimits) + queryOpt := processor.QueryOptimizationConfig{ + BatchExecution: processor.BatchExecutionConfig(dp.config.Search.QueryOptimization.BatchExecution), + } // Limit the parameter range to data boundaries to prevent histogram overflow params.From = max(params.From, dp.info.From) @@ -126,7 +129,7 @@ func (dp *sealedDataProvider) Search(params processor.SearchParams) (*seq.QPR, e t := sw.Start("total") defer t.Stop() - qpr, err := processor.IndexSearch(dp.ctx, params, dp.getSearchIndex(), aggLimits, sw) + qpr, err := processor.IndexSearch(dp.ctx, params, dp.getSearchIndex(), aggLimits, queryOpt, sw) if err != nil { return nil, err } @@ -256,6 +259,24 @@ func (ti *sealedTokenIndex) GetTIDsByTokenExpr(t parser.Token) ([]uint32, error) return tids, nil } +func (ti *sealedTokenIndex) GetFreqsByTIDs(tids []uint32, field string) []uint32 { + freqs := make([]uint32, len(tids)) + if len(tids) == 0 { + return freqs + } + + tokenTable := ti.tokenTableLoader.Load() + for i, tid := range tids { + if tid == 0 { + continue + } + entry := tokenTable.GetEntryByTID(tid, field) + block := ti.tokenBlockLoader.Load(entry.BlockIndex) + freqs[i] = block.GetFreq(entry.GetIndexInTokensBlock(tid)) + } + return freqs +} + func (ti *sealedTokenIndex) GetLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.Node { var ( getBlockIndex func(tid uint32) uint32 @@ -287,6 +308,37 @@ func (ti *sealedTokenIndex) GetLIDsFromTIDs(tids []uint32, stats lids.Counter, m return nodes } +func (ti *sealedTokenIndex) GetBatchedLIDsFromTIDs(tids []uint32, stats lids.Counter, minLID, maxLID uint32, order seq.DocsOrder) []node.BatchedNode { + var ( + getBlockIndex func(tid uint32) uint32 + getBatchedLIDsIterator func(uint32, uint32) node.BatchedNode + ) + + if order.IsReverse() { + getBlockIndex = func(tid uint32) uint32 { return ti.lidsTable.GetLastBlockIndexForTID(tid) } + getBatchedLIDsIterator = func(startIndex uint32, tid uint32) node.BatchedNode { + return lids.NewBatchedIteratorAsc(lids.NewIteratorAsc(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID)) + } + } else { + getBlockIndex = func(tid uint32) uint32 { return ti.lidsTable.GetFirstBlockIndexForTID(tid) } + getBatchedLIDsIterator = func(startIndex uint32, tid uint32) node.BatchedNode { + return lids.NewBatchedIteratorDesc(lids.NewIteratorDesc(ti.lidsTable, ti.lidsLoader, startIndex, tid, stats, minLID, maxLID)) + } + } + + startIndexes := make([]uint32, len(tids)) + for i, tid := range tids { + startIndexes[i] = getBlockIndex(tid) + } + + nodes := make([]node.BatchedNode, len(tids)) + for i, tid := range tids { + nodes[i] = getBatchedLIDsIterator(startIndexes[i], tid) + } + + return nodes +} + type sealedFetchIndex struct { fracName string idsIndex *sealedIDsIndex diff --git a/node/batch.go b/node/batch.go index 6da8ef8f5..0aaf61baf 100644 --- a/node/batch.go +++ b/node/batch.go @@ -16,7 +16,7 @@ type LIDBatch interface { Min() uint32 // Max returns max (last) value. Panics if batch is empty. Max() uint32 - CopyLIDs(desc bool, dst []LID) []LID + ManyIter(desc bool) ManyIter // Iter iterates lids in ascending way. Iter() Iter // ReverseIter iterates lids in descending way. @@ -25,6 +25,10 @@ type LIDBatch interface { Narrow(minLID, maxLID uint32) LIDBatch } +type ManyIter interface { + CopyLIDs(dst []LID, tmp []uint32) int +} + type Iter interface { Next() (uint32, bool) NextGeq(geq uint32) (uint32, bool) @@ -111,17 +115,41 @@ func (b *sliceBatch) ReverseIter() Iter { return &sliceReverseIter{lids: b.lids, idx: len(b.lids) - 1} } -func (b *sliceBatch) CopyLIDs(desc bool, dst []LID) []LID { - if desc { - for _, lid := range b.lids { - dst = append(dst, NewDescLID(lid)) - } - } else { - for i := len(b.lids) - 1; i >= 0; i-- { - dst = append(dst, NewAscLID(b.lids[i])) +func (b *sliceBatch) ManyIter(desc bool) ManyIter { + it := &sliceManyIter{lids: b.lids, desc: desc} + if !desc { + it.pos = len(b.lids) - 1 + } + return it +} + +type sliceManyIter struct { + lids []uint32 + pos int + desc bool +} + +func (it *sliceManyIter) CopyLIDs(dst []LID, tmp []uint32) int { + if len(dst) == 0 || len(tmp) == 0 { + return 0 + } + if it.desc { + n := min(len(dst), len(tmp), len(it.lids)-it.pos) + for i := 0; i < n; i++ { + dst[i] = NewDescLID(it.lids[it.pos+i]) } + it.pos += n + return n + } + if it.pos < 0 { + return 0 } - return dst + n := min(len(dst), len(tmp), it.pos+1) + for i := 0; i < n; i++ { + dst[i] = NewAscLID(it.lids[it.pos-i]) + } + it.pos -= n + return n } type sliceIter struct { @@ -214,7 +242,7 @@ func (b *bitmapBatch) Narrow(minLID, maxLID uint32) LIDBatch { out.RemoveRange(0, uint64(minLID)) } if maxLID < b.max { - out.RemoveRange(uint64(maxLID)+1, uint64(0x100000000)) + out.RemoveRange(uint64(maxLID)+1, math.MaxUint64) } return NewBitmapBatch(out) } @@ -227,19 +255,43 @@ func (b *bitmapBatch) ReverseIter() Iter { return newBitmapReverseIter(b.bm) } -func (b *bitmapBatch) CopyLIDs(desc bool, dst []LID) []LID { +func (b *bitmapBatch) ManyIter(desc bool) ManyIter { if desc { - it := b.bm.Iterator() - for it.HasNext() { - dst = append(dst, NewDescLID(it.Next())) - } - } else { - it := b.bm.ReverseIterator() - for it.HasNext() { - dst = append(dst, NewAscLID(it.Next())) - } + return &bitmapManyIterAsc{it: b.bm.ManyIterator()} } - return dst + return &bitmapManyIterDesc{it: b.bm.ReverseIterator()} +} + +type bitmapManyIterAsc struct { + it roaring.ManyIntIterable +} + +func (it *bitmapManyIterAsc) CopyLIDs(dst []LID, tmp []uint32) int { + if len(dst) == 0 || len(tmp) == 0 { + return 0 + } + n := it.it.NextMany(tmp[:min(len(dst), len(tmp))]) + for i := 0; i < n; i++ { + dst[i] = NewDescLID(tmp[i]) + } + return n +} + +type bitmapManyIterDesc struct { + it roaring.IntIterable +} + +func (it *bitmapManyIterDesc) CopyLIDs(dst []LID, tmp []uint32) int { + if len(dst) == 0 || len(tmp) == 0 { + return 0 + } + n := 0 + limit := min(len(dst), len(tmp)) + for n < limit && it.it.HasNext() { + dst[n] = NewAscLID(it.it.Next()) + n++ + } + return n } type emptyBatch struct{} @@ -261,10 +313,16 @@ func (emptyBatch) Max() uint32 { panic("Maximum called on empty batch") } -func (emptyBatch) Narrow(uint32, uint32) LIDBatch { return emptyBatchInstance } -func (emptyBatch) CopyLIDs(_ bool, dst []LID) []LID { return dst } -func (emptyBatch) Iter() Iter { return emptyIterInstance } -func (emptyBatch) ReverseIter() Iter { return emptyIterInstance } +func (emptyBatch) Narrow(uint32, uint32) LIDBatch { return emptyBatchInstance } +func (emptyBatch) ManyIter(bool) ManyIter { return emptyManyIterInstance } +func (emptyBatch) Iter() Iter { return emptyIterInstance } +func (emptyBatch) ReverseIter() Iter { return emptyIterInstance } + +type emptyManyIter struct{} + +var emptyManyIterInstance = emptyManyIter{} + +func (emptyManyIter) CopyLIDs([]LID, []uint32) int { return 0 } type emptyIter struct{} diff --git a/node/batch_ops.go b/node/batch_ops.go new file mode 100644 index 000000000..c97ae5be9 --- /dev/null +++ b/node/batch_ops.go @@ -0,0 +1,197 @@ +package node + +import ( + "math" + + "github.com/RoaringBitmap/roaring/v2" +) + +// And intersects two batches in the given document order and returns result and unprocessed parts (either left or right +// will be empty). For AND operation left and right residuals are equal to provided left or right batch, it's safe. +func And(left, right LIDBatch, desc bool) (result, leftResidual, rightResidual LIDBatch) { + empty := EmptyBatch() + if left.IsEmpty() || right.IsEmpty() { + return empty, empty, empty + } + + leftBm := toBitmapBatch(left) + rightBm := toBitmapBatch(right) + + resultBm := leftBm.bm.Clone() + resultBm.And(rightBm.bm) + result = NewBitmapBatch(resultBm) + + // If left or right are slice batches, we must return leftBm and rightBm (bitmap copies), since + // left or right might be intersected with another batch again soon. + if desc { + if leftBm.max > rightBm.max { + return result, leftBm, empty + } + if rightBm.max > leftBm.max { + return result, empty, rightBm + } + return result, empty, empty + } + + if leftBm.min < rightBm.min { + return result, leftBm, empty + } + if rightBm.min < leftBm.min { + return result, empty, rightBm + } + return result, empty, empty +} + +// AndNot finds "AND NOT" result for two batches and returns result and unprocessed parts. +func AndNot(reg, neg LIDBatch, desc bool) (result, regResidual, negResidual LIDBatch) { + empty := EmptyBatch() + if reg.IsEmpty() { + return empty, empty, neg + } + if neg.IsEmpty() { + return reg, empty, empty + } + + regBm := toBitmapBatch(reg) + negBm := toBitmapBatch(neg) + + resultBm := regBm.bm.Clone() + resultBm.AndNot(negBm.bm) + + return truncateBatches(resultBm, regBm, negBm, desc) +} + +// Or unions two batches in the given document order and returns result and unprocessed parts (either left or right +// will be empty). +func Or(left, right LIDBatch, desc bool) (result, leftResidual, rightResidual LIDBatch) { + empty := EmptyBatch() + if left.IsEmpty() { + return right, empty, empty + } + if right.IsEmpty() { + return left, empty, empty + } + + leftBm := toBitmapBatch(left) + rightBm := toBitmapBatch(right) + + resultBm := leftBm.bm.Clone() + resultBm.Or(rightBm.bm) + + return truncateBatches(resultBm, leftBm, rightBm, desc) +} + +// OrMulti unions multiple batches in the given document order and returns +// result and unprocessed parts for each input batch. +func OrMulti(batches []LIDBatch, desc bool) (result LIDBatch, residuals []LIDBatch) { + residuals = make([]LIDBatch, len(batches)) + bmBatches := make([]*bitmapBatch, len(batches)) + nonEmptyBmBatches := make([]*bitmapBatch, 0, len(batches)) + for i, b := range batches { + residuals[i] = EmptyBatch() + if b.IsEmpty() { + continue + } + bm := toBitmapBatch(b) + bmBatches[i] = bm + nonEmptyBmBatches = append(nonEmptyBmBatches, bm) + } + + if len(nonEmptyBmBatches) == 0 { + return EmptyBatch(), residuals + } + if len(nonEmptyBmBatches) == 1 { + return nonEmptyBmBatches[0], residuals + } + + bitmaps := make([]*roaring.Bitmap, len(nonEmptyBmBatches)) + for i, b := range nonEmptyBmBatches { + bitmaps[i] = b.bm + } + + resultBm := roaring.FastOr(bitmaps...) + + if desc { + minMax := nonEmptyBmBatches[0].max + for i := 1; i < len(nonEmptyBmBatches); i++ { + if nonEmptyBmBatches[i].max < minMax { + minMax = nonEmptyBmBatches[i].max + } + } + resultBm.RemoveRange(uint64(minMax)+1, math.MaxUint64) + for i, bm := range bmBatches { + if bm == nil { + continue + } + if bm.max > minMax { + residuals[i] = bm.Narrow(minMax+1, math.MaxUint32) + } + } + return NewBitmapBatch(resultBm), residuals + } + + maxMin := nonEmptyBmBatches[0].min + for i := 1; i < len(nonEmptyBmBatches); i++ { + if nonEmptyBmBatches[i].min > maxMin { + maxMin = nonEmptyBmBatches[i].min + } + } + resultBm.RemoveRange(0, uint64(maxMin)) + for i, bm := range bmBatches { + if bm == nil { + continue + } + if bm.min < maxMin { + residuals[i] = bm.Narrow(0, maxMin-1) + } + } + return NewBitmapBatch(resultBm), residuals +} + +func truncateBatches(result *roaring.Bitmap, left *bitmapBatch, right *bitmapBatch, desc bool) (LIDBatch, LIDBatch, LIDBatch) { + if desc { + if left.max > right.max { + leftRes := left.Narrow(right.max+1, math.MaxUint32) + result.RemoveRange(uint64(right.max)+1, math.MaxUint64) + return NewBitmapBatch(result), leftRes, EmptyBatch() + } + if right.max > left.max { + rightRes := right.Narrow(left.max+1, math.MaxUint32) + result.RemoveRange(uint64(left.max)+1, math.MaxUint64) + return NewBitmapBatch(result), EmptyBatch(), rightRes + } + return NewBitmapBatch(result), EmptyBatch(), EmptyBatch() + } + + if left.min < right.min { + leftRes := left.Narrow(0, right.min-1) + result.RemoveRange(0, uint64(right.min)) + return NewBitmapBatch(result), leftRes, EmptyBatch() + } + if right.min < left.min { + rightRes := right.Narrow(0, left.min-1) + result.RemoveRange(0, uint64(left.min)) + return NewBitmapBatch(result), EmptyBatch(), rightRes + } + return NewBitmapBatch(result), EmptyBatch(), EmptyBatch() +} + +func toBitmapBatch(b LIDBatch) *bitmapBatch { + if b.IsEmpty() { + panic("empty batch is not allowed to be cast to bitmap batch") + } + if bb, ok := b.(*bitmapBatch); ok { + return bb + } + slice, ok := b.(*sliceBatch) + if !ok { + panic("unsupported batch type") + } + bm := roaring.NewBitmap() + bm.AddMany(slice.lids) + return &bitmapBatch{ + bm: bm, + min: slice.Min(), + max: slice.Max(), + } +} diff --git a/node/batch_ops_test.go b/node/batch_ops_test.go new file mode 100644 index 000000000..468ddcdee --- /dev/null +++ b/node/batch_ops_test.go @@ -0,0 +1,403 @@ +package node + +import ( + "slices" + "testing" + + "github.com/stretchr/testify/assert" +) + +type batchCase struct { + name string + left []uint32 + right []uint32 + desc bool + wantResult []uint32 + wantLeftRes []uint32 + wantRightRes []uint32 +} + +type opsBatchFactory func([]uint32) LIDBatch + +var opsBatchFactories = []struct { + name string + fn opsBatchFactory +}{ + {name: "bitmap", fn: NewBitmapBatchFromLids}, + {name: "slice", fn: NewSliceBatch}, +} + +func TestLIDBatch_And(t *testing.T) { + testCases := []batchCase{ + { + name: "desc overlap left has upper tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: true, + wantResult: []uint32{1, 3, 7}, + wantLeftRes: []uint32{1, 2, 3, 7, 8, 11, 15}, + wantRightRes: nil, + }, + { + name: "desc overlap right has upper tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: true, + wantResult: []uint32{1, 3, 7}, + wantLeftRes: nil, + wantRightRes: []uint32{1, 2, 3, 7, 8, 11, 15}, + }, + { + name: "desc disjoint lower vs upper", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: true, + wantResult: nil, + wantLeftRes: nil, + wantRightRes: []uint32{10, 11}, + }, + { + name: "asc overlap left has lower tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: false, + wantResult: []uint32{1, 3, 7}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc overlap right has lower tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: false, + wantResult: []uint32{1, 3, 7}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc disjoint lower vs upper", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: false, + wantResult: nil, + wantLeftRes: []uint32{1, 2, 3}, + wantRightRes: nil, + }, + { + name: "identical inputs have no residuals", + left: []uint32{2, 4, 9}, + right: []uint32{2, 4, 9}, + desc: true, + wantResult: []uint32{2, 4, 9}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "empty left", + left: nil, + right: []uint32{5, 6}, + desc: true, + wantResult: nil, + wantLeftRes: nil, + wantRightRes: nil, + }, + } + + for _, impl := range opsBatchFactories { + t.Run(impl.name, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + left := impl.fn(tc.left) + right := impl.fn(tc.right) + + result, leftRes, rightRes := And(left, right, tc.desc) + + assertSameSet(t, tc.wantResult, toSlice(result)) + assertSameSet(t, tc.wantLeftRes, toSlice(leftRes)) + assertSameSet(t, tc.wantRightRes, toSlice(rightRes)) + }) + } + }) + } +} + +func TestLIDBatch_Or(t *testing.T) { + testCases := []batchCase{ + { + name: "desc overlap left has upper tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: true, + wantResult: []uint32{1, 2, 3, 7, 8, 10}, + wantLeftRes: []uint32{11, 15}, + wantRightRes: nil, + }, + { + name: "desc overlap right has upper tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: true, + wantResult: []uint32{1, 2, 3, 7, 8, 10}, + wantLeftRes: nil, + wantRightRes: []uint32{11, 15}, + }, + { + name: "desc disjoint lower vs upper", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: true, + wantResult: []uint32{1, 2, 3}, + wantLeftRes: nil, + wantRightRes: []uint32{10, 11}, + }, + { + name: "asc overlap left has lower tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: false, + wantResult: []uint32{1, 2, 3, 7, 8, 10, 11, 15}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc overlap right has lower tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: false, + wantResult: []uint32{1, 2, 3, 7, 8, 10, 11, 15}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc disjoint lower vs upper", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: false, + wantResult: []uint32{10, 11}, + wantLeftRes: []uint32{1, 2, 3}, + wantRightRes: nil, + }, + { + name: "empty left", + left: nil, + right: []uint32{5, 6}, + desc: false, + wantResult: []uint32{5, 6}, + wantLeftRes: nil, + wantRightRes: nil, + }, + } + + for _, impl := range opsBatchFactories { + t.Run(impl.name, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + left := impl.fn(tc.left) + right := impl.fn(tc.right) + + result, leftRes, rightRes := Or(left, right, tc.desc) + + assertSameSet(t, tc.wantResult, toSlice(result)) + assertSameSet(t, tc.wantLeftRes, toSlice(leftRes)) + assertSameSet(t, tc.wantRightRes, toSlice(rightRes)) + }) + } + }) + } +} + +func TestLIDBatch_AndNot(t *testing.T) { + testCases := []batchCase{ + { + name: "desc overlap reg has upper tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: true, + wantResult: []uint32{2, 8}, + wantLeftRes: []uint32{11, 15}, + wantRightRes: nil, + }, + { + name: "desc overlap neg has upper tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: true, + wantResult: []uint32{10}, + wantLeftRes: nil, + wantRightRes: []uint32{11, 15}, + }, + { + name: "desc disjoint lower reg vs upper neg", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: true, + wantResult: []uint32{1, 2, 3}, + wantLeftRes: nil, + wantRightRes: []uint32{10, 11}, + }, + { + name: "asc overlap reg has lower tail", + left: []uint32{1, 2, 3, 7, 8, 11, 15}, + right: []uint32{1, 3, 7, 10}, + desc: false, + wantResult: []uint32{2, 8, 11, 15}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc overlap neg has lower tail", + left: []uint32{1, 3, 7, 10}, + right: []uint32{1, 2, 3, 7, 8, 11, 15}, + desc: false, + wantResult: []uint32{10}, + wantLeftRes: nil, + wantRightRes: nil, + }, + { + name: "asc disjoint lower reg vs upper neg", + left: []uint32{1, 2, 3}, + right: []uint32{10, 11}, + desc: false, + wantResult: nil, + wantLeftRes: []uint32{1, 2, 3}, + wantRightRes: nil, + }, + { + name: "empty reg", + left: nil, + right: []uint32{5, 6}, + desc: false, + wantResult: nil, + wantLeftRes: nil, + wantRightRes: []uint32{5, 6}, + }, + { + name: "empty neg", + left: []uint32{5, 6}, + right: nil, + desc: false, + wantResult: []uint32{5, 6}, + wantLeftRes: nil, + wantRightRes: nil, + }, + } + + for _, impl := range opsBatchFactories { + t.Run(impl.name, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + reg := impl.fn(tc.left) + neg := impl.fn(tc.right) + + result, regRes, negRes := AndNot(reg, neg, tc.desc) + + assertSameSet(t, tc.wantResult, toSlice(result)) + assertSameSet(t, tc.wantLeftRes, toSlice(regRes)) + assertSameSet(t, tc.wantRightRes, toSlice(negRes)) + }) + } + }) + } +} + +func TestLIDBatch_AndMixedTypes(t *testing.T) { + left := NewSliceBatch([]uint32{1, 3, 7, 10}) + right := NewBitmapBatchFromLids([]uint32{1, 3, 7, 15}) + + result, leftRes, rightRes := And(left, right, true) + + assertSameSet(t, []uint32{1, 3, 7}, toSlice(result)) + assertSameSet(t, nil, toSlice(leftRes)) + assertSameSet(t, []uint32{1, 3, 7, 15}, toSlice(rightRes)) +} + +func TestLIDBatch_OrMixedTypes(t *testing.T) { + left := NewSliceBatch([]uint32{1, 3, 7, 10}) + right := NewBitmapBatchFromLids([]uint32{1, 3, 7, 15}) + + result, leftRes, rightRes := Or(left, right, true) + + assertSameSet(t, []uint32{1, 3, 7, 10}, toSlice(result)) + assertSameSet(t, nil, toSlice(leftRes)) + assertSameSet(t, []uint32{15}, toSlice(rightRes)) +} + +func TestLIDBatch_OrMulti(t *testing.T) { + type orMultiCase struct { + name string + desc bool + inputs [][]uint32 + wantResult []uint32 + wantResiduals [][]uint32 + } + + testCases := []orMultiCase{ + { + name: "desc overlap with one residual", + desc: true, + inputs: [][]uint32{{1, 2, 3, 7, 8, 11, 15}, {1, 3, 7, 10}, {2, 3, 5, 8, 10}}, + wantResult: []uint32{1, 2, 3, 5, 7, 8, 10}, + wantResiduals: [][]uint32{ + {11, 15}, + nil, + nil, + }, + }, + { + name: "asc overlap with one residual", + desc: false, + inputs: [][]uint32{{1, 2, 3, 7, 8, 11, 15}, {1, 3, 7, 10}, {2, 3, 5, 8, 10}}, + wantResult: []uint32{2, 3, 5, 7, 8, 10, 11, 15}, + wantResiduals: [][]uint32{ + {1}, + {1}, + nil, + }, + }, + { + name: "single non-empty behaves as pass-through", + desc: true, + inputs: [][]uint32{nil, {4, 7, 9}, nil}, + wantResult: []uint32{4, 7, 9}, + wantResiduals: [][]uint32{ + nil, + nil, + nil, + }, + }, + } + + for _, impl := range opsBatchFactories { + t.Run(impl.name, func(t *testing.T) { + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + batches := make([]LIDBatch, len(tc.inputs)) + for i, lids := range tc.inputs { + batches[i] = impl.fn(lids) + } + + result, residuals := OrMulti(batches, tc.desc) + + assertSameSet(t, tc.wantResult, toSlice(result)) + assert.Len(t, residuals, len(tc.wantResiduals)) + for i := range tc.wantResiduals { + assertSameSet(t, tc.wantResiduals[i], toSlice(residuals[i])) + } + }) + } + }) + } +} + +func assertSameSet(t *testing.T, want, got []uint32) { + t.Helper() + if len(want) == 0 { + want = nil + } + if len(got) == 0 { + got = nil + } + slices.Sort(want) + slices.Sort(got) + assert.Equal(t, want, got) +} diff --git a/node/batch_test.go b/node/batch_test.go index 1454b5389..375042c8f 100644 --- a/node/batch_test.go +++ b/node/batch_test.go @@ -240,3 +240,62 @@ func TestBatchReverseIter(t *testing.T) { }) } } + +func TestBatchManyIter(t *testing.T) { + input := []uint32{1, 5, 10, 15, 20, 25, 30} + + for _, impl := range batchFactories { + t.Run(impl.name, func(t *testing.T) { + t.Run("desc chunked", func(t *testing.T) { + b := impl.build(input) + it := b.ManyIter(true) + dst := make([]LID, 3) + tmp := make([]uint32, 3) + + var got []uint32 + for { + n := it.CopyLIDs(dst, tmp) + if n == 0 { + break + } + assert.LessOrEqual(t, n, 3) + for i := 0; i < n; i++ { + got = append(got, dst[i].Unpack()) + } + } + assert.Equal(t, input, got) + }) + + t.Run("asc chunked", func(t *testing.T) { + b := impl.build(input) + it := b.ManyIter(false) + dst := make([]LID, 3) + tmp := make([]uint32, 3) + + var got []uint32 + for { + n := it.CopyLIDs(dst, tmp) + if n == 0 { + break + } + assert.LessOrEqual(t, n, 3) + for i := 0; i < n; i++ { + got = append(got, dst[i].Unpack()) + } + } + assert.Equal(t, []uint32{30, 25, 20, 15, 10, 5, 1}, got) + }) + + t.Run("empty tmp yields zero for desc", func(t *testing.T) { + b := impl.build(input) + n := b.ManyIter(true).CopyLIDs(make([]LID, 8), nil) + assert.Equal(t, 0, n) + }) + }) + } + + t.Run("empty batch", func(t *testing.T) { + n := EmptyBatch().ManyIter(true).CopyLIDs(make([]LID, 8), make([]uint32, 8)) + assert.Equal(t, 0, n) + }) +} diff --git a/node/node.go b/node/node.go index 98cf21e3c..1fa3e3da4 100644 --- a/node/node.go +++ b/node/node.go @@ -14,9 +14,9 @@ type Node interface { type BatchedNode interface { fmt.Stringer // NextBatch returns next batch. Returns nil when exhausted. - NextBatch(need int) LIDBatch + NextBatch() LIDBatch // NextBatchGeq returns next batch (LIDs >= minLID). Returns nil when exhausted. - NextBatchGeq(need int, nextLID LID) LIDBatch + NextBatchGeq(nextID LID) LIDBatch } type Sourced interface { diff --git a/node/node_and.go b/node/node_and.go index 856e12f1e..95bad8e74 100644 --- a/node/node_and.go +++ b/node/node_and.go @@ -79,3 +79,57 @@ func (n *nodeAnd) NextGeq(nextID LID) LID { } } } + +type nodeAndBatched struct { + left BatchedNode + right BatchedNode + desc bool + + leftBatch LIDBatch + rightBatch LIDBatch +} + +// NewAndBatched returns a BatchedNode that intersects two batched iterators. +// desc is the document traversal order for NextBatch / NextBatchGeq. +func NewAndBatched(left, right BatchedNode, desc bool) BatchedNode { + return &nodeAndBatched{ + left: left, + right: right, + desc: desc, + leftBatch: EmptyBatch(), + rightBatch: EmptyBatch(), + } +} + +func (n *nodeAndBatched) String() string { + return fmt.Sprintf("(%s AND %s)", n.left.String(), n.right.String()) +} + +func (n *nodeAndBatched) NextBatch() LIDBatch { + if n.desc { + return n.NextBatchGeq(NewDescZeroLID()) + } + return n.NextBatchGeq(NewAscZeroLID()) +} + +func (n *nodeAndBatched) NextBatchGeq(nextID LID) LIDBatch { + for { + if n.leftBatch.IsEmpty() { + n.leftBatch = n.left.NextBatchGeq(nextID) + } + if n.rightBatch.IsEmpty() { + n.rightBatch = n.right.NextBatchGeq(nextID) + } + if n.leftBatch.IsEmpty() || n.rightBatch.IsEmpty() { + return EmptyBatch() + } + + inter, leftResidual, rightResidual := And(n.leftBatch, n.rightBatch, n.desc) + n.leftBatch = leftResidual + n.rightBatch = rightResidual + + if !inter.IsEmpty() { + return inter + } + } +} diff --git a/node/node_nand.go b/node/node_nand.go index 52f5ff01f..54dfa2b90 100644 --- a/node/node_nand.go +++ b/node/node_nand.go @@ -51,3 +51,60 @@ func (n *nodeNAnd) NextGeq(nextID LID) LID { } return lid } + +type nodeNAndBatched struct { + reg BatchedNode + neg BatchedNode + desc bool + + regBatch LIDBatch + negBatch LIDBatch + negDone bool +} + +func NewNAndBatched(neg, reg BatchedNode, desc bool) BatchedNode { + return &nodeNAndBatched{ + reg: reg, + neg: neg, + desc: desc, + negDone: false, + regBatch: EmptyBatch(), + negBatch: EmptyBatch(), + } +} + +func (n *nodeNAndBatched) String() string { + return fmt.Sprintf("(%s NAND %s)", n.neg.String(), n.reg.String()) +} + +func (n *nodeNAndBatched) NextBatch() LIDBatch { + if n.desc { + return n.NextBatchGeq(NewDescZeroLID()) + } + return n.NextBatchGeq(NewAscZeroLID()) +} + +func (n *nodeNAndBatched) NextBatchGeq(nextID LID) LIDBatch { + for { + if n.regBatch.IsEmpty() { + n.regBatch = n.reg.NextBatchGeq(nextID) + if n.regBatch.IsEmpty() { + return EmptyBatch() + } + } + if !n.negDone && n.negBatch.IsEmpty() { + n.negBatch = n.neg.NextBatchGeq(nextID) + if n.negBatch.IsEmpty() { + n.negDone = true + } + } + + result, regResidual, negResidual := AndNot(n.regBatch, n.negBatch, n.desc) + n.regBatch = regResidual + n.negBatch = negResidual + + if !result.IsEmpty() { + return result + } + } +} diff --git a/node/node_or.go b/node/node_or.go index ab0bf30fa..969f4cdc2 100644 --- a/node/node_or.go +++ b/node/node_or.go @@ -158,3 +158,127 @@ func (n *nodeOrAgg) NextSourcedGeq(nextID LID) (LID, uint32) { return n.NextSourced() } + +type nodeOrBatched struct { + left BatchedNode + right BatchedNode + desc bool + + leftBatch LIDBatch + rightBatch LIDBatch + leftDone bool + rightDone bool +} + +// NewOrBatched returns a BatchedNode that unions two batched iterators. +// desc is the document traversal order for NextBatch / NextBatchGeq. +func NewOrBatched(left, right BatchedNode, desc bool) BatchedNode { + return &nodeOrBatched{ + left: left, + right: right, + desc: desc, + leftBatch: EmptyBatch(), + rightBatch: EmptyBatch(), + } +} + +func (n *nodeOrBatched) String() string { + return fmt.Sprintf("(%s OR %s)", n.left.String(), n.right.String()) +} + +func (n *nodeOrBatched) NextBatch() LIDBatch { + if n.desc { + return n.NextBatchGeq(NewDescZeroLID()) + } + return n.NextBatchGeq(NewAscZeroLID()) +} + +func (n *nodeOrBatched) NextBatchGeq(nextID LID) LIDBatch { + for { + if n.leftBatch.IsEmpty() && !n.leftDone { + n.leftBatch = n.left.NextBatchGeq(nextID) + n.leftDone = n.leftBatch.IsEmpty() + } + + if n.rightBatch.IsEmpty() && !n.rightDone { + n.rightBatch = n.right.NextBatchGeq(nextID) + n.rightDone = n.rightBatch.IsEmpty() + } + + if n.leftDone && n.rightDone && n.leftBatch.IsEmpty() && n.rightBatch.IsEmpty() { + return EmptyBatch() + } + + out, leftRes, rightRes := Or(n.leftBatch, n.rightBatch, n.desc) + n.leftBatch = leftRes + n.rightBatch = rightRes + + if !out.IsEmpty() { + return out + } + } +} + +type nodeOrBatchedMulti struct { + children []BatchedNode + desc bool + + batches []LIDBatch + done []bool +} + +func NewOrBatchedMulti(children []BatchedNode, desc bool) BatchedNode { + if len(children) == 0 { + return EmptyBatched() + } + if len(children) == 1 { + return children[0] + } + batches := make([]LIDBatch, len(children)) + for i := range batches { + batches[i] = EmptyBatch() + } + return &nodeOrBatchedMulti{ + children: children, + desc: desc, + batches: batches, + done: make([]bool, len(children)), + } +} + +func (n *nodeOrBatchedMulti) String() string { + return "OR_MULTI_BATCHED" +} + +func (n *nodeOrBatchedMulti) NextBatch() LIDBatch { + if n.desc { + return n.NextBatchGeq(NewDescZeroLID()) + } + return n.NextBatchGeq(NewAscZeroLID()) +} + +func (n *nodeOrBatchedMulti) NextBatchGeq(nextID LID) LIDBatch { + for { + active := 0 + for i := range n.children { + if n.batches[i].IsEmpty() && !n.done[i] { + n.batches[i] = n.children[i].NextBatchGeq(nextID) + n.done[i] = n.batches[i].IsEmpty() + } + if !n.batches[i].IsEmpty() { + active++ + } + } + + if active == 0 { + return EmptyBatch() + } + + out, residuals := OrMulti(n.batches, n.desc) + n.batches = residuals + + if !out.IsEmpty() { + return out + } + } +} diff --git a/node/node_static.go b/node/node_static.go index baabfa37f..f8f54a67e 100644 --- a/node/node_static.go +++ b/node/node_static.go @@ -99,3 +99,91 @@ func MakeStaticNodes(data [][]uint32) []Node { } return nodes } + +type staticBatchedAsc struct { + staticCursor + batch LIDBatch +} + +type staticBatchedDesc struct { + staticCursor + batch LIDBatch +} + +func NewStaticBatched(data []uint32, reverse bool) BatchedNode { + if reverse { + return &staticBatchedDesc{staticCursor: staticCursor{ + ptr: len(data) - 1, + data: data, + }, batch: EmptyBatch()} + } + + return &staticBatchedAsc{staticCursor: staticCursor{ + ptr: 0, + data: data, + }, batch: EmptyBatch()} +} + +func (n *staticBatchedAsc) String() string { + return "STATIC_BATCHED_ASC" +} + +func (n *staticBatchedAsc) NextBatch() LIDBatch { + return n.NextBatchGeq(NewDescZeroLID()) +} + +func (n *staticBatchedAsc) NextBatchGeq(nextID LID) LIDBatch { + for { + if n.batch.IsEmpty() { + if n.ptr >= len(n.data) { + return EmptyBatch() + } + n.batch = NewSliceBatch(n.data[n.ptr:]) + n.ptr = len(n.data) + } + + if n.batch.IsEmpty() { + continue + } + if nextID.Unpack() > n.batch.Max() { + n.batch = EmptyBatch() + continue + } + + out := n.batch.Narrow(nextID.Unpack(), math.MaxUint32) + n.batch = EmptyBatch() + return out + } +} + +func (n *staticBatchedDesc) String() string { + return "STATIC_BATCHED_DESC" +} + +func (n *staticBatchedDesc) NextBatch() LIDBatch { + return n.NextBatchGeq(NewAscZeroLID()) +} + +func (n *staticBatchedDesc) NextBatchGeq(nextID LID) LIDBatch { + for { + if n.batch.IsEmpty() { + if n.ptr < 0 { + return EmptyBatch() + } + n.batch = NewSliceBatch(n.data[:n.ptr+1]) + n.ptr = -1 + } + + if n.batch.IsEmpty() { + continue + } + if nextID.Unpack() < n.batch.Min() { + n.batch = EmptyBatch() + continue + } + + out := n.batch.Narrow(0, nextID.Unpack()) + n.batch = EmptyBatch() + return out + } +} diff --git a/node/util.go b/node/util.go new file mode 100644 index 000000000..e75752fd4 --- /dev/null +++ b/node/util.go @@ -0,0 +1,81 @@ +package node + +import ( + "fmt" + "slices" +) + +const maxBatchDrain = 4 * 1024 + +// batcherNode allows to iterate over non-batched iterator batch by batch +type batcherNode struct { + source Node + desc bool + batch []uint32 +} + +func NewBatcherNode(source Node, desc bool) BatchedNode { + return &batcherNode{ + source: source, + desc: desc, + batch: make([]uint32, 0, maxBatchDrain), + } +} + +func (b *batcherNode) NextBatch() LIDBatch { + batch := b.batch[:0] + polled := 0 + for polled < maxBatchDrain { + lid := b.source.Next() + if lid.IsNull() { + break + } + batch = append(batch, lid.Unpack()) + polled++ + } + b.batch = batch[:0] + if !b.desc { + slices.Reverse(batch) + } + return NewSliceBatch(batch) +} + +func (b *batcherNode) NextBatchGeq(nextID LID) LIDBatch { + batch := b.batch[:0] + polled := 0 + for polled < maxBatchDrain { + lid := b.source.NextGeq(nextID) + if lid.IsNull() { + break + } + batch = append(batch, lid.Unpack()) + polled++ + } + b.batch = batch[:0] + if !b.desc { + slices.Reverse(batch) + } + return NewSliceBatch(batch) +} + +func (b *batcherNode) String() string { + return fmt.Sprintf("(BATCH %s)", b.source.String()) +} + +type batchedEmpty struct{} + +func EmptyBatched() BatchedNode { + return &batchedEmpty{} +} + +func (e *batchedEmpty) String() string { + return "EMPTY_BATCHED" +} + +func (e *batchedEmpty) NextBatch() LIDBatch { + return EmptyBatch() +} + +func (e *batchedEmpty) NextBatchGeq(_ LID) LIDBatch { + return EmptyBatch() +}