Support adding columns to views

NEWS.md

@@ -11,6 +11,12 @@
 * correctly handle selectors of the form `:col => AsTable` and `:col => cols`
   by expanding a single column into multiple columns
   ([#2780](https://github.com/JuliaData/DataFrames.jl/pull/2780))
+* if `sdf` is a `SubDataFrame` created with `:` as a column selector then
+  `insertcols!`, `sdf[:, col] = v`, and `sdf[:, col] .= v` where `col` is
+  a column not present in `sdf` is allowed and it creates a new column in
+  `parent(sdf)` with `missing` values stored in rows that are filtered-out
+  in `sdf`.
+  ([XXXX](https://github.com/JuliaData/DataFrames.jl/pull/XXXX))
 
 ## Bug fixes
 

docs/src/lib/indexing.md

@@ -143,6 +143,9 @@ so it is unsafe to use it afterwards (the column length correctness will be pres
 * `sdf[CartesianIndex(row, col)] = v` -> the same as `sdf[row, col] = v`;
 * `sdf[row, cols] = v` -> the same as `dfr = df[row, cols]; dfr[:] = v` in-place;
 * `sdf[rows, col] = v` -> set rows `rows` of column `col`, in-place; `v` must be an abstract vector;
+  if `sdf` was created with `:` as column selector, `rows` is `:` and `col` is a `Symbol` or `AbstractString`
+  that is not present in `df` then a new column in `df` is created and holds `v` in rows selected in `sdf`
+  and `missing` in all rows present in `parent(sdf)` but not present in `sdf`.
 * `sdf[rows, cols] = v` -> set rows `rows` of columns `cols` in-place;
                            `v` can be an `AbstractMatrix` or `v` can be `AbstractDataFrame` when column names must match;
 
@@ -171,7 +174,6 @@ The following broadcasting rules apply to `AbstractDataFrame` objects:
 Note that if broadcasting assignment operation throws an error the target data frame may be partially changed
 so it is unsafe to use it afterwards (the column length correctness will be preserved).
 
-
 Broadcasting `DataFrameRow` is currently not allowed (which is consistent with `NamedTuple`).
 
 It is possible to assign a value to `AbstractDataFrame` and `DataFrameRow` objects using the `.=` operator.
@@ -190,6 +192,10 @@ Additional rules:
   `df` is performed in-place; if `rows` is `:` and `col` is `Symbol` or `AbstractString`
   and it is missing from `df` then a new column is allocated and added;
   the length of the column is always the value of `nrow(df)` before the assignment takes place;
+* in the `sdf[:, col] .= v` if `sdf` was created with `:` as column selector
+  and `col` is a `Symbol` or `AbstractString` that is not present in `df` then a new column in `df`
+  is created and holds contents of `v` broadcasted onto rows selected in `sdf`
+  and `missing` in all rows present in `parent(sdf)` but not present in `sdf`.
 * in the `df[!, col] .= v` syntax column `col` is replaced by a freshly allocated vector;
   if `col` is `Symbol` or `AbstractString` and it is missing from `df` then a new column is allocated added;
   the length of the column is always the value of `nrow(df)` before the assignment takes place;

src/abstractdataframe/abstractdataframe.jl

@@ -2173,3 +2173,209 @@ Base.getindex(::AbstractDataFrame, ::Union{Symbol, Integer, AbstractString}) =
 
 Base.setindex!(::AbstractDataFrame, ::Any, ::Union{Symbol, Integer, AbstractString}) =
     throw(ArgumentError("syntax df[column] is not supported use df[!, column] instead"))
+
+# insertcols!
+
+"""
+    insertcols!(df::AbstractDataFrame[, col], (name=>val)::Pair...;
+                makeunique::Bool=false, copycols::Bool=true)
+
+Insert a column into a data frame in place. Return the updated data frame.
+If `col` is omitted it is set to `ncol(df)+1`
+(the column is inserted as the last column).
+
+# Arguments
+- `df` : the data frame to which we want to add columns
+- `col` : a position at which we want to insert a column, passed as an integer
+  or a column name (a string or a `Symbol`); the column selected with `col`
+  and columns following it are shifted to the right in `df` after the operation
+- `name` : the name of the new column
+- `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`;
+  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
+
+If `val` is an `AbstractRange` then the result of `collect(val)` is inserted.
+
+If `df` is a `SubDataFrame` then it must be created with `:` as column selector
+(otherwise an error is thrown). In this case the `copycols` keyword argument
+is ignored an added column is always copied and the parent data frame is
+filled with `missing` in rows that are filtered out by `df`.
+
+If `df` isa `DataFrame` that has no columns and only values
+other than `AbstractVector` are passed then it is used to create a one element
+column.
+If `df` isa `DataFrame` that has no columns and at least one `AbstractVector` is
+passed then its length is used to determine the number of elements in all
+created columns.
+In all other cases the number of rows in all created columns must match
+`nrow(df)`.
+.
+
+# Examples
+```jldoctest
+julia> df = DataFrame(a=1:3)
+3×1 DataFrame
+ Row │ a
+     │ Int64
+─────┼───────
+   1 │     1
+   2 │     2
+   3 │     3
+
+julia> insertcols!(df, 1, :b => 'a':'c')
+3×2 DataFrame
+ Row │ b     a
+     │ Char  Int64
+─────┼─────────────
+   1 │ a         1
+   2 │ b         2
+   3 │ c         3
+
+julia> insertcols!(df, 2, :c => 2:4, :c => 3:5, makeunique=true)
+3×4 DataFrame
+ Row │ b     c      c_1    a
+     │ Char  Int64  Int64  Int64
+─────┼───────────────────────────
+   1 │ a         2      3      1
+   2 │ b         3      4      2
+   3 │ c         4      5      3
+```
+"""
+function insertcols!(df::AbstractDataFrame, col::ColumnIndex, name_cols::Pair{Symbol, <:Any}...;
+                     makeunique::Bool=false, copycols::Bool=true)
+    if !(df isa DataFrame || (df isa SubDataFrame && getfield(df, :colindex) isa Index))
+        throw(ArgumentError("insertcols! is only supported for DataFrame or" *
+                            "SubDataFrame created with `:` as column selector"))
+    end
+    col_ind = Int(col isa SymbolOrString ? columnindex(df, col) : col)
+    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"))
+    end
+
+    if !makeunique
+        if !allunique(first.(name_cols))
+            throw(ArgumentError("Names of columns to be inserted into a data frame " *
+                                "must be unique when `makeunique=true`"))
+        end
+        for (n, _) in name_cols
+            if hasproperty(df, n)
+                throw(ArgumentError("Column $n is already present in the data frame " *
+                                    "which is not allowed when `makeunique=true`"))
+            end
+        end
+    end
+
+    if ncol(df) == 0 && df isa DataFrame
+        target_row_count = -1
+    else
+        target_row_count = nrow(df)
+    end
+
+    for (n, v) in name_cols
+        if v isa AbstractVector
+            if target_row_count == -1
+                target_row_count = length(v)
+            elseif length(v) != target_row_count
+                if target_row_count == nrow(df)
+                    throw(DimensionMismatch("length of new column $n which is " *
+                                            "$(length(v)) must match the number " *
+                                            "of rows in data frame ($(nrow(df)))"))
+                else
+                    throw(DimensionMismatch("all vectors passed to be inserted into " *
+                                            "a data frame must have the same length"))
+                end
+            end
+        elseif v isa AbstractArray && ndims(v) > 1
+            throw(ArgumentError("adding AbstractArray other than AbstractVector as " *
+                                "a column of a data frame is not allowed"))
+        end
+    end
+    if target_row_count == -1
+        target_row_count = 1
+    end
+
+    for (name, item) in name_cols
+        if !(item isa AbstractVector)
+            if item isa Union{AbstractArray{<:Any, 0}, Ref}
+                x = item[]
+                item_new = fill!(Tables.allocatecolumn(typeof(x), target_row_count), x)
+            else
+                @assert !(item isa AbstractArray)
+                item_new = fill!(Tables.allocatecolumn(typeof(item), target_row_count), item)
+            end
+        elseif item isa AbstractRange
+            item_new = collect(item)
+        elseif copycols
+            item_new = copy(item)
+        else
+            item_new = item
+        end
+
+        if df isa DataFrame
+            dfp = df
+        else
+            dfp = parent(df)
+            T = eltype(item_new)
+            newcol = Tables.allocatecolumn(Union{T, Missing}, nrow(dfp))
+            fill!(newcol, missing)
+            newcol[rows(df)] = item_new
+            item_new = newcol
+        end
+
+        firstindex(item_new) != 1 && _onebased_check_error()
+
+        if ncol(dfp) == 0
+            dfp[!, name] = item_new
+        else
+            if hasproperty(dfp, name)
+                @assert makeunique
+                k = 1
+                while true
+                    nn = Symbol("$(name)_$k")
+                    if !hasproperty(dfp, nn)
+                        name = nn
+                        break
+                    end
+                    k += 1
+                end
+            end
+            insert!(index(dfp), col_ind, name)
+            insert!(_columns(dfp), col_ind, item_new)
+        end
+        col_ind += 1
+    end
+    return df
+end
+
+insertcols!(df::AbstractDataFrame, col::ColumnIndex, name_cols::Pair{<:AbstractString, <:Any}...;
+                     makeunique::Bool=false, copycols::Bool=true) =
+    insertcols!(df, col, (Symbol(n) => v for (n, v) in name_cols)...,
+                makeunique=makeunique, copycols=copycols)
+
+insertcols!(df::AbstractDataFrame, name_cols::Pair{Symbol, <:Any}...;
+            makeunique::Bool=false, copycols::Bool=true) =
+    insertcols!(df, ncol(df)+1, name_cols..., makeunique=makeunique, copycols=copycols)
+
+insertcols!(df::AbstractDataFrame, name_cols::Pair{<:AbstractString, <:Any}...;
+            makeunique::Bool=false, copycols::Bool=true) =
+    insertcols!(df, (Symbol(n) => v for (n, v) in name_cols)...,
+                makeunique=makeunique, copycols=copycols)
+
+function insertcols!(df::AbstractDataFrame, col::Int=ncol(df)+1; makeunique::Bool=false, name_cols...)
+    if !(0 < col <= ncol(df) + 1)
+        throw(ArgumentError("attempt to insert a column to a data frame with " *
+                            "$(ncol(df)) columns at index $col"))
+    end
+    if !isempty(name_cols)
+        # an explicit error is thrown as keyword argument was supported in the past
+        throw(ArgumentError("inserting colums using a keyword argument is not supported, " *
+                            "pass a Pair as a positional argument instead"))
+    end
+    return df
+end