do not force unstack to create columns with missing

Author: bkamins

src/abstractdataframe/reshape.jl

@@ -198,9 +198,9 @@ wide3 = unstack(long, [:id, :a], :variable, :value)
 Note that there are some differences between the widened results above.
 """
 function unstack(df::AbstractDataFrame, rowkey::Int, colkey::Int, value::Int)
-    refkeycol = CategoricalArray{Union{eltype(df[rowkey]), Missing}}(df[rowkey])
+    refkeycol = deepcopy(categorical(df[rowkey]))
     droplevels!(refkeycol)
-    keycol = CategoricalArray{Union{eltype(df[colkey]), Missing}}(df[colkey])
+    keycol = deepcopy(categorical(df[colkey]))
     droplevels!(keycol)
     valuecol = df[value]
     _unstack(df, rowkey, colkey, value, keycol, valuecol, refkeycol)
@@ -250,7 +250,7 @@ function _unstack(df::AbstractDataFrame, rowkey::Int,
     end
     levs = levels(refkeycol)
     # we have to handle a case with missings in refkeycol as levs will skip missing
-    col = similar_missing(df[rowkey], length(levs) + hadmissing)
+    col = similar(df[rowkey], length(levs) + hadmissing)
     copy!(col, levs)
     hadmissing && (col[end] = missing)
     df2 = DataFrame(unstacked_val, map(Symbol, levels(keycol)))
@@ -279,7 +279,7 @@ function unstack(df::AbstractDataFrame, rowkeys::AbstractVector{Symbol}, colkey:
     length(rowkeys) == 0 && throw(ArgumentError("No key column found"))
     length(rowkeys) == 1 && return unstack(df, rowkeys[1], colkey, value)
     g = groupby(df, rowkeys, sort=true)
-    keycol = CategoricalArray{Union{eltype(df[colkey]), Missing}}(df[colkey])
+    keycol = deepcopy(categorical(df[colkey]))
     droplevels!(keycol)
     valuecol = df[value]
     _unstack(df, rowkeys, colkey, value, keycol, valuecol, g)
@@ -292,7 +292,7 @@ function _unstack(df::AbstractDataFrame, rowkeys::AbstractVector{Symbol},
     for i in 1:length(groupidxs)
         rowkey[groupidxs[i]] = i
     end
-    df1 = allowmissing!(df[g.idx[g.starts], g.cols], g.cols)
+    df1 = df[g.idx[g.starts], g.cols]
     Nrow = length(g)
     Ncol = length(levels(keycol))
     unstacked_val = [similar_missing(valuecol, Nrow) for i in 1:Ncol]

test/dataframe.jl

@@ -296,14 +296,15 @@ module TestDataFrame
 
     #Check the output of unstack
     df = DataFrame(Fish = CategoricalArray{Union{String, Missing}}(["Bob", "Bob", "Batman", "Batman"]),
-                   Key = CategoricalArray{Union{String, Missing}}["Mass", "Color", "Mass", "Color"],
+                   Key = CategoricalArray{Union{String, Missing}}(["Mass", "Color", "Mass", "Color"]),
                    Value = Union{String, Missing}["12 g", "Red", "18 g", "Grey"])
     # Check that reordering levels does not confuse unstack
     levels!(df[1], ["XXX", "Bob", "Batman"])
     levels!(df[2], ["YYY", "Color", "Mass"])
     #Unstack specifying a row column
     df2 = unstack(df, :Fish, :Key, :Value)
     @test levels(df[1]) == ["XXX", "Bob", "Batman"] # make sure we did not mess df[1] levels
+    @test levels(df[2]) == ["YYY", "Color", "Mass"] # make sure we did not mess df[2] levels
     #Unstack without specifying a row column
     df3 = unstack(df, :Key, :Value)
     #The expected output, XXX level should be dropped as it has no rows with this key
@@ -321,6 +322,31 @@ module TestDataFrame
     df4[1,:Mass] = missing
     @test df2 ≅ df4
 
+    #The same as above but without CategoricalArray
+    df = DataFrame(Fish = ["Bob", "Bob", "Batman", "Batman"],
+                   Key = ["Mass", "Color", "Mass", "Color"],
+                   Value = ["12 g", "Red", "18 g", "Grey"])
+    #Unstack specifying a row column
+    df2 = unstack(df, :Fish, :Key, :Value)
+    #Unstack without specifying a row column
+    df3 = unstack(df, :Key, :Value)
+    #The expected output, XXX level should be dropped as it has no rows with this key
+    df4 = DataFrame(Fish = ["Batman", "Bob"],
+                    Color = ["Grey", "Red"],
+                    Mass = ["18 g", "12 g"])
+    @test df2 ≅ df4
+    @test typeof(df2[:Fish]) <: Vector{String}
+    # first column stays as CategoricalArray in df3
+    @test df3 == df4
+    #Make sure unstack works with missing values at the start of the value column
+    allowmissing!(df, :Value)
+    df[1,:Value] = missing
+    df2 = unstack(df, :Fish, :Key, :Value)
+    #This changes the expected result
+    allowmissing!(df4, :Mass)
+    df4[2,:Mass] = missing
+    @test df2 ≅ df4
+
     # test empty set of grouping variables
     @test_throws ArgumentError unstack(df, Int[], :Key, :Value)
     @test_throws ArgumentError unstack(df, Symbol[], :Key, :Value)