streaming / low memory groupby in julia

Question

I am working on a groupby-aggregation function that will work without RAM overflow issues. Essentially, I want it to run as fast as possible, while not necessarily loading the entire data structure into RAM. The best solution would load as much as possible into RAM for increasing performance/speed, and then offload parts onto disk if required by hardware constraints.

I believe this is the exact use case from Memory mapping, and Julia has the wonderful ability to external sorting with Memory maps. Furthermore, groupby-aggregation operations should work well this way also (sort->streamfromdisk->groupby).

This is what I have here, but it seems a bit clunky. Channels would be wonderful to split one processor for simply disk IO, and then another for any of the processing of aggregation functions etc. Does anyone else have any good suggestions?

function lowmem_groupby(;mmap_fpath::AbstractString,
                        dtype::DataType,
                        groupby_symbols::Vector{Symbol},
                        out_mmap_fpath::AbstractString,
                        agg_funcs::Dict)

  # Example of agg_funcs:
  #   function cnt(x::AbstractArray)::UInt16
  #     return length(x)
  #   end
  #   function wt_med(x::AbstractArray)::Float16
  #     return median(x)
  #   end
  #   agg_funcs = Dict(count=>cnt, weight=>wt_med)

  # This Meta.parse script should produce a parameterized struct like this:
  # struct OutDType
  #   [ groupby_symbols :: groupby_fieldtypes ]
  #   foo::String
  #   bar::Int
  #   [ agg_funcs_symbols :: agg_funcs ]
  #   count::UInt16
  #   weight::Float16
  # end
  grpsymstr = join(["$(n)::$(t)" for (n, t) in zip(fieldnames(dtype), fieldtypes(dtype))
                    if n in groupby_symbols], "\n")
  grpsaggstr = join(["$(String(n))::$(Base.return_types(f)[1])\n" for (n, f) in agg_funcs], "\n")
  eval(Meta.parse("""
    struct OutDType
      $(grpsymstr)
      $(grpsaggstr)
    end
  """))


  function sort_func(a, b)
    # Get properties of the element type (a or b)
    # Make sure the groupby properties come first for sorting
    props = [groupby_symbols;
            [i for i in fieldnames(typeof(a))
             if i ∉ groupby_symbols]]
    # Property tuple generator for element
    prop_tuple_gen(el) = (getproperty(el, p) for p in props)
    return isless(prop_tuple_gen(a), prop_tuple_gen(b))
  end

  function props_eq(a, b)
    aprops = [getproperty(a, x) for x in groupby_symbols]
    bprops = [getproperty(b, x) for x in groupby_symbols]
    return aprops == bprops
  end

  open(mmap_fpath, "r+") do input_file
    Mmap.mmap(input_file, dtype) do mmap_arr
      # External sort
      sort!(mmap_arr, lt = sort_func(a, b))

      open(out_mmap_fpath, "w+") do output_file
        Mmap.mmap(output_file, out_dtype) do vec
          i = 1
          # Read while there is still lines
          while i<length(vec)
            groupitem = OutDType()
            grp = vec[i] # first line of the group
            grpitems = [grp]
            j = i # group index
            # s = 1 # count
            while (i<length(vec) && props_eq(vec[i], vec[i+1]) )
              # s += 1
              grpitems += [vec[i]]
              i += 1
            end
            for symbl in groupby_symbols
              setproperty!(groupitem, symbl, getproperty(grp, symbl))
            end
            for (symbl, func) in agg_funcs
              setproperty!(groupitem, symbl, func(grpitems))
            end

            write(output_file, groupitem)
            i += 1
          end
        end
      end
      # mmap_arr = nothing
      # close(input_file)
    end
  end
  GC.gc()
end

Answer 1

I'd separate the Mmap part from the actual streaming group operations. Here's how I'd write the latter, with some generated functions and Val instead of evaling a type. The output is just a tuple of two named tuples: the group values, and the aggregation values.

@generated function grouped_fields(x::T, ::Val{G}) where {T, G}
    fns = fieldnames(T)
    grouped_fns = G
    accessors = [:(x.$f) for f in grouped_fns]
    return :((;$(accessors...)))
end

@generated function ungrouped_fields(x::T, ::Val{G}) where {T, G}
    fns = fieldnames(T)
    grouped_fns = G
    accessors = [:(x.$f) for f in filter(∉(grouped_fns), fns)]
    return :((;$(accessors...)))
end

@generated function apply_aggregation(aggregators::NamedTuple{NT, FT}, xs) where {NT, FT}
    calls = [:(aggregators.$n(xs)) for n in NT]
    return :(NamedTuple{NT}(($(calls...),)))
end

function groupby!(output, input::AbstractArray{T}, groups::Val{G}, aggregators::NamedTuple) where {T, G}
    function grouped_isless(a, b)
        a_fields = (grouped_fields(a, groups), ungrouped_fields(a, groups))
        b_fields = (grouped_fields(b, groups), ungrouped_fields(b, groups))
        return isless(a_fields, b_fields)
    end

    function grouped_isequal(a, b)
        return isequal(grouped_fields(a, groups), grouped_fields(b, groups))
    end
    
    sort!(input, lt=grouped_isless)

    input_iterator = Iterators.Stateful(input)
    while !isempty(input_iterator)
        first, rest = Iterators.peel(input_iterator)
        group = append!([first], Iterators.takewhile(b -> grouped_isequal(first, b), input_iterator))
        group_fields = grouped_fields(first, groups)
        aggregrations = apply_aggregation(aggregators, group)
        push!(output, (group_fields, aggregrations))
    end

    return output
end

This is already pretty close to streaming code. If you assume the input iterator is already sorted, and use some package providing yield or another way to also make the output non-allocating (channels, maybe), you're there.

Example:

julia> o = Any[]; groupby!(o, testdata, Val{(:bar,)}(), (;mean = g -> mean(getproperty.(g, :foo))))
2-element Vector{Any}:
 ((bar = "a",), (mean = 5.0,))
 ((bar = "b",), (mean = 7.0,))