Merge pull request #123 from willtebbutt/wct/columns-perf

Fix performance problem for `Tables.Columns(::KeyedArray)`

Author: oxinabox

Project.toml

@@ -1,7 +1,7 @@
 name = "AxisKeys"
 uuid = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 license = "MIT"
-version = "0.2.5"
+version = "0.2.6"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/tables.jl

@@ -37,15 +37,24 @@ Tables.columnaccess(::Type{<:KeyedArray{T,N,AT}}) where {T,N,AT} =
 function Tables.columns(A::Union{KeyedArray, NdaKa})
     L = hasnames(A) ? (dimnames(A)..., :value) :
         (ntuple(d -> Symbol(:dim_,d), ndims(A))..., :value)
-    R = keys_or_axes(A)
-    G = ntuple(ndims(A)) do d
-        vec([rs[d] for rs in Iterators.product(R...)])
-        # _vec(rs[d] for rs in Iterators.product(R...))
-    end
+    G = _get_keys_columns(keys_or_axes(A))
     C = (G..., vec(parent(A)))
     NamedTuple{L}(C)
 end
 
+# indices is a tuple, the dth element of which is an index for the dth column of R.
+# By using these indices, and mapping over the columns of R, the compiler seems to
+# successfully infer the types in G, because it knows the element types of each column
+# of R, so is presumably able to unroll the call to map.
+# The previous implementation called `Iterators.product` on `R` and pulled out
+# the dth element of `indices`, whose type it could not infer.
+function _get_keys_columns(R)
+    R_inds = map(eachindex, R)
+    return map(R, ntuple(identity, length(R))) do r, d
+        vec([r[indices[d]] for indices in Iterators.product(R_inds...)])
+    end
+end
+
 function Tables.Schema(nt::NamedTuple) # 🏴‍☠️
     L = keys(nt)
     T = map(v -> typeof(first(v)), values(nt))

test/_packages.jl

@@ -49,6 +49,15 @@ end
         @test keys(first(Tables.rows(N))) == (:a, :b, :value)
 
         @test Tables.columns(N).a == [11, 12, 11, 12, 11, 12]
+
+        # Regression test: https://github.com/mcabbott/AxisKeys.jl/pull/123
+        @testset "Tables.columns(::AxisKeys) type stability" begin
+            a = randn(5)
+            b = randn(Float32, 4)
+            @inferred AxisKeys._get_keys_columns((a, b))
+            X = wrapdims(randn(5, 4); a, b)
+            @inferred Tables.columns(X)
+        end
     end
     @testset "sink" begin
         A = KeyedArray(rand(24, 11, 3); time = 0:23, loc = -5:5, id = ["a", "b", "c"])