Add error checks for constructing trees and doing queries

Co-authored-by: KristofferC <kristoffer.carlsson@juliacomputing.com>

Author: williamjsdavis

src/ball_tree.jl

@@ -38,6 +38,9 @@ function BallTree(data::AbstractVector{V},
                   parallel::Bool = Threads.nthreads() > 1) where {V <: AbstractArray}
     reorder = !isempty(reorderbuffer) || (storedata ? reorder : false)
 
+    # Reject data containing NaNs early to avoid undefined behaviour later on.
+    check_for_nan(data)
+
     tree_data = TreeData(data, leafsize)
     n_p = length(data)
 

src/brute_tree.jl

@@ -21,6 +21,7 @@ Note: `leafsize` and `reorder` parameters are ignored for BruteTree.
 """
 function BruteTree(data::AbstractVector{V}, metric::PreMetric = Euclidean();
                    reorder::Bool=false, leafsize::Int=0, storedata::Bool=true) where {V <: AbstractVector}
+    check_for_nan(data)
     if metric isa Distances.UnionMetrics
         p = parameters(metric)
         if p !== nothing && length(p) != length(V)
@@ -34,6 +35,7 @@ end
 
 function BruteTree(data::AbstractVecOrMat{T}, metric::PreMetric = Euclidean();
                    reorder::Bool=false, leafsize::Int=0, storedata::Bool=true) where {T}
+    check_for_nan(data)
     dim = size(data, 1)
     BruteTree(copy_svec(T, data, Val(dim)),
               metric, reorder = reorder, leafsize = leafsize, storedata = storedata)

src/inrange.jl

@@ -21,6 +21,7 @@ function inrange(tree::NNTree,
                  sortres=false,
                  skip::F = Returns(false)) where {T <: AbstractVector, F}
     check_input(tree, points)
+    check_for_nan_in_points(points)
     check_radius(radius)
 
     idxs = [Vector{Int}() for _ in 1:length(points)]
@@ -63,6 +64,7 @@ See also: `inrange`, `inrangecount`.
 """
 function inrange!(idxs::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, radius::Number, sortres=false, skip=Returns(false)) where {V, T <: Number}
     check_input(tree, point)
+    check_for_nan_in_points(point)
     check_radius(radius)
     length(idxs) == 0 || throw(ArgumentError("idxs must be empty"))
     inrange_point!(tree, point, radius, sortres, idxs, skip)
@@ -81,6 +83,7 @@ end
 function inrange_matrix(tree::NNTree{V}, points::AbstractMatrix{T}, radius::Number, ::Val{dim}, sortres, skip::F=Returns(false)) where {V, T <: Number, dim, F}
     # TODO: DRY with inrange for AbstractVector
     check_input(tree, points)
+    check_for_nan_in_points(points)
     check_radius(radius)
     n_points = size(points, 2)
     idxs = [Vector{Int}() for _ in 1:n_points]
@@ -107,6 +110,7 @@ Count all the points in the tree which are closer than `radius` to `points`.
 """
 function inrangecount(tree::NNTree{V}, point::AbstractVector{T}, radius::Number, skip::F=Returns(false)) where {V, T <: Number, F}
     check_input(tree, point)
+    check_for_nan_in_points(point)
     check_radius(radius)
     return inrange_point!(tree, point, radius, false, nothing, skip)
 end
@@ -115,11 +119,13 @@ function inrangecount(tree::NNTree,
         points::AbstractVector{T},
         radius::Number, skip::F=Returns(false)) where {T <: AbstractVector, F}
     check_input(tree, points)
+    check_for_nan_in_points(points)
     check_radius(radius)
     return inrange_point!.(Ref(tree), points, radius, false, nothing, skip)
 end
 
 function inrangecount(tree::NNTree{V}, point::AbstractMatrix{T}, radius::Number) where {V, T <: Number}
+    check_for_nan_in_points(point)
     dim = size(point, 1)
     npoints = size(point, 2)
     if isbitstype(T)

src/kd_tree.jl

@@ -35,6 +35,9 @@ function KDTree(data::AbstractVector{V},
                 parallel::Bool = Threads.nthreads() > 1) where {V <: AbstractArray, M <: MinkowskiMetric}
     reorder = !isempty(reorderbuffer) || (storedata ? reorder : false)
 
+    # Reject data containing NaNs early to avoid undefined behaviour later on.
+    check_for_nan(data)
+
     tree_data = TreeData(data, leafsize)
     n_p = length(data)
 

src/knn.jl

@@ -24,6 +24,7 @@ See also: `knn!`, `nn`.
 """
 function knn(tree::NNTree{V}, points::AbstractVector{T}, k::Int, sortres=false, skip::F=Returns(false)) where {V, T <: AbstractVector, F<:Function}
     check_input(tree, points)
+    check_for_nan_in_points(points)
     check_k(tree, k)
     n_points = length(points)
     dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]
