tree-ops optimizations for large trees

Project.toml

@@ -1,18 +1,28 @@
 name = "FileTrees"
 uuid = "72696420-646e-6120-6e77-6f6420746567"
-authors = ["Shashi Gowda <gowda@mit.edu>", "Julian Samaroo <jpsamaroo@jpsamaroo.me>"]
 version = "0.4.3"
+authors = ["Shashi Gowda <gowda@mit.edu>", "Julian Samaroo <jpsamaroo@jpsamaroo.me>"]
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
+TaskLocalValues = "ed4db957-447d-4319-bfb6-7fa9ae7ecf34"
+
+[weakdeps]
+Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
+
+[extensions]
+DaggerFileTreesExt = "Dagger"
 
 [compat]
 AbstractTrees = "0.3, 0.4"
 Dagger = "0.19"
 FilePathsBase = "0.9"
 Glob = "1.3"
+OrderedCollections = "1.8"
+TaskLocalValues = "0.1.3"
 julia = "1"
 
 [extras]
@@ -25,9 +35,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "Dates", "DataFrames", "CSV", "Distributed", "Dagger"]
-
-[weakdeps]
-Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
-
-[extensions]
-DaggerFileTreesExt = "Dagger"

src/FileTrees.jl

@@ -2,6 +2,8 @@ module FileTrees
 
 using FilePathsBase
 using AbstractTrees
+using TaskLocalValues
+using OrderedCollections: OrderedDict
 import AbstractTrees: children
 import Base: parent, getindex
 

src/datastructure.jl

@@ -14,17 +14,24 @@ struct NoValue end
 - `parent::Union{FileTree, Nothing}` -- The parent node. `nothing` if it's the root node.
 - `name::String` -- Name of the root.
 - `children::Vector` -- children
-- `value::Any` -- the value at the node, if no value is present, a `NoValue()` sentinal value.
+- `value::Any` -- the value at the node, if no value is present, a `NoValue()` sentinel value.
 
     FileTree(tree::FileTree; parent, name, children, value)
 
 Copy over all fields from `tree`, but use any fields provided as keyword arguments.
 
-    FileTree(dirname::String; [sort])
+    FileTree(dirname::String; [sort], [follow_symlinks], [paralleldepth])
 
 Construct a `FileTree` to reflect directory from disk in the current working directory.
 
 If `sort=true` (the default) then each level of the tree will be lexicographically sorted.
+
+If `follow_symlinks=true` (the default) then symbolic links will be followed. Setting this to `true` actually speeds up the construction, 
+even if no symbolic links are present as it elides a stat check inside `walkdir`. Only reason to set this to `false` is if there exist
+symbolic links that should not be followed.
+
+Subdirectories down to `paralleldepth` (default `0`) levels will be read in separate tasks, which can give significant speedups when the directory tree is
+wide and I/O latency is high (e.g. NFS).
 """
 struct FileTree
     parent::Union{FileTree, Nothing}
@@ -57,23 +64,27 @@ end
 
 FileTree(dir; kwargs...) = FileTree(nothing, dir; kwargs...)
 
-function FileTree(parent, dir; sort=true, root="")
+function FileTree(parent, dir; sort=true, follow_symlinks=true, root="", paralleldepth=0)
     parent′ = FileTree(parent, dir, [])
-    for (path, dirs, files) in walkdir(joinpath(root, dir)) 
-        for dir in dirs
-            push!(parent′.children, FileTree(parent′, dir; sort, root=path))
+    for (path, dirs, files) in walkdir(joinpath(root, dir); follow_symlinks)
+
+        if paralleldepth > 0 && length(dirs) > 1
+            tasks = [Threads.@spawn FileTree(parent′, d; sort, follow_symlinks, root=path, paralleldepth=paralleldepth-1) for d in dirs]
+            append!(parent′.children, fetch.(tasks))
+        else
+            for d in dirs
+                push!(parent′.children, FileTree(parent′, d; sort, follow_symlinks, root=path, paralleldepth=0))
+            end
         end
         for file in files
             push!(parent′.children, File(parent′, file))
         end
-        if sort 
+        if sort
             sort!(parent′.children; by=name)
         end
-        # this might look a bit silly, but we don't care much for walkdir recursing down the tree, only that it is
-        # much faster than readdir + isdir/isfile when it comes to separating files from directories 
         break
     end
-    parent′ 
+    parent′
 end
 
 """
