Support NetCDF input tables as Arrow alternative (#2542)

Fixes most of #2434, not all because only a subset of the tables are
supported.

This essentially gives the same treatment we currently have for Arrow
input files, but for NetCDF files. We choose between the file format
based on whether the path has a `.arrow` or `.nc` extension.

So in Python if you do:
```py
model.basin.state.set_filepath(Path("basin-state.nc"))
```
(Note you need
[this](#2039 (comment))
if you read the model from disk.)

Ribasim Python will convert the `df.to_xarray().to_netcdf()`, and if you
read it back in with Python `ds.to_dataframe()`.
In the TOML you get:
```toml
[basin]
state = "basin-state.nc"
```

The core then also reads based on the file extension. To keep things
simple, when reading data into the core we convert the N dimensional
arrays in the NetCDF to tables, so the rest of the initialization
doesn't need to handle both data structures.

We should also think about what to do for irregular tables that don't
fit in a structured array. Perhaps just throw an error?

Author: visr

core/src/Ribasim.jl

@@ -153,7 +153,7 @@ using StructArrays: StructVector
 using DataStructures: OrderedSet, OrderedDict, counter, inc!
 
 # NCDatasets is used to read and write NetCDF files.
-using NCDatasets: NCDataset, defDim, defVar
+using NCDatasets: NCDataset, defDim, defVar, dimnames
 
 using Dates: Second
 

core/src/config.jl

@@ -259,7 +259,7 @@ function Base.getproperty(config::Config, sym::Symbol)
 end
 
 "Construct a path relative to both the TOML directory and the optional `input_dir`"
-function input_path(config::Config, path::String="")
+function input_path(config::Config, path::String = "")
     return normpath(config.dir, config.input_dir, path)
 end
 
@@ -269,7 +269,7 @@ function database_path(config::Config)
 end
 
 "Construct a path relative to both the TOML directory and the optional `results_dir`"
-function results_path(config::Config, path::String="")
+function results_path(config::Config, path::String = "")
     # If the path is empty, we return the results directory.
     if !isempty(path)
         name, ext = splitext(path)

core/src/read.jl

@@ -1714,40 +1714,135 @@ DateTime. This is used to convert between the solver's inner float time, and the
 datetime_since(t::Real, t0::DateTime)::DateTime = t0 + Millisecond(round(1000 * t))
 
 """
-    load_data(db::DB, config::Config, nodetype::Symbol, kind::Symbol)::Union{Arrow.Table, Query, Nothing}
+    load_netcdf(table_path::String, table_type::Type{<:Table})::NamedTuple
 
-Load data from Arrow files if available, otherwise the database.
-Returns either an `Arrow.Table`, `SQLite.Query` or `nothing` if the data is not present.
+Load a table from a NetCDF file. The data is stored as multi-dimensional arrays, and
+converted to a table for compatibility with the rest of the internals.
+"""
+function load_netcdf(table_path::String, table_type::Type{<:Table})::NamedTuple
+    table = NCDataset(table_path) do ds
+        names = fieldnames(table_type)
+        table = OrderedDict{Symbol, AbstractVector}()
+        data_varnames = filter(x -> !(String(x) in nc_dim_names), names)
+        for data_varname in data_varnames
+            var = ds[data_varname]
+            dim_names = dimnames(var)
+            if dim_names == ("node_id",)
+                table[:node_id] = ds["node_id"][:]
+            elseif dim_names == ("node_id", "time")
+                node_id_data = ds["node_id"][:]
+                time_data = ds["time"][:]
+                ntime = length(time_data)
+                nnode = length(node_id_data)
+                table[:node_id] = repeat(node_id_data; outer = ntime)
+                table[:time] = repeat(time_data; inner = nnode)
+            else
+                error("Unsupported dimensions: $dim_names, must be (node_id, [time])")
+            end
+            table[data_varname] = vec(var[:])
+        end
+        table
+    end
+    return columntable(table)
+end
+
+"""
+    load_data(db::DB, config::Config, nodetype::Symbol, kind::Symbol)::Union{NamedTuple, Nothing}
+
+Load data from Arrow or NetCDF files if available, otherwise the database.
+Returns either a `NamedTuple` of Vectors or `nothing` if the data is not present.
 """
 function load_data(
     db::DB,
     config::Config,
     table_type::Type{<:Table},
-)::Union{Arrow.Table, Query, Nothing}
-    # TODO load_data doesn't need both config and db, use config to check which one is needed
-
+)::Union{NamedTuple, Nothing}
     toml = getfield(config, :toml)
     section_name = snake_case(node_type(table_type))
     section = getproperty(toml, section_name)
     kind = table_name(table_type)
     sql_name = sql_table_name(table_type)
 
-    path = if hasproperty(section, kind)
-        getproperty(section, kind)
-    else
-        nothing
-    end
+    path = hasproperty(section, kind) ? getproperty(section, kind) : nothing
 
-    table = if !isnothing(path)
+    if !isnothing(path)
+        # the TOML specifies a file outside the database
+        path = getproperty(section, kind)
         table_path = input_path(config, path)
-        Arrow.Table(read(table_path); convert = false)
-    elseif exists(db, sql_name)
-        execute(db, "select * from $(esc_id(sql_name))")
+        # check suffix and read with Arrow or NCDatasets
+        ext = lowercase(splitext(table_path)[2])
+        if ext == ".nc"
+            return load_netcdf(table_path, table_type)
+        elseif ext == ".arrow"
+            bytes = read(table_path)
+            arrow_table = Arrow.Table(bytes; convert = false)
+            return arrow_columntable(arrow_table, table_type)
+        else
+            error("Unsupported file format: $table_path")
+        end
     else
-        nothing
+        if exists(db, sql_name)
+            table = execute(db, "select * from $(esc_id(sql_name))")
+            return sqlite_columntable(table, db, config, table_type)
+        else
+            return nothing
+        end
+    end
+end
+
+"Faster alternative to Tables.columntable that preallocates based on the schema."
+function sqlite_columntable(
+    table::Query,
+    db::DB,
+    config::Config,
+    T::Type{<:Table},
+)::NamedTuple
+    sql_name = sql_table_name(T)
+    nrows = execute(db, "SELECT COUNT(*) FROM $(esc_id(sql_name))") |> first |> first
+
+    names = fieldnames(T)
+    types = fieldtypes(T)
+    vals = ntuple(i -> Vector{types[i]}(undef, nrows), length(names))
+    nt = NamedTuple{names}(vals)
+
+    for (i, row) in enumerate(table)
+        for name in names
+            val = row[name]
+            if name == :time
+                # time has type timestamp and is stored as a String in the database
+                # currently SQLite.jl does not automatically convert it to DateTime
+                val = if ismissing(val)
+                    DateTime(config.starttime)
+                else
+                    DateTime(
+                        replace(val, r"(\.\d{3})\d+$" => s"\1"),  # remove sub ms precision
+                        dateformat"yyyy-mm-dd HH:MM:SS.s",
+                    )
+                end
+            end
+            nt[name][i] = val
+        end
     end
+    nt
+end
 
-    return table
+"Alternative to Tables.columntable that converts time to our own to_datetime."
+function arrow_columntable(table::Query, T::Type{<:Table})::NamedTuple
+    nrows = length(first(table))
+    names = fieldnames(T)
+    types = fieldtypes(T)
+    vals = ntuple(i -> Vector{types[i]}(undef, nrows), length(names))
+    nt = NamedTuple{names}(vals)
+
+    for name in names
+        if name == :time
+            time_col = getproperty(table, name)
+            nt[name] .= [to_datetime(t) for t in time_col]
+        else
+            nt[name] .= getproperty(table, name)
+        end
+    end
+    nt
 end
 
 # alternative to convert that doesn't have warntimestamp
@@ -1762,7 +1857,7 @@ end
 """
     load_structvector(db::DB, config::Config, ::Type{T})::StructVector{T}
 
-Load data from Arrow files if available, otherwise the database.
+Load data from Arrow or NetCDF files if available, otherwise the database.
 Always returns a StructVector of the given struct type T, which is empty if the table is
 not found. This function validates the schema, and enforces the required sort order.
 """
@@ -1771,62 +1866,12 @@ function load_structvector(
     config::Config,
     ::Type{T},
 )::StructVector{T} where {T <: Table}
-    table = load_data(db, config, T)
+    nt = load_data(db, config, T)
 
-    if table === nothing
+    if nt === nothing
         return StructVector{T}(undef, 0)
     end
 
-    table_in_db = table isa Query
-
-    nt = if table_in_db
-        # faster alternative to Tables.columntable that preallocates based on the schema
-        sql_name = sql_table_name(T)
-        nrows =
-            execute(db, "SELECT COUNT(*) FROM $(esc_id(sql_name))") |> first |> first
-
-        names = fieldnames(T)
-        types = fieldtypes(T)
-        vals = ntuple(i -> Vector{types[i]}(undef, nrows), length(names))
-        nt = NamedTuple{names}(vals)
-
-        for (i, row) in enumerate(table)
-            for name in names
-                val = row[name]
-                if name == :time
-                    # time has type timestamp and is stored as a String in the database
-                    # currently SQLite.jl does not automatically convert it to DateTime
-                    val = if ismissing(val)
-                        DateTime(config.starttime)
-                    else
-                        DateTime(
-                            replace(val, r"(\.\d{3})\d+$" => s"\1"),  # remove sub ms precision
-                            dateformat"yyyy-mm-dd HH:MM:SS.s",
-                        )
-                    end
-                end
-                nt[name][i] = val
-            end
-        end
-        nt
-    else
-        nrows = length(first(table))
-        names = fieldnames(T)
-        types = fieldtypes(T)
-        vals = ntuple(i -> Vector{types[i]}(undef, nrows), length(names))
-        nt = NamedTuple{names}(vals)
-
-        for name in names
-            if name == :time
-                time_col = getproperty(table, name)
-                nt[name] .= [to_datetime(t) for t in time_col]
-            else
-                nt[name] .= getproperty(table, name)
-            end
-        end
-        nt
-    end
-
     table = StructVector{T}(nt)
     return sorted_table!(table)
 end

core/test/io_test.jl

@@ -314,3 +314,54 @@ end
     storage2_begin = current_storage
     @test storage1_end ≈ storage2_begin
 end
+
+@testitem "warm state netcdf" begin
+    # This tests that we can write Basin / state results to NetCDF, and read this in again
+    # as a warm state, such that the storages at the end of one run are equal to those
+    # at the beginning of the second run.
+
+    using IOCapture: capture
+    using Ribasim: solve!, write_results
+    import TOML
+
+    model_path_src = normpath(@__DIR__, "../../generated_testmodels/basic/")
+
+    # avoid changing the original model for other tests
+    model_path = normpath(@__DIR__, "../../generated_testmodels/basic_warm_netcdf/")
+    cp(model_path_src, model_path; force = true)
+    toml_path = normpath(model_path, "ribasim.toml")
+
+    # Configure model to use NetCDF format
+    toml_dict = TOML.parsefile(toml_path)
+    toml_dict["results"] = Dict("format" => "netcdf")
+    open(toml_path, "w") do io
+        TOML.print(io, toml_dict)
+    end
+
+    config = Ribasim.Config(toml_path)
+    model = Ribasim.Model(config)
+    (; p_independent, state_time_dependent_cache) = model.integrator.p
+    (; current_storage) = state_time_dependent_cache
+    storage1_begin = copy(current_storage)
+    solve!(model)
+    storage1_end = current_storage
+    @test storage1_begin != storage1_end
+
+    # copy state results to input
+    write_results(model)
+    state_path = Ribasim.results_path(config, Ribasim.RESULTS_FILENAME.basin_state)
+    cp(state_path, Ribasim.input_path(config, "warm_state.nc"))
+
+    # point TOML to the warm state NetCDF file
+    toml_dict = TOML.parsefile(toml_path)
+    toml_dict["basin"] = Dict("state" => "warm_state.nc")
+    open(toml_path, "w") do io
+        TOML.print(io, toml_dict)
+    end
+
+    model = Ribasim.Model(toml_path)
+    (; p_independent, state_time_dependent_cache) = model.integrator.p
+    (; current_storage) = state_time_dependent_cache
+    storage2_begin = current_storage
+    @test storage1_end ≈ storage2_begin
+end

pixi.lock

@@ -235,7 +235,7 @@ jupyter = "*"
 matplotlib = ">=3.7"
 minio = "*"
 mypy = "*"
-netcdf4 = "*"
+netcdf4 = ">=1.7.1"
 networkx = ">=3.3"
 numpy = ">=1.25, <2.2"
 packaging = ">=23.0"
@@ -264,7 +264,7 @@ teamcity-messages = "*"
 tomli = ">=2.0"
 tomli-w = ">=1.0"
 twine = "*"
-xarray = "*"
+xarray = ">=2025.8.0"
 xmipy = ">=1.3"
 xugrid = "*"
 

pixi.toml

@@ -19,6 +19,7 @@ dependencies = [
     "datacompy >=0.16",
     "geopandas >=1.0",
     "matplotlib >=3.7",
+    "netCDF4 >=1.7.1",
     "numpy >=1.25",
     "packaging >=23.0",
     "pandas >=2.0",
@@ -29,6 +30,7 @@ dependencies = [
     "shapely >=2.0",
     "tomli >=2.0",
     "tomli-w >=1.0",
+    "xarray >=2025.8.0",
 ]
 dynamic = ["version"]