Improved organization of pbench output directory to navigate easier for streams runs

Author: minhancao

stage/run_recorder.go

@@ -42,5 +42,5 @@ func (f *FileBasedRunRecorder) RecordQuery(_ context.Context, _ *Stage, result *
 }
 
 func (f *FileBasedRunRecorder) RecordRun(_ context.Context, s *Stage, _ []*QueryResult) {
-	_ = os.WriteFile(filepath.Join(s.States.OutputPath, s.Id+"_summary.csv"), []byte(f.summaryBuilder.String()), 0644)
+	_ = os.WriteFile(filepath.Join(s.States.OutputPath, s.States.RunName+"_summary.csv"), []byte(f.summaryBuilder.String()), 0644)
 }

stage/stage.go

@@ -29,6 +29,8 @@ import (
 type Stage struct {
 	// Id is used to uniquely identify a stage. It is usually the file name without its directory path and extension.
 	Id string `json:"-"`
+	// ParentStageId is used to keep track of the parent stage of this child stage.
+	ParentStageId string `json:"-"`
 	// The values in Catalog, Schema, and SessionParams are inherited by the descendant stages. Please note that if
 	// they have new values assigned in a stage, those values are NOT applied tn the Presto client until a stage
 	// creates its own client by setting StartOnNewClient = true.
@@ -233,6 +235,7 @@ func (s *Stage) run(ctx context.Context) (returnErr error) {
 	s.setDefaults()
 	s.prepareClient()
 	s.propagateStates()
+	s.createNextStagesOutputDirectories()
 	preStageErr := s.runShellScripts(ctx, s.PreStageShellScripts)
 	if preStageErr != nil {
 		return fmt.Errorf("pre-stage script execution failed: %w", preStageErr)
@@ -524,7 +527,7 @@ func (s *Stage) runQuery(ctx context.Context, query *Query) (result *QueryResult
 	)
 	if *s.SaveOutput {
 		queryOutputFile, err = os.OpenFile(
-			filepath.Join(s.States.OutputPath, querySourceStr)+".output",
+			filepath.Join(s.States.OutputPath, s.ParentStageId, s.Id, querySourceStr)+".output",
 			utils.OpenNewFileFlags, 0644)
 		if err != nil {
 			return result, err

stage/stage_utils.go

@@ -290,7 +290,7 @@ func (s *Stage) saveQueryJsonFile(result *QueryResult) {
 		querySourceStr := s.querySourceString(result)
 		{
 			queryJsonFile, err := os.OpenFile(
-				filepath.Join(s.States.OutputPath, querySourceStr)+".json",
+				filepath.Join(s.States.OutputPath, s.ParentStageId, s.Id, querySourceStr)+".json",
 				utils.OpenNewFileFlags, 0644)
 			checkErr(err)
 			if err == nil {
@@ -302,7 +302,7 @@ func (s *Stage) saveQueryJsonFile(result *QueryResult) {
 		}
 		if result.QueryError != nil {
 			queryErrorFile, err := os.OpenFile(
-				filepath.Join(s.States.OutputPath, querySourceStr)+".error.json",
+				filepath.Join(s.States.OutputPath, s.ParentStageId, s.Id, querySourceStr)+".error.json",
 				utils.OpenNewFileFlags, 0644)
 			checkErr(err)
 			if err == nil {
@@ -332,7 +332,7 @@ func (s *Stage) saveColumnMetadataFile(qr *presto.QueryResults, result *QueryRes
 		}
 	}()
 	columnMetadataFile, ioErr := os.OpenFile(
-		filepath.Join(s.States.OutputPath, querySourceStr)+".cols.json",
+		filepath.Join(s.States.OutputPath, s.ParentStageId, s.Id, querySourceStr)+".cols.json",
 		utils.OpenNewFileFlags, 0644)
 	if ioErr != nil {
 		return ioErr
@@ -372,3 +372,12 @@ func (s *Stage) querySourceString(result *QueryResult) (sourceStr string) {
 	}
 	return
 }
+
+func (s *Stage) createNextStagesOutputDirectories() {
+	parentStageId := s.Id
+	for _, nextStage := range s.NextStages {
+		nextStage.ParentStageId = parentStageId
+		stageOutputPath := filepath.Join(s.States.OutputPath, s.Id, nextStage.Id)
+		utils.PrepareOutputDirectory(stageOutputPath)
+	}
+}