Skip to content

Commit

Permalink
Return arrays with correct dims, or scalars (#603)
Browse files Browse the repository at this point in the history 
Changes to align indexing interface with Base
  • Loading branch information
@rafaqz @musm
rafaqz authored and musm committed Dec 14, 2019
1 parent ba38047 commit 7bf8632
Show file tree
Hide file tree
Showing 2 changed files with 52 additions and 26 deletions.
63 changes: 38 additions & 25 deletions src/HDF5.jl
View file
Expand Up @@ -7,7 +7,7 @@ using Base: unsafe_convert, StringVector
import Base:
close, convert, eltype, lastindex, flush, getindex, ==,
isempty, isvalid, length, names, ndims, parent, read,
setindex!, show, size, sizeof, write, isopen, iterate
setindex!, show, size, sizeof, write, isopen, iterate, eachindex, axes

import Libdl
import Mmap
Expand Down Expand Up @@ -1747,23 +1747,30 @@ write(parent::Union{HDF5File, HDF5Group}, name::String, data::Union{T, AbstractA
# For datasets, "write(dset, name, val)" means "a_write"
write(parent::HDF5Dataset, name::String, data::Union{T, AbstractArray{T}}, plists...) where {T<:ScalarOrString} = a_write(parent, name, data, plists...)

# Reading arrays using getindex: data = dset[:,:,10]
function getindex(dset::HDF5Dataset, indices::Union{AbstractRange{Int},Int}...)

# Indexing

Base.eachindex(::IndexLinear, A::HDF5Dataset) = Base.OneTo(length(A))
Base.axes(dset::HDF5Dataset) = map(Base.OneTo, size(dset))

getindex(dset::HDF5Dataset, I::Union{AbstractRange,Integer,Colon}...) =
_getindex(dset, Base.to_indices(dset, I)...)
function _getindex(dset::HDF5Dataset, I::Union{AbstractRange{Int},Int}...)
local T
dtype = datatype(dset)
try
T = hdf5_to_julia_eltype(dtype)
finally
close(dtype)
end
_getindex(dset,T, indices...)
_getindex(dset,T, I...)
end
function _getindex(dset::HDF5Dataset, T::Type, indices::Union{AbstractRange{Int},Int}...)
function _getindex(dset::HDF5Dataset, T::Type, I::Union{AbstractRange{Int},Int}...)
if !(T <: Union{HDF5Scalar, Complex{<:HDF5Scalar}})
error("Dataset indexing (hyperslab) is available only for bits types")
end
dsel_id = hyperslab(dset, indices...)
ret = Array{T}(undef,map(length, indices))
dsel_id = hyperslab(dset, I...)
ret = Array{T}(undef,map(length, I))
memtype = datatype(ret)
memspace = dataspace(ret)
try
Expand All @@ -1773,29 +1780,38 @@ function _getindex(dset::HDF5Dataset, T::Type, indices::Union{AbstractRange{Int}
close(memspace)
h5s_close(dsel_id)
end
ret
dimsize = map(length, _dropint(I...))
dimsize == () ? ret[1] : reshape(ret, dimsize...)
end

_dropint(x::Int, args...) = _dropint(args...)
_dropint(x::AbstractRange, args...) = (x, _dropint(args...)...)
_dropint() = ()


# Write to a subset of a dataset using array slices: dataset[:,:,10] = array
function setindex!(dset::HDF5Dataset, X::Array, indices::Union{AbstractRange{Int},Int}...)

setindex!(dset::HDF5Dataset, x, I::Union{AbstractRange,Integer,Colon}...) =
_setindex!(dset, x, Base.to_indices(dset, I)...)
function _setindex!(dset::HDF5Dataset, X::Array, I::Union{AbstractRange{Int},Int}...)
T = hdf5_to_julia(dset)
_setindex!(dset, T, X, indices...)
_setindex!(dset, T, X, I...)
end
function _setindex!(dset::HDF5Dataset,T::Type, X::Array, indices::Union{AbstractRange{Int},Int}...)
function _setindex!(dset::HDF5Dataset,T::Type, X::Array, I::Union{AbstractRange{Int},Int}...)
if !(T <: Array)
error("Dataset indexing (hyperslab) is available only for arrays")
end
ET = eltype(T)
if !(ET <: Union{HDF5Scalar, Complex{<:HDF5Scalar}})
error("Dataset indexing (hyperslab) is available only for bits types")
end
if length(X) != prod(map(length, indices))
if length(X) != prod(map(length, I))
error("number of elements in range and length of array must be equal")
end
if eltype(X) != ET
X = convert(Array{ET}, X)
end
dsel_id = hyperslab(dset, indices...)
dsel_id = hyperslab(dset, I...)
memtype = datatype(X)
memspace = dataspace(X)
try
Expand All @@ -1807,42 +1823,39 @@ function _setindex!(dset::HDF5Dataset,T::Type, X::Array, indices::Union{Abstract
end
X
end
function setindex!(dset::HDF5Dataset, X::AbstractArray, indices::Union{AbstractRange{Int},Int}...)
function _setindex!(dset::HDF5Dataset, X::AbstractArray, I::Union{AbstractRange{Int},Int}...)
T = hdf5_to_julia(dset)
if !(T <: Array)
error("Hyperslab interface is available only for arrays")
end
Y = convert(Array{eltype(T), ndims(X)}, X)
setindex!(dset, Y, indices...)
_setindex!(dset, Y, I...)
end

function setindex!(dset::HDF5Dataset, x::Number, indices::Union{AbstractRange{Int},Int}...)
function _setindex!(dset::HDF5Dataset, x::Number, I::Union{AbstractRange{Int},Int}...)
T = hdf5_to_julia(dset)
if !(T <: Array)
error("Hyperslab interface is available only for arrays")
end
X = fill(convert(eltype(T), x), map(length, indices))
setindex!(dset, X, indices...)
X = fill(convert(eltype(T), x), map(length, I))
_setindex!(dset, X, I...)
end

getindex(dset::HDF5Dataset, I::Union{AbstractRange{Int},Int,Colon}...) = getindex(dset, ntuple(i-> isa(I[i], Colon) ? (1:size(dset,i)) : I[i], length(I))...)
setindex!(dset::HDF5Dataset, x, I::Union{AbstractRange{Int},Int,Colon}...) = setindex!(dset, x, ntuple(i-> isa(I[i], Colon) ? (1:size(dset,i)) : I[i], length(I))...)

function hyperslab(dset::HDF5Dataset, indices::Union{AbstractRange{Int},Int}...)
function hyperslab(dset::HDF5Dataset, I::Union{AbstractRange{Int},Int}...)
local dsel_id
dspace = dataspace(dset)
try
dims, maxdims = get_dims(dspace)
n_dims = length(dims)
if length(indices) != n_dims
error("Wrong number of indices supplied, supplied length $(length(indices)) but expected $(n_dims).")
if length(I) != n_dims
error("Wrong number of indices supplied, supplied length $(length(I)) but expected $(n_dims).")
end
dsel_id = h5s_copy(dspace.id)
dsel_start = Vector{Hsize}(undef,n_dims)
dsel_stride = Vector{Hsize}(undef,n_dims)
dsel_count = Vector{Hsize}(undef,n_dims)
for k = 1:n_dims
index = indices[n_dims-k+1]
index = I[n_dims-k+1]
if isa(index, Integer)
dsel_start[k] = index-1
dsel_stride[k] = 1
Expand Down
15 changes: 14 additions & 1 deletion test/extend_test.jl
View file
Expand Up @@ -20,8 +20,21 @@ d[1, 1:5] = [1.1231, 1.313, 5.123, 2.231, 4.1231]
set_dims!(d, (1, 5))
@test size(d) == (1, 5)

# Indexing returns correct array dimensions
@test d[1, end] ≈ 4.1231
@test d[:, end] ≈ [4.1231]
@test d[end, :] == [1.1231 1.313 5.123 2.231 4.1231]
@test d[end, :] == [1.1231, 1.313, 5.123, 2.231, 4.1231]
@test d[:, :] == [1.1231 1.313 5.123 2.231 4.1231]

# Test all integer types work
@test d[UInt8(1), UInt16(1)] == 1.1231
@test d[UInt32(1), UInt128(1)] == 1.1231
@test d[Int8(1), Int16(1)] == 1.1231
@test d[Int32(1), Int128(1)] == 1.1231

# Test ranges work with steps
@test d[1, 1:2:5] == [1.1231, 5.123, 4.1231]
@test d[1:1, 2:2:4] == [1.313 2.231]

#println("d is size current $(map(int,HDF5.get_dims(d)[1])) max $(map(int,HDF5.get_dims(d)[2]))")
b = d_create(fid, "b", Int, ((1000,), (-1,)), "chunk", (100,)) #-1 is equivalent to typemax(Hsize) as far as I can tell
Expand Down

0 comments on commit 7bf8632

Please sign in to comment.