add support for NCDatasets.jl

Project.toml

@@ -15,6 +15,7 @@ ArchGDAL = "0.10"
 AxisArrays = "0.4"
 AxisKeys = "0.2"
 DimensionalData = "0.27, 0.28, 0.29, 0.30"
+NCDatasets = "0.14"
 NetCDF = "0.11, 0.12"
 Zarr = "0.10"
 
@@ -24,6 +25,7 @@ ArchGDALExt = "ArchGDAL"
 AxisArraysExt = "AxisArrays"
 AxisKeysExt = "AxisKeys"
 DimensionalDataExt = "DimensionalData"
+NCDatasetsExt = ["NCDatasets", "DiskArrays"]
 NamedDimsExt = "NamedDims"
 NetCDFExt = "NetCDF"
 ZarrExt = ["Zarr", "ZipArchives", "DiskArrays"]
@@ -35,6 +37,7 @@ AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
 NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
 Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"

ext/NCDatasetsExt.jl

@@ -0,0 +1,162 @@
+module NCDatasetsExt
+
+import YAXArrayBase: YAXArrayBase as YAB
+using NCDatasets
+
+"""
+    NCDatasetsDataset
+
+Dataset backend to read NetCDF files using NCDatasets.jl
+
+Variables are handed to YAXArrays as raw storage values (`nc[name].var`):
+YAXArrays applies `missing_value`, `scale_factor`, `add_offset` and time
+units itself, as it does for the NetCDF.jl backend.
+
+The following keyword arguments are allowed when using :ncdatasets
+as a data sink:
+
+- `compress = -1` set the deflate level for the NetCDF file
+- `zstdlevel = nothing` set the Zstandard level (requires an NCDatasets
+  with Zstandard support and the zstd HDF5 filter, e.g. H5Zzstd)
+"""
+struct NCDatasetsDataset
+    filename::String
+    mode::String
+    handle::Base.RefValue{Union{Nothing, NCDataset}}
+end
+
+function NCDatasetsDataset(filename; mode="r")
+    NCDatasetsDataset(filename, mode, Ref{Union{Nothing, NCDataset}}(nothing))
+end
+
+# Helper to execute a function block with a safe file handle
+function dsopen(f, ds::NCDatasetsDataset)
+    if ds.handle[] === nothing
+        NCDataset(ds.filename, ds.mode) do nc
+            f(nc)
+        end
+    else
+        f(ds.handle[])
+    end
+end
+
+function YAB.open_dataset_handle(f, ds::NCDatasetsDataset)
+    if ds.handle[] === nothing
+        try
+            ds.handle[] = NCDataset(ds.filename, ds.mode)
+            f(ds)
+        finally
+            ds.handle[] === nothing || close(ds.handle[])
+            ds.handle[] = nothing
+        end
+    else
+        f(ds)
+    end
+end
+
+# Implement the DiskArrays read/write interface
+import DiskArrays: AbstractDiskArray, readblock!, writeblock!, haschunks, eachchunk
+
+# raw storage variable, bypassing the CF layer of NCDatasets (see above)
+rawvar(nc, name) = nc[name].var
+
+# --- Variable representation ---
+struct NCDatasetsVariable{T,N} <: AbstractDiskArray{T,N}
+    filename::String
+    mode::String
+    varname::String
+    size::NTuple{N,Int}
+end
+
+Base.size(v::NCDatasetsVariable) = v.size
+
+readblock!(v::NCDatasetsVariable, aout, r::AbstractUnitRange...) =
+    NCDataset(nc -> (aout .= rawvar(nc, v.varname)[r...]; nothing), v.filename, v.mode)
+writeblock!(v::NCDatasetsVariable, a, r::AbstractUnitRange...) =
+    NCDataset(nc -> (rawvar(nc, v.varname)[r...] = a; nothing), v.filename, "a")
+for m in (:haschunks, :eachchunk)
+    @eval $m(v::NCDatasetsVariable) =
+        NCDataset(nc -> $m(rawvar(nc, v.varname)), v.filename, v.mode)
+end
+
+# deflate or Zstandard (the latter only with an NCDatasets that supports it)
+function YAB.iscompressed(v::NCDatasetsVariable)
+    NCDataset(v.filename, v.mode) do nc
+        rv = rawvar(nc, v.varname)
+        _, isdeflated, level = deflate(rv)
+        iszstd = isdefined(NCDatasets, :zstandard) && first(NCDatasets.zstandard(rv))
+        (isdeflated && level > 0) || iszstd
+    end
+end
+
+# --- Dataset interface implementations ---
+YAB.get_varnames(ds::NCDatasetsDataset) = dsopen(nc -> collect(keys(nc)), ds)
+YAB.get_var_dims(ds::NCDatasetsDataset, name) = dsopen(nc -> collect(dimnames(nc[name])), ds)
+YAB.get_global_attrs(ds::NCDatasetsDataset) = dsopen(nc -> Dict(nc.attrib), ds)
+Base.haskey(ds::NCDatasetsDataset, k) = dsopen(nc -> haskey(nc, k), ds)
+
+function YAB.get_var_attrs(ds::NCDatasetsDataset, name)
+    dsopen(ds) do nc
+        a = Dict{String,Any}(string(k) => v for (k, v) in nc[name].attrib)
+        # YAXArrays only honours missing_value; let _FillValue count too
+        haskey(a, "_FillValue") && !haskey(a, "missing_value") &&
+            (a["missing_value"] = a["_FillValue"])
+        a
+    end
+end
+
+function YAB.get_var_handle(ds::NCDatasetsDataset, i; persist = true)
+    if persist || ds.handle[] === nothing
+        s, et = dsopen(nc -> (size(rawvar(nc, i)), eltype(rawvar(nc, i))), ds)
+        NCDatasetsVariable{et, length(s)}(ds.filename, ds.mode, i, s)
+    else
+        rawvar(ds.handle[], i)
+    end
+end
+
+# --- Dataset creation and variable addition ---
+function YAB.create_empty(::Type{NCDatasetsDataset}, path, gatts=Dict())
+    NCDataset(path, "c", attrib = gatts) do nc
+        # Creates file structure on disk
+    end
+    NCDatasetsDataset(path, mode="a") # Return in append mode
+end
+
+function YAB.add_var(p::NCDatasetsDataset, T::Type, varname, s, dimnames, attr;
+    chunksize=s, compress = -1, zstdlevel = nothing)
+
+    dsopen(p) do nc
+        # Define dimensions if they don't exist yet
+        for (dname, dlen) in zip(dimnames, s)
+            if !haskey(nc.dim, dname)
+                defDim(nc, dname, dlen)
+            end
+        end
+
+        # only pass the storage options that are actually requested; in
+        # particular zstdlevel is not understood by every NCDatasets version
+        kw = (; chunksizes = chunksize, attrib = attr)
+        compress > -1 && (kw = (; kw..., deflatelevel = min(compress, 9), shuffle = true))
+        zstdlevel === nothing || (kw = (; kw..., zstdlevel))
+        defVar(nc, varname, T, dimnames; kw...)
+    end
+
+    NCDatasetsVariable{T, length(s)}(p.filename, p.mode, varname, s)
+end
+
+# --- Parallel and Missings support flags ---
+YAB.allow_parallel_write(::Type{<:NCDatasetsDataset}) = false
+YAB.allow_parallel_write(::NCDatasetsDataset) = false
+# NCDatasets' defVar wants a DataType; Union{Missing,T} is handled by
+# YAXArrays through missing_value, as for NetCDF.jl
+YAB.allow_missings(::Type{<:NCDatasetsDataset}) = false
+YAB.allow_missings(::NCDatasetsDataset) = false
+
+# --- Extension Registration ---
+function __init__()
+    @debug "new driver key :ncdatasets, updating backendlist."
+    YAB.backendlist[:ncdatasets] = NCDatasetsDataset
+    push!(YAB.backendregex, r".nc$" => NCDatasetsDataset)
+end
+
+end # module

