Refactor + add regression test

Author: Will Tebbutt

src/tables.jl

@@ -37,7 +37,7 @@ 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)
-    G = _get_keys_columns(A)
+    G = _get_keys_columns(keys_or_axes(A))
     C = (G..., vec(parent(A)))
     NamedTuple{L}(C)
 end
@@ -48,8 +48,7 @@ end
 # 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(A)
-    R = keys_or_axes(A)
+function _get_keys_columns(R)
     R_inds = map(eachindex, R)
     return map(
         (r, d) -> vec([r[indices[d]] for indices in Iterators.product(R_inds...)]),

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"])