diff --git a/cmd/seq-db/seq-db.go b/cmd/seq-db/seq-db.go index 5f914c79..85e6dcb8 100644 --- a/cmd/seq-db/seq-db.go +++ b/cmd/seq-db/seq-db.go @@ -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 5e454b07..f1d9d25e 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:"150000"` + } `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 5bf7852b..28fd2ea9 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, dp.info.BinaryDataVer, params, &si, aggLimits, queryOpt, sw) if err != nil { return nil, err } @@ -252,6 +255,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 { @@ -263,6 +270,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/config.go b/frac/config.go index e91392aa..e64632cf 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_concurrency_test.go b/frac/fraction_concurrency_test.go index e722ba34..0f89fbfc 100644 --- a/frac/fraction_concurrency_test.go +++ b/frac/fraction_concurrency_test.go @@ -58,7 +58,7 @@ func TestConcurrentAppendAndQuery(t *testing.T) { const numMessagesPerWriter = 5000 const bulkSize = 100 - docs, bulks, fromTime, toTime := generatesMessages(numWriters*numMessagesPerWriter, bulkSize) + docs, bulks, fromTime, toTime := generatesMessages(numWriters*numMessagesPerWriter, bulkSize, false) tmpDir := testcommon.CreateTempDir() fracPath := filepath.Join(tmpDir, "test_fraction") @@ -179,7 +179,7 @@ func TestConcurrentColdQueriesSealedFrac(t *testing.T) { const bulkSize = 100 const numIterations = 100 - docs, bulks, _, toTime := generatesMessages(numWriters*numMessagesPerWriter, bulkSize) + docs, bulks, _, toTime := generatesMessages(numWriters*numMessagesPerWriter, bulkSize, false) tmpDir := testcommon.CreateTempDir() fracPath := filepath.Join(tmpDir, "test_fraction") @@ -476,7 +476,7 @@ func readTest(t *testing.T, fraction frac.Fraction, numReaders, numQueries int, assert.NoError(t, err, "concurrent queries should complete without errors") } -func generatesMessages(numMessages, bulkSize int) ([]*testDoc, [][]string, time.Time, time.Time) { +func generatesMessages(numMessages, bulkSize int, nestedIndexes bool) ([]*testDoc, [][]string, time.Time, time.Time) { services := []string{gateway, proxy, scheduler, database, bus, kafka} messages := []string{ "request started", "request completed", "processing timed out", @@ -502,6 +502,17 @@ func generatesMessages(numMessages, bulkSize int) ([]*testDoc, [][]string, time. message += fmt.Sprintf(" %d", rand.IntN(10000000)) } + var spansJson string + + if nestedIndexes { + numSpans := 1 + rand.IntN(5) + spans := make([]string, numSpans) + for j := 0; j < numSpans; j++ { + spans[j] = fmt.Sprintf(`{"span_id":"span-%d"}`, rand.IntN(5000)) + } + spansJson = fmt.Sprintf(`, "spans":[%s]`, strings.Join(spans, ",")) + } + level := rand.IntN(6) timestamp := fromTime.Add(time.Duration(i) * time.Millisecond) id := fmt.Sprintf("id-%d", i) @@ -512,8 +523,8 @@ func generatesMessages(numMessages, bulkSize int) ([]*testDoc, [][]string, time. toTime = timestamp } - json := fmt.Sprintf(`{"timestamp":%q,"id": %q, "service":%q,"pod":%q,"client_ip":%q,"message":%q,"trace_id": %q,"level":"%d"}`, - timestamp.Format(time.RFC3339Nano), id, service, pod, clientIp, message, traceId, level) + json := fmt.Sprintf(`{"timestamp":%q,"id": %q, "service":%q,"pod":%q,"client_ip":%q,"message":%q,"trace_id": %q,"level":"%d"%s}`, + timestamp.Format(time.RFC3339Nano), id, service, pod, clientIp, message, traceId, level, spansJson) docs = append(docs, &testDoc{ json: json, diff --git a/frac/fraction_test.go b/frac/fraction_test.go index 8fd8eedc..c346fcc2 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), @@ -1223,9 +1229,18 @@ func (s *FractionTestSuite) TestSearchMultipleBulks() { s.AssertSearch(s.query("message:request"), docs, []int{6, 5, 3, 0}) } -// This test checks search on a large frac. Doc count is set to 25000 which results in ~200 kbyte docs file (3 doc blocks) +// This test checks search on a large frac func (s *FractionTestSuite) TestSearchLargeFrac() { - testDocs, bulks, fromTime, toTime := generatesMessages(25000, 1000) + s.runLargeFracTestCases(false) +} + +// This test checks search on a large frac with nested indexes +func (s *FractionTestSuite) TestSearchLargeFracNestedIndexes() { + s.runLargeFracTestCases(true) +} + +func (s *FractionTestSuite) runLargeFracTestCases(nestedIndexes bool) { + testDocs, bulks, fromTime, toTime := generatesMessages(25000, 1000, nestedIndexes) midTime := fromTime.Add(time.Duration(len(testDocs)/2) * time.Millisecond) s.insertDocuments(bulks...) @@ -1238,12 +1253,13 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { type docFilter func(doc *testDoc) bool searchTestCases := []struct { - name string - query string - filter docFilter - fromTime time.Time - toTime time.Time - limit int + name string + query string + filter docFilter + fromTime time.Time + toTime time.Time + limit int + withTotal bool }{ { name: "message:request", @@ -1259,6 +1275,14 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: midTime, }, + { + name: "message:request (time range + total)", + query: "message:request", + filter: func(doc *testDoc) bool { return strings.Contains(doc.message, "request") }, + fromTime: fromTime, + toTime: midTime, + withTotal: true, + }, { name: "message:request (time range + limit)", query: "message:request", @@ -1267,6 +1291,15 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { toTime: midTime, limit: 100, }, + { + name: "message:request (time range + limit + total)", + query: "message:request", + filter: func(doc *testDoc) bool { return strings.Contains(doc.message, "request") }, + fromTime: fromTime, + toTime: midTime, + limit: 100, + withTotal: true, + }, { name: "service:bus", query: "service:bus", @@ -1352,6 +1385,15 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { fromTime: fromTime, toTime: toTime, }, + { + name: "trace_id:trace-4999 (limit + total)", + query: "trace_id:trace-4999", + filter: func(doc *testDoc) bool { return doc.traceId == "trace-4999" }, + fromTime: fromTime, + toTime: toTime, + limit: 1, + withTotal: true, + }, { name: "trace_id:trace-2025 (time range)", query: "trace_id:trace-2025", @@ -1359,7 +1401,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", @@ -1406,6 +1448,35 @@ 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 level:5 AND message:processing (time range + limit + total)", + 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, + limit: 100, + withTotal: true, + }, { name: "service:gateway AND message:processing AND message:retry AND level:5", query: "service:gateway AND message:processing AND message:retry AND level:5", @@ -1466,7 +1537,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 { @@ -1477,6 +1548,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])", @@ -1579,6 +1669,7 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { for _, tc := range searchTestCases { s.Run(tc.name, func() { var expectedIndexes []int + var expectedTotal uint64 for i := len(testDocs) - 1; i >= 0; i-- { doc := testDocs[i] @@ -1590,10 +1681,13 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { } if tc.filter(doc) { - expectedIndexes = append(expectedIndexes, i) - if tc.limit > 0 && len(expectedIndexes) >= tc.limit { + if tc.limit == 0 || len(expectedIndexes) < tc.limit { + expectedIndexes = append(expectedIndexes, i) + } + if tc.limit > 0 && len(expectedIndexes) == tc.limit && !tc.withTotal { break } + expectedTotal++ } } @@ -1602,12 +1696,24 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { if tc.limit > 0 { options = append(options, withLimit(tc.limit)) } + if tc.withTotal { + options = append(options, withTotal()) + } - s.AssertSearch(s.query(tc.query, options...), docJsons, expectedIndexes) + // TODO(cheb) total returns fuzzy results for nested indexes + if tc.withTotal && !nestedIndexes { + s.AssertSearchWithTotal(s.query(tc.query, options...), docJsons, expectedIndexes, expectedTotal) + } else { + s.AssertSearch(s.query(tc.query, options...), docJsons, expectedIndexes) + } }) } s.Run("service:kafka | group by pod unique_count(client_ip)", func() { + // TODO(cheb0) aggregation returns fuzzy results with nested indexes enabled + if nestedIndexes { + return + } // Check both sort orders simply for aggTree to be iterated in a different order orders := []seq.DocsOrder{seq.DocsOrderDesc, seq.DocsOrderAsc} @@ -1648,6 +1754,11 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { }) s.Run("service:scheduler | group by pod avg(level)", func() { + // TODO(cheb0) aggregation returns fuzzy results with nested indexes enabled + if nestedIndexes { + return + } + // Check both sort orders simply for aggTree to be iterated in a different order orders := []seq.DocsOrder{seq.DocsOrderDesc, seq.DocsOrderAsc} @@ -1691,6 +1802,11 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { // Test large QPR with 25000 groups (all ids are unique) s.Run("_exists_:service | group by id count()", func() { + // TODO(cheb0) aggregation returns fuzzy results with nested indexes enabled + if nestedIndexes { + return + } + countById := make(map[string]int) for _, doc := range testDocs { countById[doc.id]++ @@ -1717,6 +1833,11 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { }) s.Run("NOT message:retry | group by service avg(level)", func() { + // TODO(cheb0) aggregation returns fuzzy results with nested indexes enabled + if nestedIndexes { + return + } + levelsByService := make(map[string][]int) for _, doc := range testDocs { // our query for agg will be `NOT message:retry` @@ -1754,6 +1875,11 @@ func (s *FractionTestSuite) TestSearchLargeFrac() { }) s.Run("service:database AND level:3 | hist 1s", func() { + // TODO(cheb0) histogram returns fuzzy results with nested indexes enabled + if nestedIndexes { + return + } + // Check both sort orders simply for lid tree to be iterated in a different order orders := []seq.DocsOrder{seq.DocsOrderDesc, seq.DocsOrderAsc} @@ -2028,7 +2154,7 @@ func (s *FractionTestSuite) TestSearchDownsample() { eps = 0.1 ) - _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize) + _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize, false) s.insertDocuments(bulks...) baseOpts := []searchOption{ @@ -2083,7 +2209,7 @@ func (s *FractionTestSuite) TestSearchDownsampleWithTotal() { eps = 0.1 ) - _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize) + _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize, false) s.insertDocuments(bulks...) // downsample values to test: each should return ~1/ds of total documents @@ -2123,7 +2249,7 @@ func (s *FractionTestSuite) TestSearchDownsampleZeroAndOne() { queryAll = "message:*" ) - _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize) + _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize, false) s.insertDocuments(bulks...) baseOpts := []searchOption{ @@ -2162,7 +2288,7 @@ func (s *FractionTestSuite) TestSearchDownsampleWithAggAndHist() { downsample = 3 ) - _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize) + _, bulks, fromTime, toTime := generatesMessages(totalDocs, bulkSize, false) s.insertDocuments(bulks...) commonOpts := []searchOption{ @@ -2393,9 +2519,20 @@ func mustParseTime(timeStr string) time.Time { func (s *FractionTestSuite) AssertSearch(queryObject any, originalDocs []string, expectedIndexes []int) { switch q := queryObject.(type) { case string: - s.AssertSearchWithSearchParams(s.query(q), originalDocs, expectedIndexes) + s.AssertSearchWithSearchParams(s.query(q), originalDocs, expectedIndexes, nil) case *processor.SearchParams: - s.AssertSearchWithSearchParams(q, originalDocs, expectedIndexes) + s.AssertSearchWithSearchParams(q, originalDocs, expectedIndexes, nil) + default: + s.Require().Fail("type for query object not supported") + } +} + +func (s *FractionTestSuite) AssertSearchWithTotal(queryObject any, originalDocs []string, expectedIndexes []int, expectedTotal uint64) { + switch q := queryObject.(type) { + case string: + s.AssertSearchWithSearchParams(s.query(q), originalDocs, expectedIndexes, &expectedTotal) + case *processor.SearchParams: + s.AssertSearchWithSearchParams(q, originalDocs, expectedIndexes, &expectedTotal) default: s.Require().Fail("type for query object not supported") } @@ -2405,6 +2542,7 @@ func (s *FractionTestSuite) AssertSearchWithSearchParams( params *processor.SearchParams, originalDocs []string, expectedIndexes []int, + expectedTotal *uint64, ) { sortOrders := []seq.DocsOrder{params.Order} if params.Order == seq.DocsOrderDesc && params.Limit == math.MaxInt32 { @@ -2417,6 +2555,9 @@ func (s *FractionTestSuite) AssertSearchWithSearchParams( qpr, err := s.fraction.Search(context.Background(), *params) s.Require().NoError(err, "search failed for query with order=%v", order) s.Require().Equal(len(expectedIndexes), qpr.IDs.Len(), "doc count doesn't match") + if expectedTotal != nil { + s.Require().Equal(*expectedTotal, qpr.Total, "total doesn't match") + } docs, err := s.fraction.Fetch(context.Background(), qpr.IDs.IDs(), false) s.Require().NoError(err, "failed to fetch docs") diff --git a/frac/processor/aggregator_test.go b/frac/processor/aggregator_test.go index ae074afc..644932d5 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 00000000..e5a8d98c --- /dev/null +++ b/frac/processor/batch_eval_tree.go @@ -0,0 +1,227 @@ +package processor + +import ( + "errors" + "fmt" + + "github.com/ozontech/seq-db/config" + "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, + fracVer config.BinaryDataVersion, + 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 fracVer < config.BinaryDataV6 { + // block batching for earlier versions (avoid delta-encoded posting lists converted to bitmaps) + 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 v := root.Value.(type) { + case *parser.Literal: + return len(v.Terms) == 1 && v.Terms[0].Kind == parser.TermText + case *parser.Range: + return true + case *parser.Logical: + if v.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 + } + + // Currently we do not support batches for queries with deep trees. For example, queries like 'service:abc* AND level:*' + 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 00000000..77b15dc6 --- /dev/null +++ b/frac/processor/batch_eval_tree_test.go @@ -0,0 +1,259 @@ +package processor + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/ozontech/seq-db/config" + "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) GetTIDsByField(field string) ([]uint32, error) { + 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, config.BinaryDataV6, 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, config.BinaryDataV6, 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, config.BinaryDataV6, 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, config.BinaryDataV6, 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, config.BinaryDataV6, 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, config.BinaryDataV6, 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 2db0ef8c..34bfb06b 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 24e09ba5..975826bb 100644 --- a/frac/processor/search.go +++ b/frac/processor/search.go @@ -8,6 +8,7 @@ import ( "sync" "time" + "github.com/ozontech/seq-db/config" "github.com/ozontech/seq-db/consts" "github.com/ozontech/seq-db/frac/sealed/lids" "github.com/ozontech/seq-db/metric/stopwatch" @@ -35,7 +36,9 @@ type tokenIndex interface { GetValByTID(tid uint32, field string) []byte GetTIDsByField(field string) ([]uint32, error) 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 { @@ -48,6 +51,7 @@ type searchBuffers struct { mids []seq.MID rids []seq.RID lids []node.LID + tmp []uint32 } var searchBuffersPool = sync.Pool{ @@ -56,19 +60,20 @@ 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, + fracVer config.BinaryDataVersion, params SearchParams, index searchIndex, aggLimits AggLimits, + queryOpt QueryOptimizationConfig, sw *stopwatch.Stopwatch, ) (qpr *seq.QPR, err error) { stats := &searchStats{} @@ -77,18 +82,55 @@ 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, fracVer, 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()) + } else { + batchExecutionFracsTotal.Inc() + } + defer func(start time.Time) { stats.TreeDuration += time.Since(start) }(time.Now()) if util.IsCancelled(ctx) { @@ -117,22 +159,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() @@ -176,35 +202,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) { @@ -227,8 +229,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") @@ -239,71 +241,95 @@ func iterateEvalTree( sample := sampler(params.Downsample) var aggs []Aggregator - for { + for (params.Limit-len(ids)) > 0 || needScanAllRange { if util.IsCancelled(ctx) { return total, ids, hist, aggs, ctx.Err() } - needIDs := params.Limit - len(ids) - if needIDs < 1 && !needScanAllRange { - break - } - - maxBatchSize := 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 - } - 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 - } + // Process batch part by part (batches can be quite large currently) + for (params.Limit-len(ids)) > 0 || needScanAllRange { + 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() + needIDs := params.Limit - len(ids) - if hasHist { - timerHist.Start() - hist.Update(mids) - timerHist.Stop() + // Estimate how many LIDs we want in the next part to keep the balance between unneeded work and batch part size. + var toProcessLIDs int + if needIDs > 0 { + // We have IDs to fill for search - iterate batch by supposedly smaller parts with length equal to count of IDs needed + // This allows fetching MIDs for the entire batch part to serve for IDs creation and hist + toProcessLIDs = min(needIDs, cap(lidsBuf)) + } else if needScanAllRange || params.Downsample > 1 { + // We don't have IDs to fill for search. We now operate on larger parts (size is capped by tmp buff). + // If it's a hist request, then we fetch MIDs for entire part which means no unneeded work is done. + toProcessLIDs = cap(lidsBuf) } - if needIDs > 0 { - needLIDs := min(needIDs, len(lidsBatch)) + timerEval.Start() + n := iter.CopyLIDs(lidsBuf[:toProcessLIDs], tmpBuf[:toProcessLIDs]) + timerEval.Stop() + + // no more LIDs left in the current batch + if n == 0 { + break + } - timerRID.Start() - rids = idsIndex.GetRIDs(lidsBatch[:needLIDs], rids[:0]) - timerRID.Stop() + // get the copied LIDs part of batch + lidsBatch := lidsBuf[:n] + + // TODO(cheb0) not correct for nested indexes + total += n + + lidsBatch = sample(lidsBatch) + + if len(lidsBatch) == 0 { + continue + } - // 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 hasHist || needIDs > 0 { + timerMID.Start() + mids = idsIndex.GetMIDs(lidsBatch, mids[:0]) + timerMID.Stop() + + if hasHist { + timerHist.Start() + hist.Update(mids) + timerHist.Stop() + } + + if needIDs > 0 { + timerRID.Start() + rids = idsIndex.GetRIDs(lidsBatch, rids[:0]) + timerRID.Stop() + + // fill IDs for search + for i := 0; i < len(lidsBatch) && params.Limit-len(ids) > 0; i++ { + id := seq.ID{MID: mids[i], RID: rids[i]} + if len(ids) == 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 + } } } } @@ -353,17 +379,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/processor/search_stats.go b/frac/processor/search_stats.go index 31788b95..6afdbcb9 100644 --- a/frac/processor/search_stats.go +++ b/frac/processor/search_stats.go @@ -49,6 +49,12 @@ var ( Buckets: prometheus.ExponentialBuckets(1, 5, 20), Help: "Number of document hits per search", }) + batchExecutionFracsTotal = promauto.NewCounter(prometheus.CounterOpts{ + Namespace: metricsNamespace, + Subsystem: metricsSubsystem, + Name: "batch_execution_fracs_total", + Help: "Number of fractions queried with batch execution", + }) ) type searchStats struct { diff --git a/frac/processor/search_test.go b/frac/processor/search_test.go new file mode 100644 index 00000000..e20af701 --- /dev/null +++ b/frac/processor/search_test.go @@ -0,0 +1,67 @@ +package processor + +import ( + "context" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/ozontech/seq-db/metric/stopwatch" + "github.com/ozontech/seq-db/node" + "github.com/ozontech/seq-db/seq" +) + +type stubIDsIndex struct { + ids []seq.ID +} + +func (s *stubIDsIndex) LessOrEqual(lid seq.LID, id seq.ID) bool { panic("not used") } +func (s *stubIDsIndex) GetMID(lid seq.LID) seq.MID { return s.ids[lid].MID } +func (s *stubIDsIndex) GetRID(lid seq.LID) seq.RID { return s.ids[lid].RID } +func (s *stubIDsIndex) Len() int { return len(s.ids) } + +func (s *stubIDsIndex) GetMIDs(lids []node.LID, out []seq.MID) []seq.MID { + for _, lid := range lids { + out = append(out, s.ids[lid.Unpack()].MID) + } + return out +} + +func (s *stubIDsIndex) GetRIDs(lids []node.LID, out []seq.RID) []seq.RID { + for _, lid := range lids { + out = append(out, s.ids[lid.Unpack()].RID) + } + return out +} + +func TestIterateEvalTreeDuplicateIDsAtBatchBoundary(t *testing.T) { + // LID 0 is unused; IDs descend as LID grows; LIDs 1 and 2 share one seq.ID. + idx := &stubIDsIndex{ids: []seq.ID{ + 1: {MID: 600, RID: 1}, + 2: {MID: 600, RID: 1}, // duplicate of LID 1 + 3: {MID: 500, RID: 1}, + 4: {MID: 400, RID: 1}, + 5: {MID: 300, RID: 1}, + 6: {MID: 200, RID: 1}, + }} + + // Single batch with 6 LIDs, longer than the limit. + evalTree := node.NewStaticBatched([]uint32{1, 2, 3, 4, 5, 6}, seq.DocsOrderDesc.IsReverse()) + + params := SearchParams{ + Limit: 3, + Order: seq.DocsOrderDesc, + } + + total, ids, _, _, err := iterateEvalTree(context.Background(), params, idx, evalTree, nil, stopwatch.New()) + require.NoError(t, err) + + // 6 LIDs exist and the limit is not yet satisfied after the duplicate, + // so the search must keep scanning: expected IDs are 600, 500, 400. + got := make([]seq.MID, 0, len(ids)) + for _, id := range ids { + got = append(got, id.ID.MID) + } + require.Equal(t, []seq.MID{600, 500, 400}, got) + require.Equal(t, 4, total) // 4 LIDs scanned to produce 3 distinct IDs +} diff --git a/frac/sealed/lids/block_test.go b/frac/sealed/lids/block_test.go index a429c16e..06293d7f 100644 --- a/frac/sealed/lids/block_test.go +++ b/frac/sealed/lids/block_test.go @@ -197,10 +197,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 f2634553..0bfaca08 100644 --- a/frac/sealed/lids/iterator_batched_asc.go +++ b/frac/sealed/lids/iterator_batched_asc.go @@ -65,11 +65,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 c6465c6c..5c007f14 100644 --- a/frac/sealed/lids/iterator_batched_desc.go +++ b/frac/sealed/lids/iterator_batched_desc.go @@ -65,11 +65,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 b079743e..da302958 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, dp.info.BinaryDataVer, params, dp.getSearchIndex(), aggLimits, queryOpt, sw) if err != nil { return nil, err } @@ -275,6 +278,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 @@ -306,6 +327,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 6da8ef8f..0aaf61ba 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 00000000..5ce0056e --- /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, 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 00000000..468ddcde --- /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 1454b538..375042c8 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 98cf21e3..27d3b111 100644 --- a/node/node.go +++ b/node/node.go @@ -13,10 +13,10 @@ type Node interface { type BatchedNode interface { fmt.Stringer - // NextBatch returns next batch. Returns nil when exhausted. - NextBatch(need int) LIDBatch - // NextBatchGeq returns next batch (LIDs >= minLID). Returns nil when exhausted. - NextBatchGeq(need int, nextLID LID) LIDBatch + // NextBatch returns next batch. Returns empty batch when exhausted. + NextBatch() LIDBatch + // NextBatchGeq returns next batch (LIDs >= minLID). Returns empty batch when exhausted. + NextBatchGeq(nextID LID) LIDBatch } type Sourced interface { diff --git a/node/node_and.go b/node/node_and.go index 856e12f1..95bad8e7 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 52f5ff01..54dfa2b9 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 ab0bf30f..969f4cdc 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 baabfa37..f8f54a67 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 00000000..a9b69ec6 --- /dev/null +++ b/node/util.go @@ -0,0 +1,77 @@ +package node + +import ( + "fmt" + "slices" +) + +const maxBatchDrain = 4 * 1024 + +// batcherNode allows to iterate over non-batched iterator batch by batch. +// A caller must immediately consume a yielded batch after calling NextBatch, since +// the underlying slice is reused. +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 { + b.batch = b.batch[:0] + for len(b.batch) < maxBatchDrain { + lid := b.source.Next() + if lid.IsNull() { + break + } + b.batch = append(b.batch, lid.Unpack()) + } + if !b.desc { + slices.Reverse(b.batch) + } + return NewSliceBatch(b.batch) +} + +func (b *batcherNode) NextBatchGeq(nextID LID) LIDBatch { + b.batch = b.batch[:0] + for len(b.batch) < maxBatchDrain { + lid := b.source.NextGeq(nextID) + if lid.IsNull() { + break + } + b.batch = append(b.batch, lid.Unpack()) + } + if !b.desc { + slices.Reverse(b.batch) + } + return NewSliceBatch(b.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() +}