From 92766f75120b5e6ecfeee4d9478732788626f262 Mon Sep 17 00:00:00 2001 From: Kristoffer Date: Wed, 22 Nov 2023 10:06:47 +0100 Subject: [PATCH] use ArrayOfArrays for return value to reduce the number of allocated arrays A long standing gripe for me has been that indices and distances are returned as standard nested `Array`s. Typically, each inner array hold quite a small number of neighbors so it means that we allocate a large number of small arrays. Using ArrayOfArrays, these are stored contigously in one large flat array instead. The difference in allocations can be readily seen: ```julia julia> input = rand(3, 10^6); julia> tree = KDTree(rand(3, 10^6)); julia> @time knn(tree, input, 5); 1.538003 seconds (2.00 M allocations: 221.253 MiB, 10.03% gc time) julia> @time knn(tree, input, 5); 1.489310 seconds (98 allocations: 189.884 MiB, 0.29% gc time) ``` --- Project.toml | 2 ++ src/NearestNeighbors.jl | 1 + src/inrange.jl | 14 ++++++++++---- src/knn.jl | 25 +++++++++++++++++++------ test/test_knn.jl | 4 ++-- 5 files changed, 34 insertions(+), 12 deletions(-) diff --git a/Project.toml b/Project.toml index d0071f5..9d71bda 100644 --- a/Project.toml +++ b/Project.toml @@ -3,6 +3,7 @@ uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce" version = "0.4.18" [deps] +ArraysOfArrays = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -10,6 +11,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Distances = "0.9, 0.10" StaticArrays = "0.9, 0.10, 0.11, 0.12, 1.0" julia = "1.6" +ArraysOfArrays = "0.6" [extras] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/NearestNeighbors.jl b/src/NearestNeighbors.jl index 28d998d..aa5772e 100644 --- a/src/NearestNeighbors.jl +++ b/src/NearestNeighbors.jl @@ -4,6 +4,7 @@ using Distances import Distances: PreMetric, Metric, result_type, eval_reduce, eval_end, eval_op, eval_start, evaluate, parameters using StaticArrays +using ArraysOfArrays import Base.show export NNTree, BruteTree, KDTree, BallTree, DataFreeTree diff --git a/src/inrange.jl b/src/inrange.jl index d271639..52a306c 100644 --- a/src/inrange.jl +++ b/src/inrange.jl @@ -15,10 +15,13 @@ function inrange(tree::NNTree, check_input(tree, points) check_radius(radius) - idxs = [Vector{Int}() for _ in 1:length(points)] + idxs = VectorOfArrays{Int, 1}() + idx = Int[] for i in 1:length(points) - inrange_point!(tree, points[i], radius, sortres, idxs[i]) + inrange_point!(tree, points[i], radius, sortres, idx) + push!(idxs, idx) + resize!(idx, 0) end return idxs end @@ -66,11 +69,14 @@ function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Numb check_input(tree, points) check_radius(radius) n_points = size(points, 2) - idxs = [Vector{Int}() for _ in 1:n_points] + idxs = VectorOfArrays{Int, 1}() + idx = Int[] for i in 1:n_points point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim))) - inrange_point!(tree, point, radius, sortres, idxs[i]) + inrange_point!(tree, point, radius, sortres, idx) + push!(idxs, idx) + resize!(idx, 0) end return idxs end diff --git a/src/knn.jl b/src/knn.jl index 775f7eb..1136179 100644 --- a/src/knn.jl +++ b/src/knn.jl @@ -18,10 +18,17 @@ function knn(tree::NNTree{V}, points::AbstractVector{T}, k::Int, sortres=false, check_input(tree, points) check_k(tree, k) n_points = length(points) - dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points] - idxs = [Vector{Int}(undef, k) for _ in 1:n_points] + dists = VectorOfArrays{get_T(eltype(V)), 1}() + idxs = VectorOfArrays{Int, 1}() + dist = zeros(get_T(eltype(V)), k) + idx = zeros(Int, k) + for i in 1:n_points - knn_point!(tree, points[i], sortres, dists[i], idxs[i], skip) + knn_point!(tree, points[i], sortres, dist, idx, skip) + push!(dists, dist) + push!(idxs, idx) + fill!(dist, 0) + fill!(idx, 0) end return idxs, dists end @@ -77,12 +84,18 @@ function knn_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, ::Val{di check_input(tree, points) check_k(tree, k) n_points = size(points, 2) - dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points] - idxs = [Vector{Int}(undef, k) for _ in 1:n_points] + dists = VectorOfArrays{Float64, 1}() + idxs = VectorOfArrays{Int, 1}() + dist = zeros(Float64, k) + idx = zeros(Int, k) for i in 1:n_points point = SVector{dim,T}(ntuple(j -> points[j, i], Val(dim))) - knn_point!(tree, point, sortres, dists[i], idxs[i], skip) + knn_point!(tree, point, sortres, dist, idx, skip) + push!(dists, dist) + push!(idxs, idx) + fill!(dist, 0) + fill!(idx, 0) end return idxs, dists end diff --git a/test/test_knn.jl b/test/test_knn.jl index 99d96a5..c8b923c 100644 --- a/test/test_knn.jl +++ b/test/test_knn.jl @@ -133,8 +133,8 @@ end points = rand(SVector{3, Float64}, 100) kdtree = KDTree(points) idxs, dists = knn(kdtree, view(points, 1:10), 3) - @test idxs isa Vector{Vector{Int}} - @test dists isa Vector{Vector{Float64}} + @test eltype(idxs) <: AbstractVector{Int} + @test eltype(dists) <: AbstractVector{Float64} end @testset "mutating" begin