test/Project.toml

@@ -3,7 +3,9 @@ ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
 AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
+DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 NamedDims = "356022a1-0364-5f58-8944-0da4b18d706f"
 NetCDF = "30363a11-5582-574a-97bb-aa9a979735b9"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

test/datasets.jl

@@ -3,7 +3,7 @@ using YAXArrayBase, Test
   @test_throws "No backend found." YAXArrayBase.backendfrompath("test.zarr")
 end
 
-using NetCDF, Zarr
+using NetCDF, Zarr, NCDatasets, DiskArrays
 
 using Pkg.Artifacts
 import Downloads
@@ -63,6 +63,57 @@ YAXArrayBase.open_dataset_handle(ds_nc2) do ds_nc
 end
 end
 
+@testset "Reading NCDatasets" begin
+  ds_nc = YAXArrayBase.to_dataset(p2, driver=:ncdatasets)
+  vn = get_varnames(ds_nc)
+  @test sort(vn) == ["area", "lat", "lat_bnds", "lon", "lon_bnds", "msk_rgn",
+   "plev", "pr", "tas", "time", "time_bnds", "ua"]
+  @test get_var_dims(ds_nc, "tas") == ["lon", "lat", "time"]
+  @test get_var_dims(ds_nc, "area") == ["lon", "lat"]
+  @test get_var_dims(ds_nc, "time") == ["time"]
+  @test get_var_dims(ds_nc, "time_bnds") == ["bnds", "time"]
+  @test get_var_attrs(ds_nc,"tas")["long_name"] == "air_temperature"
+  h = get_var_handle(ds_nc, "tas")
+  @test !YAXArrayBase.iscompressed(h)
+  @test all(isapprox.(h[1:2,1:2], [215.893 217.168; 215.805 217.03]))
+  @test allow_parallel_write(ds_nc) == false
+  @test allow_missings(ds_nc) == false
+  # handles carry the file's chunking
+  @test DiskArrays.eachchunk(h) == DiskArrays.eachchunk(get_var_handle(YAXArrayBase.to_dataset(p2), "tas"))
+  # a failed open reports the real error, not close(nothing)
+  @test_throws NCDatasets.NetCDFError YAXArrayBase.open_dataset_handle(identity,
+    YAXArrayBase.to_dataset(tempname() * ".nc", driver=:ncdatasets))
+end
+
+@testset "NCDatasets hands over raw storage values" begin
+  # scale_factor/add_offset/_FillValue are applied by YAXArrays, not the backend
+  p = tempname() * ".nc"
+  NCDataset(p, "c") do nc
+    defDim(nc, "x", 4)
+    v = defVar(nc, "v", Int16, ("x",); fillvalue=Int16(-1),
+               attrib=["scale_factor"=>0.5, "add_offset"=>100.0])
+    v.var[:] = Int16[2, 4, -1, 8]
+    t = defVar(nc, "time", Float64, ("x",); attrib=["units"=>"days since 2000-01-01"])
+    t.var[:] = [0.0, 1.0, 2.0, 3.0]
+    c = defVar(nc, "c", Float32, ("x",); chunksizes=(2,), deflatelevel=1)
+    c[:] = 1:4
+  end
+  ds = YAXArrayBase.to_dataset(p, driver=:ncdatasets)
+  h = get_var_handle(ds, "v")
+  @test eltype(h) == Int16
+  @test h[:] == Int16[2, 4, -1, 8]
+  a = get_var_attrs(ds, "v")
+  @test a["scale_factor"] == 0.5 && a["add_offset"] == 100.0
+  @test a["_FillValue"] == Int16(-1) && a["missing_value"] == Int16(-1)
+  @test eltype(get_var_handle(ds, "time")) == Float64
+  @test get_var_handle(ds, "time")[:] == [0.0, 1.0, 2.0, 3.0]
+  c = get_var_handle(ds, "c")
+  @test DiskArrays.eachchunk(c) == DiskArrays.GridChunks(c, (2,))
+  @test YAXArrayBase.iscompressed(c)
+  @test !YAXArrayBase.iscompressed(h)
+  rm(p)
+end
+
 @testset "Reading Zarr" begin
   p = "https://s3.bgc-jena.mpg.de:9000/esdl-esdc-v3.0.2/esdc-16d-2.5deg-46x72x1440-3.0.2.zarr"
   for ds_zarr in [to_dataset(p,driver=:zarr), to_dataset(zopen(p))]
@@ -136,6 +187,10 @@ end
   test_write(YAXArrayBase.backendlist[:netcdf])
 end
 
+@testset "Writing NCDatasets" begin
+  test_write(YAXArrayBase.backendlist[:ncdatasets])
+end
+
 @testset "Writing Zarr" begin
   test_write(YAXArrayBase.backendlist[:zarr])
 end