Flexible vcat

src/abstractdataframe/abstractdataframe.jl

@@ -1017,20 +1017,41 @@ Base.hcat(df1::AbstractDataFrame, df2::AbstractDataFrame, dfn::AbstractDataFrame
           makeunique=makeunique, copycols=copycols)
 
 """
-    vcat(dfs::AbstractDataFrame...)
+    vcat(dfs::AbstractDataFrame...; cols::Union{Symbol, AbstractVector{Symbol}}=:equal)
 
-Vertically concatenate `AbstractDataFrames`.
+Vertically concatenate `AbstractDataFrame`s.
 
-Column names in all passed data frames must be the same, but they can have
-different order. In such cases the order of names in the first passed
-`DataFrame` is used.
+The `cols` keyword argument determines the columns of the returned data frame:
+
+* `:equal` (the default): require all data frames to have the same column names.
+  If they appear in different orders, the order of the first provided data frame is used.
+* `:intersect`: only the columns present in *all* provided data frames are kept.
+  If the intersection is empty, an empty data frame is returned.
+* `:union`: columns present in *at least one* of the provided data frames are kept.
+  Columns not present in some data frames are filled with `missing` where necessary.
+* A vector of `Symbol`s: only listed columns are kept.
+  Columns not present in some data frames are filled with `missing` where necessary.
+
+The order of columns is determined by the order they appear in the included
+data frames, searching through the header of the first data frame, then the 
+second, etc. 
+
+The element types of columns are determined using `promote_type`,
+as with `vcat` for `AbstractVector`s.
+
+`vcat` ignores empty data frames, making it possible to initialize an empty 
+data frame at the beginning of a loop and `vcat` onto it. 
 
 # Example
 ```jldoctest
 julia> df1 = DataFrame(A=1:3, B=1:3);
 
 julia> df2 = DataFrame(A=4:6, B=4:6);
 
+julia> df3 = DataFrame(A=7:9, C=7:9);
+
+julia> d4 = DataFrame();
+
 julia> vcat(df1, df2)
 6×2 DataFrame
 │ Row │ A     │ B     │
@@ -1042,47 +1063,105 @@ julia> vcat(df1, df2)
 │ 4   │ 4     │ 4     │
 │ 5   │ 5     │ 5     │
 │ 6   │ 6     │ 6     │
+
+julia> vcat(df1, df3, cols=:union)
+6×3 DataFrame
+│ Row │ A     │ B       │ C       │
+│     │ Int64 │ Int64⍰  │ Int64⍰  │
+├─────┼───────┼─────────┼─────────┤
+│ 1   │ 1     │ 1       │ missing │
+│ 2   │ 2     │ 2       │ missing │
+│ 3   │ 3     │ 3       │ missing │
+│ 4   │ 7     │ missing │ 7       │
+│ 5   │ 8     │ missing │ 8       │
+│ 6   │ 9     │ missing │ 9       │
+
+julia> vcat(df1, df3, cols=:intersect)
+6×1 DataFrame
+│ Row │ A     │
+│     │ Int64 │
+├─────┼───────┤
+│ 1   │ 1     │
+│ 2   │ 2     │
+│ 3   │ 3     │
+│ 4   │ 7     │
+│ 5   │ 8     │
+│ 6   │ 9     │
+
+julia> vcat(d4, df1)
+3×2 DataFrame
+│ Row │ A     │ B     │
+│     │ Int64 │ Int64 │
+├─────┼───────┼───────┤
+│ 1   │ 1     │ 1     │
+│ 2   │ 2     │ 2     │
+│ 3   │ 3     │ 3     │
 ```
+
 """
-Base.vcat(df::AbstractDataFrame) = DataFrame(df)
-Base.vcat(dfs::AbstractDataFrame...) = _vcat(collect(dfs))
-function _vcat(dfs::AbstractVector{<:AbstractDataFrame})
+Base.vcat(dfs::AbstractDataFrame...;
+          cols::Union{Symbol, AbstractVector{Symbol}}=:equal) =
+    _vcat([df for df in collect(dfs) if ncol(df) != 0]; cols=cols)
+
+function _vcat(dfs::AbstractVector{<:AbstractDataFrame}; 
+               cols::Union{Symbol, AbstractVector{Symbol}}=:equal)
+
     isempty(dfs) && return DataFrame()
