Merge pull request #88 from VEZY/optimizations

Optimize traversal and queries

Author: VEZY

benchmark/benchmarks.jl

@@ -6,6 +6,8 @@ using Tables
 const SUITE = BenchmarkGroup()
 const HAS_EXPLICIT_ATTRIBUTE_API = isdefined(MultiScaleTreeGraph, :attribute) && isdefined(MultiScaleTreeGraph, :attribute!)
 const HAS_TABLE_VIEWS_API = isdefined(MultiScaleTreeGraph, :symbol_table) && isdefined(MultiScaleTreeGraph, :mtg_table)
+const HAS_ATTRIBUTE_POSITIONAL_DEFAULT = HAS_EXPLICIT_ATTRIBUTE_API &&
+                                         hasmethod(attribute, Tuple{MultiScaleTreeGraph.Node,Symbol,Any})
 const DEFAULT_ATTR_KEY_IS_SYMBOL = true # always true
 
 const SIZE_TIERS = (
@@ -37,6 +39,15 @@ end
     return sym_is_symbol ? syms : Tuple(String(s) for s in syms)
 end
 
+@inline function _attribute_default(node, key, default)
+    HAS_EXPLICIT_ATTRIBUTE_API || return node[key]
+    if HAS_ATTRIBUTE_POSITIONAL_DEFAULT
+        return attribute(node, key, default)
+    else
+        return attribute(node, key; default=default)
+    end
+end
+
 function synthetic_mtg(; n_nodes::Int=10_000, seed::Int=42)
     rng = MersenneTwister(seed)
     mass_key = _default_attr_key(:mass)
@@ -204,7 +215,7 @@ end
 
 function traverse_update_one_explicit_api!(root, key_mass)
     traverse!(root) do node
-        m = HAS_EXPLICIT_ATTRIBUTE_API ? attribute(node, key_mass, default=0.0) : node[key_mass]
+        m = _attribute_default(node, key_mass, 0.0)
         m === nothing && (m = 0.0)
         HAS_EXPLICIT_ATTRIBUTE_API ? attribute!(node, key_mass, m + 0.1) : (node[key_mass] = m + 0.1)
     end
@@ -225,8 +236,8 @@ end
 
 function traverse_update_multi_leaf_explicit_api!(root, key_width, key_area, symbol_leaf)
     traverse!(root, symbol=symbol_leaf) do node
-        width = HAS_EXPLICIT_ATTRIBUTE_API ? attribute(node, key_width, default=0.0) : node[key_width]
-        area = HAS_EXPLICIT_ATTRIBUTE_API ? attribute(node, key_area, default=0.0) : node[key_area]
+        width = _attribute_default(node, key_width, 0.0)
+        area = _attribute_default(node, key_area, 0.0)
         width === nothing && (width = 0.0)
         area === nothing && (area = 0.0)
         if HAS_EXPLICIT_ATTRIBUTE_API
@@ -254,9 +265,9 @@ end
 
 function traverse_update_multi_mixed_explicit_api!(root, key_mass, key_counter, symbol_leaf_internode)
     traverse!(root, symbol=symbol_leaf_internode) do node
-        m = HAS_EXPLICIT_ATTRIBUTE_API ? attribute(node, key_mass, default=0.0) : node[key_mass]
+        m = _attribute_default(node, key_mass, 0.0)
         m === nothing && (m = 0.0)
-        counter = HAS_EXPLICIT_ATTRIBUTE_API ? attribute(node, key_counter, default=0) : node[key_counter]
+        counter = _attribute_default(node, key_counter, 0)
         isnothing(counter) && (counter = 0)
         if HAS_EXPLICIT_ATTRIBUTE_API
             attribute!(node, key_mass, m * 0.999 + 0.0001)

docs/Project.toml

@@ -1,3 +1,6 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 MultiScaleTreeGraph = "dd4a991b-8a45-4075-bede-262ee62d5583"
+
+[sources]
+MultiScaleTreeGraph = { path = ".." }

docs/src/tutorials/0.read_write.md

@@ -22,18 +22,18 @@ The file given in input can be either a `.mtg`, `.csv`, `.xlsx` or `.xlsm` file.
 
 ### Options
 
-The function has two optional arguments to set the type used for the attributes, and the type used for the MTG field (see next section for more details). It also has a keyword argument to choose the sheet name in case you're reading an `xlsx` or `xlsm` file.
+The function has one optional argument to set the type used for the MTG field (see next section for more details). It also has a keyword argument to choose the sheet name in case you're reading an `xlsx` or `xlsm` file.
 
 #### Attributes type
 
-The type used for the attributes should be a `NamedTuple`-alike or a `Dict`-alike type. Here is a more in-depth recommendation, use:
+Attributes are stored using the columnar backend (`ColumnarAttrs`) by default.
 
-- `NamedTuple` if you don't plan to modify the attributes of the MTG, *e.g.* to use them for plotting or computing statistics...
-- `MutableNamedTuple` if you plan to modify the attributes values but not adding new attributes very often, *e.g.* recompute an attribute value...
-- `Dict` or similar (*e.g.* `OrderedDict`) if you plan to heavily modify the attributes, *e.g.* adding/removing attributes a lot
+This backend keeps one typed table per symbol and is optimized for repeated traversal and attribute retrieval. In practice:
 
-!!! note
-    If you don't know what to use, just use the default.
+- users can still create nodes with `Dict`/`NamedTuple`-like attribute inputs;
+- those inputs are converted to the columnar representation when nodes are attached to an MTG.
+
+So for most workflows, there is no extra option to choose here: use the default.
 
 #### MTG encoding type
 

docs/src/tutorials/2.descendants_ancestors_filters.md

@@ -54,7 +54,9 @@ The previous notations are both equivalent to:
 node_attributes(node_5)[:Length]
 ```
 
-But we strongly recommend to avoid this last notation. First, in our case the attributes are stored in a Dictionary (Dict, the default), so we access their values using the Dict notation: `node_attributes(node_5)[:Length]`. But if the attributes are stored as a NamedTuple-alike structure, we must use the dot notation instead: `node_attributes(node_5).Length` (see [Attributes type](@ref) for more details). Second, the attributes of the nodes are not stored in the structure returned by `node_attributes`, so updating the value of an attribute using this notation does not update the value in the MTG. This is a common mistake that can lead to very hard-to-debug issues.
+`node_attributes(node_5)` returns the node attribute container view for this node (columnar backend by default), so reading and writing through it does update the MTG.
+
+However, we still recommend the higher-level APIs (`node[:attr]`, `attribute(node, :attr)`, and `attribute!(node, :attr, value)`) because they are clearer and keep user code backend-agnostic.
 
  That is why the package implements the more generic `node_5[:Length]` notation that works with any structure used for the attributes, which helps develop more generic code.
 

src/compute_MTG/ancestors.jl

@@ -1,3 +1,48 @@
+@inline function _ancestors_values_no_filter_ignore!(node, key::Symbol, val, recursivity_level, ignore_nothing::Bool)
+    current = node
+    remaining = recursivity_level
+
+    while !isroot(current) && remaining != 0
+        parent_ = parent(current)
+        v = unsafe_getindex(parent_, key)
+        ignore_nothing && v === nothing || push!(val, v)
+        remaining -= 1
+        current = parent_
+    end
+    return val
+end
+
+@inline function _ancestors_values_no_filter_collect(node, key::Symbol, recursivity_level, ignore_nothing::Bool)
+    vals = Any[]
+    current = node
+    remaining = recursivity_level
+
+    while !isroot(current) && remaining != 0
+        parent_ = parent(current)
+        v = unsafe_getindex(parent_, key)
+        ignore_nothing && v === nothing || push!(vals, v)
+        remaining -= 1
+        current = parent_
+    end
+    return vals
+end
+
+@inline function _typed_from_any(vals::Vector{Any}, fallback_type=Any)
+    if isempty(vals)
+        return Array{fallback_type,1}()
+    end
+    T = typeof(vals[1])
+    @inbounds for i in 2:length(vals)
+        T = Union{T,typeof(vals[i])}
+    end
+    out = Vector{T}(undef, length(vals))
+    @inbounds for i in eachindex(vals)
+        out[i] = vals[i]
+    end
+    return out
+end
+
+
 """
     ancestors(node::Node,key,<keyword arguments>)
     ancestors(node::Node,<keyword arguments>)
@@ -60,6 +105,13 @@ function ancestors(
     recursivity_level=-1,
     ignore_nothing=false,
     type::Union{Union,DataType}=Any)
+    if no_node_filters(scale, symbol, link, filter_fun) && all && !self && type === Any
+        key_ = Symbol(key)
+        vals = _ancestors_values_no_filter_collect(node, key_, recursivity_level, ignore_nothing)
+        fallback_T = infer_columnar_attr_type(node, key_, nothing, ignore_nothing)
+        return _typed_from_any(vals, fallback_T)
+    end
+
     symbol = normalize_symbol_filter(symbol)
     link = normalize_link_filter(link)
     _maybe_depwarn_traversal_type_kw(:ancestors, type)
@@ -216,6 +268,13 @@ function ancestors!(
     ignore_nothing=false,
     type::Union{Union,DataType}=Any,
 )
+    if no_node_filters(scale, symbol, link, filter_fun) && all && !self && type === Any
+        key_ = Symbol(key)
+        empty!(out)
+        _ancestors_values_no_filter_ignore!(node, key_, out, recursivity_level, ignore_nothing)
+        return out
+    end
+
     symbol = normalize_symbol_filter(symbol)
     link = normalize_link_filter(link)
     _maybe_depwarn_traversal_type_kw(:ancestors!, type)

src/compute_MTG/descendants.jl

@@ -26,6 +26,33 @@ end
     return mask
 end
 
+@inline function _single_bucket_id(store::MTGAttributeStore, symbol_filter)
+    symbol_filter === nothing && return 0
+    bids = _symbol_bucket_ids(store, symbol_filter)
+    return length(bids) == 1 ? bids[1] : 0
+end
+
+@inline function _collect_descendant_values_indexed_single_bucket!(
+    out::AbstractVector,
+    idx::SubtreeIndexCache,
+    store::MTGAttributeStore,
+    bid::Int,
+    col::Column{T},
+    left::Int,
+    right::Int,
+    ignore_nothing::Bool,
+) where {T}
+    sizehint!(out, length(out) + (right - left + 1))
+    @inbounds for i in left:right
+        nid = idx.dfs_order[i]
+        store.node_bucket[nid] == bid || continue
+        v = col.data[store.node_row[nid]]
+        ignore_nothing && v === nothing && continue
+        push!(out, v)
+    end
+    return true
+end
+
 function _collect_descendant_values_indexed!(
     out::AbstractVector,
     node,
@@ -49,6 +76,34 @@ function _collect_descendant_values_indexed!(
     left > right && return true
 
     allow_mask = _bucket_allow_mask(store, symbol_filter)
+    single_bid = _single_bucket_id(store, symbol_filter)
+
+    if single_bid != 0
+        col_idx = key_plan.col_idx_by_bucket[single_bid]
+        if col_idx == 0
+            ignore_nothing && return true
+            sizehint!(out, length(out) + (right - left + 1))
+            @inbounds for i in left:right
+                nid = idx.dfs_order[i]
+                store.node_bucket[nid] == single_bid || continue
+                push!(out, nothing)
+            end
+            return true
+        end
+
+        col = store.buckets[single_bid].columns[col_idx]
+        return _collect_descendant_values_indexed_single_bucket!(
+            out,
+            idx,
+            store,
+            single_bid,
+            col,
+            left,
+            right,
+            ignore_nothing,
+        )
+    end
+
     sizehint!(out, length(out) + (right - left + 1))
 
     @inbounds for i in left:right