keep attributes for reduced grouped variables

Author: Alexander-Barth

src/groupby.jl

@@ -44,6 +44,7 @@ end
 struct ReducedGroupedVariable{T,N,TGV,TF}  <: AbstractVariable{T,N}
     gv::TGV
     reduce_fun::TF
+    _attrib::OrderedDict{String,Any}
 end
 
 
@@ -141,6 +142,11 @@ function variable(gds::ReducedGroupedDataset,varname::SymbolOrString)
     end
 end
 
+
+attribnames(gds::ReducedGroupedDataset) = attribnames(gds.ds)
+attrib(gds::ReducedGroupedDataset,attribname::SymbolOrString) =
+    attrib(gds.ds,attribname)
+
 #
 # methods with GroupedVariable as main argument
 #
@@ -396,7 +402,8 @@ function ReducedGroupedVariable(gv::GroupedVariable,reduce_fun)
     T = eltype(gv.v)
     @debug "inference " T reduce_fun  Base.return_types(reduce_fun, (Vector{T},))
     N = ndims(gv.v)
-    ReducedGroupedVariable{T,N,typeof(gv),typeof(reduce_fun)}(gv,reduce_fun)
+    _attrib = OrderedDict(gv.v.attrib)
+    ReducedGroupedVariable{T,N,typeof(gv),typeof(reduce_fun)}(gv,reduce_fun,_attrib)
 end
 
 function ReducedGroupedDataset(gds::GroupedDataset,reduce_fun)
@@ -440,6 +447,12 @@ end
 dimnames(gr::ReducedGroupedVariable) = dimnames(gr.gv.v)
 name(gr::ReducedGroupedVariable) = name(gr.gv.v)
 
+
+attribnames(gr::ReducedGroupedVariable) = collect(keys(gr._attrib))
+attrib(gr::ReducedGroupedVariable,attribname::SymbolOrString) = gr._attrib[attribname]
+defAttrib(gr::ReducedGroupedVariable,attribname::SymbolOrString,value) =
+    gr._attrib[attribname] = value
+
 struct ReducedGroupedVariableStyle <: BroadcastStyle end
 Base.BroadcastStyle(::Type{<:ReducedGroupedVariable}) = ReducedGroupedVariableStyle()
 

test/test_groupby.jl

@@ -6,14 +6,18 @@ using Statistics
 using CommonDataModel
 using CommonDataModel:
     @groupby,
-    name,
     GroupedVariable,
+    MemoryDataset,
     ReducedGroupedVariable,
     _array_selectdim_indices,
     _dest_indices,
     _dim_after_getindex,
     _indices,
-    groupby
+    dataset,
+    defVar,
+    groupby,
+    name,
+    variable
 
 
 #include("memory_dataset.jl");
@@ -53,13 +57,14 @@ data3 .= data
 data3[1,1,:] .= missing
 data3[1,2,1] = missing
 
-TDS(fname,"c") do ds
+TDS(fname,"c",attrib = ["title" => "test"]) do ds
     defVar(ds,"lon",1:size(data,1),("lon",))
     defVar(ds,"lat",1:size(data,2),("lat",))
     defVar(ds,"time",time,("time",))
-    defVar(ds,"data",data,("lon","lat","time"))
+    defVar(ds,"data",data,("lon","lat","time"),attrib = ["foo" => "bar"])
     defVar(ds,"data2",data .+ 1,("lon","lat","time"))
     defVar(ds,"data3",data3,("lon","lat","time"))
+    defVar(ds,"data4",data,("lon","lat","time"),attrib = ["scale_factor" => 2])
 end
 
 ds = TDS(fname)
@@ -108,6 +113,9 @@ gd = groupby(ds[:data],:time => Dates.Month);
 d_sum = cat([sum(ds[varname][:,:,:][:,:,findall(Dates.month.(ds[coordname][:]) .== m)],dims=3)
                  for m in 1:12]...,dims=3)
 
+d_mean = cat([mean(ds[varname][:,:,:][:,:,findall(Dates.month.(ds[coordname][:]) .== m)],dims=3)
+                 for m in 1:12]...,dims=3)
+
 
 gd = groupby(ds["data"],"time" => Dates.Month)
 month_sum = sum(gd);
@@ -118,11 +126,25 @@ gd = groupby(ds[:data],:time => Dates.Month)
 month_sum = sum(gd);
 @test month_sum[:,:,:] == d_sum
 
-# group dataset
-gds = sum(groupby(ds,:time => Dates.Month))
-@test gds["data"][:,:,:] == d_sum
+# group dataset function
+gds = mean(groupby(ds,:time => Dates.Month))
+@test gds["data"][:,:,:] == d_mean
 @test gds["lon"][:] == ds["lon"][:]
 @test gds["lat"][:] == ds["lat"][:]
+@test gds["data4"][:,:,:] == d_mean
+@test gds.attrib["title"] == "test"
+@test gds["data"].attrib["foo"] == "bar"
+@test collect(keys(gds.attrib)) == ["title"]
+@test collect(keys(gds["data"].attrib)) == ["foo"]
+
+# group dataset macro
+gds = mean(@groupby(ds,Dates.Month(time)))
+@test gds["data"][:,:,:] == d_mean
+@test gds["lon"][:] == ds["lon"][:]
+@test gds["lat"][:] == ds["lat"][:]
+
+gr = mean(groupby(ds["data4"],:time => Dates.Month))
+@test gr[:,:,:] == d_mean
 
 
 gr = month_sum