Add support for KeyedArrays as a Tables sink.

Author: rofinn

src/tables.jl

@@ -115,3 +115,65 @@ end
 
 #     end
 # end
+
+"""
+    populate!(A, table, value)
+
+Populate A with the contents of the `value` column in a provided table.
+The provided table contain columns corresponding to the keys in A and support row iteration.
+"""
+function populate!(A, table, value::Symbol)
+    cols = [value, dimnames(A)...]
+    for r in Tables.rows(table)
+        setkey!(A, (Tables.getcolumn(r, c) for c in cols)...)
+    end
+    return A
+end
+
+"""
+    KeyedArray(table, value, keys...; default=undef)
+    NamedDimsArray(table, value, keys...; default=undef)
+
+Construct a KeyedArray/NamedDimsArray from a `table` supporting column and row.
+`keys` columns are extracted as is, without sorting, and are assumed to uniquely identify
+each `value`. The `default` value is used in cases where no value is identified for a given
+keypair.
+"""
+function KeyedArray(table, value::Symbol, keys::Symbol...; kwargs...)
+   return _construct_from_table(KeyedArray, table, value, keys...; kwargs...)
+end
+
+function NamedDimsArray(table, value::Symbol, keys::Symbol...; kwargs...)
+   return _construct_from_table(NamedDimsArray, table, value, keys...; kwargs...)
+end
+
+
+# Internal function for constructing the KeyedArray or NamedDimsArray.
+# This code doesn't care which type we produce so we just pass that along.
+function _construct_from_table(
+    T::Type, table, value::Symbol, keys::Symbol...;
+    default=undef, issorted=false
+)
+    # get columns of the input table source
+    cols = Tables.columns(table)
+
+    # Extract key columns
+    kw = Tuple(k => unique(Tables.getcolumn(cols, k)) for k in keys)
+
+    # Extract data/value column
+    vals = Tables.getcolumn(cols, value)
+
+    # Initialize the KeyedArray
+    sz = length.(last.(kw))
+
+    A = if default === undef
+        data = similar(vals, sz)
+    else
+        data = similar(vals, Union{eltype(vals), typeof(default)}, sz)
+        fill!(data, default)
+    end
+
+    A = T(data; kw...)
+    populate!(A, table, value)
+    return A
+end

test/_packages.jl

@@ -29,14 +29,39 @@ end
 @testset "tables" begin
     using Tables
 
-    R = wrapdims(rand(2,3), 11:12, 21:23)
-    N = wrapdims(rand(2,3), a=[11, 12], b=[21, 22, 23.0])
-
-    @test keys(first(Tables.rows(R))) == (:dim_1, :dim_2, :value)
-    @test keys(first(Tables.rows(N))) == (:a, :b, :value)
-
-    @test Tables.columns(N).a == [11, 12, 11, 12, 11, 12]
-
+    @testset "source" begin
+        R = wrapdims(rand(2,3), 11:12, 21:23)
+        N = wrapdims(rand(2,3), a=[11, 12], b=[21, 22, 23.0])
+
+        @test keys(first(Tables.rows(R))) == (:dim_1, :dim_2, :value)
+        @test keys(first(Tables.rows(N))) == (:a, :b, :value)
+
+        @test Tables.columns(N).a == [11, 12, 11, 12, 11, 12]
+    end
+    @testset "sink" begin
+        A = KeyedArray(rand(24, 11, 3); :time => 0:23, :loc => -5:5, :id => ["a", "b", "c"])
+        table = Tables.columntable(A)
+
+        # Test fully constructing from a table
+        # Common when working with adhoc data
+        B = KeyedArray(table, :value, :time, :loc, :id)
+        @test B == A
+
+        # Test populating an existing array (e.g., expected data based on calculated targets/offsets)
+        C = KeyedArray(
+            zeros(Float64, size(A));
+            time = unique(table.time),
+            loc = unique(table.loc),
+            id = unique(table.id),
+        )
+        @test C != A
+        AxisKeys.populate!(C, table, :value)
+        @test C == A
+
+        # Constructing a NamedDimsArray with different default value and table type
+        D = NamedDimsArray(Tables.rowtable(A), :value, :time, :loc, :id; default=missing)
+        @test D == A
+    end
 end
 @testset "stack" begin
     using LazyStack