Add broadcasting of AbstractDataFrame

docs/src/lib/indexing.md

@@ -47,6 +47,7 @@ For performance reasons, accessing, via `getindex` or `view`, a single `row` and
 * `df[col]` -> the vector contained in column `col`;
 * `df[cols]` -> a freshly allocated `DataFrame` containing the copies of vectors contained in columns `cols`;
 * `df[row, col]` -> the value contained in row `row` of column `col`, the same as `df[col][row]`;
+* `df[CartesianIndex(row, col)]` -> the same as `df[row,col]`;
 * `df[row, cols]` -> a `DataFrameRow` with parent `df` if `cols` is a colon and `df[cols]` otherwise;
 * `df[rows, col]` -> a copy of the vector `df[col]` with only the entries corresponding to `rows` selected, the same as `df[col][rows]`;
 * `df[rows, cols]` -> a `DataFrame` containing copies of columns `cols` with only the entries corresponding to `rows` selected.
@@ -83,6 +84,14 @@ Under construction
 
 ## Broadcasting
 
+The following broadcasting rules apply to `AbstractDataFrame` objects:
+* `AbstractDataFrame` behaves in broadcasting like a two-dimensional collection compatible with matrices.
+* If an `AbstractDataFrame` takes part in broadcasting then a `DataFrame` is always produced as a result.
+  In this case the requested broadcasting operation produce an object with exactly two dimensions.
+  An exception is when an `AbstractDataFrame` is used only as a source of broadcast assignment into an object
+  of dimensionality higher than two.
+* If multiple `AbstractDataFrame` objects take part in broadcasting then they have to have identical column names.
+
 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.
 

src/other/broadcasting.jl

@@ -1,3 +1,78 @@
+### Broadcasting
+
+Base.getindex(df::AbstractDataFrame, idx::CartesianIndex{2}) = df[idx[1], idx[2]]
+Base.setindex!(df::AbstractDataFrame, val, idx::CartesianIndex{2}) =
+    (df[idx[1], idx[2]] = val)
+
+Base.broadcastable(df::AbstractDataFrame) = df
+
+struct DataFrameStyle <: Base.Broadcast.BroadcastStyle end
+
+Base.Broadcast.BroadcastStyle(::Type{<:AbstractDataFrame}) =
+    DataFrameStyle()
+
+Base.Broadcast.BroadcastStyle(::DataFrameStyle, ::Base.Broadcast.BroadcastStyle) = DataFrameStyle()
+Base.Broadcast.BroadcastStyle(::Base.Broadcast.BroadcastStyle, ::DataFrameStyle) = DataFrameStyle()
+Base.Broadcast.BroadcastStyle(::DataFrameStyle, ::DataFrameStyle) = DataFrameStyle()
+
+function copyto_widen!(res::AbstractVector{T},
+                       bc::Base.Broadcast.Broadcasted{DataFrameStyle},
+                       pos, col) where T
+    for i in pos:length(axes(bc)[1])
+        val = bc[CartesianIndex(i, col)]
+        S = typeof(val)
+        if S <: T || promote_type(S, T) <: T
+            res[i] = val
+        else
+            newres = similar(Vector{promote_type(S, T)}, length(res))
+            copyto!(newres, 1, res, 1, i-1)
+            newres[i] = val
+            return copyto_widen!(newres, bc, i + 1, 2)
+        end
+    end
+    return res
+end
+
+function getcolbc(bcf::Base.Broadcast.Broadcasted{Style}, colind) where {Style}
+    # we assume that bcf is already flattened and unaliased
+    newargs = map(bcf.args) do x
+        Base.Broadcast.extrude(x isa AbstractDataFrame ? x[colind] : x)
+    end
+    Base.Broadcast.Broadcasted{Style}(bcf.f, newargs, bcf.axes)
+end
+
+function Base.copy(bc::Base.Broadcast.Broadcasted{DataFrameStyle})
+    ndim = length(axes(bc))
+    if ndim != 2
+        throw(DimensionMismatch("cannot broadcast a data frame into $ndim dimensions"))
+    end
+    bcf = Base.Broadcast.flatten(bc)
+    colnames = unique([_names(df) for df in bcf.args if df isa AbstractDataFrame])
+    if length(colnames) != 1
+        wrongnames = setdiff(union(colnames...), intersect(colnames...))
+        msg = join(wrongnames, ", ", " and ")
+        throw(ArgumentError("Column names in broadcasted data frames must match. " *
+                            "Non matching column names are $msg"))
+    end
+    nrows = length(axes(bcf)[1])
+    df = DataFrame()
+    for i in axes(bcf)[2]
+        if nrows == 0
+            col = Any[]
+        else
+            bcf′ = getcolbc(bcf, i)
+            v1 = bcf′[CartesianIndex(1, i)]
+            startcol = similar(Vector{typeof(v1)}, nrows)
+            startcol[1] = v1
+            col = copyto_widen!(startcol, bcf′, 2, i)
+        end
+        df[colnames[1][i]] = col
+    end
+    return df
+end
+
+### Broadcasting assignment
+
 struct LazyNewColDataFrame
     df::DataFrame
     col::Symbol
