From 82a798b5a1ce60f6f608aef5d690c5b6841e1867 Mon Sep 17 00:00:00 2001 From: Minhan Cao Date: Thu, 30 Apr 2026 15:51:41 -0700 Subject: [PATCH] feat: Add ability to have the formatted text plan be populated and inserted for presto_query_plans MySQL table from the loadjson command --- cmd/loadjson/main.go | 14 +++++- go.mod | 2 +- go.sum | 4 +- prestoapi/plan_formatter.go | 85 ++++++++++++++++++++++++++++++++ prestoapi/plan_formatter_test.go | 52 +++++++++++++++++++ 5 files changed, 153 insertions(+), 4 deletions(-) create mode 100644 prestoapi/plan_formatter.go create mode 100644 prestoapi/plan_formatter_test.go diff --git a/cmd/loadjson/main.go b/cmd/loadjson/main.go index f613ba0a..f52b570c 100644 --- a/cmd/loadjson/main.go +++ b/cmd/loadjson/main.go @@ -9,6 +9,7 @@ import ( "os/signal" "path/filepath" "pbench/log" + "pbench/prestoapi" "pbench/stage" "pbench/utils" "reflect" @@ -17,7 +18,6 @@ import ( "time" "github.com/prestodb/presto-go-client/v2/queryjson" - "github.com/spf13/cobra" ) @@ -205,6 +205,18 @@ func processFile(ctx context.Context, path string) { log.Error().Err(err).Str("path", path).Msg("failed to pre-process query info JSON") return } + + // Generate the text plan from the JSON plan + if queryInfo.AssembledQueryPlanJson != "" { + textPlan, err := prestoapi.FormatQueryPlanAsText(queryInfo.AssembledQueryPlanJson) + if err != nil { + // Log the error but don't fail the entire operation + // The json_plan will still be available + log.Error().Err(err).Str("path", path).Msg("failed to generate text plan") + } else { + queryInfo.TextPlan = textPlan + } + } } if ExtractPlanJson { planJsonFilePath := filepath.Join(OutputPath, fileName[0:len(fileName)-len(filepath.Ext(fileName))]+".plan.json") diff --git a/go.mod b/go.mod index 6f738d66..4074b8f4 100644 --- a/go.mod +++ b/go.mod @@ -8,7 +8,7 @@ require ( github.com/go-sql-driver/mysql v1.9.3 github.com/influxdata/influxdb-client-go/v2 v2.14.0 github.com/pkg/errors v0.9.1 - github.com/prestodb/presto-go-client/v2 v2.1.1 + github.com/prestodb/presto-go-client/v2 v2.1.2 github.com/rs/zerolog v1.34.0 github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 diff --git a/go.sum b/go.sum index bf17eeee..bba98cce 100644 --- a/go.sum +++ b/go.sum @@ -61,8 +61,8 @@ github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prestodb/presto-go-client/v2 v2.1.1 h1:3mIaMitux4x/kV5CCpT0XQEcyPuPxi7eIfdXtB9VGJg= -github.com/prestodb/presto-go-client/v2 v2.1.1/go.mod h1:xXanlGM7ptbiRDX9rn4G7GXlITtR+pWX5MoR5G7t8cI= +github.com/prestodb/presto-go-client/v2 v2.1.2 h1:UQK1ipQAHoQl2uShbBjhTXjo5HCkoMHOY8E2dIzlNy0= +github.com/prestodb/presto-go-client/v2 v2.1.2/go.mod h1:xXanlGM7ptbiRDX9rn4G7GXlITtR+pWX5MoR5G7t8cI= github.com/rs/xid v1.6.0/go.mod h1:7XoLgs4eV+QndskICGsho+ADou8ySMSjJKDIan90Nz0= github.com/rs/zerolog v1.34.0 h1:k43nTLIwcTVQAncfCw4KZ2VY6ukYoZaBPNOE8txlOeY= github.com/rs/zerolog v1.34.0/go.mod h1:bJsvje4Z08ROH4Nhs5iH600c3IkWhwp44iRc54W6wYQ= diff --git a/prestoapi/plan_formatter.go b/prestoapi/plan_formatter.go new file mode 100644 index 00000000..a5bcb69e --- /dev/null +++ b/prestoapi/plan_formatter.go @@ -0,0 +1,85 @@ +package prestoapi + +import ( + "encoding/json" + "fmt" + "strings" +) + +// PlanNode represents a node in the query execution plan +type PlanNode struct { + ID string `json:"id"` + Name string `json:"name"` + Identifier string `json:"identifier"` + Details string `json:"details"` + Children []PlanNode `json:"children"` + RemoteSources []interface{} `json:"remoteSources"` + Estimates []map[string]interface{} `json:"estimates"` +} + +// StagePlanWrapper wraps the plan for a stage +type StagePlanWrapper struct { + Plan PlanNode `json:"plan"` +} + +// FormatQueryPlanAsText converts a JSON query plan to human-readable text format +// The input should be the AssembledQueryPlanJson which is a map of stage IDs to plan wrappers +func FormatQueryPlanAsText(jsonPlan string) (string, error) { + if jsonPlan == "" { + return "", nil + } + + // Parse the JSON plan - it's a map of stage IDs to plan wrappers + var stages map[string]StagePlanWrapper + if err := json.Unmarshal([]byte(jsonPlan), &stages); err != nil { + return "", fmt.Errorf("failed to parse query plan JSON: %w", err) + } + + var result strings.Builder + + // Process each stage (fragment) in order + // Note: stages are typically numbered 0, 1, 2, etc. + for stageID, stageWrapper := range stages { + result.WriteString(fmt.Sprintf("Fragment %s\n", stageID)) + result.WriteString(formatPlanNode(&stageWrapper.Plan, 0)) + result.WriteString("\n") + } + + return result.String(), nil +} + +// formatPlanNode recursively formats a plan node and its children +func formatPlanNode(node *PlanNode, depth int) string { + if node == nil { + return "" + } + + var result strings.Builder + indent := strings.Repeat(" ", depth) + + // Format the node header with identifier if present + if node.Identifier != "" { + result.WriteString(fmt.Sprintf("%s- %s[PlanNodeId %s]%s\n", + indent, node.Name, node.ID, node.Identifier)) + } else { + result.WriteString(fmt.Sprintf("%s- %s[PlanNodeId %s]\n", + indent, node.Name, node.ID)) + } + + // Add details if present (already formatted with proper indentation in the JSON) + if node.Details != "" { + detailLines := strings.Split(strings.TrimSpace(node.Details), "\n") + for _, line := range detailLines { + if line != "" { + result.WriteString(fmt.Sprintf("%s %s\n", indent, line)) + } + } + } + + // Recursively format children with increased indentation + for i := range node.Children { + result.WriteString(formatPlanNode(&node.Children[i], depth+1)) + } + + return result.String() +} diff --git a/prestoapi/plan_formatter_test.go b/prestoapi/plan_formatter_test.go new file mode 100644 index 00000000..ff1deaed --- /dev/null +++ b/prestoapi/plan_formatter_test.go @@ -0,0 +1,52 @@ +package prestoapi + +import ( + "strings" + "testing" +) + +func TestFormatQueryPlanAsText(t *testing.T) { + // Sample JSON plan from the user + jsonPlan := `{"0":{"plan":{"id":"22","name":"Output","identifier":"[n_name, revenue]","details":"revenue := sum (4:5)\n","children":[{"id":"1311","name":"TopN","identifier":"[1 by (sum DESC_NULLS_LAST)]","details":"","children":[{"id":"1310","name":"TopNPartial","identifier":"[1 by (sum DESC_NULLS_LAST)]","details":"","children":[{"id":"14","name":"Aggregate(FINAL)[n_name]","identifier":"","details":"sum := \"presto.default.sum\"((sum_13)) (4:5)\n","children":[{"id":"1705","name":"LocalExchange","identifier":"[SINGLE] ()","details":"","children":[{"id":"1703","name":"Aggregate(PARTIAL)[n_name]","identifier":"","details":"sum_13 := \"presto.default.sum\"((expr)) (4:5)\n","children":[{"id":"340","name":"Project","identifier":"[projectLocality = LOCAL]","details":"expr := (l_extendedprice) * ((DOUBLE'1.0') - (l_discount)) (8:6)\n","children":[{"id":"1410","name":"InnerJoin","identifier":"[(\"l_suppkey\" = \"s_suppkey\") AND (\"c_nationkey\" = \"s_nationkey\")]","details":"Distribution: REPLICATED\n","children":[{"id":"1409","name":"InnerJoin","identifier":"[(\"l_orderkey\" = \"o_orderkey\")]","details":"Distribution: REPLICATED\n","children":[{"id":"3","name":"TableScan","identifier":"[TableHandle {connectorId='tpch', connectorHandle='lineitem:sf1.0', layout='Optional[lineitem:sf1.0]'}]","details":"l_orderkey := tpch:l_orderkey (8:5)\nl_extendedprice := tpch:l_extendedprice (8:5)\nl_suppkey := tpch:l_suppkey (8:5)\nl_discount := tpch:l_discount (8:5)\n","children":[],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}],"remoteSources":[],"estimates":[]}}}` + + result, err := FormatQueryPlanAsText(jsonPlan) + if err != nil { + t.Fatalf("FormatQueryPlanAsText failed: %v", err) + } + + // Verify the output contains expected elements + expectedStrings := []string{ + "Fragment 0", + "- Output[PlanNodeId 22][n_name, revenue]", + "revenue := sum (4:5)", + "- TopN[PlanNodeId 1311][1 by (sum DESC_NULLS_LAST)]", + "- TableScan[PlanNodeId 3]", + "l_orderkey := tpch:l_orderkey", + } + + for _, expected := range expectedStrings { + if !strings.Contains(result, expected) { + t.Errorf("Expected output to contain %q, but it didn't.\nGot:\n%s", expected, result) + } + } + + // Print the result for manual inspection + t.Logf("Formatted plan:\n%s", result) +} + +func TestFormatQueryPlanAsText_EmptyInput(t *testing.T) { + result, err := FormatQueryPlanAsText("") + if err != nil { + t.Fatalf("FormatQueryPlanAsText with empty input failed: %v", err) + } + if result != "" { + t.Errorf("Expected empty result for empty input, got: %q", result) + } +} + +func TestFormatQueryPlanAsText_InvalidJSON(t *testing.T) { + _, err := FormatQueryPlanAsText("invalid json") + if err == nil { + t.Error("Expected error for invalid JSON, got nil") + } +}