@@ -119,13 +130,25 @@ Return a copy of node with `val` set as the value.
 setvalue(x::FileTree, val) = FileTree(x; value=val)
 
 """
-    setparent(node::Union{FileTree, File}, parent)
+    setparent(node::Union{FileTree, File}, parent; [deep])
 
 Return a copy of node with `parent` set as the parent.
 """
 function setparent(x::FileTree, parent=parent(x))
-    p = FileTree(x, parent=parent, children=copy(x.children))
-    copy!(p.children, setparent.(x.children, (p,)))
+    p = FileTree(x, parent=parent, children=similar(x.children))
+    @inbounds for i in eachindex(p.children, x.children)
+        p.children[i] =  setparent(x.children[i], p)
+    end
+    p
+end
+
+# Update children without copying them. Only for when we are dead certain we have
+# created the whole tree from scratch
+function setparent!(x::FileTree, parent=parent(x))
+    p = FileTree(x, parent=parent, children=x.children)
+    @inbounds for i in eachindex(p.children, x.children)
+        p.children[i] =  setparent!(x.children[i], p)
+    end
     p
 end
 
@@ -138,7 +161,7 @@ Base.show(io::IO, d::FileTree) = AbstractTrees.print_tree(io, d)
 
 - `parent::Union{FileTree, Nothing}` -- The parent node. `nothing` if it's the root node.
 - `name::String` -- Name of the root.