@@ -50,9 +125,88 @@ function _copyto_helper!(dfcol::AbstractVector, bc::Base.Broadcast.Broadcasted,
     end
 end
 
+function Base.Broadcast.broadcast_unalias(dest::AbstractDataFrame, src)
+    for col in eachcol(dest)
+        src = Base.Broadcast.unalias(col, src)
+    end
+    src
+end
+
+function Base.Broadcast.broadcast_unalias(dest, src::AbstractDataFrame)
+    wascopied = false
+    for (i, col) in enumerate(eachcol(src))
+        if Base.mightalias(dest, col)
+            if src isa SubDataFrame
+                if !wascopied
+                    src = SubDataFrame(copy(parent(src), copycols=false),
+                                       index(src), rows(src))
+                end
+                parentidx = parentcols(index(src), i)
+                parent(src)[parentidx] = Base.unaliascopy(parent(src)[parentidx])
+            else
+                if !wascopied
+                    src = copy(src, copycols=false)
+                end
+                src[i] = Base.unaliascopy(col)
+            end
+            wascopied = true
+        end
+    end
+    src
+end
+
+function _broadcast_unalias_helper(dest::AbstractDataFrame, scol::AbstractVector,
+                                   src::AbstractDataFrame, col2::Int, wascopied::Bool)
+    # col1 can be checked till col2 point as we are writing broadcasting
+    # results from 1 to ncol
+    # we go downwards because aliasing when col1 == col2 is most probable
+    for col1 in col2:-1:1
+        dcol = dest[col1]
+        if Base.mightalias(dcol, scol)
+            if src isa SubDataFrame
+                if !wascopied
+                    src =SubDataFrame(copy(parent(src), copycols=false),
+                                      index(src), rows(src))
+                end
+                parentidx = parentcols(index(src), col2)
+                parent(src)[parentidx] = Base.unaliascopy(parent(src)[parentidx])
+            else
+                if !wascopied
+                    src = copy(src, copycols=false)
+                end
+                src[col2] = Base.unaliascopy(scol)
+            end
+            return src, true
+        end
+    end
+    return src, wascopied
+end
+
+function Base.Broadcast.broadcast_unalias(dest::AbstractDataFrame, src::AbstractDataFrame)
+    if size(dest, 2) != size(src, 2)
+        throw(ArgumentError("Dimension mismatch in broadcasting."))
+    end
+    wascopied = false
+    for col2 in axes(dest, 2)
+        scol = src[col2]
+        src, wascopied = _broadcast_unalias_helper(dest, scol, src, col2, wascopied)
+    end
+    src
+end
+
 function Base.copyto!(df::AbstractDataFrame, bc::Base.Broadcast.Broadcasted)
-    for col in axes(df, 2)
-        _copyto_helper!(df[col], bc, col)
+    bcf = Base.Broadcast.flatten(bc)
+    colnames = unique([_names(df) for df in bcf.args if df isa AbstractDataFrame])
+    if length(colnames) > 1 || (length(colnames) == 1 && _names(df) != colnames[1])
+        wrongnames = setdiff(union(colnames...), intersect(colnames...))
+        msg = join(wrongnames, ", ", " and ")
+        throw(ArgumentError("Column names in broadcasted data frames must match. " *
+                            "Non matching column names are $msg"))
+    end
+
+    bcf′ = Base.Broadcast.preprocess(df, bcf)
+    for i in axes(df, 2)
+        _copyto_helper!(df[i], getcolbc(bcf′, i), i)
     end
     df
 end
@@ -65,13 +219,17 @@ function Base.copyto!(df::AbstractDataFrame, bc::Base.Broadcast.Broadcasted{<:Ba
         end
         df
     else
-        copyto!(df, convert(Broadcasted{Nothing}, bc))
+        copyto!(df, convert(Base.Broadcast.Broadcasted{Nothing}, bc))
     end
 end
 
+Base.Broadcast.broadcast_unalias(dest::DataFrameRow, src) =
+    Base.Broadcast.broadcast_unalias(parent(dest), src)
+
 function Base.copyto!(dfr::DataFrameRow, bc::Base.Broadcast.Broadcasted)
-    for I in eachindex(bc)
-        dfr[I] = bc[I]
+    bc′ = Base.Broadcast.preprocess(dfr, bc)
+    for I in eachindex(bc′)
+        dfr[I] = bc′[I]
     end
     dfr
 end