Add ZarrDataset extension

This is a first draft for opening zarr data with ZarrDatasets in Rasters

Project.toml

@@ -32,6 +32,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 RasterDataSources = "3cb90ccd-e1b6-4867-9617-4276c8b2ca36"
+ZarrDatasets = "519a4cdf-1362-424a-9ea1-b1d782dbb24b"
 
 [extensions]
 RastersArchGDALExt = "ArchGDAL"
@@ -41,6 +42,7 @@ RastersHDF5Ext = "HDF5"
 RastersMakieExt = "Makie"
 RastersNCDatasetsExt = "NCDatasets"
 RastersRasterDataSourcesExt = "RasterDataSources"
+RastersZarrDatasetsExt = "ZarrDatasets"
 
 [compat]
 Aqua = "0.8"
@@ -75,6 +77,7 @@ Setfield = "0.6, 0.7, 0.8, 1"
 Shapefile = "0.10, 0.11"
 Statistics = "1"
 Test = "1"
+ZarrDatasets = "0.1"
 julia = "1.10"
 
 [extras]

ext/RastersZarrDatasetsExt/RastersZarrDatasetsExt.jl

@@ -0,0 +1,34 @@
+module RastersZarrDatasetsExt
+
+@static if isdefined(Base, :get_extension) # julia < 1.9
+    using Rasters, ZarrDatasets, CommonDataModel
+else    
+    using ..Rasters, ..GRIBDatasets, ..CommonDataModel
+end
+
+import DiskArrays,
+    FillArrays,
+    Extents,
+    GeoInterface,
+    Missings
+
+using Dates, 
+    DimensionalData,
+    GeoFormatTypes
+
+using Rasters.Lookups
+using Rasters.Dimensions
+using Rasters: Zarrsource
+
+using ZarrDatasets
+using CommonDataModel: AbstractDataset
+
+const RA = Rasters
+const DD = DimensionalData
+const DA = DiskArrays
+const GI = GeoInterface
+const LA = Lookups
+
+include("zarrdatasets_source.jl")
+
+end

ext/RastersZarrDatasetsExt/zarrdatasets_source.jl

@@ -0,0 +1,21 @@
+const GDS = ZarrDatasets
+
+function RA.OpenStack(fs::RA.FileStack{Zarrsource,K}) where K
+    RA.OpenStack{Zarrsource,K}(ZD.ZarrDataset(RA.filename(fs)))
+end
+
+# In ZarrDatasets, the file is open for reading the values and closed afterwards. 
+Base.close(os::RA.OpenStack{Zarrsource}) = nothing
+
+function RA._open(f, ::Zarrsource, filename::AbstractString; write=false, kw...)
+    #isfile(filename) || RA._filenotfound_error(filename)
+    ds = ZarrDatasets.ZarrDataset(filename)
+    RA._open(f, Zarrsource(), ds; kw...)
+end
+
+# Hack to get the inner DiskArrays chunks as they are not exposed at the top level
+RA._get_eachchunk(var::ZD.ZarrVariable) = DiskArrays.eachchunk(var.zarray)
+RA._get_haschunks(var::ZD.ZarrVariable) = DiskArrays.haschunks(var.zarray)
+
+RA._sourcetrait(::ZD.ZarrVariable) = Zarrsource()
+RA._sourcetrait(::ZD.ZarrDataset) = Zarrsource()

src/sources/commondatamodel.jl

@@ -185,9 +185,10 @@ function _layers(ds::AbstractDataset, ::Nothing=nothing)
         end
         union(dimkeys, boundskeys)::Vector{String}
     else
-        dimkeys::Vector{String}
+        collect(dimkeys)::Vector{String}
     end
-    nondim = setdiff(keys(ds), toremove)
+    # Maybe this should be fixed in ZarrDatasets but it works with this patch.
+    nondim = collect(setdiff(keys(ds), toremove))
     grid_mapping = String[]
     vars = map(k -> ds[k], nondim)
     attrs = map(CDM.attribs, vars)

src/sources/sources.jl

@@ -4,6 +4,7 @@ abstract type Source end
 abstract type CDMsource <: Source end
 struct NCDsource <: CDMsource end
 struct GRIBsource <: CDMsource end
+struct Zarrsource <: CDMsource end
 struct GRDsource <: Source end
 struct GDALsource <: Source end
 struct SMAPsource <: Source end
@@ -22,6 +23,7 @@ const SYMBOL2SOURCE = Dict(
     :netcdf => NCDsource(), 
     :grib => GRIBsource(), 
     :smap => SMAPsource(),
+    :zarr => Zarrsource(),
 )
 
 const SOURCE2SYMBOL = Dict(map(reverse, collect(pairs(SYMBOL2SOURCE))))
@@ -32,12 +34,14 @@ const SOURCE2EXT = Dict(
     NCDsource() => (".nc",), 
     GRIBsource() => (".grib",), 
     SMAPsource() => (".h5",),
+    Zarrsource() => (".zarr", ".zarr/"),
 )
 const SOURCE2PACKAGENAME = Dict(
     GDALsource() => "ArchGDAL",
     NCDsource() => "NCDatasets",
     GRIBsource() => "GRIBDatasets",
     SMAPsource() => "HDF5",
+    Zarrsource() => "ZarrDataset",
 )
 
 const EXT2SOURCE = Dict(
@@ -46,6 +50,7 @@ const EXT2SOURCE = Dict(
     ".nc" => NCDsource(), 
     ".grib" => GRIBsource(), 
     ".h5" => SMAPsource(),
+    ".zarr" => Zarrsource(),
 )
 
 # exception to be raised when backend extension is not satisfied

test/sources/zarr.jl

@@ -0,0 +1,2 @@
+using Rasters, Zarr
+using ZarrDatasets