+    # Array of all headers
     allheaders = map(names, dfs)
+    # Array of unique headers across all data frames
     uniqueheaders = unique(allheaders)
+    # All symbols present across all headers
     unionunique = union(uniqueheaders...)
+    # List of symbols present in all dataframes
     intersectunique = intersect(uniqueheaders...)
-    coldiff = setdiff(unionunique, intersectunique)
-
-    if !isempty(coldiff)
-        # if any DataFrames are a full superset of names, skip them
-        filter!(u -> Set(u) != Set(unionunique), uniqueheaders)
-        estrings = Vector{String}(undef, length(uniqueheaders))
-        for (i, u) in enumerate(uniqueheaders)
-            matching = findall(h -> u == h, allheaders)
-            headerdiff = setdiff(coldiff, u)
-            cols = join(headerdiff, ", ", " and ")
-            args = join(matching, ", ", " and ")
-            estrings[i] = "column(s) $cols are missing from argument(s) $args"
-        end
+
+    if cols === :equal
+        header = unionunique
+        coldiff = setdiff(unionunique, intersectunique)
+
+        if !isempty(coldiff) 
+            # if any DataFrames are a full superset of names, skip them
+            filter!(u -> !issetequal(u, header), uniqueheaders)
+            estrings = map(enumerate(uniqueheaders)) do (i, head)
+                matching = findall(h -> head == h, allheaders)
+                headerdiff = setdiff(coldiff, head)
+                cols = join(headerdiff, ", ", " and ")
+                args = join(matching, ", ", " and ")
+                return "column(s) $cols are missing from argument(s) $args"
+            end
         throw(ArgumentError(join(estrings, ", ", ", and ")))
+        end
+
+    elseif cols === :intersect 
+        header = intersectunique
+    elseif cols === :union
+        header = unionunique
+    else
+        header = cols 
     end
 
-    header = allheaders[1]
     length(header) == 0 && return DataFrame()
-    cols = Vector{AbstractVector}(undef, length(header))
+    all_cols = Vector{AbstractVector}(undef, length(header))
     for (i, name) in enumerate(header)
-        data = [df[name] for df in dfs]
-        lens = map(length, data)
-        T = mapreduce(eltype, promote_type, data)
-        cols[i] = Tables.allocatecolumn(T, sum(lens))
+        newcols = map(dfs) do df
+            if haskey(df, name)
+                return df[name]
+            else
+                Iterators.repeated(missing, nrow(df))
+            end
+        end
+
+        lens = map(length, newcols)
+        T = mapreduce(eltype, promote_type, newcols)
+        all_cols[i] = Tables.allocatecolumn(T, sum(lens))
         offset = 1
-        for j in 1:length(data)
-            copyto!(cols[i], offset, data[j])
+        for j in 1:length(newcols)
+            copyto!(all_cols[i], offset, newcols[j])
             offset += lens[j]
         end
     end
-    return DataFrame(cols, header, copycols=false)
+    return DataFrame(all_cols, header, copycols=false)
 end
 
 function Base.reduce(::typeof(vcat), dfs::AbstractVector{<:AbstractDataFrame})

test/cat.jl

@@ -269,16 +269,19 @@ end
 
     @test vcat(missing_df) == DataFrame()
     @test vcat(missing_df, missing_df) == DataFrame()
-    @test_throws ArgumentError vcat(missing_df, df)
-    @test_throws ArgumentError vcat(df, missing_df)
+    @test vcat(missing_df) == DataFrame()
+    @test vcat(missing_df, missing_df) == DataFrame()
+    @test vcat(missing_df, df) == df
+    @test vcat(df, missing_df) == df
     @test eltypes(vcat(df, df)) == Type[Float64, Float64, Int]
     @test size(vcat(df, df)) == (size(df, 1) * 2, size(df, 2))
     res = vcat(df, df)
     @test res[1:size(df, 1), :] == df
