Is there any way to find K-nearest neighbors allocation free?

It turns out this is incomplete, but I'll post my attempt nevertheless.

"Reinterpreting" a matrix into a vector of vectors without allocation is conceptually simple, but requires implementing a new array type. Such a type is provided by JuliennedArrays.jl's Sliced.

The simplest implementation IHMO is this one:

mapslices(y, dims=2) do row
    partialsortperm(Slices(x, 2), 1:5, by=x -> norm(x - row))
end 

Which still allocates something; this must at least be the indexing vector used by partialsortperm and the intermediate rows.

I tried to get rid of that in this function:

function knnslice!(result, x, y, k)
    result_sliced = Slices(result, 2)
    x_sliced = Slices(x, 2)
    y_sliced = Slices(y, 2)
    indices = collect(axes(x, 1))
    for i in eachindex(result_sliced, y_sliced)
        result_sliced[i] .= partialsortperm!(indices, x_sliced, 1:k, by=x -> norm(x - y_sliced[i]))
    end
    return result
end
knnslice(x, y, k) = knnslice!(similar(x, Int, size(y, 1), k), x, y, k)

But the results are hardly an improvement, at least when compared on arrays of the size of your example data. I'm not sure how this can be brought down any more with this kind of implementation.

The missing piece would be a sortperm implementation which directly works on slices. For small k, this should be doable with one iteration over x and maintaining the results row as buffer (or even small heap) of that size instead of performing a partial sort. Something like:

function knnslice!(result, x, y, k)
    for (i_r, i_y) in zip(axes(result, 1), axes(y, 1))
        result_row = @view(result[i_r, :])
        fill!(result_row, 1)
        f(r) = norm(@view(x[r, :]) - @view(y[i_y, :]))
        for j_x in axes(x, 1)
            heappush!(result_row, j_x; by=f)
        end
    end
    return result
end

where heappush! should be insertion into a bounded minimum heap with order by (something like what heapq in Python provides, but keeping the size of the queue fixed).