avoid type piracy in reduce with vcat

JuliaData · Sep 11, 2024 · 34d8736 · 34d8736
1 parent e4b7627
commit 34d8736
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 3 deletions.
diff --git a/src/abstractdataframe/iteration.jl b/src/abstractdataframe/iteration.jl
@@ -598,6 +598,12 @@ function mapcols!(f::Union{Function,Type}, df::DataFrame; cols=All())
     return df
 end
 
+##############################################################################
+##
+## Reduction
+##
+##############################################################################
+
 """
     reduce(::typeof(vcat),
            dfs::Union{AbstractVector{<:AbstractDataFrame},
@@ -686,7 +692,10 @@ julia> reduce(vcat, [df1, df2, df3], cols=:union, source=:source)
 ```
 """
 function Base.reduce(::typeof(vcat),
-    dfs::Union{AbstractVector{<:AbstractDataFrame},
+    dfs::Union{AbstractVector{AbstractDataFrame},
+        AbstractVector{DataFrame},
+        AbstractVector{SubDataFrame},
+        AbstractVector{Union{DataFrame,SubDataFrame}},
         Tuple{AbstractDataFrame,Vararg{AbstractDataFrame}}};
     cols::Union{Symbol,AbstractVector{Symbol},
         AbstractVector{<:AbstractString}}=:setequal,
@@ -741,8 +750,10 @@ end
 
 # definition needed to avoid dispatch ambiguity
 Base.reduce(::typeof(vcat),
-    dfs::SentinelArrays.ChainedVector{T,A} where {T<:AbstractDataFrame,
-        A<:AbstractVector{T}};
+    dfs::Union{SentinelArrays.ChainedVector{AbstractDataFrame,<:AbstractVector{AbstractDataFrame}},
+        SentinelArrays.ChainedVector{DataFrame,<:AbstractVector{DataFrame}},
+        SentinelArrays.ChainedVector{SubDataFrame,<:AbstractVector{SubDataFrame}},
+        SentinelArrays.ChainedVector{Union{DataFrame,SubDataFrame},<:AbstractVector{Union{DataFrame,SubDataFrame}}}};
     cols::Union{Symbol,AbstractVector{Symbol},
         AbstractVector{<:AbstractString}}=:setequal,
     source::Union{Nothing,SymbolOrString,

diff --git a/test/cat.jl b/test/cat.jl
@@ -477,4 +477,19 @@ end
     @test reduce(vcat, (df1, df2)) == DataFrame(a=[1, 1], b=[2, 2])
 end
 
+@testset "vcat type piracy" begin
+    x = Int[]
+    @test reduce(vcat, Union{}[], init=x) === x
+
+    @test reduce(vcat, AbstractDataFrame[DataFrame(a=1), DataFrame(a=2)]) ==
+          DataFrame(a=[1, 2])
+    @test reduce(vcat, Union{DataFrame, SubDataFrame}[DataFrame(a=1), DataFrame(a=2)]) ==
+          DataFrame(a=[1, 2])
+    @test reduce(vcat, AbstractDataFrame[DataFrame(a=1), DataFrame(a=2)]; source=:source) ==
+          DataFrame(a=[1, 2], source=[1, 2])
+    @test reduce(vcat, Union{DataFrame,SubDataFrame}[DataFrame(a=1), DataFrame(a=2)]; source=:source) ==
+          DataFrame(a=[1, 2], source=[1, 2])
+end
+
+
 end # module