@@ -73,6 +74,7 @@ Useful to avoid allocations or specify the element type of the output vectors.
 See also: `knn`, `nn`.
 """
 function knn!(idxs::AbstractVector{<:Integer}, dists::AbstractVector, tree::NNTree{V}, point::AbstractVector{T}, k::Int, sortres=false, skip::F=Returns(false)) where {V, T <: Number, F<:Function}
+    check_for_nan_in_points(point)
     check_k(tree, k)
     length(idxs) == k || throw(ArgumentError("idxs must be of length k"))
     length(dists) == k || throw(ArgumentError("dists must be of length k"))
@@ -91,6 +93,7 @@ function knn(tree::NNTree{V}, points::AbstractMatrix{T}, k::Int, sortres=false,
 
     # TODO: DRY with knn for AbstractVector
     check_input(tree, points)
+    check_for_nan_in_points(points)
     check_k(tree, k)
     n_points = size(points, 2)
     dists = [Vector{get_T(eltype(V))}(undef, k) for _ in 1:n_points]

src/utilities.jl

@@ -91,3 +91,28 @@ end
 # Instead of ReinterpretArray wrapper, copy an array, interpreting it as a vector of SVectors
 copy_svec(::Type{T}, data, ::Val{dim}) where {T, dim} =
         [SVector{dim,T}(ntuple(i -> data[n+i], Val(dim))) for n in 0:dim:(length(data)-1)]::Vector{SVector{dim,T}}
+
+# Check for NaN values in data; throw if any are present
+function check_for_nan(data)
+    @inbounds for p in data
+        if any(isnan, p)
+            throw(ArgumentError("Tree cannot be constructed from data containing NaN values"))
+        end
+    end
+    return
+end
+
+# Check for NaN values in input points; throw if any are present
+function check_for_nan_in_points(points::Union{AbstractVector, AbstractMatrix})
+    if any(isnan, points)
+        throw(ArgumentError("Tree cannot be queried with points containing NaN values"))
+    end
+    return
+end
+
+function check_for_nan_in_points(points::AbstractVector{<:AbstractVector})
+    for p in points
+        check_for_nan_in_points(p)
+    end
+    return
+end

test/runtests.jl

@@ -35,6 +35,7 @@ include("test_monkey.jl")
 include("test_datafreetree.jl")
 include("test_tree_data.jl")
 include("test_periodic.jl")
+include("test_tree_nan.jl")
 
 @testset "views of SVector" begin
     x = [rand(SVector{3}) for i in 1:20]

test/test_tree_nan.jl

@@ -0,0 +1,65 @@
+# Tests for KDTree, BallTree, BruteTree that reject data containing NaNs
+
+@testset "Trees reject NaNs" begin
+    data_vec = [SVector{2,Float64}(NaN, 0.0), SVector{2,Float64}(1.0, 1.0)]
+    data_mat = [NaN 0.0; 1.0 1.0]
+
+    for TreeType in (KDTree, BallTree, BruteTree)
+        @test_throws ArgumentError TreeType(data_vec)
+        @test_throws ArgumentError TreeType(data_mat)
+    end
+end
+
+@testset "knn rejects NaNs" begin
+    for TreeType in (KDTree, BallTree, BruteTree)
+        data = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(1.0, 1.0)]
+        tree = TreeType(data)
+
+        # Single query point (vector) containing NaN
+        @test_throws ArgumentError knn(tree, [NaN, 0.0], 1)
+
+        # Vector-of-vectors query containing NaN
+        query_vec = [SVector{2,Float64}(NaN, 0.0)]
+        @test_throws ArgumentError knn(tree, query_vec, 1)
+
+        # Matrix query containing NaN
+        query_mat = [NaN 0.0; 0.0 1.0]
+        @test_throws ArgumentError knn(tree, query_mat, 1)
+    end
+end
+
+@testset "inrange rejects NaNs" begin
+    for TreeType in (KDTree, BallTree, BruteTree)
+        data = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(1.0, 1.0)]
+        tree = TreeType(data)
+
+        # Single query point (vector) containing NaN
+        @test_throws ArgumentError inrange(tree, [NaN, 0.0], 1.0)
+
+        # Vector-of-vectors query containing NaN
+        query_vec = [SVector{2,Float64}(NaN, 0.0)]
+        @test_throws ArgumentError inrange(tree, query_vec, 1.0)
+
+        # Matrix query containing NaN
+        query_mat = [NaN 0.0; 0.0 1.0]
+        @test_throws ArgumentError inrange(tree, query_mat, 1.0)
+    end
+end
+
+@testset "inrangecount rejects NaNs" begin
+    for TreeType in (KDTree, BallTree, BruteTree)
+        data = [SVector{2,Float64}(0.0, 0.0), SVector{2,Float64}(1.0, 1.0)]
+        tree = TreeType(data)
+
+        # Single query point (vector) containing NaN
+        @test_throws ArgumentError inrangecount(tree, [NaN, 0.0], 1.0)
+
+        # Vector-of-vectors query containing NaN
+        query_vec = [SVector{2,Float64}(NaN, 0.0)]
+        @test_throws ArgumentError inrangecount(tree, query_vec, 1.0)
+
+        # Matrix query containing NaN
+        query_mat = [NaN 0.0; 0.0 1.0]
+        @test_throws ArgumentError inrangecount(tree, query_mat, 1.0)
+    end
+end