add source option to outerjoin

Author: sl-solution

src/join/join.jl

@@ -699,8 +699,33 @@ function _fill_oncols_left_table_left_outer!(res, x, notinleft, en, total)
     end
 end
 
+function _fill_source_for_outer!(res, ranges, notinleft, lval, rval, en, total)
+    cnt = 0
+    for i in 1:length(ranges)
+        if length(ranges[i]) == 0
+            cnt += 1
+            res[cnt] = lval
+        else
+            cnt += length(ranges[i])
+        end
+    end
+    for i in en[end]+1:total
+        res[i] = rval
+    end
+end
+
+
+
+function _create_source_for_outer(ranges, notinleft, total_length, en)
+    res = allowmissing(PooledArray(["left", "right", "both"]))
+    resize!(res.refs, total_length)
+    fill!(res.refs, get(res.invpool, "both", missing))
+    _fill_source_for_outer!(res.refs, ranges, notinleft, get(res.invpool, "left", missing), get(res.invpool, "right", missing), en, total_length)
+    res
+end
+
 
-function _join_outer(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeunique = false, mapformats = [true, true], stable = false, alg = HeapSort, check = true, accelerate = false, method = :sort, threads = true) where T
+function _join_outer(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeunique = false, mapformats = [true, true], stable = false, alg = HeapSort, check = true, accelerate = false, method = :sort, threads = true, source::Bool = false, source_col_name = :source) where T
     isempty(dsl) || isempty(dsr) && throw(ArgumentError("in `outerjoin` both left and right tables must be non-empty"))
     oncols_left = onleft
     oncols_right = onright
