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..608a2ea --- /dev/null +++ b/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 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