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.12, 0.13, 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,145 @@
+module NCDatasetsExt
+
+import YAXArrayBase: YAXArrayBase as YAB
+using NCDatasets
+
+"""
+    NCDatasetsDataset
+
+Dataset backend to read NetCDF files using NCDatasets.jl
+"""
+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
+            close(ds.handle[])
+            ds.handle[] = nothing
+        end
+    else
+        f(ds)
+    end
+end
+
+# Implement the DiskArrays read/write interface
+import DiskArrays: AbstractDiskArray
+import NCDatasets: readblock!, writeblock!, haschunks, eachchunk
+
+# --- 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
+
+function readblock!(v::NCDatasetsVariable, aout, r::AbstractUnitRange...)
+    NCDataset(v.filename, v.mode) do nc
+        aout .= nc[v.varname][r...]
+    end
+end
+
+function writeblock!(v::NCDatasetsVariable, a, r::AbstractUnitRange...)
+    NCDataset(v.filename, "a") do nc
+        nc[v.varname][r...] = a
+    end
+end
+
+# NCDatasets supports querying the deflate level from the variable
+function YAB.iscompressed(v::NCDatasetsVariable)
+    NCDataset(v.filename, v.mode) do nc
+        cfvar = nc[v.varname]
+        rawvar = hasproperty(cfvar, :var) ? cfvar.var : cfvar
+        isshuffled, isdeflated, deflate_level = deflate(rawvar)
+        return isdeflated && deflate_level > 0
+    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 -> dimnames(nc[name]), ds)...]
+YAB.get_var_attrs(ds::NCDatasetsDataset, name) = dsopen(nc -> Dict(nc[name].attrib), 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_handle(ds::NCDatasetsDataset, i; persist = true)
+    if persist || ds.handle[] === nothing
+        s, et = dsopen(nc -> (size(nc[i]), eltype(nc[i])), ds)
+        NCDatasetsVariable{et, length(s)}(ds.filename, ds.mode, i, s)
+    else
+        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)
+
+    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
+
+        # Set up compression options
+        deflatelevel = compress > -1 ? min(compress, 9) : 0
+        shuffle = deflatelevel > 0
+
+        # Define the compressed variable natively in NCDatasets
+        defVar(nc, varname, T, dimnames, 
+               attrib = attr, 
+               chunksizes = chunksize, 
+               deflatelevel = deflatelevel, 
+               shuffle = shuffle)
+    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
+YAB.allow_missings(::Type{<:NCDatasetsDataset}) = true
+YAB.allow_missings(::NCDatasetsDataset) = true
+
+# --- 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

@@ -4,6 +4,7 @@ AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
 AxisKeys = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 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
 
 using Pkg.Artifacts
 import Downloads
@@ -63,6 +63,23 @@ 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) == true
+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 +153,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