Add transform and combine from DataFrames (#22)

krynju · web-flow · commit 2bc883fcc083 · 2023-02-26T12:01:05.000+01:00
diff --git a/Project.toml b/Project.toml
@@ -9,6 +9,7 @@ DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
 SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TableOperations = "ab02a1b2-a7df-11e8-156e-fb1833f50b87"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
diff --git a/src/DTables.jl b/src/DTables.jl
@@ -16,6 +16,7 @@ using DataFrames:
     make_pair_concrete
 using InvertedIndices: Not
 using SentinelArrays: ChainedVector
+using Statistics: mean
 using TableOperations: TableOperations
 using Tables:
     columnindex,
@@ -54,7 +55,7 @@ import Base:
 import DataAPI: leftjoin, ncol, nrow, innerjoin
 import Tables:
     columnaccess, columnnames, columns, getcolumn, istable, partitions, rowaccess, rows, schema
-import DataFrames: broadcast_pair, select, index
+import DataFrames: broadcast_pair, combine, groupby, select, index, transform
 
 ############################################################################################
 # Export
@@ -64,6 +65,7 @@ export All,
     AsTable,
     Between,
     ByRow,
+    combine,
     Cols,
     DTable,
     DTableColumn,
@@ -73,8 +75,10 @@ export All,
     Not,
     nrow,
     select,
+    groupby,
     tabletype,
     tabletype!,
+    transform,
     trim,
     trim!
 ############################################################################################
diff --git a/src/operations/dataframes_interface.jl b/src/operations/dataframes_interface.jl
@@ -54,14 +54,16 @@ function _manipulate(df::DTable, normalized_cs::Vector{Any}, copycols::Bool, kee
     # STAGE 1: Spawning full column thunks - also multicolumn when needed (except identity)
     # These get saved later and used in last stages.
     #########
-    colresults = Dict{Int,Any}()
-    for (i, (colidx, (f, _))) in enumerate(normalized_cs)
-        if !(colidx isa AsTable) && !(f isa ByRow) && f != identity
-            if length(colidx) > 0
-                cs = DTableColumn.(Ref(df), [colidx...])
-                colresults[i] = Dagger.@spawn f(cs...)
+    normalized_cs_results = Dict{Int,Dagger.EagerThunk}()
+    for (idx, (column_index, (fun, result_column_symbol))) in enumerate(normalized_cs)
+        if (!(column_index isa AsTable) && !(fun isa ByRow) && fun != identity)
+            if length(column_index) > 0
+                normalized_cs_results[idx] = Dagger.@spawn fun(
+                    DTableColumn.(Ref(df), [column_index...])...
+                )
             else
-                colresults[i] = Dagger.@spawn f() # case of select(d, [] => fun)
+                # case of select(d, [] => fun) where there are no input columns
+                normalized_cs_results[idx] = Dagger.@spawn fun()
             end
         end
     end
@@ -71,51 +73,78 @@ function _manipulate(df::DTable, normalized_cs::Vector{Any}, copycols::Bool, kee
     # These will be just injected as values in the mapping, because it's a vector full of these values
     #########
 
-    colresults = Dict{Int,Any}(
-        k => fetch(Dagger.spawn(length, v)) == 1 ? fetch(v) : v for (k, v) in colresults
+    mappable_part_of_normalized_cs = filter(
+        x -> !haskey(normalized_cs_results, x[1]), collect(enumerate(normalized_cs))
     )
 
-    mapmask = [
-        haskey(colresults, x) && colresults[x] isa Dagger.EagerThunk for
-        (x, _) in enumerate(normalized_cs)
-    ]
-
-    mappable_part_of_normalized_cs = filter(x -> !mapmask[x[1]], collect(enumerate(normalized_cs)))
-
     #########
     # STAGE 3: Mapping function (need to ensure this is compiled only once)
     # It's awful right now, but it covers all cases
     # Essentially we skip all the non-mappable stuff here
     #########
 
-    rd = map(x -> select_rowfunction(x, mappable_part_of_normalized_cs, colresults), df)
+    has_any_mappable = length(mappable_part_of_normalized_cs) > 0
+
+    rd = if has_any_mappable || keeprows
+        map(x -> select_rowfunction(x, mappable_part_of_normalized_cs), df)
+    else
+        nothing # in case there's nothing mappable we just go ahead with an empty dtable (just nothing)
+    end
 
     #########
     # STAGE 4: Preping for last stage - getting all the full column thunks with not 1 lengths
     #########
-    cpcolresults = Dict{Int,Any}()
 
-    for (k, v) in colresults
-        if v isa Dagger.EagerThunk
-            cpcolresults[k] = v
-        end
+    fullcolumn_ops_result_lengths = Int[
+        fetch(Dagger.spawn(length, v)) for v in values(normalized_cs_results)
+    ]
+
+    collength_to_compare_against = if has_any_mappable || keeprows
+        length(df)
+    else
+        maximum(fullcolumn_ops_result_lengths)
     end
 
-    for (_, v) in colresults
-        if v isa Dagger.EagerThunk
-            if fetch(Dagger.spawn(length, v)) != length(df)
-                throw("result column is not the size of the table")
-            end
-        end
+    if !all(map(x -> x == 1 || x == collength_to_compare_against, fullcolumn_ops_result_lengths))
+        throw(ArgumentError("New columns must have the same length as old columns"))
     end
+
     #########
     # STAGE 5: Fill columns - meaning the previously omitted full column tasks
     # will be now merged into the final DTable
     #########
 
-    rd = fillcolumns(rd, cpcolresults, normalized_cs, chunk_lengths(df))
+    new_chunk_lengths = if has_any_mappable || keeprows
+        chunk_lengths(df)
+    elseif maximum(fullcolumn_ops_result_lengths) == 1
+        z = zeros(Int, nchunks(df))
+        z[1] = 1
+        z
+    else
+        b = maximum(fullcolumn_ops_result_lengths)
+        a = zeros(Int, nchunks(df))
+        avg_chunk_length = floor(Int, mean(chunk_lengths(df)))
+        for (i, c) in enumerate(chunk_lengths(df))
+            if b >= c
+                a[i] += c
+                b -= c
+            else
+                a[i] += b
+                b = 0
+            end
+        end
+        while b > 0
+            bm = min(b, avg_chunk_length)
+            push!(a, bm)
+            b -= bm
+        end
+        a
+    end
 
-    return rd
+    rd2 = fillcolumns(
+        rd, normalized_cs_results, normalized_cs, new_chunk_lengths, fullcolumn_ops_result_lengths
+    )
+    return rd2
 end
 
 """
@@ -147,3 +176,23 @@ function select(
         renamecols=renamecols,
     )
 end
+
+function transform(
+    df::DTable,
+    @nospecialize(args...);
+    copycols::Bool=true,
+    renamecols::Bool=true,
+    threads::Bool=true,
+)
+    return select(df, :, args...; copycols=copycols, renamecols=renamecols, threads=threads)
+end
+
+function combine(df::DTable, @nospecialize(args...); renamecols::Bool=true, threads::Bool=true)
+    return manipulate(
+        df,
+        map(x -> broadcast_pair(df, x), args)...;
+        copycols=true,
+        keeprows=false,
+        renamecols=renamecols,
+    )
+end
diff --git a/src/operations/dataframes_interface_utils.jl b/src/operations/dataframes_interface_utils.jl
@@ -1,5 +1,5 @@
 
-function select_rowfunction(row, mappable_part_of_normalized_cs, colresults)
+function select_rowfunction(row, mappable_part_of_normalized_cs)
     _cs = [
         begin
             kk = result_colname === AsTable ? Symbol("AsTable$(i)") : result_colname
@@ -21,8 +21,6 @@ function select_rowfunction(row, mappable_part_of_normalized_cs, colresults)
                     f.fun(args)
                 elseif f == identity
                     args
-                elseif length(colresults[i]) == 1
-                    colresults[i]
                 else
                     throw(ErrorException("Weird unhandled stuff"))
                 end
@@ -33,62 +31,106 @@ function select_rowfunction(row, mappable_part_of_normalized_cs, colresults)
     return (; _cs...)
 end
 
-function fillcolumns(
-    dt::DTable, ics::Dict{Int,Any}, normalized_cs, chunk_lengths_of_original_dt::Vector{Int}
+function fillcolumn(
+    chunk,
+    csymbols::Union{Vector{DataType},Vector{Symbol},Vector{Union{DataType,Symbol}}},
+    colfragments::Union{Vector{Dagger.EagerThunk},Vector{Any}},
+    expected_chunk_length::Int,
+    normalized_cs::Vector{Any},
 )
-    col_keys_indices = collect(keys(ics))::Vector{Int}
-    col_vecs = map(x -> ics[x], col_keys_indices)::Union{Vector{Any},Vector{Dagger.EagerThunk}}
-
-    f =
-        (ch, csymbols, colfragments) -> begin
-            col_vecs_fetched = fetch.(colfragments)
-            colnames = Vector{Symbol}()
-            cols = Vector{Any}()
-            last_astable = 0
+    col_vecs_fetched = fetch.(colfragments)
+    colnames = Vector{Symbol}()
+    cols = Vector{Any}()
+    last_astable = 0
 
-            for (idx, (_, (_, sym))) in enumerate(normalized_cs)
-                if sym !== AsTable
-                    col = if sym in csymbols
-                        index = something(indexin(csymbols, [sym])...)
-                        col_vecs_fetched[index]
-                    else
-                        getcolumn(ch, sym)
-                    end
-                    push!(colnames, sym)
-                    push!(cols, col)
-                elseif sym === AsTable
-                    i = findfirst(x -> x === AsTable, csymbols[(last_astable + 1):end])
-                    if i === nothing
-                        c = getcolumn(ch, Symbol("AsTable$(idx)"))
-                    else
-                        last_astable = i
-                        c = col_vecs_fetched[i]
-                    end
-
-                    push!.(Ref(colnames), columnnames(columns(c)))
-                    push!.(Ref(cols), getcolumn.(Ref(columns(c)), columnnames(columns(c))))
+    for (idx, (_, (_, sym))) in enumerate(normalized_cs)
+        if sym !== AsTable
+            col = if sym in csymbols
+                index = findfirst(x -> x === sym, csymbols)
+                if col_vecs_fetched[index] isa AbstractVector
+                    col_vecs_fetched[index]
+                else
+                    repeat([col_vecs_fetched[index]], expected_chunk_length)
+                end
+            else
+                getcolumn(chunk, sym)
+            end
+            push!(colnames, sym)
+            push!(cols, col)
+        elseif sym === AsTable
+            i = findfirst(x -> x === AsTable, csymbols[(last_astable + 1):end])
+            c = if i === nothing
+                getcolumn(chunk, Symbol("AsTable$(idx)"))
+            else
+                last_astable = i
+                if col_vecs_fetched[i] isa AbstractVector
+                    col_vecs_fetched[i]
                 else
-                    throw(ErrorException("something is off"))
+                    repeat([col_vecs_fetched[i]], expected_chunk_length)
                 end
             end
-            materializer(ch)(
-                merge(NamedTuple(), (; [e[1] => e[2] for e in zip(colnames, cols)]...))
-            )
+
+            push!.(Ref(colnames), columnnames(columns(c)))
+            push!.(Ref(cols), getcolumn.(Ref(columns(c)), columnnames(columns(c))))
+        else
+            throw(ErrorException("something is off"))
         end
+    end
+    return materializer(chunk)(
+        merge(NamedTuple(), (; [e[1] => e[2] for e in zip(colnames, cols)]...))
+    )
+end
+
+function fillcolumns(
+    dt::Union{Nothing,DTable},
+    normalized_cs_results::Dict{Int,Dagger.EagerThunk},
+    normalized_cs::Vector{Any},
+    new_chunk_lengths::Vector{Int},
+    fullcolumn_ops_result_lengths::Vector{Int},
+)
+    fullcolumn_ops_indices_in_normalized_cs = collect(keys(normalized_cs_results))::Vector{Int}
+    fullcolumn_ops_results_ordered = map(
+        x -> normalized_cs_results[x], fullcolumn_ops_indices_in_normalized_cs
+    )::Union{Vector{Any},Vector{Dagger.EagerThunk}}
 
     colfragment = (column, s, e) -> Dagger.@spawn getindex(column, s:e)
-    clenghts = chunk_lengths_of_original_dt
-    result_column_symbols = getindex.(Ref(map(x -> x[2][2], normalized_cs)), col_keys_indices)
+    result_column_symbols =
+        getindex.(Ref(map(x -> x[2][2], normalized_cs)), fullcolumn_ops_indices_in_normalized_cs)
 
-    chunks = [
-        begin
-            cfrags = [
-                colfragment(column, 1 + sum(clenghts[1:(i - 1)]), sum(clenghts[1:i])) for
-                column in col_vecs
-            ]
-            Dagger.@spawn f(ch, result_column_symbols, cfrags)
-        end for (i, ch) in enumerate(dt.chunks)
+    dtchunks = if dt === nothing
+        [Dagger.spawn(() -> nothing) for _ in 1:length(new_chunk_lengths)]
+    else
+        dt.chunks
+    end
+    dtchunks_filled = [
+        x <= length(dtchunks) ? dtchunks[x] : Dagger.spawn(() -> nothing) for
+        x in 1:length(new_chunk_lengths)
     ]
 
-    return DTable(chunks, dt.tabletype)
+    chunks = Dagger.EagerThunk[
+        Dagger.spawn(
+            fillcolumn,
+            chunk,
+            result_column_symbols,
+            [
+                if len > 1
+                    colfragment(
+                        column, 1 + sum(new_chunk_lengths[1:(i - 1)]), sum(new_chunk_lengths[1:i])
+                    )
+                else
+                    column
+                end for
+                (column, len) in zip(fullcolumn_ops_results_ordered, fullcolumn_ops_result_lengths)
+            ],
+            new_chunk_length,
+            normalized_cs,
+        ) for
+        (i, (chunk, new_chunk_length)) in enumerate(zip(dtchunks_filled, new_chunk_lengths)) if
+        new_chunk_length > 0
+    ]
+    if dt === nothing
+        return DTable(chunks, nothing)
+    else
+        return DTable(chunks, dt.tabletype)
+    end
 end
diff --git a/test/table_dataframes.jl b/test/table_dataframes.jl
