Implement support for Zarr read/write

Squashed commit of the following:

commit 649aa0a
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Tue Aug 26 15:11:46 2025 +0200

    enable unit tests for Zarr and make them pass

commit 8ed25ed
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Mon Aug 25 18:17:58 2025 +0200

    make tests pass again

commit 17bb1b2
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Mon Aug 25 16:09:11 2025 +0200

    update JuliaFormatter config and reformat everything

commit a702865
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Mon Aug 25 16:07:06 2025 +0200

    implement support for writing Zarr files

commit db9abf9
Merge: ef4d948 eefb2c6
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Fri Aug 22 15:20:40 2025 +0200

    Merge branch 'main' into zarr

commit ef4d948
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Fri Apr 28 15:00:31 2023 +0200

    split reading into separate functions for hdf5 and zarr

commit 22bde93
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Fri Apr 28 14:19:33 2023 +0200

    Zarr for mudata, some bugfixes for backed storage

commit c61faba
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Fri Apr 28 13:35:12 2023 +0200

    initial Zarr-backed AnnData support

commit 735cd1a
Author: Ilia Kats <ilia-kats@gmx.net>
Date:   Fri Apr 28 10:09:29 2023 +0200

    add basic Zarr support

    Only non-backed reading works currently. This requires
    JuliaIO/Zarr.jl#116 and
    JuliaIO/Zarr.jl#117

Author: ilia-kats

.JuliaFormatter.toml

@@ -1,9 +1,10 @@
 indent = 4
 margin = 100
 always_for_in = true
+for_in_replacement = "∈"
 whitespace_typedefs = true
 whitespace_ops_in_indices = true
 remove_extra_newlines = true
 whitespace_in_kwargs = false
 annotate_untyped_fields_with_any = false
-#normalize_line_endings = "unix"
+normalize_line_endings = "unix"

Project.toml

@@ -16,6 +16,7 @@ PooledArrays = "2dfb63ee-cc39-5dd5-95bd-886bf059d720"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 
 [compat]
 CategoricalArrays = "0.9, 0.10, 1"
@@ -28,6 +29,7 @@ OrderedCollections = "1.6"
 PooledArrays = "1"
 StructArrays = "0.6.4"
 julia = "1.5"
+Zarr = "0.9"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

src/Muon.jl

@@ -5,6 +5,7 @@ using SparseArrays
 import LinearAlgebra: Adjoint
 
 using HDF5
+using Zarr
 using DataFrames
 using CategoricalArrays
 using StructArrays
@@ -13,7 +14,18 @@ import CompressHashDisplace: FrozenDict
 import OrderedCollections: OrderedDict
 using FileIO
 
-export readh5mu, readh5ad, writeh5mu, writeh5ad, isbacked, update_obs!, update_var!, update!
+export readh5mu,
+    readh5ad,
+    readzarrmu,
+    readzarrad,
+    writeh5mu,
+    writeh5ad,
+    writezarrmu,
+    writezarrad,
+    isbacked,
+    update_obs!,
+    update_var!,
+    update!
 export AnnData, MuData
 export var_names_make_unique!, obs_names_make_unique!
 
@@ -27,11 +39,13 @@ end
 MUDATAVERSION = v"0.1.0"
 ANNDATAVERSION = v"0.1.0"
 
-
+include("typedefs.jl")
 include("index.jl")
 include("sparsedataset.jl")
 include("transposeddataset.jl")
+include("common_io.jl")
 include("hdf5_io.jl")
+include("zarr_io.jl")
 include("alignedmapping.jl")
 include("anndata.jl")
 include("mudata.jl")

src/alignedmapping.jl

@@ -20,41 +20,37 @@ struct AlignedMapping{T <: Tuple, K, R} <: AbstractAlignedMapping{
     }
 
     function AlignedMapping{T, K}(r, d::AbstractDict{K}) where {T <: Tuple, K}
-        for (k, v) in d
+        for (k, v) ∈ d
             checkdim(T, v, r, k)
         end
         return new{T, K, typeof(r)}(r, d)
     end
 end
 
-mutable struct BackedAlignedMapping{T <: Tuple, R} <:
+mutable struct BackedAlignedMapping{T <: Tuple, G <: Group, R} <:
                AbstractAlignedMapping{T, String, AbstractArray{<:Number}}
     ref::R
