From 8c20d0ab5ab8cebf8594ee924fd8352f4e7ecd35 Mon Sep 17 00:00:00 2001 From: Felix Cremer Date: Thu, 23 Apr 2026 17:13:38 +0200 Subject: [PATCH 1/2] Enable reading of geozarr data --- ext/ZarrExt.jl | 87 ++++++++++++++++++++++++++++++++++++++++++++++++ test/datasets.jl | 18 ++++++++++ 2 files changed, 105 insertions(+) diff --git a/ext/ZarrExt.jl b/ext/ZarrExt.jl index 79b2ae3..f141f59 100644 --- a/ext/ZarrExt.jl +++ b/ext/ZarrExt.jl @@ -10,6 +10,7 @@ function __init__() @debug "new driver key :zarr, updating backendlist." YAB.backendlist[:zarr] = ZarrDataset push!(YAB.backendregex, r"(.zarr$)|(.zarr/$)|(zarr.zip$)" => ZarrDataset) + YAB.backendlist[:geozarr] = GeoZarrDataset end struct ZarrDataset @@ -106,6 +107,92 @@ function DiskArrays.eachchunk(a::SimpleFileDiskArray) end end +# Start GeoZarrDataset +struct GeoZarrDataset + g::ZGroup + axes_ranges::Dict{String,Any} # Mapping from dimension name to dimension + axes::Dict{String,Any} # Mapping from array name to axis name +end +function GeoZarrDataset(g::Union{String,ZGroup}; mode="r", path="", kwargs...) + zg = if g isa ZGroup + g + else + store = if endswith(g, "zip") + ZipStore(ZipReader(SimpleFileDiskArray(g))) + else + g + end + (zopen(store, mode, fill_as_missing=false, path=path)) + end + zg + axis_ranges = Dict{String,Any}() + axes = Dict{String,Any}() + for (name, arr) in zg.arrays + if haskey(arr.attrs, "spatial:dimensions") + transformtype = get(arr.attrs, "spatial:transform", "affine") + transform = arr.attrs["spatial:transform"] + axnames = arr.attrs["spatial:dimensions"] + @show transform + firstrange = if transform[2] == 0 + range(transform[3], length = size(arr, 1), step = transform[1]) + else + throw(ArgumentError("Rotation is currently not supported")) + end + secrange = if transform[4] == 0 + range(transform[6], length = size(arr, 2), step = transform[5]) + else + throw(ArgumentError("Rotation is currently not supported")) + end + @show axis_ranges + firstaxis = filter(x->last(x) == firstrange, axis_ranges) + if isempty(firstaxis) + push!(axis_ranges, axnames[1] => firstrange) + else + axnames[1] = first(firstaxis) + end + secaxis = (filter(x->last(x) == secrange, axis_ranges)) + if isempty(secaxis) + push!(axis_ranges, axnames[2] => secrange) + else + axnames[2] = first(secaxis) + end + + push!(axes, name => axnames) + @show arr.attrs + @show arr.attrs["spatial:dimensions"] + end + end + @show axis_ranges, axes + GeoZarrDataset(zg, axis_ranges, axes) +end + + +YAB.get_var_dims(ds::GeoZarrDataset, name) = ds.axes[name] +YAB.get_varnames(ds::GeoZarrDataset) = collect(keys(ds.g.arrays)) +YAB.get_global_attrs(ds::GeoZarrDataset) = ds.g.attrs +Base.getindex(ds::GeoZarrDataset, i) = ds.g[i] +Base.haskey(ds::GeoZarrDataset, k) = haskey(ds.g, k) || haskey(ds.axes_ranges, k) + +function YAB.get_var_attrs(ds::GeoZarrDataset, name) + #We add the fill value to the attributes to be consistent with NetCDF + haskey(ds.g, name) || return Dict() + 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 + +function YAB.get_var_handle(ds::GeoZarrDataset, name; persist=true) + @show name + @show keys(ds.axes_ranges) + if haskey(ds.axes_ranges, name) + ds.axes_ranges[name] + else + ds.g[name] + end +end end \ No newline at end of file diff --git a/test/datasets.jl b/test/datasets.jl index e4b4f9e..78b8933 100644 --- a/test/datasets.jl +++ b/test/datasets.jl @@ -78,6 +78,24 @@ end @test allow_missings(ds_zarr) == false end end + +@testset "Reading GeoZarr" begin + + path = "geozarr-spec/test-datasets/valid/spatial/spatial-registration-default.zarr" + ds_geozarr = to_dataset(path, driver=:geozarr) + vn = get_varnames(ds_geozarr) + @test sort(vn) == ["ar"] +end + +@testset "Writing and loading zerodim Zarr" begin + path = tempname() * "testdatazerodim.zarr" + ds = create_empty(YAXArrayBase.backendlist[:zarr], path) + add_var(ds, fill(1), "layer", (), Dict{String, Any}()) + ds_loaded = to_dataset(path, driver=:zarr) + @test ds_loaded["layer"][] == 1 + @test ds_loaded["layer"].attrs["_ARRAY_DIMENSIONS"] == [] +end + @testset "Reading ArchGDAL" begin using ArchGDAL import Downloads From c46a6f35cc8ed89a1cb1056e0930021ab1efae47 Mon Sep 17 00:00:00 2001 From: Felix Cremer Date: Wed, 24 Jun 2026 16:28:24 +0200 Subject: [PATCH 2/2] Add writing for geozarr --- ext/ZarrExt.jl | 54 ++++++++++++++++++++++++++++++++++++++++++++++-- test/datasets.jl | 27 +++++++++++++++++++----- 2 files changed, 74 insertions(+), 7 deletions(-) diff --git a/ext/ZarrExt.jl b/ext/ZarrExt.jl index f141f59..127ba20 100644 --- a/ext/ZarrExt.jl +++ b/ext/ZarrExt.jl @@ -125,13 +125,13 @@ function GeoZarrDataset(g::Union{String,ZGroup}; mode="r", path="", kwargs...) end (zopen(store, mode, fill_as_missing=false, path=path)) end - zg + @show zg axis_ranges = Dict{String,Any}() axes = Dict{String,Any}() for (name, arr) in zg.arrays if haskey(arr.attrs, "spatial:dimensions") transformtype = get(arr.attrs, "spatial:transform", "affine") - transform = arr.attrs["spatial:transform"] + transform = get(arr.attrs, "spatial:transform", [1,0,0,1,0,0]) axnames = arr.attrs["spatial:dimensions"] @show transform firstrange = if transform[2] == 0 @@ -195,4 +195,54 @@ function YAB.get_var_handle(ds::GeoZarrDataset, name; persist=true) end end +coords(d) = Dict("type" => "array") + +function YAB.add_var(p::GeoZarrDataset, T::Type, varname, s, dimnames, attr; + chunksize=s, fill_as_missing=false, kwargs...) + # This follows the zarr-coords specification for handling the coordinates + # See https://github.com/christophenoel/zarr-coords + dims = reverse(collect(dimnames)) + attr2 = merge(attr, Dict("_ARRAY_DIMENSIONS" => dims)) + @show attr, varname + @show T + coordsattr = Dict(d => coords(d) for d in dims) + merge!(attr2, Dict("coords:coordinates" => coordsattr)) + 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 + +function create_dataset(T::GeoZarrDataset, path, gatts, dimnames, dimvals, dimattrs, vartypes, varnames, vardims, varattrs, varchunks; kwargs...) + @show dimnames, dimvals, dimattr + ds = create_empty(T, path, gatts) + axlengths = Dict{String, Int}() + for (dname, dval, dattr) in zip(dimnames, dimvals, dimattrs) + add_var(ds, dval, dname, (dname,), dattr) + axlengths[dname] = length(dval) + end + for (T, vn, vd, va, vc) in zip(vartypes, varnames, vardims, varattrs, varchunks) + s = getindex.(Ref(axlengths),vd) + add_var(ds, T, vn, (s...,), vd, va; chunksize = vc, kwargs...) + end + ds +end + + +#Special case for init with Arrays +function YAB.add_var(p::GeoZarrDataset, a::AbstractArray, varname, dimnames, attr; + kwargs...) + T = to_zarrtype(a) + b = add_var(p, T, varname, size(a), dimnames, attr; kwargs...) + b .= a + a +end + +YAB.create_empty(::Type{GeoZarrDataset}, path, gatts=Dict()) = GeoZarrDataset(zgroup(path, attrs=gatts)) + +YAB.allow_parallel_write(::GeoZarrDataset) = true +YAB.allow_missings(::GeoZarrDataset) = false + end \ No newline at end of file diff --git a/test/datasets.jl b/test/datasets.jl index 78b8933..7044127 100644 --- a/test/datasets.jl +++ b/test/datasets.jl @@ -80,11 +80,28 @@ end end @testset "Reading GeoZarr" begin - - path = "geozarr-spec/test-datasets/valid/spatial/spatial-registration-default.zarr" - ds_geozarr = to_dataset(path, driver=:geozarr) - vn = get_varnames(ds_geozarr) - @test sort(vn) == ["ar"] + @testset "Registration default" begin + path = "geozarr-spec/test-datasets/valid/spatial/spatial-registration-default.zarr" + ds_geozarr = to_dataset(path, driver=:geozarr) + vn = get_varnames(ds_geozarr) + @test sort(vn) == ["ar"] + @test get_var_handle(ds_geozarr, "x") == 8.0:-1.0:1.0 + @test get_var_handle(ds_geozarr, "y") == 0:7 + end + + @testset "Rotation broken" begin + path_rotation = "geozarr-spec/test-datasets/valid/spatial/spatial-rotated.zarr" + @test_throws ArgumentError ds_rot = to_dataset(path_rotation, driver=:geozarr) + end + + @testset "crs" begin + path_crs = "geozarr-spec/test-datasets/valid/crs/crs-epsg-4326.zarr" + ds_crs = to_dataset(path_crs, driver=:geozarr) + + end + + + end @testset "Writing and loading zerodim Zarr" begin