migrate Zarr backend to ZarrCore and add ZarrZip extension support

Project.toml

@@ -1,7 +1,7 @@
 name = "YAXArrayBase"
 uuid = "90b8fcef-0c2d-428d-9c56-5f86629e9d14"
 authors = ["Fabian Gans <fgans@bgc-jena.mpg.de>"]
-version = "0.7.11"
+version = "0.8.0"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -15,8 +15,10 @@ ArchGDAL = "0.10"
 AxisArrays = "0.4"
 AxisKeys = "0.2"
 DimensionalData = "0.27, 0.28, 0.29, 0.30"
+DiskArrays = "0.4"
 NetCDF = "0.11, 0.12"
-Zarr = "0.10"
+ZarrCore = "0.11.1"
+ZarrZip = "0.11"
 
 
 [extensions]
@@ -26,16 +28,16 @@ AxisKeysExt = "AxisKeys"
 DimensionalDataExt = "DimensionalData"
 NamedDimsExt = "NamedDims"
 NetCDFExt = "NetCDF"
-ZarrExt = ["Zarr", "ZipArchives", "DiskArrays"]
+ZarrCoreExt = "ZarrCore"
+ZarrZipExt = ["ZarrCore", "ZarrZip", "DiskArrays"]
 
 [weakdeps]
 ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
-Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
 NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
-Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
-ZipArchives = "49080126-0e18-4c2a-b176-c102e4b3760c"
+ZarrCore = "77f5b75c-4c08-499f-ba13-550b0a0af171"
+ZarrZip = "54bcf1c8-e61d-4e4f-b375-463581846cca"

README.md

@@ -83,3 +83,13 @@ dimvals(a,1)
 ````
 
 Look at `src/axisinterface` to get a full description of the interface.
+
+## Zarr backend
+
+The Zarr dataset backend is enabled by loading `ZarrCore`. ZarrCore on its own
+only handles local stores and writes uncompressed data. Load the subpackages
+you need, or simply `using Zarr` to get all of them:
+
+- `ZarrHTTP`, `ZarrS3`, `ZarrGCS` for remote stores
+- `ZarrBlosc`, `ZarrZlib`, `ZarrZstd` for compression (`ZarrBlosc` makes Blosc the default compressor)
+- `ZarrZip` for reading `.zarr.zip` files

ext/NetCDFExt.jl

@@ -70,7 +70,14 @@ function readblock!(v::NetCDFVariable, aout, r::AbstractUnitRange...)
     aout .= aouttemp
   end
 end
-YAB.iscompressed(v::NetCDFVariable) = NetCDF.open(v->v.compress > 0, v.filename, v.varname)
+# NcVar.compress is not read back from the file, so ask the library for the deflate setting
+function YAB.iscompressed(v::NetCDFVariable)
+  NetCDF.open(v.filename, v.varname) do nv
+    shuffle, deflate, level = Ref{Cint}(0), Ref{Cint}(0), Ref{Cint}(0)
+    NetCDF.nc_inq_var_deflate(nv.ncid, nv.varid, shuffle, deflate, level)
+    deflate[] != 0
+  end
+end
 
 Base.size(v::NetCDFVariable) = v.size
 