-    d::Union{HDF5.Group, Nothing}
-    parent::Union{HDF5.File, HDF5.Group, Nothing}
+    d::Union{G, Nothing}
+    parent::Union{G, Nothing}
     path::Union{String, Nothing}
 
-    function BackedAlignedMapping{T}(r, g::HDF5.Group) where {T <: Tuple}
-        for k in keys(g)
+    function BackedAlignedMapping{T}(r, g) where {T <: Tuple}
+        for k ∈ keys(g)
             checkdim(T, backed_matrix(g[k]), r, k)
         end
-        return new{T, typeof(r)}(r, g, nothing, nothing)
+        return new{T, typeof(g), typeof(r)}(r, g, nothing, nothing)
     end
-    function BackedAlignedMapping{T}(
-        r,
-        parent::Union{HDF5.File, HDF5.Group},
-        path::String,
-    ) where {T <: Tuple}
+    function BackedAlignedMapping{T}(r, parent::Group, path::String) where {T <: Tuple}
         if haskey(parent, path)
             return BackedAlignedMapping{T}(r, parent[path])
         else
-            return new{T, typeof(r)}(r, nothing, parent, path)
+            return new{T, typeof(parent), typeof(r)}(r, nothing, parent, path)
         end
     end
 end
 
 function checkdim(::Type{T}, v, ref, k) where {T <: Tuple}
-    for (vdim, refdim) in T.parameters
+    for (vdim, refdim) ∈ T.parameters
         vsize = size(v, vdim)
         rsize = size(ref, refdim)
         if vsize != rsize
@@ -102,13 +98,12 @@ Base.sizehint!(d::AlignedMapping, n) = sizehint!(d.d, n)
 AlignedMapping{T}(r, d::AbstractDict) where {T <: Tuple} = AlignedMapping{T, keytype(d)}(r, d)
 AlignedMapping{T, K}(ref) where {T, K} = AlignedMapping{T}(ref, Dict{K, AbstractMatrix{<:Number}}())
 AlignedMapping{T, K}(ref, ::Nothing) where {T, K} = AlignedMapping{T, K}(ref)
-AlignedMapping{T}(r, d::HDF5.Group) where {T <: Tuple} =
-    AligedMapping{T}(ref, read_dict_of_mixed(d))
+AlignedMapping{T}(r, d::Group) where {T <: Tuple} = AligedMapping{T}(ref, read_dict_of_mixed(d))
 
 Base.delete!(d::BackedAlignedMapping, k) = !isnothing(d.d) && delete_object(d.d, k)
 function Base.empty!(d::BackedAlignedMapping)
     if !isnothing(d.d)
-        for k in keys(d.d)
+        for k ∈ keys(d.d)
             delete_object(d.d, k)
         end
     end
@@ -121,24 +116,27 @@ Base.get(default::Base.Callable, d::BackedAlignedMapping, key) =
     isnothing(d.d) || !haskey(d.d, key) ? default() : backed_matrix(d.d[key])
 Base.haskey(d::BackedAlignedMapping, key) = isnothing(d.d) ? false : haskey(d.d, key)
 Base.isempty(d::BackedAlignedMapping) = isnothing(d.d) ? true : isempty(d.d)
-function Base.iterate(d::BackedAlignedMapping)
+function Base.iterate(
+    d::BackedAlignedMapping{T, G},
+    i=nothing,
+) where {T, G <: Union{HDF5.File, HDF5.Group}}
     if isnothing(d.d)
         return nothing
     else
-        next = iterate(d.d)
+        next = iterate(d.d, i)
         return isnothing(next) ? next :
                (hdf5_object_name(next[1]) => backed_matrix(next[1]), next[2])
     end
 end
-function Base.iterate(d::BackedAlignedMapping, i)
-    if isnothing(d.d)
+function Base.iterate(d::BackedAlignedMapping{T, G}, i=nothing) where {T, G <: ZGroup}
+    if (isnothing(d.d))
         return nothing
     else
         next = iterate(d.d, i)
-        return isnothing(next) ? next :
-               (hdf5_object_name(next[1]) => backed_matrix(next[1]), next[2])
+        return isnothing(next) ? next : (next[1][1] => backed_matrix(next[1][2]), next[2])
     end
 end
+
 Base.length(d::BackedAlignedMapping) = isnothing(d.d) ? 0 : length(d.d)
 function Base.pop!(d::BackedAlignedMapping)
     if isnothing(d.d)
