diff --git a/cmd/pyroscope/help-all.txt.tmpl b/cmd/pyroscope/help-all.txt.tmpl index 3470e41c5c..d20783e9a0 100644 --- a/cmd/pyroscope/help-all.txt.tmpl +++ b/cmd/pyroscope/help-all.txt.tmpl @@ -935,6 +935,8 @@ Usage of ./pyroscope: Port to advertise to query-scheduler and querier (defaults to -server.http-listen-port). -query-frontend.max-async-query-concurrency int Maximum number of concurrent async queries per tenant. 0 to disable async queries. (default 5) + -query-frontend.query-planner-strategy string + Sets the query planner strategy, options: classic, balanced (default "classic") -query-frontend.scheduler-worker-concurrency int Number of concurrent workers forwarding queries to single query-scheduler. (default 5) -query-scheduler.grpc-client-config.backoff-max-period duration diff --git a/cmd/pyroscope/help.txt.tmpl b/cmd/pyroscope/help.txt.tmpl index e1ae759017..eab246e977 100644 --- a/cmd/pyroscope/help.txt.tmpl +++ b/cmd/pyroscope/help.txt.tmpl @@ -199,6 +199,8 @@ Usage of ./pyroscope: Include profiles that were sampled out and stored with stacktraces stripped (marked __sampled__) in query results. -query-frontend.max-async-query-concurrency int Maximum number of concurrent async queries per tenant. 0 to disable async queries. (default 5) + -query-frontend.query-planner-strategy string + Sets the query planner strategy, options: classic, balanced (default "classic") -query-scheduler.max-outstanding-requests-per-tenant int Maximum number of outstanding requests per tenant per query-scheduler. In-flight requests above this limit will fail with HTTP response status code 429. (default 100) -query-scheduler.ring.consul.hostname string diff --git a/pkg/frontend/frontend.go b/pkg/frontend/frontend.go index 6d7818451b..7b33dc1b01 100644 --- a/pkg/frontend/frontend.go +++ b/pkg/frontend/frontend.go @@ -54,6 +54,13 @@ type Config struct { // on the request is rejected with Unimplemented. AsyncQueriesEnabled bool `yaml:"async_queries_enabled" category:"experimental"` + // QueryPlannerStrategy sets the query planner strategy. By default this is + // "classic" which provides the legacy query planner behavior. + // + // Optionally this can be "balanced" which uses the balanced query planner + // algorithm. + QueryPlannerStrategy string `yaml:"query_planner_strategy" doc:"hidden"` + // Used to find local IP address, that is sent to scheduler and querier-worker. InfNames []string `yaml:"instance_interface_names" category:"advanced" doc:"default=[]"` Addr string `yaml:"instance_addr" category:"advanced"` @@ -79,6 +86,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet, logger log.Logger) { f.BoolVar(&cfg.EnableIPv6, "query-frontend.instance-enable-ipv6", false, "Enable using a IPv6 instance address. (default false)") f.IntVar(&cfg.Port, "query-frontend.instance-port", 0, "Port to advertise to query-scheduler and querier (defaults to -server.http-listen-port).") f.BoolVar(&cfg.AsyncQueriesEnabled, "query-frontend.async-queries-enabled", false, "Enable the experimental asynchronous query path on SelectMergeStacktraces (default false)") + f.StringVar(&cfg.QueryPlannerStrategy, "query-frontend.query-planner-strategy", "classic", "Sets the query planner strategy, options: classic, balanced") cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-frontend.grpc-client-config", f) } @@ -87,6 +95,14 @@ func (cfg *Config) Validate() error { return fmt.Errorf("scheduler address cannot be specified when query-scheduler service discovery mode is set to '%s'", cfg.QuerySchedulerDiscovery.Mode) } + switch cfg.QueryPlannerStrategy { + case "": + cfg.QueryPlannerStrategy = "classic" + case "classic", "balanced": + default: + return fmt.Errorf("unknown query planner strategy: %q", cfg.QueryPlannerStrategy) + } + return cfg.GRPCClientConfig.Validate() } diff --git a/pkg/frontend/readpath/queryfrontend/query_frontend.go b/pkg/frontend/readpath/queryfrontend/query_frontend.go index eb2796c2ae..ec2b5abd90 100644 --- a/pkg/frontend/readpath/queryfrontend/query_frontend.go +++ b/pkg/frontend/readpath/queryfrontend/query_frontend.go @@ -66,6 +66,7 @@ type QueryFrontend struct { symbolizer Symbolizer diagnosticsStore DiagnosticsStore now func() time.Time + queryPlanType string metrics *queryFrontendMetrics } @@ -132,6 +133,7 @@ func newQueryFrontendMetrics(reg prometheus.Registerer) *queryFrontendMetrics { func NewQueryFrontend( logger log.Logger, limits frontend.Limits, + cfg frontend.Config, metadataQueryClient metastorev1.MetadataQueryServiceClient, tenantServiceClient metastorev1.TenantServiceClient, querybackendClient QueryBackend, @@ -148,6 +150,7 @@ func NewQueryFrontend( symbolizer: sym, diagnosticsStore: diagnosticsStore, now: time.Now, + queryPlanType: cfg.QueryPlannerStrategy, metrics: newQueryFrontendMetrics(reg), } return qf @@ -228,7 +231,7 @@ func (q *QueryFrontend) doQuery( endTime := time.UnixMilli(req.EndTime) queryWindow := endTime.Sub(startTime).Round(time.Second) traceID, _ := tracing.ExtractTraceID(ctx) - logArgs := []interface{}{ + logArgs := []any{ "msg", "query weight", "trace_id", traceID, "tenant", strings.Join(tenants, ","), @@ -255,8 +258,14 @@ func (q *QueryFrontend) doQuery( blocks[i], blocks[j] = blocks[j], blocks[i] }) xrandMutex.Unlock() - // TODO(kolesnikovae): Should be dynamic. - p := queryplan.Build(blocks, 4, 20) + + var p *queryv1.QueryPlan + switch q.queryPlanType { + case "balanced": + p = queryplan.BuildBalanced(blocks, 4, 20) + default: // Always fallback to the classic planner + p = queryplan.Build(blocks, 4, 20) + } backend := q.querybackend if backendC != nil { diff --git a/pkg/frontend/readpath/queryfrontend/query_frontend_test.go b/pkg/frontend/readpath/queryfrontend/query_frontend_test.go index 66f0a1ee71..9b749d3b92 100644 --- a/pkg/frontend/readpath/queryfrontend/query_frontend_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_frontend_test.go @@ -18,6 +18,7 @@ import ( typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1" "github.com/grafana/pyroscope/v2/pkg/block/metadata" "github.com/grafana/pyroscope/v2/pkg/featureflags" + "github.com/grafana/pyroscope/v2/pkg/frontend" "github.com/grafana/pyroscope/v2/pkg/tenant" "github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend" "github.com/grafana/pyroscope/v2/pkg/test/mocks/mockmetastorev1" @@ -186,6 +187,7 @@ func Test_QueryFrontend_LabelNames_WithFiltering(t *testing.T) { qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, @@ -329,6 +331,7 @@ func Test_QueryFrontend_Series_WithLabelNameFiltering(t *testing.T) { qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go index 7b446e9cf2..d20b11622c 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go @@ -19,6 +19,7 @@ import ( queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1" typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1" "github.com/grafana/pyroscope/v2/pkg/block/metadata" + "github.com/grafana/pyroscope/v2/pkg/frontend" phlaremodel "github.com/grafana/pyroscope/v2/pkg/model" "github.com/grafana/pyroscope/v2/pkg/pprof" "github.com/grafana/pyroscope/v2/pkg/tenant" @@ -41,6 +42,7 @@ func newSMPQueryFrontend( return NewQueryFrontend( log.NewNopLogger(), limits, + frontend.Config{}, metaClient, nil, // tenantServiceClient backend, @@ -682,6 +684,7 @@ func TestSelectMergeProfiles_Symbolization(t *testing.T) { qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go index 9c22bff2d0..47dc57bd04 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go @@ -17,6 +17,7 @@ import ( querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1" queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1" "github.com/grafana/pyroscope/v2/pkg/block/metadata" + "github.com/grafana/pyroscope/v2/pkg/frontend" "github.com/grafana/pyroscope/v2/pkg/pprof" "github.com/grafana/pyroscope/v2/pkg/tenant" "github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend" @@ -172,6 +173,7 @@ func TestSelectMergeSpanProfile_Symbolization(t *testing.T) { qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go index 0a9c6df7b6..4aeefca28e 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go @@ -16,6 +16,7 @@ import ( querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1" queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1" "github.com/grafana/pyroscope/v2/pkg/block/metadata" + "github.com/grafana/pyroscope/v2/pkg/frontend" phlaremodel "github.com/grafana/pyroscope/v2/pkg/model" "github.com/grafana/pyroscope/v2/pkg/pprof" "github.com/grafana/pyroscope/v2/pkg/tenant" @@ -280,6 +281,7 @@ func TestSelectMergeStacktrace_Symbolization(t *testing.T) { qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, @@ -354,7 +356,7 @@ func TestSelectMergeStacktraces_DotFormat(t *testing.T) { }}, }, nil) - qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, nil, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, nil, nil, nil) ctx := tenant.InjectTenantID(context.Background(), "tenant1") start, end := smpValidTimeRange() diff --git a/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go b/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go index 19c359665c..770b9f2777 100644 --- a/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go @@ -10,6 +10,7 @@ import ( "github.com/stretchr/testify/require" querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1" + "github.com/grafana/pyroscope/v2/pkg/frontend" "github.com/grafana/pyroscope/v2/pkg/tenant" "github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend" ) @@ -24,7 +25,7 @@ func TestSelectSeries_RejectsSubMillisecondStep(t *testing.T) { limits.On("MaxQueryLookback", "test-tenant").Return(time.Duration(0)).Maybe() limits.On("MaxQueryLength", "test-tenant").Return(time.Duration(0)).Maybe() - qf := NewQueryFrontend(log.NewNopLogger(), limits, nil, nil, nil, nil, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), limits, frontend.Config{}, nil, nil, nil, nil, nil, nil) ctx := tenant.InjectTenantID(context.Background(), "test-tenant") _, err := qf.SelectSeries(ctx, connect.NewRequest(&querierv1.SelectSeriesRequest{ @@ -46,7 +47,7 @@ func TestSelectHeatmap_RejectsSubMillisecondStep(t *testing.T) { for _, step := range []float64{0, 0.0001, 0.0005, 0.0009999} { t.Run("step="+formatStep(step), func(t *testing.T) { limits := mockfrontend.NewMockLimits(t) - qf := NewQueryFrontend(log.NewNopLogger(), limits, nil, nil, nil, nil, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), limits, frontend.Config{}, nil, nil, nil, nil, nil, nil) ctx := tenant.InjectTenantID(context.Background(), "test-tenant") _, err := qf.SelectHeatmap(ctx, connect.NewRequest(&querierv1.SelectHeatmapRequest{ diff --git a/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go b/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go index 89f46bc316..6be58c037e 100644 --- a/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go +++ b/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go @@ -14,6 +14,7 @@ import ( metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1" querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1" typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1" + "github.com/grafana/pyroscope/v2/pkg/frontend" phlaremodel "github.com/grafana/pyroscope/v2/pkg/model" "github.com/grafana/pyroscope/v2/pkg/tenant" "github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend" @@ -136,6 +137,7 @@ func Test_QueryFrontend_Series_ProfileTypeQueryServedFromMetadata(t *testing.T) qf := NewQueryFrontend( log.NewNopLogger(), mockLimits, + frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, diff --git a/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go b/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go index 9c6259a87f..c039744d64 100644 --- a/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go +++ b/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go @@ -16,6 +16,7 @@ import ( querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1" queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1" "github.com/grafana/pyroscope/lidia" + "github.com/grafana/pyroscope/v2/pkg/frontend" phlaremodel "github.com/grafana/pyroscope/v2/pkg/model" "github.com/grafana/pyroscope/v2/pkg/model/symbolref" "github.com/grafana/pyroscope/v2/pkg/tenant" @@ -177,7 +178,7 @@ func TestSelectMergeStacktracesTree_SymbolRefFlagOn(t *testing.T) { }}, }, nil).Once() - qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil) ctx := tenant.InjectTenantID(context.Background(), "tenant1") start, end := smpValidTimeRange() @@ -235,7 +236,7 @@ func TestSelectMergeStacktracesTree_SymbolRefResolution(t *testing.T) { Blocks: []*metastorev1.BlockMeta{{Id: "block_id"}}, }, nil).Once() - qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil) before := testutil.ToFloat64(qf.metrics.symbolRefLocationsTotal.WithLabelValues(symbolRefLocationResolved)) @@ -269,7 +270,7 @@ func TestSelectMergeStacktracesTree_SymbolRefResolution(t *testing.T) { // path (see TestRebuildInlineChainExpansionOrder for the Rebuild-side // contract). func TestBuildLookup_ReversesLidiaFrameOrder(t *testing.T) { - qf := NewQueryFrontend(log.NewNopLogger(), nil, nil, nil, nil, nil, nil, nil) + qf := NewQueryFrontend(log.NewNopLogger(), nil, frontend.Config{}, nil, nil, nil, nil, nil, nil) lookup := qf.buildLookup([]binaryResolution{{ binary: symbolref.UnresolvedBinary{BuildID: "build-a", BinaryName: "libfoo.so", Addresses: []uint64{0x100}}, frames: [][]lidia.SourceInfoFrame{{ diff --git a/pkg/pyroscope/modules_experimental.go b/pkg/pyroscope/modules_experimental.go index 4b5e43b8e1..1278af6594 100644 --- a/pkg/pyroscope/modules_experimental.go +++ b/pkg/pyroscope/modules_experimental.go @@ -87,6 +87,7 @@ func (f *Pyroscope) initQueryFrontendV2() (services.Service, error) { f.queryFrontend = queryfrontend.NewQueryFrontend( queryFrontendLogger, f.Overrides, + f.Cfg.Frontend, f.metastoreClient, f.metastoreClient, f.queryBackendClient, @@ -145,6 +146,7 @@ func (f *Pyroscope) initQueryFrontendV12() (services.Service, error) { f.queryFrontend = queryfrontend.NewQueryFrontend( queryFrontendLogger, f.Overrides, + f.Cfg.Frontend, f.metastoreClient, f.metastoreClient, f.queryBackendClient, diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go index 475585d9d0..7e7582e94a 100644 --- a/pkg/querybackend/queryplan/query_plan.go +++ b/pkg/querybackend/queryplan/query_plan.go @@ -36,10 +36,7 @@ func Build( leafNodeCount := (len(blocks) + maxReads - 1) / maxReads nodes := allocateContiguous[queryv1.QueryNode](leafNodeCount) for start, idx := 0, 0; start < len(blocks); start, idx = start+maxReads, idx+1 { - end := start + maxReads - if end > len(blocks) { - end = len(blocks) - } + end := min(start+maxReads, len(blocks)) nodes[idx].Type = queryv1.QueryNode_READ nodes[idx].Blocks = blocks[start:end] } @@ -50,10 +47,7 @@ func Build( mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount) for start, idx := 0, 0; start < len(nodes); start, idx = start+maxMerges, idx+1 { - end := start + maxMerges - if end > len(nodes) { - end = len(nodes) - } + end := min(start+maxMerges, len(nodes)) mergeNodes[idx].Type = queryv1.QueryNode_MERGE mergeNodes[idx].Children = nodes[start:end:end] } @@ -66,6 +60,167 @@ func Build( } } +// BuildBalanced builds a balanced query tree where each node of the tree has a +// similar number of blocks it needs to process compared to its siblings. +func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int) *queryv1.QueryPlan { + if len(blocks) == 0 || maxReads < 1 || maxMerges < 2 { + return new(queryv1.QueryPlan) + } + + leafNodeCount := (len(blocks) + maxReads - 1) / maxReads + nodes := allocateContiguous[queryv1.QueryNode](leafNodeCount) + + // Uniformly assign blocks to leaf nodes. + var start int + for idx := range leafNodeCount { + end := (idx + 1) * len(blocks) / leafNodeCount + nodes[idx].Type = queryv1.QueryNode_READ + nodes[idx].Blocks = blocks[start:end:end] + start = end + } + + // Recursively build a balanced tree of merge nodes. + root := buildMergeTree(nodes, maxMerges) + return &queryv1.QueryPlan{ + Root: root, + } +} + +// buildMergeTree will recursively create a query tree of merge nodes with at +// most maxMerges number of children. At each level, the tree will maintain a +// similar number of blocks assigned to each merge node compared to its +// siblings. +// +// The algorithm is straightforward, but it has nuances. Since len(nodes) > 1, +// we know we need to create a merge node and place nodes underneath it. +// +// When creating a merge node we want to partition the nodes evenly into groups +// such that we have no more than maxMerges groups (since a given merge node +// cannot exceed maxMerges children). Once we select a number of groups, we use +// balanceGroupItems to evenly (as evenly as possible) spread the nodes across +// all the groups. +// +// It's important to note that we want to select a groupSize which is a power of +// maxMerges. This ensures that each subtree has the same depth as its siblings. +// +// As an example, consider the following input: +// +// maxMerges = 3 +// nodes = [ n0 n1 n2 n3 n4 n5 n6 n7 n8 n9 ] +// +// We want to partition nodes such that each group has an even number of +// nodes itself. Naively we could compute: +// +// groupSize = ceil(len(nodes) / maxMerges) +// +// However, this would result in an imbalanced tree: +// +// [ M0 ] +// [ n0 n1 n2 n3 ] [ n4 n5 n6 ] [ n7 n8 n9 ] +// +// From here, the first partition was given 4 nodes, which exceeds maxMerges, +// so it needs to be broken down further. The other two partitions have 3 nodes, +// so they do not need to be branched further. After splitting the first +// partition, the tree becomes imbalanced. +// +// [ M0 ] +// [ M1 ] [ n4 n5 n6 ] [ n7 n8 n9 ] +// [ n0 n1 ] [ n2 n3 ] +// +// Instead, we select the largest power of maxMerges K that's less than +// len(nodes). In this case: +// +// K = maxMerges^N +// for max(N) and maxMerges^N < len(nodes) +// if N = 2, then 3^2 < 10 so K = 9 +// +// If we allow K to represent the largest size of a partition, we can calculate +// how many partitions we need: +// +// # of partitions = ceil(len(nodes) / K) +// = ceil(10 / 5) +// = 2 +// +// Now we can balance the nodes across 2 partitions: +// +// [ M0 ] +// [ n0 n1 n2 n3 n4 ] [ n5 n6 n7 n8 n9 ] +// +// Both partitions exceed maxMerges, so we repeat the algorithm: +// +// [ M0 ] +// [ M1 ] [ M2 ] +// [ n0 n1 n2 ] [ n3 n4 ] [ n5 n6 n7 ] [ n8 n9 ] +// +// At this point, each subtree does the same amount of work as its siblings. +// +// As a theoretical note, this algorithm may produce an imbalanced depth tree +// if the len(nodes) is not a pwer of maxMerges and maxMerges = 2. This case is +// unlikely in production--as maxMerges = 20 is common--but also irrelevant to +// performance overall since such a low maxMerges value would produce other +// inefficiencies. +func buildMergeTree(nodes []*queryv1.QueryNode, maxMerges int) *queryv1.QueryNode { + if len(nodes) == 1 { + // The base case, this is a leaf node. + return nodes[0] + } + + // We have len(nodes) number of nodes. We need to partition them into groups + // such that we have no more than maxMerges number of groups. Importantly, we + // want a group size such that we don't place all the nodes into a single + // group. + groupSize := 1 + for groupSize*maxMerges < len(nodes) { + groupSize *= maxMerges + } + + // Given groupSize as the maximum number of nodes that can be assigned to + // each subtree, we calculate + // + // ceil(len(nodes) / groupSize) + // + // to determine how many children this node will have. + childCount := (len(nodes) + groupSize - 1) / groupSize + + parent := &queryv1.QueryNode{ + Type: queryv1.QueryNode_MERGE, + Children: make([]*queryv1.QueryNode, childCount), + } + + // Evenly distribute all the nodes to each child of this merge node. + var start int + for idx := range childCount { + end := start + balanceGroupItems(len(nodes), childCount, idx) + parent.Children[idx] = buildMergeTree(nodes[start:end], maxMerges) + start = end + } + + return parent +} + +// balanceGroupItems will take numItems and distribute them evenly across +// numGroups groups. If there is a remainder R, that remainder is then spread +// evenly across the first R groups. It returns the number of items groupIdx +// should have for all groups to remain balanced. +// +// For example, given these parameters: +// +// numItems = 5 +// numGroups = 3 +// +// We would get the following (for various values of groupIdx): +// +// Group 0: 2 +// Group 1: 2 +// Group 2: 1 +func balanceGroupItems(numItems int, numGroups int, groupIdx int) int { + groupSize := numItems / numGroups + if groupIdx < numItems%numGroups { + groupSize++ + } + return groupSize +} + // allocateContiguous returns a []*T of length size where every element points // into a single backing []T allocation. This avoids the per-element heap // allocations from N separate &T{} expressions. diff --git a/pkg/querybackend/queryplan/query_plan_test.go b/pkg/querybackend/queryplan/query_plan_test.go index 318238bbdc..240b05ec7f 100644 --- a/pkg/querybackend/queryplan/query_plan_test.go +++ b/pkg/querybackend/queryplan/query_plan_test.go @@ -20,9 +20,14 @@ import ( var update = flag.Bool("update", false, "rewrite golden files in testdata/ from the current plan output") -// Test_Build verifies the shape of query plans produced by Build against -// golden files in testdata/. Each subtest's golden file is named after the -// subtest. E.g. Test_Build/single_block reads testdata/single_block.txt. +// Test_Build verifies the shape of query plans produced by Build and +// BuildBalanced against golden files in testdata/, running both builders over +// the same input table so the two algorithms are easy to compare. Each +// subtest's golden file is named after the subtest: Build's golden file is +// named after the case (e.g. Test_Build/single_block reads +// testdata/single_block.txt), and BuildBalanced's is named after the case +// with a "balanced_" prefix (e.g. Test_Build/balanced_single_block reads +// testdata/balanced_single_block.txt). // // To regenerate all golden files: // @@ -30,7 +35,7 @@ var update = flag.Bool("update", false, "rewrite golden files in testdata/ from // // To regenerate a specific golden file: // -// go test ./pkg/querybackend/queryplan/ -run Test_Build/ -update +// go test ./pkg/querybackend/queryplan/ -run 'Test_Build/$' -update func Test_Build(t *testing.T) { tests := []struct { name string @@ -47,35 +52,49 @@ func Test_Build(t *testing.T) { {name: "full_depth_2", blocks: 6, maxReads: 2, maxMerges: 3}, {name: "just_over_depth_2", blocks: 7, maxReads: 2, maxMerges: 3}, {name: "twenty_five_blocks", blocks: 25, maxReads: 2, maxMerges: 3}, + {name: "full_merge_vs_single_leaf_merge", blocks: 33, maxReads: 4, maxMerges: 8}, + {name: "forced_equal_split", blocks: 16, maxReads: 1, maxMerges: 8}, + {name: "three_way_split", blocks: 20, maxReads: 1, maxMerges: 8}, + } + + builders := []struct { + name string + build func(blocks []*metastorev1.BlockMeta, maxReads, maxMerges int) *queryv1.QueryPlan + prefix string + }{ + {name: "Build", build: Build, prefix: "build_"}, + {name: "BuildBalanced", build: BuildBalanced, prefix: "balanced_"}, } for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - blocks := makeBlocks(tt.blocks) - p := Build(blocks, tt.maxReads, tt.maxMerges) + for _, b := range builders { + t.Run(b.prefix+tt.name, func(t *testing.T) { + blocks := makeBlocks(tt.blocks) + p := b.build(blocks, tt.maxReads, tt.maxMerges) - var buf bytes.Buffer - writePlan(t, &buf, "", p.Root) + var buf bytes.Buffer + writePlan(t, &buf, "", p.Root) - // Ensure that the plan has not been modified during traversal. - assert.Equal(t, Build(blocks, tt.maxReads, tt.maxMerges), p) + // Ensure that the plan has not been modified during traversal. + assert.Equal(t, b.build(blocks, tt.maxReads, tt.maxMerges), p) - if *update { - require.NoError(t, os.WriteFile(goldenFile(t), buf.Bytes(), 0o644)) - return - } + if *update { + require.NoError(t, os.WriteFile(goldenFile(t), buf.Bytes(), 0o644)) + return + } - expected, err := os.ReadFile(goldenFile(t)) - require.NoError(t, err) - assert.Equal(t, string(expected), buf.String()) - }) + expected, err := os.ReadFile(goldenFile(t)) + require.NoError(t, err) + assert.Equal(t, string(expected), buf.String()) + }) + } } } // makeBlocks creates n BlockMeta with sequential string IDs starting at "1". func makeBlocks(n int) []*metastorev1.BlockMeta { blocks := make([]*metastorev1.BlockMeta, n) - for i := 0; i < n; i++ { + for i := range n { blocks[i] = &metastorev1.BlockMeta{Id: strconv.Itoa(i + 1)} } return blocks @@ -90,6 +109,42 @@ func goldenFile(t *testing.T) string { return filepath.Join("testdata", parts[len(parts)-1]+".txt") } +// Test_balanceGroupItems verifies that items are spread evenly across groups, +// with any remainder distributed one-per-group starting from group 0. +func Test_balanceGroupItems(t *testing.T) { + tests := []struct { + name string + numItems int + numGroups int + want []int // want[i] is the expected size of group i + }{ + {name: "zero_items", numItems: 0, numGroups: 3, want: []int{0, 0, 0}}, + {name: "single_group", numItems: 5, numGroups: 1, want: []int{5}}, + {name: "even_split", numItems: 6, numGroups: 3, want: []int{2, 2, 2}}, + {name: "remainder_spread", numItems: 5, numGroups: 3, want: []int{2, 2, 1}}, + {name: "remainder_almost_full", numItems: 8, numGroups: 3, want: []int{3, 3, 2}}, + {name: "one_item_per_group", numItems: 3, numGroups: 3, want: []int{1, 1, 1}}, + {name: "more_groups_than_items", numItems: 2, numGroups: 5, want: []int{1, 1, 0, 0, 0}}, + {name: "no_items", numItems: 0, numGroups: 1, want: []int{0}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Len(t, tt.want, tt.numGroups, "test case is malformed: want must have numGroups elements") + + got := make([]int, tt.numGroups) + sum := 0 + for idx := range tt.numGroups { + got[idx] = balanceGroupItems(tt.numItems, tt.numGroups, idx) + sum += got[idx] + } + + assert.Equal(t, tt.want, got) + assert.Equal(t, tt.numItems, sum, "group sizes must sum back to numItems") + }) + } +} + // writePlan writes an indented textual representation of the plan rooted at // n to w. A nil root produces no output. The test fails on malformed nodes. func writePlan(t *testing.T, w io.Writer, pad string, n *queryv1.QueryNode) { diff --git a/pkg/querybackend/queryplan/testdata/empty.txt b/pkg/querybackend/queryplan/testdata/balanced_empty.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/empty.txt rename to pkg/querybackend/queryplan/testdata/balanced_empty.txt diff --git a/pkg/querybackend/queryplan/testdata/exact_one_leaf.txt b/pkg/querybackend/queryplan/testdata/balanced_exact_one_leaf.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/exact_one_leaf.txt rename to pkg/querybackend/queryplan/testdata/balanced_exact_one_leaf.txt diff --git a/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt b/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt new file mode 100644 index 0000000000..08aff02ba7 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt @@ -0,0 +1,35 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 1} + id:"2" + READ {children: 0, blocks: 1} + id:"3" + READ {children: 0, blocks: 1} + id:"4" + READ {children: 0, blocks: 1} + id:"5" + READ {children: 0, blocks: 1} + id:"6" + READ {children: 0, blocks: 1} + id:"7" + READ {children: 0, blocks: 1} + id:"8" + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"9" + READ {children: 0, blocks: 1} + id:"10" + READ {children: 0, blocks: 1} + id:"11" + READ {children: 0, blocks: 1} + id:"12" + READ {children: 0, blocks: 1} + id:"13" + READ {children: 0, blocks: 1} + id:"14" + READ {children: 0, blocks: 1} + id:"15" + READ {children: 0, blocks: 1} + id:"16" diff --git a/pkg/querybackend/queryplan/testdata/full_depth_2.txt b/pkg/querybackend/queryplan/testdata/balanced_full_depth_2.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/full_depth_2.txt rename to pkg/querybackend/queryplan/testdata/balanced_full_depth_2.txt diff --git a/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt new file mode 100644 index 0000000000..44384abc9b --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt @@ -0,0 +1,45 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 5, blocks: 0} + READ {children: 0, blocks: 3} + id:"1" + id:"2" + id:"3" + READ {children: 0, blocks: 4} + id:"4" + id:"5" + id:"6" + id:"7" + READ {children: 0, blocks: 4} + id:"8" + id:"9" + id:"10" + id:"11" + READ {children: 0, blocks: 3} + id:"12" + id:"13" + id:"14" + READ {children: 0, blocks: 4} + id:"15" + id:"16" + id:"17" + id:"18" + MERGE {children: 4, blocks: 0} + READ {children: 0, blocks: 4} + id:"19" + id:"20" + id:"21" + id:"22" + READ {children: 0, blocks: 3} + id:"23" + id:"24" + id:"25" + READ {children: 0, blocks: 4} + id:"26" + id:"27" + id:"28" + id:"29" + READ {children: 0, blocks: 4} + id:"30" + id:"31" + id:"32" + id:"33" diff --git a/pkg/querybackend/queryplan/testdata/invalid_max_merges.txt b/pkg/querybackend/queryplan/testdata/balanced_invalid_max_merges.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/invalid_max_merges.txt rename to pkg/querybackend/queryplan/testdata/balanced_invalid_max_merges.txt diff --git a/pkg/querybackend/queryplan/testdata/invalid_max_reads.txt b/pkg/querybackend/queryplan/testdata/balanced_invalid_max_reads.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/invalid_max_reads.txt rename to pkg/querybackend/queryplan/testdata/balanced_invalid_max_reads.txt diff --git a/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt new file mode 100644 index 0000000000..53feaac55b --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt @@ -0,0 +1,14 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 2, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 2} + id:"2" + id:"3" + MERGE {children: 2, blocks: 0} + READ {children: 0, blocks: 2} + id:"4" + id:"5" + READ {children: 0, blocks: 2} + id:"6" + id:"7" diff --git a/pkg/querybackend/queryplan/testdata/single_block.txt b/pkg/querybackend/queryplan/testdata/balanced_single_block.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/single_block.txt rename to pkg/querybackend/queryplan/testdata/balanced_single_block.txt diff --git a/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt new file mode 100644 index 0000000000..6bf750c031 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt @@ -0,0 +1,44 @@ +MERGE {children: 3, blocks: 0} + MERGE {children: 7, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 1} + id:"2" + READ {children: 0, blocks: 1} + id:"3" + READ {children: 0, blocks: 1} + id:"4" + READ {children: 0, blocks: 1} + id:"5" + READ {children: 0, blocks: 1} + id:"6" + READ {children: 0, blocks: 1} + id:"7" + MERGE {children: 7, blocks: 0} + READ {children: 0, blocks: 1} + id:"8" + READ {children: 0, blocks: 1} + id:"9" + READ {children: 0, blocks: 1} + id:"10" + READ {children: 0, blocks: 1} + id:"11" + READ {children: 0, blocks: 1} + id:"12" + READ {children: 0, blocks: 1} + id:"13" + READ {children: 0, blocks: 1} + id:"14" + MERGE {children: 6, blocks: 0} + READ {children: 0, blocks: 1} + id:"15" + READ {children: 0, blocks: 1} + id:"16" + READ {children: 0, blocks: 1} + id:"17" + READ {children: 0, blocks: 1} + id:"18" + READ {children: 0, blocks: 1} + id:"19" + READ {children: 0, blocks: 1} + id:"20" diff --git a/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt new file mode 100644 index 0000000000..ad8ac16d21 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt @@ -0,0 +1,46 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 3, blocks: 0} + MERGE {children: 3, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 2} + id:"2" + id:"3" + READ {children: 0, blocks: 2} + id:"4" + id:"5" + MERGE {children: 2, blocks: 0} + READ {children: 0, blocks: 2} + id:"6" + id:"7" + READ {children: 0, blocks: 2} + id:"8" + id:"9" + MERGE {children: 2, blocks: 0} + READ {children: 0, blocks: 2} + id:"10" + id:"11" + READ {children: 0, blocks: 2} + id:"12" + id:"13" + MERGE {children: 2, blocks: 0} + MERGE {children: 3, blocks: 0} + READ {children: 0, blocks: 2} + id:"14" + id:"15" + READ {children: 0, blocks: 2} + id:"16" + id:"17" + READ {children: 0, blocks: 2} + id:"18" + id:"19" + MERGE {children: 3, blocks: 0} + READ {children: 0, blocks: 2} + id:"20" + id:"21" + READ {children: 0, blocks: 2} + id:"22" + id:"23" + READ {children: 0, blocks: 2} + id:"24" + id:"25" diff --git a/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt b/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt new file mode 100644 index 0000000000..a5313fd486 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt @@ -0,0 +1,6 @@ +MERGE {children: 2, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 2} + id:"2" + id:"3" diff --git a/pkg/querybackend/queryplan/testdata/build_empty.txt b/pkg/querybackend/queryplan/testdata/build_empty.txt new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt b/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt new file mode 100644 index 0000000000..a680a080ca --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt @@ -0,0 +1,3 @@ +READ {children: 0, blocks: 2} + id:"1" + id:"2" diff --git a/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt b/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt new file mode 100644 index 0000000000..08aff02ba7 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt @@ -0,0 +1,35 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 1} + id:"2" + READ {children: 0, blocks: 1} + id:"3" + READ {children: 0, blocks: 1} + id:"4" + READ {children: 0, blocks: 1} + id:"5" + READ {children: 0, blocks: 1} + id:"6" + READ {children: 0, blocks: 1} + id:"7" + READ {children: 0, blocks: 1} + id:"8" + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"9" + READ {children: 0, blocks: 1} + id:"10" + READ {children: 0, blocks: 1} + id:"11" + READ {children: 0, blocks: 1} + id:"12" + READ {children: 0, blocks: 1} + id:"13" + READ {children: 0, blocks: 1} + id:"14" + READ {children: 0, blocks: 1} + id:"15" + READ {children: 0, blocks: 1} + id:"16" diff --git a/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt b/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt new file mode 100644 index 0000000000..bdb06b694c --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt @@ -0,0 +1,10 @@ +MERGE {children: 3, blocks: 0} + READ {children: 0, blocks: 2} + id:"1" + id:"2" + READ {children: 0, blocks: 2} + id:"3" + id:"4" + READ {children: 0, blocks: 2} + id:"5" + id:"6" diff --git a/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt new file mode 100644 index 0000000000..e8316ea5e7 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt @@ -0,0 +1,45 @@ +MERGE {children: 2, blocks: 0} + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 4} + id:"1" + id:"2" + id:"3" + id:"4" + READ {children: 0, blocks: 4} + id:"5" + id:"6" + id:"7" + id:"8" + READ {children: 0, blocks: 4} + id:"9" + id:"10" + id:"11" + id:"12" + READ {children: 0, blocks: 4} + id:"13" + id:"14" + id:"15" + id:"16" + READ {children: 0, blocks: 4} + id:"17" + id:"18" + id:"19" + id:"20" + READ {children: 0, blocks: 4} + id:"21" + id:"22" + id:"23" + id:"24" + READ {children: 0, blocks: 4} + id:"25" + id:"26" + id:"27" + id:"28" + READ {children: 0, blocks: 4} + id:"29" + id:"30" + id:"31" + id:"32" + MERGE {children: 1, blocks: 0} + READ {children: 0, blocks: 1} + id:"33" diff --git a/pkg/querybackend/queryplan/testdata/build_invalid_max_merges.txt b/pkg/querybackend/queryplan/testdata/build_invalid_max_merges.txt new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pkg/querybackend/queryplan/testdata/build_invalid_max_reads.txt b/pkg/querybackend/queryplan/testdata/build_invalid_max_reads.txt new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pkg/querybackend/queryplan/testdata/just_over_depth_2.txt b/pkg/querybackend/queryplan/testdata/build_just_over_depth_2.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/just_over_depth_2.txt rename to pkg/querybackend/queryplan/testdata/build_just_over_depth_2.txt diff --git a/pkg/querybackend/queryplan/testdata/build_single_block.txt b/pkg/querybackend/queryplan/testdata/build_single_block.txt new file mode 100644 index 0000000000..e0e4cf04c2 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_single_block.txt @@ -0,0 +1,2 @@ +READ {children: 0, blocks: 1} + id:"1" diff --git a/pkg/querybackend/queryplan/testdata/build_three_way_split.txt b/pkg/querybackend/queryplan/testdata/build_three_way_split.txt new file mode 100644 index 0000000000..780dc17603 --- /dev/null +++ b/pkg/querybackend/queryplan/testdata/build_three_way_split.txt @@ -0,0 +1,44 @@ +MERGE {children: 3, blocks: 0} + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"1" + READ {children: 0, blocks: 1} + id:"2" + READ {children: 0, blocks: 1} + id:"3" + READ {children: 0, blocks: 1} + id:"4" + READ {children: 0, blocks: 1} + id:"5" + READ {children: 0, blocks: 1} + id:"6" + READ {children: 0, blocks: 1} + id:"7" + READ {children: 0, blocks: 1} + id:"8" + MERGE {children: 8, blocks: 0} + READ {children: 0, blocks: 1} + id:"9" + READ {children: 0, blocks: 1} + id:"10" + READ {children: 0, blocks: 1} + id:"11" + READ {children: 0, blocks: 1} + id:"12" + READ {children: 0, blocks: 1} + id:"13" + READ {children: 0, blocks: 1} + id:"14" + READ {children: 0, blocks: 1} + id:"15" + READ {children: 0, blocks: 1} + id:"16" + MERGE {children: 4, blocks: 0} + READ {children: 0, blocks: 1} + id:"17" + READ {children: 0, blocks: 1} + id:"18" + READ {children: 0, blocks: 1} + id:"19" + READ {children: 0, blocks: 1} + id:"20" diff --git a/pkg/querybackend/queryplan/testdata/twenty_five_blocks.txt b/pkg/querybackend/queryplan/testdata/build_twenty_five_blocks.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/twenty_five_blocks.txt rename to pkg/querybackend/queryplan/testdata/build_twenty_five_blocks.txt diff --git a/pkg/querybackend/queryplan/testdata/two_leaves.txt b/pkg/querybackend/queryplan/testdata/build_two_leaves.txt similarity index 100% rename from pkg/querybackend/queryplan/testdata/two_leaves.txt rename to pkg/querybackend/queryplan/testdata/build_two_leaves.txt