Skip to content

Commit

Permalink
Reduce number of changes, use mapreduce for T
Browse files Browse the repository at this point in the history
  • Loading branch information
pdeffebach committed Sep 17, 2018
1 parent 01dc729 commit 649e78f
Showing 1 changed file with 15 additions and 12 deletions.
27 changes: 15 additions & 12 deletions src/DataFramesMeta.jl
Original file line number Diff line number Diff line change
Expand Up @@ -388,30 +388,33 @@ end

function transform(g::GroupedDataFrame; kwargs...)
result = DataFrame(g)
starts = getproperty(g, :starts)
ends = getproperty(g, :ends)

function spread_scalar(x::Vector, obs_in_group::Int)
function spread_scalar(x::AbstractVector, obs_in_group)
if length(x) == obs_in_group
return x
else
throw(DimensionError("If a function returns a vector, the result " *
throw("If a function returns a vector, the result " *
"must have the same length as the groups it " *
"operates on"))
"operates on")
end
return x
end

function spread_scalar(x::CategoricalArrays.CategoricalValue, obs_in_group::Int)
vec = CategoricalArray(fill(x, obs_in_group))
levels!(vec, CategoricalArrays.index(CategoricalArrays.pool(x)))
end

function spread_scalar(x::Any, obs_in_group::Int)
fill(x, obs_in_group)
function spread_scalar(x, obs_in_group)
return fill(x, obs_in_group)
end

for (k, v) in kwargs
result[k] = reduce(vcat, spread_scalar(v(ig), size(ig, 1)) for ig in g)
output_iterator = (spread_scalar(v(ig), size(ig, 1)) for ig in g)
T = mapreduce(eltype, promote_type, g)
result[k] = Vector{T}(undef, size(result, 1))
for (i, output) in enumerate(output_iterator)
result[starts[i]:ends[i], k] = output
end
end
return result
result
end

function transform_helper(x, args...)
Expand Down

0 comments on commit 649e78f

Please sign in to comment.