@@ -95,7 +102,7 @@ function YAB.add_var(p::NetCDFDataset, T::Type, varname, s, dimnames, attr;
   NetCDFVariable{T,length(s)}(p.filename,varname,(s...,))
 end
 
-function YAB.create_empty(::Type{NetCDFDataset}, path, gatts=Dict())
+function YAB.create_empty(::Type{NetCDFDataset}, path, gatts=Dict(); kwargs...)
   NetCDF.create(_->nothing, path, NcVar[], gatts = gatts)
   NetCDFDataset(path)
 end

ext/ZarrCoreExt.jl

@@ -0,0 +1,110 @@
+module ZarrCoreExt
+using YAXArrayBase
+using ZarrCore: ZarrCore, ZArray, ZGroup, zgroup, zcreate, zopen, NoCompressor, CRC32cV3Codec
+import YAXArrayBase: YAXArrayBase as YAB
+import ZarrCore: zarr_format
+export ZarrDataset
+
+function __init__()
+  @debug "new driver key :zarr, updating backendlist."
+  YAB.backendlist[:zarr] = ZarrDataset
+  push!(YAB.backendregex, r"(.zarr$)|(.zarr/$)|(zarr.zip$)" => ZarrDataset)
+end
+
+struct ZarrDataset
+  g::ZGroup
+end
+function ZarrDataset(g::Union{String,ZGroup}; mode="r", path="", kwargs...)
+  if g isa ZGroup
+    return ZarrDataset(g)
+  end
+  store = if endswith(g, "zip")
+    YAB.zarr_zipstore(g)
+  else
+    g
+  end
+  ZarrDataset(zopen(store, mode, fill_as_missing=false, path=path))
+end
+
+function YAB.get_var_dims(ds::ZarrDataset, name)
+  a = ds[name]
+  # Zarr v3 stores dimension names in the array metadata, already in Julia order
+  dn = ZarrCore.dimension_names(a)
+  if dn !== nothing && all(!isnothing, dn)
+    return collect(String, dn)
+  end
+  if haskey(a.attrs, "_ARRAY_DIMENSIONS")
+    return reverse(a.attrs["_ARRAY_DIMENSIONS"])
+  end
+  # Scalars have no dimensions to name, Zarr v3 allows omitting dimension_names for them
+  ndims(a) == 0 && return String[]
+  throw(ArgumentError("Zarr array $name has no dimension names"))
+end
+YAB.get_varnames(ds::ZarrDataset) = collect(keys(ds.g.arrays))
+function YAB.get_var_attrs(ds::ZarrDataset, name)
+  #We add the fill value to the attributes to be consistent with NetCDF
+  a = ds[name]
+  if a.metadata.fill_value !== nothing
+    merge(ds[name].attrs, Dict("_FillValue" => a.metadata.fill_value))
+  else
+    ds[name].attrs
+  end
+end
+YAB.get_global_attrs(ds::ZarrDataset) = ds.g.attrs
+Base.getindex(ds::ZarrDataset, i) = ds.g[i]
+Base.haskey(ds::ZarrDataset, k) = haskey(ds.g, k)
+
+# function add_var(p::ZarrDataset, T::Type{>:Missing}, varname, s, dimnames, attr; kwargs...)
+#   S = Base.nonmissingtype(T)
+#   add_var(p,S, varname, s, dimnames, attr; fill_value = defaultfillval(S), fill_as_missing=true, kwargs...)
+# end
+
+function YAB.add_var(p::ZarrDataset, T::Type, varname, s, dimnames, attr;
+  chunksize=s, fill_as_missing=false, format=nothing, kwargs...)
+  # The format is set by the group in create_empty
+  if format !== nothing && format != zarr_format(p)
+    throw(ArgumentError("Can not create a Zarr v$format array in a Zarr v$(zarr_format(p)) group"))
+  end
+  if !haskey(kwargs, :compressor) && ZarrCore.default_compressor() isa NoCompressor
+    @info "No Zarr compressor package is loaded, so data will be written uncompressed. Load e.g. ZarrBlosc or Zarr to enable compression." maxlog = 1
+  end
+  if zarr_format(p) == 2
+    attr2 = merge(attr, Dict("_ARRAY_DIMENSIONS" => reverse(collect(String, dimnames))))
+  else
+    attr2 = attr
+    kwargs = (; kwargs..., dimension_names=Tuple(collect(String, dimnames)))
+  end
+  fv = get(attr, "_FillValue", get(attr, "missing_value", YAB.defaultfillval(T)))
+  attr3 = filter(attr2) do (k, v)
+    !isa(v, AbstractFloat) || !isnan(v)
+  end
+  za = zcreate(T, p.g, varname, s...; fill_value=fv, fill_as_missing, attrs=attr3, chunks=chunksize, kwargs...)
+  za
+end
+
+#Special case for init with Arrays
+function YAB.add_var(p::ZarrDataset, a::AbstractArray, varname, dimnames, attr;
+  kwargs...)
+  # to_zarrtype is not public in ZarrCore
+  T = ZarrCore.to_zarrtype(a)
+  b = add_var(p, T, varname, size(a), dimnames, attr; kwargs...)
+  b .= a
+  a
+end
+
+YAB.create_empty(::Type{ZarrDataset}, path, gatts=Dict(); format=2, kwargs...) =
+  ZarrDataset(zgroup(ZarrCore.storefromstring(path, true)..., format; attrs=gatts))
+zarr_format(ds::ZarrDataset) = zarr_format(ds.g.metadata)
+
+
+
+YAB.allow_parallel_write(::ZarrDataset) = true
+YAB.allow_missings(::ZarrDataset) = false
+YAB.to_dataset(g::ZGroup; kwargs...) = ZarrDataset(g; kwargs...)
+# get_pipeline, V2Pipeline and V3Pipeline are not public in ZarrCore
+YAB.iscompressed(a::ZArray) = _iscompressed(ZarrCore.get_pipeline(a.metadata))
+_iscompressed(p::ZarrCore.V2Pipeline) = !isa(p.compressor, NoCompressor)
+# Checksum codecs are bytes-to-bytes codecs too, but do not compress
+_iscompressed(p::ZarrCore.V3Pipeline) = any(c -> !isa(c, CRC32cV3Codec), p.bytes_bytes)
+
+end

ext/ZarrZipExt.jl

@@ -0,0 +1,35 @@
+module ZarrZipExt
+using ZarrZip: ZipStore
+import DiskArrays: AbstractDiskArray, DiskArrays, Unchunked, Chunked, GridChunks
+import YAXArrayBase: YAXArrayBase as YAB
+
+#Add ability to read zipped zarrs
+YAB.zarr_zipstore(path::AbstractString) = ZipStore(SimpleFileDiskArray(path))
+
+struct SimpleFileDiskArray{C<:Union{Int,Nothing}} <: AbstractDiskArray{UInt8,1}
+  file::String
+  s::Int
+  chunksize::C
+end
+Base.size(s::SimpleFileDiskArray) = (s.s,)
+function SimpleFileDiskArray(filename; chunksize=nothing)
+  isfile(filename) || throw(ArgumentError("File $filename does not exist"))
+  s = filesize(filename)
+  SimpleFileDiskArray(filename, s, chunksize)
+end
+function DiskArrays.readblock!(a::SimpleFileDiskArray, aout, i::AbstractUnitRange)
+  open(a.file) do f
+    seek(f, first(i) - 1)
+    read!(f, aout)
+  end
+end
+DiskArrays.haschunks(a::SimpleFileDiskArray) = a.chunksize === nothing ? Unchunked() : Chunked()
+function DiskArrays.eachchunk(a::SimpleFileDiskArray)
+  if a.chunksize === nothing
+    DiskArrays.estimate_chunksize(a)
+  else
+    GridChunks((a.s,), (a.chunksize,))
+  end
+end
+
+end

src/datasets/datasetinterface.jl

@@ -15,7 +15,12 @@ function get_var_attrs end
 function get_global_attrs end
 
 #Functions to be implemented for Dataset sinks
-"Initialize and return a handle to a new empty dataset"
+"""
+    create_empty(T, path, gatts; kwargs...)
+
+Initialize and return a handle to a new empty dataset. Backends ignore keyword
+arguments they do not support.
+"""
 function create_empty end
 
 "Apply a function `f` on a dataset `ds` while keeping possible file handles open during the operations"
@@ -56,10 +61,13 @@ function add_var(ds,x::AbstractArray,name,dimlist,atts;kwargs...)
 end
 
 function create_dataset(T::Type, path, gatts, dimnames, dimvals, dimattrs, vartypes, varnames, vardims, varattrs, varchunks; kwargs...)
-  ds = create_empty(T, path, gatts)
+  ds = create_empty(T, path, gatts; kwargs...)
   axlengths = Dict{String, Int}()
+  # Coordinates share the compression settings (Zarr compressor, NetCDF compress),
+  # other options like filters may be type specific
+  coordkwargs = (; (k => v for (k, v) in pairs(kwargs) if k in (:compressor, :compress))...)
   for (dname, dval, dattr) in zip(dimnames, dimvals, dimattrs)
-    add_var(ds, dval, dname, (dname,), dattr)
+    add_var(ds, dval, dname, (dname,), dattr; coordkwargs...)
     axlengths[dname] = length(dval)
   end
   for (T, vn, vd, va, vc) in zip(vartypes, varnames, vardims, varattrs, varchunks)
@@ -81,6 +89,14 @@ backendlist = OrderedDict{Symbol, Any}(
 
 backendregex = Pair[]
 
+"""
+    zarr_zipstore(path)
+
+Open the zipped Zarr store at `path`. Used by the Zarr backend and implemented
+when `ZarrZip` is loaded.
+"""
+zarr_zipstore(path) = throw(ArgumentError("Reading zipped Zarr stores requires ZarrZip. Run `using ZarrZip` (or `using Zarr`) first."))
+
 function backendfrompath(g::String; driver = :all)
   if driver == :all
     isempty(backendregex) && throw("No backend found. Load a backend by using the corresponding package.")