-    @test res[1+size(df, 1):end, :] == df
-    @test eltypes(vcat(df, df, df)) == Type[Float64, Float64, Int]
-    @test size(vcat(df, df, df)) == (size(df, 1) * 3, size(df, 2))
+    @test res[(1+size(df, 1)):end, :] == df
     res = vcat(df, df, df)
+    @test eltypes(res) == Type[Float64, Float64, Int]
+    @test size(res) == (size(df, 1) * 3, size(df, 2))
+
     s = size(df, 1)
     for i in 1:3
         @test res[1+(i-1)*s:i*s, :] == df
@@ -348,24 +351,26 @@ end
 @testset "vcat out of order" begin
     df1 = DataFrame(A = 1:3, B = 4:6, C = 7:9)
     df2 = DataFrame([2x for x in eachcol(df1)], reverse(names(df1)))
-    @test vcat(df1, df2) == DataFrame([[1, 2, 3, 14, 16, 18],
-                                       [4, 5, 6, 8, 10, 12],
-                                       [7, 8, 9, 2, 4, 6]], [:A, :B, :C])
-    @test vcat(df1, df1, df2) == DataFrame([[1, 2, 3, 1, 2, 3, 14, 16, 18],
-                                            [4, 5, 6, 4, 5, 6, 8, 10, 12],
-                                            [7, 8, 9, 7, 8, 9, 2, 4, 6]], [:A, :B, :C])
-    @test vcat(df1, df2, df2) == DataFrame([[1, 2, 3, 14, 16, 18, 14, 16, 18],
-                                            [4, 5, 6, 8, 10, 12, 8, 10, 12],
-                                            [7, 8, 9, 2, 4, 6, 2, 4, 6]], [:A, :B, :C])
-    @test vcat(df2, df1, df2) == DataFrame([[2, 4, 6, 7, 8, 9, 2, 4, 6],
-                                            [8, 10, 12, 4, 5, 6, 8, 10, 12],
-                                            [14, 16, 18, 1, 2, 3, 14, 16, 18]] ,[:C, :B, :A])
-
+    @test vcat(df1, df2) == DataFrame(A = [1, 2, 3, 14, 16, 18],
+                                      B = [4, 5, 6, 8, 10, 12],
+                                      C = [7, 8, 9, 2, 4, 6])
+    # test with cols keyword argument
+    @test vcat(df1, df2, cols = :equal) == DataFrame(A = [1, 2, 3, 14, 16, 18],
+                                                     B = [4, 5, 6, 8, 10, 12],
+                                                     C = [7, 8, 9, 2, 4, 6])
+    @test vcat(df1, df1, df2) == DataFrame(A = [1, 2, 3, 1, 2, 3, 14, 16, 18],
+                                           B = [4, 5, 6, 4, 5, 6, 8, 10, 12],
+                                           C = [7, 8, 9, 7, 8, 9, 2, 4, 6])
+    @test vcat(df1, df2, df2) == DataFrame(A = [1, 2, 3, 14, 16, 18, 14, 16, 18],
+                                           B = [4, 5, 6, 8, 10, 12, 8, 10, 12],
+                                           C = [7, 8, 9, 2, 4, 6, 2, 4, 6])
+    @test vcat(df2, df1, df2) == DataFrame(C = [2, 4, 6, 7, 8, 9, 2, 4, 6],
+                                           B = [8, 10, 12, 4, 5, 6, 8, 10, 12],
+                                           A = [14, 16, 18, 1, 2, 3, 14, 16, 18])
     @test size(vcat(df1, df1, df1, df2, df2, df2)) == (18, 3)
     df3 = df1[[1, 3, 2]]
     res = vcat(df1, df1, df1, df2, df2, df2, df3, df3, df3, df3)
     @test res == reduce(vcat, [df1, df1, df1, df2, df2, df2, df3, df3, df3, df3])
-
     @test size(res) == (30, 3)
     @test res[1:3,:] == df1
     @test res[4:6,:] == df1
@@ -382,11 +387,46 @@ end
     @test [df1; df2] == df3 == reduce(vcat, [df1, df2])
 end
 