@@ -192,10 +190,10 @@ function copy_subset(
             J
         else
             (:)
-        end for (vdim, refdim) in T.parameters
+        end for (vdim, refdim) ∈ T.parameters
     )
-    for (k, v) in src
-        dst[k] = v[idx..., ((:) for i in 1:(ndims(v) - length(idx)))...]
+    for (k, v) ∈ src
+        dst[k] = v[idx..., ((:) for i ∈ 1:(ndims(v) - length(idx)))...]
     end
 end
 
@@ -208,7 +206,7 @@ end
 function aligned_view(d::AlignedMappingView{T}, A) where {T <: Tuple}
     idx = Vector{Union{Colon, typeof(d.indices).parameters...}}(undef, ndims(A))
     idx .= (:)
-    for ((vdim, refdim), cidx) in zip(T.parameters, d.indices)
+    for ((vdim, refdim), cidx) ∈ zip(T.parameters, d.indices)
         idx[vdim] = cidx
     end
     return @inbounds view(A, idx...)
@@ -249,7 +247,8 @@ Base.pop!(d::AlignedMappingView, k, default) =
 Base.parent(d::AlignedMappingView) = d.parent
 Base.parentindices(d::AlignedMappingView) = d.indices
 
-Base.setindex!(d::AlignedMappingView, v::AbstractArray, k) = throw(ArgumentError("Replacing or adding elements of an AlignedMappingView is not supported."))
+Base.setindex!(d::AlignedMappingView, v::AbstractArray, k) =
+    throw(ArgumentError("Replacing or adding elements of an AlignedMappingView is not supported."))
 
 function Base.view(parent::AbstractAlignedMapping{T}, indices...) where {T <: Tuple}
     @boundscheck if length(T.parameters) != length(indices)