@@ -713,7 +738,7 @@ function _join_outer(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeu
         end
         ranges = Vector{UnitRange{T}}(undef, nrow(dsl))
         if length(oncols_left) == 1 && nrow(dsr)>1
-            success, result = _join_outer_dict(dsl, dsr, ranges, oncols_left, oncols_right, oncols_left, oncols_right, right_cols, Val(T); makeunique = makeunique, mapformats = mapformats, check = check, threads = threads)
+            success, result = _join_outer_dict(dsl, dsr, ranges, oncols_left, oncols_right, oncols_left, oncols_right, right_cols, Val(T); makeunique = makeunique, mapformats = mapformats, check = check, threads = threads, source = source, source_col_name = source_col_name)
             if success
                 return result
             end
@@ -727,6 +752,9 @@ function _join_outer(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeu
     notinleft = _find_right_not_in_left(ranges, nrow(dsr), idx)
     cumsum!(new_ends, new_ends)
     total_length = new_ends[end] + length(notinleft)
+    if source
+        source_col = _create_source_for_outer(ranges, notinleft, total_length, new_ends)
+    end
     if check
         @assert total_length < 10*nrow(dsl) "the output data set will be very large ($(total_length)×$(ncol(dsl)+length(right_cols))) compared to the left data set size ($(nrow(dsl))×$(ncol(dsl))), make sure that the `on` keyword is selected properly, alternatively, pass `check = false` to ignore this error."
     end
@@ -765,6 +793,9 @@ function _join_outer(dsl, dsr::AbstractDataset, ::Val{T}; onleft, onright, makeu
         push!(index(newds), new_var_name)
         setformat!(newds, index(newds)[new_var_name], getformat(dsr, _names(dsr)[right_cols[j]]))
     end
+    if source
+        insertcols!(newds, source_col_name => source_col)
+    end
     newds
 
 end

src/join/join_dict.jl

@@ -353,7 +353,7 @@ function _join_inner_dict(dsl, dsr, ranges, onleft, onright, right_cols, ::Val{T
 
 end
 
-function _join_outer_dict(dsl, dsr, ranges, onleft, onright, oncols_left, oncols_right, right_cols, ::Val{T}; makeunique = makeunique, mapformats = mapformats, check = check, threads = true) where T
+function _join_outer_dict(dsl, dsr, ranges, onleft, onright, oncols_left, oncols_right, right_cols, ::Val{T}; makeunique = makeunique, mapformats = mapformats, check = check, threads = true, source::Bool = false, source_col_name = :source) where T
     _fl = _date_value∘identity
     _fr = _date_value∘identity
     if mapformats[1]
@@ -372,6 +372,10 @@ function _join_outer_dict(dsl, dsr, ranges, onleft, onright, oncols_left, oncols
     notinleft = _find_right_not_in_left(ranges, nrow(dsr), 1:nrow(dsr))
     cumsum!(new_ends, new_ends)
     total_length = new_ends[end] + length(notinleft)
+    if source
+        source_col = _create_source_for_outer(ranges, notinleft, total_length, new_ends)
+    end
+
     if check
         @assert total_length < 10*nrow(dsl) "the output data set will be very large ($(total_length)×$(ncol(dsl)+length(right_cols))) compared to the left data set size ($(nrow(dsl))×$(ncol(dsl))), make sure that the `on` keyword is selected properly, alternatively, pass `check = false` to ignore this error."
     end
@@ -410,6 +414,9 @@ function _join_outer_dict(dsl, dsr, ranges, onleft, onright, oncols_left, oncols
         push!(index(newds), new_var_name)
         setformat!(newds, index(newds)[new_var_name], getformat(dsr, _names(dsr)[right_cols[j]]))
     end
+    if source
+        insertcols!(newds, source_col_name => source_col)
+    end
     true, newds
 
 end

src/join/main.jl

@@ -461,7 +461,7 @@ julia> outerjoin(dsl, dsr, on = :year, mapformats = true) # Use formats for data
    4 │ 2012        true  missing
 ```
 """
-function DataAPI.outerjoin(dsl::AbstractDataset, dsr::AbstractDataset; on = nothing, makeunique = false,  mapformats::Union{Bool, Vector{Bool}} = true, stable = false, alg = HeapSort, check = true, accelerate = false, method = :sort, threads::Bool = true)
+function DataAPI.outerjoin(dsl::AbstractDataset, dsr::AbstractDataset; on = nothing, makeunique = false,  mapformats::Union{Bool, Vector{Bool}} = true, stable = false, alg = HeapSort, check = true, accelerate = false, method = :sort, threads::Bool = true, source::Bool = false, source_name = :source)
     !(method in (:hash, :sort)) && throw(ArgumentError("method must be :hash or :sort"))
     on === nothing && throw(ArgumentError("`on` keyword must be specified"))
     if !(on isa AbstractVector)
@@ -485,7 +485,7 @@ function DataAPI.outerjoin(dsl::AbstractDataset, dsr::AbstractDataset; on = noth
     else
         throw(ArgumentError("`on` keyword must be a vector of column names or a vector of pairs of column names"))
     end
-    _join_outer(dsl, dsr, nrow(dsr) < typemax(Int32) ? Val(Int32) : Val(Int64), onleft = onleft, onright = onright, makeunique = makeunique, mapformats = mapformats, stable = stable, alg = alg, check = check, accelerate = accelerate, method = method, threads = threads)
+    _join_outer(dsl, dsr, nrow(dsr) < typemax(Int32) ? Val(Int32) : Val(Int64), onleft = onleft, onright = onright, makeunique = makeunique, mapformats = mapformats, stable = stable, alg = alg, check = check, accelerate = accelerate, method = method, threads = threads, source = source, source_col_name = source_name)
 end
 
 """

test/join.jl

@@ -1303,6 +1303,31 @@ closefinance_tol10ms_noexact = Dataset([Union{Missing, DateTime}[DateTime("2016-
     @test closejoin(dsl, dsr, on = :x1, direction = :nearest, border = :none, threads = false) == Dataset(x1=[.3,.74,.53,.30, .65,1], y = [missing,3,3,missing,3, missing])
     @test closejoin(dsl, dsr, on = :x1, direction = :nearest, border = :none, method = :hash, threads = false) == Dataset(x1=[.3,.74,.53,.30, .65,1], y = [missing,3,3,missing,3, missing])
 
+    dsl = Dataset(x1 = [1,2,3,4], y = [100,200,300,400])
+    dsr = Dataset(x1 = [2,1,5,6], y1 = [-100,-200,-300,-400])
+    out1 = outerjoin(dsl, dsr, on = :x1, source = true)
+    out2 = outerjoin(dsr, dsl, on = :x1, source = true)
+    out1_t = Dataset(AbstractVector[Union{Missing, Int64}[1, 2, 3, 4, 5, 6], Union{Missing, Int64}[100, 200, 300, 400, missing, missing], Union{Missing, Int64}[-200, -100, missing, missing, -300, -400], Union{Missing, String}["both", "both", "left", "left", "right", "right"]], ["x1", "y", "y1", "source"])
+    out2_t = Dataset(AbstractVector[Union{Missing, Int64}[2, 1, 5, 6, 3, 4], Union{Missing, Int64}[-100, -200, -300, -400, missing, missing], Union{Missing, Int64}[200, 100, missing, missing, 300, 400], Union{Missing, String}["both", "both", "left", "left", "right", "right"]], ["x1", "y1", "y", "source"])
+    @test out1 == out1_t
+    @test out2 == out2_t
+    dsl = Dataset(x1 = [1,2,3,4], y = [100,200,300,400])
+    dsr = Dataset(x1 = [2,1], y1 = [-100,missing])
+    out1 = outerjoin(dsl, view(dsr,[2,1],:), on = :x1, source = true)
+    out2 = outerjoin(view(dsr,[2,1],:), dsl, on = :x1, source = true)
+    out1_t = Dataset(AbstractVector[Union{Missing, Int64}[1, 2, 3, 4], Union{Missing, Int64}[100, 200, 300, 400], Union{Missing, Int64}[missing, -100, missing, missing], Union{Missing, String}["both", "both", "left", "left"]], ["x1", "y", "y1", "source"])
+    out2_t = Dataset(AbstractVector[Union{Missing, Int64}[1, 2, 3, 4], Union{Missing, Int64}[missing, -100, missing, missing], Union{Missing, Int64}[100, 200, 300, 400], Union{Missing, String}["both", "both", "right", "right"]], ["x1", "y1", "y", "source"])
+    @test out1 == out1_t
+    @test out2 == out2_t
+    dsl = Dataset(x1 = [1,2,3,4],x2=[1,1,1,1], y = [100,200,300,400])
+    dsr = Dataset(x1 = [2,1,5,6],x2= [1,1,1,1], y1 = [-100,-200,-300,-400])
+    out1 = outerjoin(dsl, dsr, on = [:x1, :x2], source = true)
+    out2 = outerjoin(dsr, dsl, on = [:x1, :x2], source = true)
+    out1_t = Dataset(AbstractVector[Union{Missing, Int64}[1, 2, 3, 4, 5, 6], Union{Missing, Int64}[1, 1, 1, 1, 1, 1], Union{Missing, Int64}[100, 200, 300, 400, missing, missing], Union{Missing, Int64}[-200, -100, missing, missing, -300, -400], Union{Missing, String}["both", "both", "left", "left", "right", "right"]], ["x1", "x2", "y", "y1", "source"])
+    out2_t = Dataset(AbstractVector[Union{Missing, Int64}[2, 1, 5, 6, 3, 4], Union{Missing, Int64}[1, 1, 1, 1, 1, 1], Union{Missing, Int64}[-100, -200, -300, -400, missing, missing], Union{Missing, Int64}[200, 100, missing, missing, 300, 400], Union{Missing, String}["both", "both", "left", "left", "right", "right"]], ["x1", "x2", "y1", "y", "source"])
+    @test out1 == out1_t
+    @test out2 == out2_t
+
 end
 
 @testset "Test empty inputs 1" begin