Broadcasting in DataFrames

docs/src/lib/indexing.md

@@ -79,3 +79,20 @@ For performance reasons, accessing, via `getindex` or `view`, a single `row` and
 ## `setindex!`
 
 Under construction
+
+## Broadcasting
+
+It is possible to assign a value to `AbstractDataFrame` and `DataFrameRow` objects using the `.=` operator.
+In such an operation `AbstractDataFrame` is considered as two-dimensional and `DataFrameRow` as single-dimensional.
+
+!!! note
+
+    The rule above means that, similar to single-dimensional objects in Base (e.g. vectors),
+    `DataFrameRow` is considered to be column-oriented.
+
+If column indexing using `Symbol` names is performed the order of columns in the operation is specified
+by the order of names.
+
+`df[col] .= value` is allowed when `col` is a `Symbol` even if `col` is not present in the `DataFrame`
+under the condition that `df` is not empty: a new column will be created.
+On the contrary, `df.col .= value` is not allowed if `col` is not present in `df`.

src/DataFrames.jl

@@ -83,6 +83,8 @@ include("dataframerow/dataframerow.jl")
 include("groupeddataframe/grouping.jl")
 include("dataframerow/utils.jl")
 
+include("other/broadcasting.jl")
+
 include("abstractdataframe/iteration.jl")
 include("abstractdataframe/join.jl")
 include("abstractdataframe/reshape.jl")

src/other/broadcasting.jl

@@ -0,0 +1,72 @@
+struct LazyNewColDataFrame
+    df::DataFrame
+    col::Symbol
+end
+
+# we allow LazyNewColDataFrame only for data frames with at least one column
+Base.axes(x::LazyNewColDataFrame) = axes(x.df[1])
+
+Base.maybeview(df::AbstractDataFrame, idxs...) = view(df, idxs...)
+
+function Base.maybeview(df::AbstractDataFrame, idxs)
+    if ncol(df) == 0
+        throw(ArgumentError("Broadcasting into a data frame with no columns is not allowed"))
+    end
+    if idxs isa Symbol
+        if !haskey(df, idxs)
+            if !(df isa DataFrame)
+                # this will throw an appropriate error message
+                df[idxs]
+            end
+            return LazyNewColDataFrame(df, idxs)
+        end
+    end
+    view(df, idxs)
+end
+
+function Base.copyto!(lazydf::LazyNewColDataFrame, bc::Base.Broadcast.Broadcasted)
+    if isempty(lazydf.df)
+        throw(ArgumentError("creating a column via broadcasting is not allowed on empty data frames"))
+    end
+    T = mapreduce(i -> typeof(bc[i]), promote_type, eachindex(bc); init=Union{})
+    col = Tables.allocatecolumn(T, nrow(lazydf.df))
+    copyto!(col, bc)
+    lazydf.df[lazydf.col] = col
+end
+
+function _copyto_helper!(dfcol::AbstractVector, bc::Base.Broadcast.Broadcasted, col::Int)
+    if axes(dfcol, 1) != axes(bc)[1]
+        # this should never happen unless data frame is corrupted (has unequal column lengths)
+        throw(ArgumentError("Dimension mismatch in broadcasting. " *
+                            "The updated data frame is invalid and should not be used"))
+    end
+    @inbounds for row in eachindex(dfcol)
+        dfcol[row] = bc[CartesianIndex(row, col)]
+    end
+end
+
+function Base.copyto!(df::AbstractDataFrame, bc::Base.Broadcast.Broadcasted)
+    for col in axes(df, 2)
+        _copyto_helper!(df[col], bc, col)
+    end
+    df
+end
+
+function Base.copyto!(df::AbstractDataFrame, bc::Base.Broadcast.Broadcasted{<:Base.Broadcast.AbstractArrayStyle{0}})
+    # special case of fast approach when bc is providing an untransformed scalar
+    if bc.f === identity && bc.args isa Tuple{Any} && Base.Broadcast.isflat(bc)
+        for col in axes(df, 2)
+            fill!(df[col], bc.args[1][])
+        end
+        df
+    else
+        copyto!(df, convert(Broadcasted{Nothing}, bc))
+    end
+end
+
+function Base.copyto!(dfr::DataFrameRow, bc::Base.Broadcast.Broadcasted)
+    for I in eachindex(bc)
+        dfr[I] = bc[I]
+    end
+    dfr
+end