From 65d860b6e00fd5870c5745600ec1c4f4524ce698 Mon Sep 17 00:00:00 2001 From: Maarten Plonk <114733774+repro-code@users.noreply.github.com> Date: Sat, 10 Jan 2026 23:06:32 +0000 Subject: [PATCH] Implement SpatialTreeInterface for packed R-tree Add PackedRTree and PackedRTreeNode types that provide a hierarchical view over the flat-packed R-tree array. This implements the SpatialTreeInterface from GeometryOps.jl, enabling generic spatial tree algorithms. Key changes: - New packedrtree.jl with PackedRTree wrapper types - Implements isspatialtree, isleaf, getchild, nchild, child_indices_extents, node_extent interface methods - Includes generic depth_first_search and query functions - Rewrite filter!/search to use the new interface - Backwards-compatible legacy search function preserved The interface is designed to be generic over node element types, not just NodeItem, allowing reuse for other packed R-tree implementations. Co-Authored-By: Claude Opus 4.5 --- src/FlatGeobuf.jl | 1 + src/index.jl | 71 ++++++++-------- src/packedrtree.jl | 201 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 239 insertions(+), 34 deletions(-) create mode 100644 src/packedrtree.jl diff --git a/src/FlatGeobuf.jl b/src/FlatGeobuf.jl index 33431f5..2b7d451 100644 --- a/src/FlatGeobuf.jl +++ b/src/FlatGeobuf.jl @@ -7,6 +7,7 @@ include("schema/header.jl") include("schema/feature.jl") include("flatgeobuffer.jl") +include("packedrtree.jl") include("index.jl") include("io.jl") include("table.jl") diff --git a/src/index.jl b/src/index.jl index ab717cf..bc5626b 100644 --- a/src/index.jl +++ b/src/index.jl @@ -1,4 +1,5 @@ using Extents + function Base.read(io::IO, ::Type{NodeItem}) NodeItem( Base.read(io, Float64), @@ -8,61 +9,63 @@ function Base.read(io::IO, ::Type{NodeItem}) Base.read(io, UInt64) ) end + Base.convert(::Type{Extent}, node::NodeItem) = Extent(X=(node.min_x, node.max_x), Y=(node.min_y, node.max_y)) Base.convert(::Type{NodeItem}, ex::Extent) = NodeItem(ex.X[1], ex.Y[1], ex.X[2], ex.Y[2], 0) -function intersects(a::NodeItem, b::NodeItem) +""" + intersects(a::Extent, b::Extent) + +Check if two extents intersect (AABB intersection test). +""" +function intersects(a::Extent, b::Extent) !( - a.min_x > b.max_x || - a.min_y > b.max_y || - a.max_x < b.min_x || - a.max_y < b.min_y + a.X[1] > b.X[2] || + a.Y[1] > b.Y[2] || + a.X[2] < b.X[1] || + a.Y[2] < b.Y[1] ) end -function search(nodes::Vector{NodeItem}, bbox::NodeItem, nfeatures::UInt64, node_size::UInt16) - # Queue stores both order, as the nodes - queue = Vector{Tuple{Int64,NodeItem}}([(1, nodes[1])]) - offsets = Vector{Int}() +""" + search(tree::PackedRTree, bbox::Extent) - # First n nodes are part of tree and have offsets to other nodes - # while afterwards the nodes are terminal and the offsets point to the features - n_non_leaves = length(nodes) - nfeatures - while (length(queue) != 0) - (i, node) = pop!(queue) - if intersects(bbox, node) - if i <= n_non_leaves - # Still part of tree, add all child nodes to queue - leaves = Int(node.offset + 1):Int(node.offset + 1)+Int(node_size)-1 - append!(queue, zip(leaves, nodes[leaves])) - else - # Or terminal node - push!(offsets, node.offset) # Depth first traversal - end - end - end - offsets +Search the packed R-tree for all features intersecting the bounding box. +Returns a vector of feature offsets. + +This uses the SpatialTreeInterface's depth_first_search internally. +""" +function search(tree::PackedRTree, bbox::Extent) + predicate = ext -> intersects(ext, bbox) + query(tree, predicate) end +# Legacy interface for backwards compatibility +function search(nodes::Vector{NodeItem}, bbox::NodeItem, nfeatures::UInt64, node_size::UInt16) + tree = PackedRTree(nodes, node_size, nfeatures) + bbox_extent = convert(Extent, bbox) + search(tree, bbox_extent) +end """Find all features within a given bounding box.""" function Base.filter(fgb::FlatGeobuffer, bboxv::Vector{<:Real}) - bbox = NodeItem(bboxv[1], bboxv[2], bboxv[3], bboxv[4], 0) - results = search(fgb.rtree, bbox, fgb.header.features_count, fgb.header.index_node_size) + bbox = Extent(X=(bboxv[1], bboxv[3]), Y=(bboxv[2], bboxv[4])) + tree = PackedRTree(fgb) + results = search(tree, bbox) # Results has offsets into file, but we already parsed the file # So we use the offsets to find the relative location of features - offsets = sort(map(x -> x.offset, nodes[end-nfeatures+1:end])) - fgb.features[findfirst.(isequal.(results), Ref(offsets))] + leaf_offsets = sort(map(x -> Int(x.offset), fgb.rtree[end-Int(fgb.header.features_count)+1:end])) + fgb.features[findfirst.(isequal.(results), Ref(leaf_offsets))] end function Base.filter!(fgb::FlatGeobuffer, bboxv::Vector{<:Real}) - bbox = NodeItem(bboxv[1], bboxv[2], bboxv[3], bboxv[4], 0) + bbox = Extent(X=(bboxv[1], bboxv[3]), Y=(bboxv[2], bboxv[4])) Base.filter!(fgb, bbox) end -Base.filter!(fgb::FlatGeobuffer, ex::Extent) = Base.filter!(fgb, convert(NodeItem, ex)) -function Base.filter!(fgb::FlatGeobuffer, bbox::NodeItem) - fgb.offsets = search(fgb.rtree, bbox, fgb.header.features_count, fgb.header.index_node_size) +Base.filter!(fgb::FlatGeobuffer, ex::Extent) = begin + tree = PackedRTree(fgb) + fgb.offsets = search(tree, ex) fgb.filtered = true fgb end diff --git a/src/packedrtree.jl b/src/packedrtree.jl new file mode 100644 index 0000000..76c2af0 --- /dev/null +++ b/src/packedrtree.jl @@ -0,0 +1,201 @@ +# PackedRTree - A view over a flat-packed R-tree array +# Implements the SpatialTreeInterface from GeometryOps.jl + +using Extents + +# Interface trait for packed rtree node elements +# Implement these for your node element type to use PackedRTree + +""" + node_extent(node) + +Return the Extent of a packed rtree node element. +""" +function node_extent end + +""" + node_offset(node) + +Return the offset stored in a packed rtree node element. +For internal nodes, this is the index of the first child. +For leaf nodes, this is the feature identifier/offset. +""" +function node_offset end + +# Implementation for NodeItem +node_extent(node::NodeItem) = Extent(X=(node.min_x, node.max_x), Y=(node.min_y, node.max_y)) +node_offset(node::NodeItem) = node.offset + +""" + PackedRTree{T, V} + +Wrapper around a flat-packed R-tree stored as a contiguous array. +This provides a hierarchical view over the flat array structure. + +The flat array layout follows the FlatGeobuf specification: +- Nodes are stored level by level, root first +- Internal nodes store child offset in their offset field +- Leaf nodes store feature identifiers in their offset field +""" +struct PackedRTree{T, V <: AbstractVector{T}} + nodes::V # The flat array of nodes + node_size::Int # Branching factor (max children per internal node) + n_features::Int # Number of features (= number of leaf nodes) + n_non_leaves::Int # nodes[1:n_non_leaves] are internal nodes +end + +""" + PackedRTree(nodes::AbstractVector, node_size::Integer, n_features::Integer) + +Construct a PackedRTree from a flat array of nodes. +""" +function PackedRTree(nodes::V, node_size::Integer, n_features::Integer) where {T, V <: AbstractVector{T}} + n_non_leaves = length(nodes) - n_features + PackedRTree{T, V}(nodes, Int(node_size), Int(n_features), n_non_leaves) +end + +""" + PackedRTree(fgb::FlatGeobuffer) + +Construct a PackedRTree view over a FlatGeobuffer's spatial index. +""" +function PackedRTree(fgb::FlatGeobuffer) + PackedRTree(fgb.rtree, fgb.header.index_node_size, fgb.header.features_count) +end + +""" + PackedRTreeNode{T, V} + +A view into a single node of a PackedRTree. +Carries reference to tree metadata to compute children and leaf status. +""" +struct PackedRTreeNode{T, V <: AbstractVector{T}} + tree::PackedRTree{T, V} + index::Int # 1-based index into tree.nodes +end + +# Get the underlying node element +@inline nodeitem(node::PackedRTreeNode) = node.tree.nodes[node.index] + +# === SpatialTreeInterface implementation === + +# Predicate to identify spatial trees +isspatialtree(::Type{<:PackedRTree}) = true +isspatialtree(::Type{<:PackedRTreeNode}) = true +isspatialtree(::T) where T = isspatialtree(T) + +# Root node accessor for tree +root(tree::PackedRTree) = PackedRTreeNode(tree, 1) + +# node_extent - bounding box of a node +node_extent(tree::PackedRTree) = node_extent(root(tree)) +node_extent(node::PackedRTreeNode) = node_extent(nodeitem(node)) + +# isleaf - check if node is a leaf (stores feature offsets, not child nodes) +isleaf(tree::PackedRTree) = isleaf(root(tree)) +isleaf(node::PackedRTreeNode) = node.index > node.tree.n_non_leaves + +# nchild - number of children +function nchild(tree::PackedRTree) + isempty(tree.nodes) && return 0 + nchild(root(tree)) +end + +function nchild(node::PackedRTreeNode) + if isleaf(node) + # Leaf nodes have exactly 1 "child" - their feature reference + return 1 + else + # Internal nodes: children start at offset+1, up to node_size children + # But we need to not exceed the array bounds + first_child = Int(node_offset(nodeitem(node))) + 1 + max_children = min(node.tree.node_size, length(node.tree.nodes) - first_child + 1) + return max_children + end +end + +# getchild - get children iterator or specific child +function getchild(tree::PackedRTree) + isempty(tree.nodes) && return PackedRTreeNode{eltype(tree.nodes), typeof(tree.nodes)}[] + getchild(root(tree)) +end +getchild(tree::PackedRTree, i) = getchild(root(tree), i) + +function getchild(node::PackedRTreeNode) + if isleaf(node) + # Leaf node returns itself as the single "child" for iteration + return (node,) + else + tree = node.tree + first_child = Int(node_offset(nodeitem(node))) + 1 + n = nchild(node) + return (PackedRTreeNode(tree, first_child + i - 1) for i in 1:n) + end +end + +function getchild(node::PackedRTreeNode, i) + if isleaf(node) + i == 1 || throw(BoundsError(node, i)) + return node + else + first_child = Int(node_offset(nodeitem(node))) + 1 + return PackedRTreeNode(node.tree, first_child + i - 1) + end +end + +# child_indices_extents - iterator over (index, extent) for leaf children +# For leaf nodes, returns the feature offset and extent +function child_indices_extents(node::PackedRTreeNode) + if !isleaf(node) + error("child_indices_extents can only be called on leaf nodes") + end + item = nodeitem(node) + # Return single (offset, extent) pair + # The offset is the feature identifier/file offset + return ((Int(node_offset(item)), node_extent(item)),) +end + +# === Depth-first search implementation === +# This is a local implementation compatible with SpatialTreeInterface + +""" + depth_first_search(f, predicate, tree::PackedRTree) + depth_first_search(f, predicate, node::PackedRTreeNode) + +Perform depth-first search over the tree, calling `f(index)` for each +leaf node whose extent satisfies `predicate(extent)`. + +Returns nothing; results should be accumulated via the callback `f`. +""" +function depth_first_search(f::F, predicate::P, tree::PackedRTree) where {F, P} + isempty(tree.nodes) && return nothing + depth_first_search(f, predicate, root(tree)) +end + +function depth_first_search(f::F, predicate::P, node::PackedRTreeNode) where {F, P} + if isleaf(node) + for (idx, ext) in child_indices_extents(node) + if predicate(ext) + f(idx) + end + end + else + for child in getchild(node) + if predicate(node_extent(child)) + depth_first_search(f, predicate, child) + end + end + end + return nothing +end + +""" + query(tree::PackedRTree, predicate) + +Return a vector of feature indices/offsets that satisfy the predicate. +""" +function query(tree::PackedRTree, predicate) + results = Int[] + depth_first_search(Base.Fix1(push!, results), predicate, tree) + return results +end