feat: allow writing compressed files (#123)

Author: radioflash

README.md

@@ -23,7 +23,7 @@ To load the module:
 using MAT
 ```
 
-To read a single variable from a MAT file:
+To read a single variable from a MAT file (compressed files are detected and handled automatically):
 
 ```julia
 file = matopen("matfile.mat")
@@ -45,13 +45,14 @@ To read all variables from a MAT file as a Dict:
 vars = matread("matfile.mat")
 ```
 
-To write a Dict to a MAT file, using its keys as variable names:
+To write a Dict to a MAT file, using its keys as variable names.
+The `compress` argument is optional, and compression is off by default:
 
 ```julia
 matwrite("matfile.mat", Dict(
 	"myvar1" => 0,
 	"myvar2" => 1
-))
+); compress = true)
 ```
 
 To get a list of variable names in a MAT file:

src/MAT.jl

@@ -35,11 +35,11 @@ export matopen, matread, matwrite, names, exists, @read, @write
 
 # Open a MATLAB file
 const HDF5_HEADER = UInt8[0x89, 0x48, 0x44, 0x46, 0x0d, 0x0a, 0x1a, 0x0a]
-function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Bool, ff::Bool)
+function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Bool, ff::Bool, compress::Bool)
     # When creating new files, create as HDF5 by default
     fs = filesize(filename)
     if cr && (tr || fs == 0)
-        return MAT_HDF5.matopen(filename, rd, wr, cr, tr, ff)
+        return MAT_HDF5.matopen(filename, rd, wr, cr, tr, ff, compress)
     elseif fs == 0
         error("File \"$filename\" does not exist and create was not specified")
     end
@@ -76,28 +76,28 @@ function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Boo
         seek(rawfid, offset)
         if read!(rawfid, Vector{UInt8}(undef, 8)) == HDF5_HEADER
             close(rawfid)
-            return MAT_HDF5.matopen(filename, rd, wr, cr, tr, ff)
+            return MAT_HDF5.matopen(filename, rd, wr, cr, tr, ff, compress)
         end
     end
 
     close(rawfid)
     error("\"$filename\" is not a MAT file")
 end
 
-function matopen(fname::AbstractString, mode::AbstractString)
-    mode == "r"  ? matopen(fname, true , false, false, false, false) :
-    mode == "r+" ? matopen(fname, true , true , false, false, false) :
-    mode == "w"  ? matopen(fname, false, true , true , true , false) :
-#     mode == "w+" ? matopen(fname, true , true , true , true , false) :
-#     mode == "a"  ? matopen(fname, false, true , true , false, true ) :
-#     mode == "a+" ? matopen(fname, true , true , true , false, true ) :
+function matopen(fname::AbstractString, mode::AbstractString; compress::Bool = false)
+    mode == "r"  ? matopen(fname, true , false, false, false, false, false) :
+    mode == "r+" ? matopen(fname, true , true , false, false, false, compress) :
+    mode == "w"  ? matopen(fname, false, true , true , true , false, compress) :
+#     mode == "w+" ? matopen(fname, true , true , true , true , false, compress) :
+#     mode == "a"  ? matopen(fname, false, true , true , false, true, compress) :
+#     mode == "a+" ? matopen(fname, true , true , true , false, true, compress) :
     error("invalid open mode: ", mode)
 end
 
-matopen(fname::AbstractString) = matopen(fname, "r")
+matopen(fname::AbstractString; kwargs...) = matopen(fname, "r"; kwargs...)
 
-function matopen(f::Function, args...)
-    fid = matopen(args...)
+function matopen(f::Function, args...; kwargs...)
+    fid = matopen(args...; kwargs...)
     try
         f(fid)
     finally
@@ -106,12 +106,15 @@ function matopen(f::Function, args...)
 end
 
 """
-    matopen(filename [, mode]) -> handle
-    matopen(f::Function, filename [, mode]) -> f(handle)
+    matopen(filename [, mode]; compress = false) -> handle
+    matopen(f::Function, filename [, mode]; compress = false) -> f(handle)
 
 Mode defaults to "r" for read.  It can also be "w" for write, or "r+" for
 read or write without creation or truncation.
 
+Compression on reading is detected/handled automatically; the compress
+keyword argument only affects write operations.
+
 Use with `read`, `write`, `close`, `names`, and `exists`.
 """
 matopen
@@ -136,13 +139,13 @@ end
 
 # Write a dict to a MATLAB file
 """
-    matwrite(filename, d::Dict)
+    matwrite(filename, d::Dict; compress::Bool = false)
 
 Write a dictionary containing variable names as keys and values as values
 to a Matlab file, opening and closing it automatically.
 """
-function matwrite(filename::AbstractString, dict::AbstractDict{S, T}) where {S, T}
-    file = matopen(filename, "w")
+function matwrite(filename::AbstractString, dict::AbstractDict{S, T}; compress::Bool = false) where {S, T}
+    file = matopen(filename, "w"; compress = compress)
     try
         for (k, v) in dict
             local kstring
@@ -158,3 +161,4 @@ function matwrite(filename::AbstractString, dict::AbstractDict{S, T}) where {S,
     end
 end
 end
+

src/MAT_HDF5.jl

@@ -41,9 +41,10 @@ mutable struct MatlabHDF5File <: HDF5.DataFile
     toclose::Bool
     writeheader::Bool
     refcounter::Int
+    compress::Bool
 
-    function MatlabHDF5File(plain, toclose::Bool=true, writeheader::Bool=false, refcounter::Int=0)
-        f = new(plain, toclose, writeheader, refcounter)
+    function MatlabHDF5File(plain, toclose::Bool=true, writeheader::Bool=false, refcounter::Int=0, compress::Bool=false)
+        f = new(plain, toclose, writeheader, refcounter, compress)
         if toclose
             finalizer(close, f)
         end
@@ -79,7 +80,7 @@ function close(f::MatlabHDF5File)
     nothing
 end
 
-function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Bool, ff::Bool)
+function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Bool, ff::Bool, compress::Bool)
     local f
     if ff && !wr
         error("Cannot append to a write-only file")
@@ -101,7 +102,7 @@ function matopen(filename::AbstractString, rd::Bool, wr::Bool, cr::Bool, tr::Boo
         writeheader = false
     end
     close(pa)
-    fid = MatlabHDF5File(HDF5File(f, filename), true, writeheader, 0)
+    fid = MatlabHDF5File(HDF5File(f, filename), true, writeheader, 0, compress)
     pathrefs = "/#refs#"
     if exists(fid.plain, pathrefs)
         g = fid.plain[pathrefs]
@@ -326,8 +327,8 @@ function m_writeempty(parent::HDF5Parent, name::String, data::AbstractArray)
 end
 
 # Write an array to a dataset in a MATLAB file, returning the dataset
-function m_writearray(parent::HDF5Parent, name::String, adata::AbstractArray{T}) where {T<:HDF5BitsOrBool}
-    dset, dtype = d_create(parent, name, adata)
+function m_writearray(parent::HDF5Parent, name::String, adata::AbstractArray{T}, compress::Bool) where {T<:HDF5BitsOrBool}
+    dset, dtype = compress ? d_create(parent, name, adata, "chunk",HDF5.heuristic_chunk(adata), "compress",3) : d_create(parent, name, adata)
     try
         HDF5.writearray(dset, dtype.id, adata)
         dset
@@ -338,11 +339,19 @@ function m_writearray(parent::HDF5Parent, name::String, adata::AbstractArray{T})
         close(dtype)
     end
 end
-function m_writearray(parent::HDF5Parent, name::String, adata::AbstractArray{Complex{T}}) where {T<:HDF5BitsOrBool}
+function m_writearray(parent::HDF5Parent, name::String, adata::AbstractArray{Complex{T}}, compress::Bool) where {T<:HDF5BitsOrBool}
     dtype = build_datatype_complex(T)
     try
         stype = dataspace(adata)
-        obj_id = HDF5.h5d_create(parent.id, name, dtype.id, stype.id)
+        if compress
+            p = p_create(HDF5.H5P_DATASET_CREATE)
+            p["compress"] = 3
+            p["chunk"] = HDF5.heuristic_chunk(adata)
+            obj_id = HDF5.h5d_create(parent.id, name, dtype.id, stype.id,
+                                     HDF5._link_properties(name),p.id,HDF5.H5P_DEFAULT)
+        else
+            obj_id = HDF5.h5d_create(parent.id, name, dtype.id, stype.id)
+        end
         dset = HDF5Dataset(obj_id, file(parent))
         try
             arr = reshape(reinterpret(T, adata), tuple(2, size(adata)...))
@@ -365,7 +374,7 @@ function m_write(mfile::MatlabHDF5File, parent::HDF5Parent, name::String, data::
         m_writeempty(parent, name, data)
         return
     end
-    dset = m_writearray(parent, name, toarray(data))
+    dset = m_writearray(parent, name, toarray(data), mfile.compress)
     try
         m_writetypeattr(dset, T)
     finally
@@ -380,10 +389,10 @@ function m_write(mfile::MatlabHDF5File, parent::HDF5Parent, name::String, data::
         m_writetypeattr(g, T)
         a_write(g, sparse_attr_matlab, UInt64(size(data, 1)))
         if !isempty(data.nzval)
-            close(m_writearray(g, "data", toarray(data.nzval)))
-            close(m_writearray(g, "ir", add!(isa(data.rowval, Vector{UInt64}) ? copy(data.rowval) : convert(Vector{UInt64}, data.rowval), typemax(UInt64))))
+            close(m_writearray(g, "data", toarray(data.nzval), mfile.compress))
+            close(m_writearray(g, "ir", add!(isa(data.rowval, Vector{UInt64}) ? copy(data.rowval) : convert(Vector{UInt64}, data.rowval), typemax(UInt64)), mfile.compress))
         end
-        close(m_writearray(g, "jc", add!(isa(data.colptr, Vector{UInt64}) ? copy(data.colptr) : convert(Vector{UInt64}, data.colptr), typemax(UInt64))))
+        close(m_writearray(g, "jc", add!(isa(data.colptr, Vector{UInt64}) ? copy(data.colptr) : convert(Vector{UInt64}, data.colptr), typemax(UInt64)), mfile.compress))
     finally
         close(g)
     end

test/write.jl

@@ -2,8 +2,8 @@ using MAT
 
 tmpfile = string(tempname(), ".mat")
 
-function test_write(data)
-    matwrite(tmpfile, data)
+function test_write(data; kwargs...)
+    matwrite(tmpfile, data; kwargs...)
 
     fid = matopen(tmpfile, "r")
     local result
@@ -18,6 +18,22 @@ function test_write(data)
     end
 end
 
+function test_write(data)
+    test_write(data; compress = false)
+    test_write(data; compress = true)
+end
+
+function test_compression_effective(data)
+    test_write(data; compress = false)
+    sizeUncompressed = stat(tmpfile).size
+    test_write(data; compress = true)
+    sizeCompressed = stat(tmpfile).size
+
+    if sizeCompressed >= sizeUncompressed
+        error("Compression was not effective")
+    end
+end
+
 test_write(Dict(
     "int8" => Int8(1),
     "uint8" => UInt8(1),
@@ -109,3 +125,7 @@ sd = SortedDict(Dict(
     "sparse_empty" => sparse(Matrix{Float64}(undef, 0, 0))
 ))
 test_write(sd)
+
+# note: compression is NOT effective when the dict contains many duplicate entries
+# which are not compressible by themselves!
+test_compression_effective(Dict("data" => fill(1.0, 1000)))