Add content

Author: jakobnissen

Project.toml

@@ -1,16 +1,14 @@
 name = "KWayMerges"
 uuid = "f29e91c7-719d-4dbc-8870-0ce36bf055b7"
-authors = ["Jakob Nybo Nissen <[email protected]>"]
 version = "0.1.0"
+authors = ["Jakob Nybo Nissen <[email protected]>"]
 
-[deps]
+[compat]
+Test = "1.11"
+julia = "1.11"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
-[compat]
-Test = "1.10"
-julia = "1.10"
-
 [targets]
 test = ["Test"]

README.md

@@ -7,10 +7,17 @@
 
 Implementation of k-way merge.
 
-This package implements `kway_merge(x::AbstractVector{<:AbstractVector})`, which returns a `KWayMerger`.
-This type is a lazy iterator of the elements in the inner vectors. If the inner vectors are sorted, the output of the `KWayMerger` is also guaranteed to be sorted.
+This package implements the `KWayMerger` type.
+It is a stateful, lazy iterator of the elements in an iterator of iterators, similar to `Iterators.flatten`. However, the elements of the inner iterators will be yielded in an order given by a predicate optionally passed to `KWayMerger` (default: `isless`).
+If the inner iterators are sorted by the predicate, the output of the `KWayMerger` is also guaranteed to be sorted.
 
-The function `peek` can be used to check the next element without advancing the iterator. 
+The primary purpose of `KWayMerger` is to efficiently merge N sorted iterables into one sorted stream.
+
+The function `peek` can be used to check the next element without advancing the iterator.
+
+## Documentation
+This package's two public functions are the `KWayMerger` constructor, and its `Base.peek` method.
+See their docstrings for more details.
 
 ## Contributing
 We appreciate contributions from users including reporting bugs, fixing

docs/Project.toml

@@ -1,43 +1,161 @@
 module KWayMerges
 
-export my_function, MyType
+export KWayMerger
+
+include("heap.jl")
 
 """
-    MyStruct{T <: Real}
+    KWayMerger{F, T, I, S}
+
+Stateful iterator which does a K-way merge between multiple
+iterators of the same type.
+This iterator will yield the elements in every contained
+iterator. At each iteration, it will choose the element from
+the iterator with the lowest precedence according to the order
+determined by `f::F` (default: `isless`).
 
-This struct represents a point in 2D-space.
+If all inner iterators are sorted by `f`, the yielded elements
+will be in sorted order.
+A `KWayMerger` is typically used to combined multiple sorted arrays
+into one sorted array.
 
 # Examples
 ```jldoctest
-julia> x = MyType(1.0, 2.0)
-MyType{Float64}(1.0, 2.0)
+julia> arrs = [[1,4,10], [2,3], [6,8], [5,7,9]];
+
+julia> it = KWayMerger(arrs);
 
-julia> my_function(x)
-2.23606797749979
+julia> print(collect(it))
+[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
 ```
 
-See also: [`my_function`](@ref)
+# Extended help
+The type parameters are:
+* `F`: Type of function used to compare the elements. It defaults
+  to `typeof(Base.isless)`
+* `T`: Element type of iterators
+* `I`: Iterator type
+* `S`: Type of state of iterators
+
+All iterators must be of the same type. `Base.eltype` is used
+to determine `T` and `I`; since its default implementation
+returns `Any`, these type parameters might need to be explicitly
+passed to get good performance for some iterators.
+
+`S` is derived automatically, but this must be a fixed type;
+iterators that use states of multiple different types may
+not be supported by `KWayMerger`.
 """
-struct MyType{T <: Real}
-    x::T
-    y::T
-end 
+struct KWayMerger{F, T, I, S}
+    f::F
+    iterators::Memory{I}
+    states::Memory{S}
+    heap::Vector{Tuple{T, Int}}
+end
+
+ord(x::KWayMerger) = ord(x.f)
+ord(f) = (i, j) -> f(first(i), first(j))
+
+collect_to_memory(x) = collect_to_memory(Base.IteratorSize(typeof(x)), x)
+
+function collect_to_memory(::Union{Base.HasShape, Base.HasLength}, x)
+    mem = Memory{eltype(x)}(undef, length(x))
+    i = 0
+    for it in x
+        i += 1
+        mem[i] = it
+    end
+    i == length(mem) || error("Implementation error: More elements than reported")
+    return mem
+end
+
+function collect_to_memory(::Base.SizeUnknown, x)
+    T = eltype(typeof(x))
+    v = collect(T, x)
+    return copy!(Memory{T}(undef, length(v)), v)
+end
+
+function KWayMerger{F, T, I}(f::F, iterators) where {F, T, I}
+    iters = collect_to_memory(iterators)
+    states = nothing
+    things = Tuple{T, Int}[]
+    for i in eachindex(iters)
+        it = iterate(iters[i])
+        isnothing(it) && continue
+        (thing::T, state) = it
+        if isnothing(states)
+            states = Memory{typeof(state)}(undef, length(iters))
+        end
+        push!(things, (thing, i))
+        states[i] = state
+    end
+    heapify!(ord(f), things)
+    states = if isnothing(states)
+        Memory{Union{}}(undef, length(iters))
+    else
+        states
+    end
+    return KWayMerger{F, T, I, eltype(states)}(f, iters, states, things)
+end
+
+function KWayMerger{T, I}(iterators) where {T, I}
+    return KWayMerger{typeof(isless), T, I}(isless, iterators)
+end
+
+KWayMerger(iterators) = KWayMerger(isless, iterators)
+
+function KWayMerger(f::F, iterators) where {F}
+    I = eltype(typeof(iterators))
+    T = eltype(I)
+    return KWayMerger{F, T, I}(f, iterators)
+end
+
+# We could technically know this, but KWayMerger is stateful, and
+# Julia's iterator length works badly with stateful iterators.
+Base.IteratorSize(::Type{<:KWayMerger}) = Base.SizeUnknown()
+Base.eltype(::Type{<:KWayMerger{F, T}}) where {F, T} = T
+
+function Base.iterate(x::KWayMerger, ::Nothing = nothing)
+    isempty(x.heap) && return nothing
+    (item, i) = @inbounds x.heap[1]
+    iterator = @inbounds x.iterators[i]
+    state = @inbounds x.states[i]
+    it = iterate(iterator, state)
+    if it === nothing
+        @inbounds heappop!(ord(x.f), x.heap)
+    else
+        (new_item, new_state) = it
+        @inbounds x.states[i] = new_state
+        @inbounds heapreplace!(ord(x.f), x.heap, (new_item, i))
+    end
+    return (item, nothing)
+end
+
+Base.isempty(x::KWayMerger) = isempty(x.heap)
+Base.isdone(x::KWayMerger) = isempty(x.heap)
 
 """
-    my_function(x::MyType) -> Real
+    peek(x::KWayMerger)::Union{Some, Nothing}
+
+Get the first element of `x` without advancing the iterator, or `nothing` if the
+iterator is empty.
 
-Computes the L2 norm of `x`.
 # Examples
 ```jldoctest
-julia> x = MyType(1.0, 2.0)
-MyType{Float64}(1.0, 2.0)
+julia> it = KWayMerger([[3, 4], [2, 7]]);
 
-julia> my_function(x)
-2.23606797749979
-```
+julia> peek(it)
+Some(2)
 
-See also: [`MyType`](@ref)
+julia> collect(it); # empty stateful iterator
+
+julia> peek(it) === nothing
+true
+```
 """
