@@ -33,34 +33,28 @@ type Downloader struct {
3333 mode core.IOStrategy_DownloadMode
3434}
3535
36- // By default, blobs (FlyteFiles) were not and still are not written with a file extension.
37- // For example, a data: FlyteFile["csv"] blob should be written to "data", even though
38- // Format="csv".
36+ // By default, blobs (FlyteFiles) were not and still are not written with a
37+ // file extension. For example, a data: FlyteFile["csv"] blob is written
38+ // to "inputs/data", even though Format="csv".
3939//
4040// When FileExtension="" (the default), this old behavior is preserved.
4141//
42- // However, a data: Annotated[FlyteFile["csv"], FileDownloadConfig(file_extension="csv")]
43- // blob should be written to "data.csv" - both Format="csv" and FileExtension="csv" (new behavior).
44- //
45- // So when e.g. FileExtension="csv", the file is written to "data.csv".
46- // Also, when e.g. FileExtension="csv" and EnableLegacyFilename=true, the file is written to
47- // "data" and "data.csv" (partially new behavior, bridges the gap of backward compatibility).
48- func resolveVarFilenames (vars * core.VariableMap ) map [string ][]string {
49- varFilenames := make (map [string ][]string , len (vars .GetVariables ()))
42+ // However, an input blob
43+ // `data: Annotated[FlyteFile["csv"], FileDownloadConfig(file_extension="csv")]`
44+ // should be written to "inputs/data.csv" (when FileExtension="csv" - new behavior).
45+ func resolveVarFilenames (vars * core.VariableMap ) map [string ]string {
46+ varFilenames := make (map [string ]string , len (vars .GetVariables ()))
5047 for varName , variable := range vars .GetVariables () {
5148 varType := variable .GetType ()
5249 switch varType .GetType ().(type ) {
5350 case * core.LiteralType_Blob :
5451 if varType .GetBlob ().GetFileExtension () == "" {
55- varFilenames [varName ] = append ( varFilenames [ varName ], varName )
52+ varFilenames [varName ] = varName
5653 } else {
57- varFilenames [varName ] = append (varFilenames [varName ], varName + "." + varType .GetBlob ().GetFileExtension ())
58- if varType .GetBlob ().GetEnableLegacyFilename () {
59- varFilenames [varName ] = append (varFilenames [varName ], varName )
60- }
54+ varFilenames [varName ] = varName + "." + varType .GetBlob ().GetFileExtension ()
6155 }
6256 default :
63- varFilenames [varName ] = append ( varFilenames [ varName ], varName )
57+ varFilenames [varName ] = varName
6458 }
6559 }
6660 return varFilenames
@@ -465,7 +459,7 @@ func (d Downloader) handleLiteral(ctx context.Context, lit *core.Literal, filePa
465459 if err != nil {
466460 return nil , nil , errors .Wrapf (err , "failed to create directory [%s]" , filePath )
467461 }
468- v , m , err := d .RecursiveDownload (ctx , lit .GetMap (), filePath , make (map [string ][] string ), writeToFile )
462+ v , m , err := d .RecursiveDownload (ctx , lit .GetMap (), filePath , make (map [string ]string ), writeToFile )
469463 if err != nil {
470464 return nil , nil , err
471465 }
@@ -501,7 +495,7 @@ type downloadedResult struct {
501495 v interface {}
502496}
503497
504- func (d Downloader ) RecursiveDownload (ctx context.Context , inputs * core.LiteralMap , dir string , varFilenames map [string ][] string , writePrimitiveToFile bool ) (VarMap , * core.LiteralMap , error ) {
498+ func (d Downloader ) RecursiveDownload (ctx context.Context , inputs * core.LiteralMap , dir string , varFilenames map [string ]string , writePrimitiveToFile bool ) (VarMap , * core.LiteralMap , error ) {
505499 childCtx , cancel := context .WithCancel (ctx )
506500 defer cancel ()
507501 if inputs == nil || len (inputs .GetLiterals ()) == 0 {
@@ -519,26 +513,18 @@ func (d Downloader) RecursiveDownload(ctx context.Context, inputs *core.LiteralM
519513 }
520514 logger .Infof (ctx , "read object at location [%s]" , offloadedMetadataURI )
521515 }
516+ filename := variable
517+ if varFilename , ok := varFilenames [variable ]; ok {
518+ filename = varFilename
519+ }
520+ varPath := path .Join (dir , filename )
522521 lit := literal
523522 f [variable ] = futures .NewAsyncFuture (childCtx , func (ctx2 context.Context ) (interface {}, error ) {
524- var filenames []string
525- var resultLit * core.Literal
526- var resultV interface {}
527- var err error
528- if len (varFilenames [variable ]) == 0 {
529- filenames = []string {variable }
530- } else {
531- filenames = varFilenames [variable ]
532- }
533- for _ , filename := range filenames {
534- varPath := path .Join (dir , filename )
535- // TODO: Refactor handleLiteral to accept a list of file paths and return a list of downloaded results
536- resultV , resultLit , err = d .handleLiteral (ctx2 , lit , varPath , writePrimitiveToFile )
537- if err != nil {
538- return nil , err
539- }
523+ v , lit , err := d .handleLiteral (ctx2 , lit , varPath , writePrimitiveToFile )
524+ if err != nil {
525+ return nil , err
540526 }
541- return downloadedResult {lit : resultLit , v : resultV }, nil
527+ return downloadedResult {lit : lit , v : v }, nil
542528 })
543529 }
544530
0 commit comments