Edited to have parent stage id

Author: minhancao

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.
@@ -525,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, s.Id, 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, s.Id, 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, s.Id, 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, s.Id, querySourceStr)+".cols.json",
+		filepath.Join(s.States.OutputPath, s.ParentStageId, s.Id, querySourceStr)+".cols.json",
 		utils.OpenNewFileFlags, 0644)
 	if ioErr != nil {
 		return ioErr
@@ -374,11 +374,10 @@ func (s *Stage) querySourceString(result *QueryResult) (sourceStr string) {
 }
 
 func (s *Stage) createNextStagesOutputDirectories() {
-	parentOutputPath := s.States.OutputPath
-	utils.PrepareOutputDirectory(filepath.Join(parentOutputPath, s.Id))
+	parentStageId := s.Id
 	for _, nextStage := range s.NextStages {
-		nextStage.States.OutputPath = filepath.Join(parentOutputPath, nextStage.Id)
-		utils.PrepareOutputDirectory(nextStage.States.OutputPath)
+		nextStage.ParentStageId = parentStageId
+		stageOutputPath := filepath.Join(s.States.OutputPath, s.Id, nextStage.Id)
+		utils.PrepareOutputDirectory(stageOutputPath)
 	}
-	s.States.OutputPath = parentOutputPath
 }