-my_function(x::MyType) = sqrt(x.x^2 + x.y^2)
+function Base.peek(x::KWayMerger{F, T})::Union{Some{T}, Nothing} where {F, T}
+    isempty(x.heap) && return nothing
+    return @inbounds Some(x.heap[1][1])
+end
 
 end # module KWayMerges

docs/make.jl

@@ -0,0 +1,48 @@
+@inline function heapify!(f, xs::Vector)
+    for i in div(length(xs), 2):-1:1
+        percolate_down!(f, xs, i, xs[i])
+    end
+    return xs
+end
+
+@inline function percolate_down!(
+        f::F,
+        xs::Vector,
+        i::Integer,
+        x,
+    ) where {F}
+    len = length(xs)
+    @inbounds while (l = 2i) <= len
+        r = 2i + 1
+        j = if r > len || f(@inbounds(xs[l]), @inbounds(xs[r]))
+            l
+        else
+            r
+        end
+        if f(@inbounds(xs[j]), x)
+            @inbounds xs[i] = xs[j]
+            i = j
+        else
+            break
+        end
+    end
+    return @inbounds xs[i] = x
+end
+
+@noinline function heappop!(f, xs::Vector)
+    isempty(xs) && throw(BoundsError(xs, 1))
+    x = @inbounds xs[1]
+    y = @inbounds pop!(xs)
+    if !isempty(xs)
+        percolate_down!(f, xs, 1, y)
+    end
+    return x
+end
+
+@inline function heapreplace!(f, xs::Vector, x)
+    @boundscheck isempty(xs) && throw(BoundsError(xs, 1))
+    res = @inbounds xs[1]
+    @inbounds xs[1] = x
+    percolate_down!(f, xs, 1, x)
+    return res
+end

docs/src/index.md

@@ -1,12 +1,53 @@
 module KWayMergesTest
 
 using KWayMerges
-
 using Test
 
-@testset "Approximation" begin
-    x = MyType(3, 4)
-    @test isapprox(my_function(x), 5)
+@testset "Construction" begin
+    v = [1:3, 4:6, 9:11]
+
+    T = KWayMerger{typeof(isless), Int, UnitRange{Int}, Int}
+
+    @test KWayMerger{typeof(isless), Int, UnitRange{Int}}(isless, v) isa T
+    @test KWayMerger{Int, UnitRange{Int}}(v) isa T
+    @test KWayMerger(v) isa T
+    @test KWayMerger(isless, v) isa T
+    @test KWayMerger(<, v) isa KWayMerger{typeof(<), Int, UnitRange{Int}}
+end
+
+@testset "Forward sorting" begin
+    v = [[3, 5, 8], [1, 1], [10, 11], [1, 2, 7], Int[]]
+    r = sort(collect(Iterators.flatten(v)))
+
+    @test collect(KWayMerger(v)) == r
+end
+
+@testset "Using a predicate" begin
+    v = [["de", "abc"], [""], ["xysa", "dsakljdwe"]]
+    r = sort(collect(Iterators.flatten(v)); by = length)
+
+    @test collect(KWayMerger((i, j) -> isless(length(i), length(j)), v)) == r
+end
+
+@testset "Reverse sorting" begin
+    v = [[3, 5, 8], [1, 1], [10, 11], [1, 2, 7], Int[]]
+    for i in v
+        sort!(i; rev = true)
+    end
+    r = sort(collect(Iterators.flatten(v)); rev = true)
+
+    @test collect(KWayMerger((i, j) -> isless(j, i), v)) == r
+end
+
+@testset "Some edge cases" begin
+    it = KWayMerger([])
+    @test it isa KWayMerger{typeof(isless)}
+    @test collect(it) == Any[]
+
+    it = KWayMerger([1:0, 11:10])
+    @test it isa KWayMerger{typeof(isless), Int}
+    @test collect(it) == Int[]
+    @test typeof(collect(it)) == Vector{Int}
 end
 
 end # module KWayMergesTest