-- `value::Any` -- the value at the node, if no value is present, a `NoValue()` sentinal value.
+- `value::Any` -- the value at the node, if no value is present, a `NoValue()` sentinel value.
 """
 struct File
     parent::Union{Nothing, FileTree}
@@ -184,6 +207,9 @@ setvalue(x::File, val) = File(x, value=val)
 
 setparent(x::File, parent=parent(x)) = File(x, parent=parent)
 
+setparent!(f::File, args...;kws...) = setparent(f, args...)
+
+
 Base.getindex(tree::FileTree, i::Int) = tree.children[i]
 
 function Base.getindex(tree::FileTree, ix::Vector)
@@ -279,22 +305,80 @@ maketree(node::Vector) = maketree("."=>node)
 
 Base.basename(d::Node) = Path(d.name)
 
+# Internal struct for collecting the path from a Node with minimal overhead
+# We reuse it so we don't need to allocate the data array each time we use it
+# TaskLocalValue is used to thread safety
+struct PathStringBuffer
+    data::Vector{UInt8}
+end
+PathStringBuffer() = PathStringBuffer(Vector{UInt8}[])
+const _PATH_STRING_BUFFER = TaskLocalValue{PathStringBuffer}(PathStringBuffer)
+
+
+function _rpath(buf::PathStringBuffer, f::FileTrees.Node, delim)
+    pos = _rpath!(buf, f, 0, delim)
+    res = GC.@preserve buf unsafe_string(pointer(buf.data), pos)
+    empty!(buf.data)
+    res
+end
+
+function _rpath!(buf, f::FileTrees.Node, pos::Int, delim::Vector{UInt8})
+    if !isnothing(parent(f))
+        pos = _rpath!(buf, parent(f), pos, delim)  
+        append!(buf.data, delim)
+        pos += length(delim)
+    end
+    n = codeunits(name(f))
+    append!(buf.data, n)
+    return pos + length(n)
+end
+
+# Internal struct for collecting the Path segments from a Node with minimal overhead
+# We reuse it so we don't need to allocate the path array each time we use it
+# TaskLocalValue is used to thread safety
+struct PathStringArray
+    path::Vector{String}
+end
+PathStringArray() = PathStringArray(sizehint!(String[], 4))
+const _PATH_STRING_ARRAY = TaskLocalValue{PathStringArray}(PathStringArray)
+
+function _rpath(arr::PathStringArray, f::FileTrees.Node)
+    len = _rpath!(arr, f, 1)-1
+    ntuple(i -> arr.path[i], len)
+end
+function _rpath!(arr::PathStringArray, f::FileTrees.Node, n)
+    if isnothing(parent(f))
+        if length(arr.path) < n
+            resize!(arr.path, n)
+        end
+        arr.path[1] = name(f)
+        return 2
+    end
+
+    pos = _rpath!(arr, parent(f), n+1) 
+    arr.path[pos] = name(f)
+    pos+1
+end
+
 """
     Path(file::Union{File, FileTree)
 
 Returns an [`AbstractPath`](https://rofinn.github.io/FilePathsBase.jl/stable/design/#Path-Types-1) object which is the Path of the file from the
 root node leading up to this file.
 """
-function Path(d::Node)
-    parent(d) === nothing && return Path(d.name)
-    # Prevent that empty names become ".". Not 100 that this is always right, but at least is prevents mv from silently deleting them
-    isempty(d.name) && return Path(parent(d)) 
-    Path(parent(d)) / Path(d.name)
+function Path(d::Node) 
+    arr =_PATH_STRING_ARRAY[]
+    len = _rpath!(arr, d, 1)-1
+    Path(ntuple(i -> arr.path[i], len))
 end
 
-path(x::Node) = string(Path(x))
+const _PATH_SEPARATOR = Vector{UInt8}(Base.Filesystem.path_separator)
+const _CANONICAL_PATH_SEPARATOR = [UInt8('/')]
+
+path(x::Node) =  _rpath(_PATH_STRING_BUFFER[], x, _PATH_SEPARATOR)
+canonical_path(x::Node) = _rpath(_PATH_STRING_BUFFER[], x, _CANONICAL_PATH_SEPARATOR)
 
-Base.dirname(d::Node) = dirname(Path(d))
+Base.dirname(d::Node) = Path(parent(d))
 
 Base.getindex(d::Node) = d.value
 

src/parallelism.jl

@@ -96,7 +96,7 @@ baremodule Executor
 
     Use Julia's standard library `Threads` to spawn each computation in a separate task when passed as first argument to `exec`.
 
-    If the keyword argument `unwrap_exceptions` is set to `true`` (the default), any `TaskFailedExceptions` will be unwrapped,
+    If the keyword argument `unwrap_exceptions` is set to `true` (the default), any `TaskFailedExceptions` will be unwrapped,
     which typically results in less visual noise in case the computation throws an exception.
 
     The keyword argument `pool` (default `:default`) is given as first argument to `Threads.@spawn`. 

src/patterns.jl

@@ -57,7 +57,7 @@ function Base.getindex(x::Union{FileTree, File}, regex::Regex)
 end
 
 function _regex_map(yes, no, t::FileTree, regex::Regex, toplevel=true)
-    if !toplevel && !isnothing(match(regex, canonical_path(Path(t))))
+    if !toplevel && !isnothing(match(regex, canonical_path(t)))
         return yes(t)
     end
 
@@ -78,7 +78,7 @@ function _regex_map(yes, no, t::FileTree, regex::Regex, toplevel=true)
 end
 
 function _regex_map(yes, no, t::File, regex::Regex, toplevel=false)
-    !isnothing(match(regex, canonical_path(Path(t)))) ? yes(t) : no(t)
+    !isnothing(match(regex, canonical_path(t))) ? yes(t) : no(t)
 end