diff --git a/src/functions.jl b/src/functions.jl index 3d152a7..f1758f8 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -94,7 +94,13 @@ function Base.permutedims(A::KeyedArray, perm) KeyedArray(data, new_keys)#, copy(A.meta)) end -if VERSION >= v"1.1" +@static if VERSION > v"1.9-DEV" + function Base.eachslice(A::KeyedArray; dims) + dims_ix = AxisKeys.dim(A, dims) |> Tuple + data = @invoke eachslice(A::AbstractArray; dims=dims_ix) + return KeyedArray(NamedDimsArray(data, map(d -> dimnames(A, d), dims_ix)), map(d -> axiskeys(A, d), dims_ix)) + end +elseif VERSION >= v"1.1" # This copies the implementation from Base, except with numerical_dims: @inline function Base.eachslice(A::KeyedArray; dims) numerical_dims = NamedDims.dim(A, dims) @@ -107,6 +113,19 @@ if VERSION >= v"1.1" end end +@static if VERSION > v"1.9-DEV" + # TODO: this will ERROR if given dims, instead of falling back to Base + # TODO: ideally it would dispatch on the element type, for e.g. a generator of KeyedArrays + function Base.stack(A::KeyedArray; dims::Colon=:) + data = @invoke stack(A::AbstractArray; dims) + if !allequal(named_axiskeys(a) for a in A) + throw(DimensionMismatch("stack expects uniform axiskeys for all arrays")) + end + akeys = (; named_axiskeys(first(A))..., named_axiskeys(A)...) + KeyedArray(data; akeys...) + end +end + function Base.mapslices(f, A::KeyedArray; dims) numerical_dims = NamedDims.dim(A, dims) data = mapslices(f, parent(A); dims=numerical_dims) diff --git a/src/selectors.jl b/src/selectors.jl index 3963d8d..62f5917 100644 --- a/src/selectors.jl +++ b/src/selectors.jl @@ -208,8 +208,14 @@ end @inline Base.to_indices(A::Union{KeyedArray,NdaKa}, ax, inds::Tuple{Function, Vararg}) = select_to_indices(A, ax, inds) -using Base: to_indices, tail, safe_tail, uncolon +using Base: to_indices, tail, safe_tail + +if VERSION > v"1.9-DEV" + const uncolon = Base.uncolon +else + uncolon(inds) = Base.uncolon(inds, (:,)) +end @inline Base.to_indices(A::Union{KeyedArray,NdaKa}, inds, I::Tuple{Colon, Vararg{Any}}) = - (uncolon(inds, I), to_indices(A, safe_tail(inds), tail(I))...) + (uncolon(inds), to_indices(A, safe_tail(inds), tail(I))...) diff --git a/src/stack.jl b/src/stack.jl index 0d035dc..1849591 100644 --- a/src/stack.jl +++ b/src/stack.jl @@ -1,5 +1,5 @@ -using LazyStack +import LazyStack # for stack_iter LazyStack.no_wraps(a::KeyedArray) = LazyStack.no_wraps(NamedDims.unname(parent(a))) @@ -23,14 +23,14 @@ stack_keys(xs::Tuple{Vararg{KeyedArray}}) = # array of arrays: first strip off outer containers... function LazyStack.stack(xs::KeyedArray{<:AbstractArray}) - KeyedArray(stack(parent(xs)), stack_keys(xs)) + KeyedArray(LazyStack.stack(parent(xs)), stack_keys(xs)) end function LazyStack.stack(xs::KeyedArray{<:AbstractArray,N,<:NamedDimsArray{L}}) where {L,N} - data = stack(parent(parent(xs))) + data = LazyStack.stack(parent(parent(xs))) KeyedArray(LazyStack.ensure_named(data, LazyStack.getnames(xs)), stack_keys(xs)) end function LazyStack.stack(xs::NamedDimsArray{L,<:AbstractArray,N,<:KeyedArray}) where {L,N} - data = stack(parent(parent(xs))) + data = LazyStack.stack(parent(parent(xs))) LazyStack.ensure_named(KeyedArray(data, stack_keys(xs)), LazyStack.getnames(xs)) end diff --git a/test/_basic.jl b/test/_basic.jl index de90d1d..59db1ca 100644 --- a/test/_basic.jl +++ b/test/_basic.jl @@ -151,7 +151,7 @@ end N2 = NamedDimsArray(KeyedArray(data, axiskeys(N1)), dimnames(N1)) N3 = KeyedArray(NamedDimsArray(data, dimnames(N1)), axiskeys(N1)) - @testset "with $(typeof(N).name) outside" for N in [N2, N3] + @testset "with $(typeof(N).name) outside" for N in (N2, N3) @test axiskeys(N) == (['a', 'b', 'c'], 10:10:40) @test axiskeys(N, :iter) == 10:10:40 @test dimnames(N) == (:obs, :iter) diff --git a/test/_functions.jl b/test/_functions.jl index 83410bc..189bf1e 100644 --- a/test/_functions.jl +++ b/test/_functions.jl @@ -46,6 +46,17 @@ A3 = wrapdims(rand(Int8, 3,4,2), r='a':'c', c=2:5, p=[10.0, 20.0]) if VERSION >= v"1.1" @test axiskeys(first(eachslice(M, dims=:r))) === (2:5,) end + if VERSION >= v"1.9-DEV" + arr = KeyedArray([ + KeyedArray([1, 2], a=[:x, :y]), + KeyedArray([3, 4], a=[:x, :y]), + KeyedArray([5, 6], a=[:x, :y]), + ], b=10:12) + + sk = stack(arr)::KeyedArray + @test sk == [1 3 5; 2 4 6] + @test named_axiskeys(sk) == (a=[:x, :y], b=10:12) + end @test axiskeys(selectdim(M, :r, [true, false, true])) == (['a', 'c'], 2:5) diff --git a/test/_packages.jl b/test/_packages.jl index 50facda..35bfa08 100644 --- a/test/_packages.jl +++ b/test/_packages.jl @@ -129,25 +129,25 @@ end end end @testset "stack" begin - using LazyStack + using LazyStack: stack as lstack rin = [wrapdims(1:3, a='a':'c') for i=1:4] - @test axiskeys(stack(rin), :a) == 'a':'c' - @test axiskeys(stack(:b, rin...), :a) == 'a':'c' # tuple - @test axiskeys(stack(z for z in rin), :a) == 'a':'c' # generator + @test axiskeys(lstack(rin), :a) == 'a':'c' + @test axiskeys(lstack(:b, rin...), :a) == 'a':'c' # tuple + @test axiskeys(lstack(z for z in rin), :a) == 'a':'c' # generator rout = wrapdims([[1,2], [3,4]], b=10:11) - @test axiskeys(stack(rout), :b) == 10:11 + @test axiskeys(lstack(rout), :b) == 10:11 rboth = wrapdims(rin, b=10:13) - @test axiskeys(stack(rboth), :a) == 'a':'c' - @test axiskeys(stack(rboth), :b) == 10:13 + @test axiskeys(lstack(rboth), :a) == 'a':'c' + @test axiskeys(lstack(rboth), :b) == 10:13 nts = [(i=i, j="j", k=33) for i=1:3] - @test axiskeys(stack(nts), 1) == [:i, :j, :k] - @test axiskeys(stack(:z, nts...), 1) == [:i, :j, :k] - @test axiskeys(stack(n for n in nts), 1) == [:i, :j, :k] + @test axiskeys(lstack(nts), 1) == [:i, :j, :k] + @test axiskeys(lstack(:z, nts...), 1) == [:i, :j, :k] + @test axiskeys(lstack(n for n in nts), 1) == [:i, :j, :k] end @testset "dates" begin diff --git a/test/runtests.jl b/test/runtests.jl index 9cd94c5..c8bdc9b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,5 @@ using Test, AxisKeys, NamedDims -using Statistics, OffsetArrays, Tables, UniqueVectors, LazyStack +using Statistics, OffsetArrays, Tables, UniqueVectors using ChainRulesCore: ProjectTo, NoTangent using ChainRulesTestUtils: test_rrule using FiniteDifferences