@@ -262,7 +261,7 @@ function Base.view(parent::AbstractAlignedMapping{T}, indices...) where {T <: Tu
     return AlignedMappingView(parent, indices)
 end
 
-function Base.view(parentview::AlignedMappingView{T}, indices...) where T <: Tuple
+function Base.view(parentview::AlignedMappingView{T}, indices...) where {T <: Tuple}
     @boundscheck if length(T.parameters) != length(indices)
         throw(
             DimensionMismatch(

src/anndata.jl

@@ -1,7 +1,7 @@
 abstract type AbstractAnnData end
 
 mutable struct AnnData <: AbstractAnnData
-    file::Union{HDF5.File, HDF5.Group, Nothing}
+    file::Union{HDF5.File, HDF5.Group, ZGroup, Nothing}
 
     X::Union{AbstractMatrix{<:Number}, Nothing}
 
@@ -21,13 +21,13 @@ mutable struct AnnData <: AbstractAnnData
 
     uns::Dict{<:AbstractString, <:Any}
 
-    function AnnData(file::Union{HDF5.File, HDF5.Group}, backed=false, checkversion=true)
+    function AnnData(file::Union{HDF5.File, HDF5.Group, ZGroup}, backed=false, checkversion=true)
         if checkversion
             attrs = attributes(file)
             if !haskey(attrs, "encoding-type")
-                @warn "The HDF5 file was not created by muon, we can't guarantee that everything will work correctly"
-            elseif attrs["encoding-type"] != "AnnData"
-                error("This HDF5 file does not appear to hold an AnnData object")
+                @warn "This file was not created by muon, we can't guarantee that everything will work correctly"
+            elseif attrs["encoding-type"] != "anndata"
+                error("This file does not appear to hold an AnnData object")
             end
         end
 
@@ -158,9 +158,29 @@ function readh5ad(filename::AbstractString; backed=false)
     return adata
 end
 
+function readzarrad(filename::AbstractString; backed=false)
+    filename = abspath(filename)
+    if !backed
+        fid = zopen(filename, "r")
+    else
+        fid = zopen(filename, "r+")
+    end
+    local adata
+    try
+        adata = AnnData(fid, backed, true)
+    catch e
+        close(fid)
+        rethrow()
+    end
+    if !backed
+        close(fid)
+    end
+    return adata
+end
+
 function writeh5ad(filename::AbstractString, adata::AbstractAnnData; compress::UInt8=0x9)
     filename = abspath(filename)
-    if file(adata) === nothing || filename != HDF5.filename(file(adata))
+    if isnothing(file(adata)) || filename != HDF5.filename(file(adata))
         hfile = h5open(filename, "w", userblock=512)
         try
             write(hfile, adata, compress=compress)
@@ -179,22 +199,43 @@ function writeh5ad(filename::AbstractString, adata::AbstractAnnData; compress::U
     return nothing
 end
 
-function Base.write(
-    parent::Union{HDF5.File, HDF5.Group},
-    name::AbstractString,
-    adata::AbstractAnnData;
-    compress::UInt8=0x9
-)
+function writezarrad(filename::AbstractString, adata::AbstractAnnData; compress::UInt8=0x9)
+    filename = abspath(filename)
+    if isnothing(file(adata)) || filename != zarr_filename(file(adata))
+        rm(filename, force=true, recursive=true)
+        zfile = zgroup(filename)
+        write(zfile, adata, compress=compress)
+    else
+        write(adata, compress=compress)
+    end
+    return nothing
+end
+
+# HDF5.jl defines Base.write(::Union{HDF5.File, HDF5.Group}, ::Union{Nothing, AbstractString}, ::Any; kwargs...)
+# Using the below as Base.write leads to ambiguity: The HDF5.jl definition is more specific in the first argument,
+# ours is more specific in the third argument. Thus the ugly workaround.
+function _write(parent::Group, name::AbstractString, adata::AbstractAnnData; compress::UInt8=0x9)
     g = create_group(parent, name)
     write(g, adata, compress=compress)
 end
-
-function Base.write(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnnData; compress::UInt8=0x9)
-    attrs = attributes(parent)
-    attrs["encoding-type"] = "anndata"
-    attrs["encoding-version"] = string(ANNDATAVERSION)
-    attrs["encoder"] = NAME
-    attrs["encoder-version"] = string(VERSION)
+Base.write(
+    parent::Union{HDF5.File, HDF5.Group},
+    name::AbstractString,
+    adata::AbstractAnnData;
+    compress::UInt8=0x9,
+)=_write(parent, name, adata, compress=compress)
+Base.write(
+    parent::ZGroup,
+    name::AbstractString,
+    adata::AbstractAnnData;
+    compress::UInt8=0x9,
+)=_write(parent, name, adata, compress=compress)
+
+function Base.write(parent::Group, adata::AbstractAnnData; compress::UInt8=0x9)
+    write_attribute(parent, "encoding-type", "anndata")
+    write_attribute(parent, "encoding-version", string(ANNDATAVERSION))
+    write_attribute(parent, "encoder", NAME)
+    write_attribute(parent, "encoder-version", string(VERSION))
     if parent === file(adata)
         write(adata, compress=compress)
     else
@@ -211,7 +252,7 @@ function Base.write(adata; compress::UInt8=0x9)
     write_metadata(file(adata), adata, compress=compress)
 end
 
-function write_metadata(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnnData; compress::UInt8=0x9)
+function write_metadata(parent::Group, adata::AbstractAnnData; compress::UInt8=0x9)
     write_attr(parent, "obs", adata.obs, index=adata.obs_names, compress=compress)
     write_attr(parent, "obsm", adata.obsm, index=adata.obs_names, compress=compress)
     write_attr(parent, "obsp", adata.obsp, compress=compress)
@@ -222,7 +263,8 @@ function write_metadata(parent::Union{HDF5.File, HDF5.Group}, adata::AbstractAnn
 end
 # FileIO support
 load(f::File{format"h5ad"}; backed::Bool=false) = readh5ad(filename(f), backed=backed)
-save(f::File{format"h5ad"}, data::AbstractAnnData; compress::UInt8=0x9) = writeh5ad(filename(f), data, compress=compress)
+save(f::File{format"h5ad"}, data::AbstractAnnData; compress::UInt8=0x9) =
+    writeh5ad(filename(f), data, compress=compress)
 
 Base.size(adata::AbstractAnnData) = (length(adata.obs_names), length(adata.var_names))
 Base.size(adata::AbstractAnnData, d::Integer) = size(adata)[d]
@@ -247,15 +289,15 @@ end
 function Base.getindex(
     adata::AbstractAnnData,
     I::Union{
-        AbstractUnitRange,
+        OrdinalRange,
         Colon,
         AbstractVector{<:Integer},
         AbstractVector{<:AbstractString},
         Integer,
         AbstractString,
     },
     J::Union{
-        AbstractUnitRange,
+        OrdinalRange,
         Colon,
         AbstractVector{<:Integer},
         AbstractVector{<:AbstractString},