+@testset "vcat with cols=:union" begin
+    df1 = DataFrame(A = 1:3, B = 4:6)
+    df2 = DataFrame(A = 7:9)
+    df3 = DataFrame(B = 4:6, A = 1:3)
+
+    @test vcat(df1, df2; cols = :union) ≅
+        DataFrame(A = [1, 2, 3, 7, 8, 9],
+                  B = [4, 5, 6, missing, missing, missing])
+    @test vcat(df1, df2, df3; cols = :union) ≅
+        DataFrame(A = [1, 2, 3, 7, 8, 9, 1, 2, 3],
+                  B = [4, 5, 6, missing, missing, missing, 4, 5, 6])
+end
+
+@testset "vcat with cols=:intersect" begin
+    df1 = DataFrame(A = 1:3, B = 4:6)
+    df2 = DataFrame(A = 7:9)
+    df3 = DataFrame(A = 10:12, C = 13:15)
+
+    @test vcat(df1, df2; cols = :intersect) ≅ DataFrame(A = [1, 2, 3, 7, 8, 9])
+    @test vcat(df1, df2, df3; cols = :intersect) ≅ DataFrame(A = [1, 2, 3, 7, 8, 9,
+                                                                     10, 11, 12])
+end
+
+@testset "vcat with cols::Vector" begin
+    df1 = DataFrame(A = 1:3, B = 4:6)
+    df2 = DataFrame(A = 7:9)
+    df3 = DataFrame(A = 10:12, C = 13:15)
+
+    @test vcat(df1, df2; cols = [:A, :B, :C]) ≅
+        DataFrame(A = [1, 2, 3, 7, 8, 9],
+                  B = [4, 5, 6, missing, missing, missing],
+                  C = [missing, missing, missing, missing, missing, missing])
+
+    @test vcat(df1, df2, df3; cols = [:A, :B, :C]) ≅
+        DataFrame(A = [1, 2, 3, 7, 8, 9, 10, 11, 12],
+                  B = [4, 5, 6, missing, missing, missing, missing, missing, missing],
+                  C = [missing, missing, missing, missing, missing, missing, 13, 14, 15])
+end
+
 @testset "vcat errors" begin
-    err = @test_throws ArgumentError vcat(DataFrame(), DataFrame(), DataFrame(x=[]))
-    @test err.value.msg == "column(s) x are missing from argument(s) 1 and 2"
-    err = @test_throws ArgumentError vcat(DataFrame(), DataFrame(), DataFrame(x=[1]))
-    @test err.value.msg == "column(s) x are missing from argument(s) 1 and 2"
     df1 = DataFrame(A = 1:3, B = 1:3)
     df2 = DataFrame(A = 1:3)
     # right missing 1 column
@@ -396,10 +436,9 @@ end
     err = @test_throws ArgumentError vcat(df2, df1)
     @test err.value.msg == "column(s) B are missing from argument(s) 1"
     # multiple missing 1 column
-    err = @test_throws ArgumentError vcat(df1, df2, df2, df2, df2, df2)
+    err1 = @test_throws ArgumentError vcat(df1, df2, df2, df2, df2, df2)
     err2 = @test_throws ArgumentError reduce(vcat, [df1, df2, df2, df2, df2, df2])
-    @test err == err2
-    @test err.value.msg == "column(s) B are missing from argument(s) 2, 3, 4, 5 and 6"
+    @test err1.value.msg == err2.value.msg == "column(s) B are missing from argument(s) 2, 3, 4, 5 and 6"
     # argument missing >1 columns
     df1 = DataFrame(A = 1:3, B = 1:3, C = 1:3, D = 1:3, E = 1:3)
     err = @test_throws ArgumentError vcat(df1, df2)
@@ -438,6 +477,7 @@ end
     err = @test_throws ArgumentError vcat(df1, df2, df3, df4, df1, df2, df3, df4, df1, df2, df3, df4)
     @test err.value.msg == "column(s) E and F are missing from argument(s) 1, 5 and 9, column(s) B are missing from argument(s) 2, 6 and 10, and column(s) F are missing from argument(s) 3, 7 and 11"
 end
+
 x = view(DataFrame(A = Vector{Union{Missing, Int}}(1:3)), 2:2, :)
 y = DataFrame(A = 4:5)
 @test vcat(x, y) == DataFrame(A = [2, 4, 5]) == reduce(vcat, [x, y])