-
Notifications
You must be signed in to change notification settings - Fork 15
/
binsparse.jl
115 lines (106 loc) · 3.37 KB
/
binsparse.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
using HDF5
using JSON
bswrite_type_lookup = Dict(
UInt32 => "uint32",
UInt64 => "uint64",
Int32 => "int32",
Int64 => "int64",
Float32 => "float32",
Float64 => "float64",
)
bsread_type_lookup =
Dict(value => key for (key, value) in bswrite_type_lookup)
function bswrite_data(f, desc, key, data)
desc["data_types"]["$(key)_type"] = bswrite_type_lookup[eltype(data)]
f[key] = data
end
function bsread_data(f, desc, key)
data = read(f[key])
T = bsread_type_lookup[desc["data_types"]["$(key)_type"]]
convert(Vector{T}, reinterpret(T, data))
end
function bswrite(fname, fbr::Fiber, attrs = Dict())
h5open(fname, "w") do f
desc = Dict(
"format" => Dict(),
"swizzle" => collect(ndims(fbr)-1:-1:0),
"shape" => size(fbr),
"data_types" => Dict(),
"attrs" => attrs,
)
bswrite_level(f, desc, desc["format"], fbr.lvl)
f["descriptor"] = json(desc, 4)
end
fname
end
function bsread(fname)
h5open(fname, "r") do f
desc = JSON.parse(read(f["descriptor"]))
Fiber(bsread_level(f, desc, desc["format"]))
end
end
bsread_level(f, desc, fmt) = bsread_level(f, desc, fmt, Val(Symbol(fmt["level"])))
function bswrite_level(f, desc, fmt, lvl::ElementLevel{D}) where {D}
fmt["level"] = "element"
@assert D == zero(eltype(lvl.val)) #TODO what about arbitrary fill?
bswrite_data(f, desc, "values", lvl.val)
end
function bsread_level(f, desc, fmt, ::Val{:element})
val = bsread_data(f, desc, "values")
D = zero(eltype(val)) #TODO what about arbitrary fill?
ElementLevel(D, val)
end
function bswrite_level(f, desc, fmt, lvl::DenseLevel{D}) where {D}
fmt["level"] = "dense"
fmt["rank"] = 1
fmt["subformat"] = Dict()
bswrite_level(f, desc, fmt["subformat"], lvl.lvl)
end
function bsread_level(f, desc, fmt, ::Val{:dense})
lvl = bsread_level(f, desc, fmt["subformat"])
R = fmt["rank"]
for r = 1:R
n = level_ndims(typeof(lvl))
shape = desc["shape"][end - n]
lvl = DenseLevel(lvl, shape)
end
lvl
end
function bswrite_level(f, desc, fmt, lvl::SparseListLevel)
fmt["level"] = "sparse"
fmt["rank"] = 1
n = level_ndims(typeof(lvl))
N = length(desc["shape"])
bswrite_data(f, desc, "pointers_$(N - n - 1)", lvl.ptr)
bswrite_data(f, desc, "indices_$(N - n)", lvl.idx)
fmt["subformat"] = Dict()
bswrite_level(f, desc, fmt["subformat"], lvl.lvl)
end
function bswrite_level(f, desc, fmt, lvl::SparseCOOLevel{R}) where {R}
fmt["level"] = "sparse"
fmt["rank"] = R
n = level_ndims(typeof(lvl))
N = length(desc["shape"])
bswrite_data(f, desc, "pointers_$(N - n - 1)", lvl.ptr)
for r = 1:R
bswrite_data(f, desc, "indices_$(N - n + r - 1)", lvl.tbl[r])
end
fmt["subformat"] = Dict()
bswrite_level(f, desc, fmt["subformat"], lvl.lvl)
end
function bsread_level(f, desc, fmt, ::Val{:sparse})
R = fmt["rank"]
lvl = bsread_level(f, desc, fmt["subformat"])
n = level_ndims(typeof(lvl)) + R
N = length(desc["shape"])
shape = (desc["shape"][n - R + 1: n]...,)
ptr = bsread_data(f, desc, "pointers_$(N - n - 1)")
tbl = (map(1:R) do r
bsread_data(f, desc, "indices_$(N - n + r - 1)")
end...,)
if R == 1
SparseListLevel(lvl, shape[1], ptr, tbl[1])
else
SparseCOOLevel{R}(lvl, shape, tbl, ptr)
end
end