From 46f92d7823caa3d769e8ec0d777117cc38c2cb43 Mon Sep 17 00:00:00 2001
From: Five Grant <5@fivegrant.com>
Date: Thu, 15 Jun 2023 13:26:29 -0500
Subject: [PATCH] Add optional column mapping (#51)

---
 paths.yaml                  | 2 ++
 src/service/ArgIO.jl        | 2 +-
 src/service/AssetManager.jl | 8 +++++---
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/paths.yaml b/paths.yaml
index f1c509a..f7ab5e2 100644
--- a/paths.yaml
+++ b/paths.yaml
@@ -75,6 +75,8 @@
                   filename:
                     type: string
                     example: "dataset.csv"
+                  mappings:
+                    type: object
               timespan:
                 type: object
                 properties:
diff --git a/src/service/ArgIO.jl b/src/service/ArgIO.jl
index 6fc8f32..df45f7f 100644
--- a/src/service/ArgIO.jl
+++ b/src/service/ArgIO.jl
@@ -31,7 +31,7 @@ function prepare_input(args; context...)
         args[:model] = fetch_model(args[:model_config_id])
     end
     if in(:dataset, keys(args)) && !isa(args[:dataset], String)
-        args[:dataset] = fetch_dataset(args[:dataset]["id"], args[:dataset]["filename"])
+        args[:dataset] = fetch_dataset(args[:dataset]["id"], args[:dataset]["filename"], get(args[:dataset], "mappings", Dict()))
     end
     if in(:model_config_ids, keys(args))
         args[:models] = fetch_model.(map(string, args[:model_ids]))
diff --git a/src/service/AssetManager.jl b/src/service/AssetManager.jl
index 29bc34f..2d9bdcf 100644
--- a/src/service/AssetManager.jl
+++ b/src/service/AssetManager.jl
@@ -3,7 +3,7 @@ Asset fetching from TDS
 """
 module AssetManager
 
-import DataFrames: DataFrame
+import DataFrames: rename!, DataFrame
 import CSV, Downloads, HTTP
 import OpenAPI.Clients: Client
 import JSON3 as JSON
@@ -31,7 +31,7 @@ end
 """
 Return csv from TDS by ID
 """
-function fetch_dataset(dataset_id::String, filename::String)
+function fetch_dataset(dataset_id::String, filename::String, mappings::Dict=Dict())
     # TODO(five): Select name dynamicially
     url = "$(settings["TDS_URL"])/datasets/$dataset_id/download-url?filename=$filename"
     response = HTTP.get(url, ["Content-Type" => "application/json"])
@@ -39,7 +39,9 @@ function fetch_dataset(dataset_id::String, filename::String)
     io = IOBuffer()
     Downloads.download(body.url, io)
     seekstart(io)
-    CSV.read(io, DataFrame)
+    dataframe = CSV.read(io, DataFrame)
+    for (from, to) in mappings rename!(dataframe, Symbol(from)=>Symbol(to)) end
+    dataframe
 end
 
 """