feat: insertcols! receives kwarg after (#2829)

Author: pstorozenko

NEWS.md

@@ -21,6 +21,10 @@
 * `Tables.materializer` when passed the following types or their subtypes:
   `AbstractDataFrame`, `DataFrameRows`, `DataFrameColumns` returns `DataFrame`.
   ([#2839](https://github.com/JuliaData/DataFrames.jl/pull/2839))
+* the `insertcols!` function receives new keyword argument `after`
+  (with `false` default) that specifies if columns should be inserted after
+  or before `col`.
+  ([#2829](https://github.com/JuliaData/DataFrames.jl/pull/2829))
 
 ## Bug fixes
 

src/dataframe/dataframe.jl

@@ -748,7 +748,7 @@ end
 
 """
     insertcols!(df::DataFrame[, col], (name=>val)::Pair...;
-                makeunique::Bool=false, copycols::Bool=true)
+                after::Bool=false, makeunique::Bool=false, copycols::Bool=true)
 
 Insert a column into a data frame in place. Return the updated `DataFrame`.
 If `col` is omitted it is set to `ncol(df)+1`
@@ -763,8 +763,9 @@ If `col` is omitted it is set to `ncol(df)+1`
 - `val` : an `AbstractVector` giving the contents of the new column or a value of any
   type other than `AbstractArray` which will be repeated to fill a new vector;
   As a particular rule a values stored in a `Ref` or a `0`-dimensional `AbstractArray`
-  are unwrapped and treated in the same way.
-- `makeunique` : Defines what to do if `name` already exists in `df`;
+  are unwrapped and treated in the same way
+- `after` : if `true` columns are inserted after `col`
+- `makeunique` : defines what to do if `name` already exists in `df`;
   if it is `false` an error will be thrown; if it is `true` a new unique name will
   be generated by adding a suffix
 - `copycols` : whether vectors passed as columns should be copied
@@ -799,11 +800,33 @@ julia> insertcols!(df, 2, :c => 2:4, :c => 3:5, makeunique=true)
    1 │ a         2      3      1
    2 │ b         3      4      2
    3 │ c         4      5      3
+
+julia> insertcols!(df, :b, :d => 7:9, after=true)
+3×5 DataFrame
+ Row │ b     d      c      c_1    a
+     │ Char  Int64  Int64  Int64  Int64
+─────┼──────────────────────────────────
+   1 │ a         7      2      3      1
+   2 │ b         8      3      4      2
+   3 │ c         9      4      5      3
 ```
 """
 function insertcols!(df::DataFrame, col::ColumnIndex, name_cols::Pair{Symbol, <:Any}...;
-                     makeunique::Bool=false, copycols::Bool=true)
-    col_ind = Int(col isa SymbolOrString ? columnindex(df, col) : col)
+                     after::Bool=false, makeunique::Bool=false, copycols::Bool=true)
+
+    if col isa SymbolOrString
+        col_ind = Int(columnindex(df, col))
+        if col_ind == 0
+            throw(ArgumentError("column $col does not exist in data frame"))
+        end
+    else
+        col_ind = Int(col)
+    end
+
+    if after
+        col_ind += 1
+    end
+
     if !(0 < col_ind <= ncol(df) + 1)
         throw(ArgumentError("attempt to insert a column to a data frame with " *
                             "$(ncol(df)) columns at index $col_ind"))
@@ -894,18 +917,19 @@ function insertcols!(df::DataFrame, col::ColumnIndex, name_cols::Pair{Symbol, <:
 end
 
 insertcols!(df::DataFrame, col::ColumnIndex, name_cols::Pair{<:AbstractString, <:Any}...;
-                     makeunique::Bool=false, copycols::Bool=true) =
+            after::Bool=false, makeunique::Bool=false, copycols::Bool=true) =
     insertcols!(df, col, (Symbol(n) => v for (n, v) in name_cols)...,
-                makeunique=makeunique, copycols=copycols)
+                after=after, makeunique=makeunique, copycols=copycols)
 
 insertcols!(df::DataFrame, name_cols::Pair{Symbol, <:Any}...;
-            makeunique::Bool=false, copycols::Bool=true) =
-    insertcols!(df, ncol(df)+1, name_cols..., makeunique=makeunique, copycols=copycols)
+            after::Bool=false, makeunique::Bool=false, copycols::Bool=true) =
+    insertcols!(df, ncol(df)+1, name_cols..., after=after,
+                makeunique=makeunique, copycols=copycols)
 
 insertcols!(df::DataFrame, name_cols::Pair{<:AbstractString, <:Any}...;
-            makeunique::Bool=false, copycols::Bool=true) =
+            after::Bool=false, makeunique::Bool=false, copycols::Bool=true) =
     insertcols!(df, (Symbol(n) => v for (n, v) in name_cols)...,
-                makeunique=makeunique, copycols=copycols)
+                after=after, makeunique=makeunique, copycols=copycols)
 
 function insertcols!(df::DataFrame, col::Int=ncol(df)+1; makeunique::Bool=false, name_cols...)
     if !(0 < col <= ncol(df) + 1)

test/dataframe.jl

@@ -349,6 +349,32 @@ end
     @test_throws ArgumentError insertcols!(df, 2, y=2:3)
 end
 
+@testset "insertcols! after" begin
+    df = DataFrame(a=1:3)
+    insertcols!(df, :a, "b" => 2:4, after=true)
+    @test df == DataFrame(a=1:3, b=2:4)
+    insertcols!(df, :b, "b" => 2:4, after=true, makeunique=true)
+    @test df == DataFrame(a=1:3, b=2:4, b_1=2:4)
+    insertcols!(df, 0, :e => 1:3, after=true)
+    @test df == DataFrame(e=1:3, a=1:3, b=2:4, b_1=2:4)
+
+    @test_throws ArgumentError insertcols!(df, :a, "b" => 2:4, after=true)
+    @test_throws DimensionMismatch insertcols!(df, :a, :c => 2:5, after=true)
+    @test_throws ArgumentError insertcols!(df, :c, :b => 2:4, after=true, makeunique=true)
+    @test_throws ArgumentError insertcols!(df, ncol(df)+1, :d => 1:3, after=true)
+    @test_throws ArgumentError insertcols!(df, -1, :d => 1:3, after=true)
+
+    df = DataFrame(a=1:3, b=2:4)
+    insertcols!(df, 1, :c => 7:9, after=true)
+    @test df == DataFrame(a=1:3, c=7:9, b=2:4)
+
+    df = DataFrame(a=1:3)
+    insertcols!(df, 1, :b => 2:4, :c => 7:9, after=true)
+    @test df == DataFrame(a=1:3, b=2:4, c=7:9)
+    insertcols!(df, "b", :b => 2:4, :c => 7:9, after=true, makeunique=true)
+    @test df == DataFrame(a=1:3, b=2:4, b_1=2:4, c_1=7:9, c=7:9)
+end
+
 @testset "DataFrame constructors" begin
     @test DataFrame([Union{Int, Missing}[1, 2, 3], Union{Float64, Missing}[2.5, 4.5, 6.5]],
                     [:A, :B]) ==