From 2dab34cd6667bc13b037fc44b586fee1344c8d6c Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Wed, 22 Jul 2026 13:19:28 -0400 Subject: [PATCH 1/2] add support for NCDatasets.jl --- Project.toml | 3 + ext/NCDatasetsExt.jl | 146 +++++++++++++++++++++++++++++++++++++++++++ test/Project.toml | 1 + test/datasets.jl | 23 ++++++- 4 files changed, 172 insertions(+), 1 deletion(-) create mode 100644 ext/NCDatasetsExt.jl diff --git a/Project.toml b/Project.toml index a7501ad..8136562 100644 --- a/Project.toml +++ b/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" diff --git a/ext/NCDatasetsExt.jl b/ext/NCDatasetsExt.jl new file mode 100644 index 0000000..e5bccf7 --- /dev/null +++ b/ext/NCDatasetsExt.jl @@ -0,0 +1,146 @@ +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, 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 + + # 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, + zstdlevel = zstdlevel) + 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 diff --git a/test/Project.toml b/test/Project.toml index e845daa..2bda035 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/datasets.jl b/test/datasets.jl index f89ae04..38244fe 100644 --- a/test/datasets.jl +++ b/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 From b61bd6e2716df976ad6675adbef313f20bbdc168 Mon Sep 17 00:00:00 2001 From: Ben Arthur Date: Fri, 18 Sep 2026 19:58:17 +0000 Subject: [PATCH 2/2] NCDatasetsExt: hand YAXArrays raw storage values, forward chunks The extension read and wrote through NCDatasets' CFVariable, so the CF decoding (scale_factor, add_offset, fill values, time units) was applied twice: once by NCDatasets and again by YAXArrays' CFDiskArray/timedecode. A variable with scale_factor 0.5 and add_offset 100 came back as 150.5 instead of 101, and a "days since" time axis got a corrupt element type. Wrap the raw Variable (nc[name].var) instead, as the NetCDF.jl backend does, and copy _FillValue into missing_value so YAXArrays honours it. Also: - forward haschunks/eachchunk to the file (they were imported but never defined, so every handle fell back to a single whole-array chunk) - iscompressed recognises Zstandard when NCDatasets provides zstandard() - open_dataset_handle no longer masks a failed open with close(nothing) - add_var passes deflate/zstd keywords only when requested, so zstdlevel never reaches a defVar that does not know it - allow_missings is false: defVar wants a DataType, and YAXArrays handles Union{Missing,T} through missing_value as for NetCDF.jl - NCDatasets compat is "0.14"; 0.12 predates the DiskArrays integration - tests: chunking, failed-open and raw-value/attribute round trip Co-Authored-By: Claude Fable 5.1 --- Project.toml | 2 +- ext/NCDatasetsExt.jl | 86 ++++++++++++++++++++++++++------------------ test/Project.toml | 1 + test/datasets.jl | 38 ++++++++++++++++++-- 4 files changed, 89 insertions(+), 38 deletions(-) diff --git a/Project.toml b/Project.toml index 8136562..5a1d583 100644 --- a/Project.toml +++ b/Project.toml @@ -15,7 +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" +NCDatasets = "0.14" NetCDF = "0.11, 0.12" Zarr = "0.10" diff --git a/ext/NCDatasetsExt.jl b/ext/NCDatasetsExt.jl index e5bccf7..9bf0916 100644 --- a/ext/NCDatasetsExt.jl +++ b/ext/NCDatasetsExt.jl @@ -7,6 +7,17 @@ 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 @@ -35,7 +46,7 @@ function YAB.open_dataset_handle(f, ds::NCDatasetsDataset) ds.handle[] = NCDataset(ds.filename, ds.mode) f(ds) finally - close(ds.handle[]) + ds.handle[] === nothing || close(ds.handle[]) ds.handle[] = nothing end else @@ -44,8 +55,10 @@ function YAB.open_dataset_handle(f, ds::NCDatasetsDataset) end # Implement the DiskArrays read/write interface -import DiskArrays: AbstractDiskArray -import NCDatasets: readblock!, writeblock!, haschunks, eachchunk +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} @@ -57,41 +70,47 @@ 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 +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 -# NCDatasets supports querying the deflate level from the variable +# deflate or Zstandard (the latter only with an NCDatasets that supports it) 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 + 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 -> dimnames(nc[name]), ds)...] -YAB.get_var_attrs(ds::NCDatasetsDataset, name) = dsopen(nc -> Dict(nc[name].attrib), 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(nc[i]), eltype(nc[i])), ds) + s, et = dsopen(nc -> (size(rawvar(nc, i)), eltype(rawvar(nc, i))), ds) NCDatasetsVariable{et, length(s)}(ds.filename, ds.mode, i, s) else - ds.handle[][i] + rawvar(ds.handle[], i) end end @@ -114,17 +133,12 @@ function YAB.add_var(p::NCDatasetsDataset, T::Type, varname, s, dimnames, attr; 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, - zstdlevel = zstdlevel) + # 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) @@ -133,8 +147,10 @@ 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 +# 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__() diff --git a/test/Project.toml b/test/Project.toml index 2bda035..9864af0 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,7 @@ 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" diff --git a/test/datasets.jl b/test/datasets.jl index 38244fe..a3007f3 100644 --- a/test/datasets.jl +++ b/test/datasets.jl @@ -3,7 +3,7 @@ using YAXArrayBase, Test @test_throws "No backend found." YAXArrayBase.backendfrompath("test.zarr") end -using NetCDF, Zarr, NCDatasets +using NetCDF, Zarr, NCDatasets, DiskArrays using Pkg.Artifacts import Downloads @@ -77,7 +77,41 @@ end @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 + @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