diff --git a/docs/src/pythoncall-reference.md b/docs/src/pythoncall-reference.md index f09d5f12..7b39f351 100644 --- a/docs/src/pythoncall-reference.md +++ b/docs/src/pythoncall-reference.md @@ -185,6 +185,7 @@ PySet PyDict PyIterable PyArray +PyDenseArray PyIO PyTable PyPandasDataFrame diff --git a/docs/src/pythoncall.md b/docs/src/pythoncall.md index 7c604b9a..8966318c 100644 --- a/docs/src/pythoncall.md +++ b/docs/src/pythoncall.md @@ -254,6 +254,24 @@ Python: array('i', [0, 4, 5]) It directly wraps the underlying data buffer, so array operations such as indexing are about as fast as for an ordinary `Array`. +If the data is contiguous in memory then [`PyDenseArray`](@ref) can wrap it as a +`DenseArray` instead, so that it works with code specialised for dense or strided arrays, +such as BLAS. Julia arrays are column-major whereas numpy arrays are row-major by default, so +the dimensions are reversed when the data is row-major: + +```julia-repl +julia> x = pyimport("numpy").arange(6.0).reshape(2, 3) +Python: +array([[0., 1., 2.], + [3., 4., 5.]]) + +julia> PyDenseArray(x) +3×2 PyDenseArray{Float64, 2}: + 0.0 3.0 + 1.0 4.0 + 2.0 5.0 +``` + The [`PyIO`](@ref) wrapper type views a Python file object as a Julia IO object: ```julia-repl diff --git a/src/API/exports.jl b/src/API/exports.jl index 3d476182..730b6194 100644 --- a/src/API/exports.jl +++ b/src/API/exports.jl @@ -113,6 +113,7 @@ export pyconvert_unconverted # Wrap export PyArray +export PyDenseArray export PyDict export PyIO export PyIterable diff --git a/src/API/types.jl b/src/API/types.jl index 55a4211e..5d160bc4 100644 --- a/src/API/types.jl +++ b/src/API/types.jl @@ -90,6 +90,43 @@ struct PyArray{T,N,M,L,R} <: AbstractArray{T,N} end end +""" + PyDenseArray{T,N,M}(x; copy=true, array=true, buffer=true) <: DenseArray + +Wrap the Python array `x` as a Julia `DenseArray{T,N}`. + +This is like [`PyArray`](@ref) but requires the data to be contiguous in memory, so that the +result can be used wherever a `DenseArray` or `StridedArray` is expected, such as BLAS +routines. + +Julia arrays are column-major but most Python arrays (including `numpy.ndarray` by default) +are row-major. If the data is row-major then the dimensions are reversed, so a numpy array +of shape `(2, 3)` becomes a `PyDenseArray` of size `(3, 2)`. Column-major arrays +keep their shape. + +The type parameters are all optional, and are identical to the `T`, `N` and `M` +parameters of `PyArray`. The element type `T` is always the element type of the +underlying buffer. +""" +struct PyDenseArray{T,N,M} <: DenseArray{T,N} + ptr::Ptr{T} # pointer to the data + size::NTuple{N,Int} # size of the array (reversed if the data is row-major) + py::Py # underlying python object + handle::Py # the data in this array is valid as long as this handle is alive + function PyDenseArray{T,N,M}( + ::Val{:new}, + ptr::Ptr{T}, + size::NTuple{N,Int}, + py::Py, + handle::Py, + ) where {T,N,M} + T isa DataType || error("T must be a DataType") + N isa Int || error("N must be an Int") + M isa Bool || error("M must be a Bool") + new{T,N,M}(ptr, size, py, handle) + end +end + """ PyDict{K=Py,V=Py}([x]) diff --git a/src/Compat/serialization.jl b/src/Compat/serialization.jl index 41883a90..4664ba79 100644 --- a/src/Compat/serialization.jl +++ b/src/Compat/serialization.jl @@ -46,18 +46,21 @@ end Serialization.deserialize(s::AbstractSerializer, ::Type{PyException}) = PyException(deserialize_py(s)) -### PyArray +### PyArray and PyDenseArray # -# This type holds a pointer and a handle (usually a python memoryview or capsule) which are +# These types hold a pointer and a handle (usually a python memoryview or capsule) which are # not serializable by default, and even if they were would not be consistent after # serializing each field independently. So we just serialize the wrapped Python object. -function Serialization.serialize(s::AbstractSerializer, x::PyArray) +function Serialization.serialize(s::AbstractSerializer, x::Union{PyArray,PyDenseArray}) Serialization.serialize_type(s, typeof(x), false) serialize_py(s, x.py) end -function Serialization.deserialize(s::AbstractSerializer, ::Type{T}) where {T<:PyArray} +function Serialization.deserialize( + s::AbstractSerializer, + ::Type{T}, +) where {T<:Union{PyArray,PyDenseArray}} # TODO: set buffer and array args too? T(deserialize_py(s); copy = false) end diff --git a/src/Wrap/PyArray.jl b/src/Wrap/PyArray.jl index dc11c085..8ed5a0f1 100644 --- a/src/Wrap/PyArray.jl +++ b/src/Wrap/PyArray.jl @@ -729,3 +729,110 @@ function pyarray_check_T(::Type{T}, ::Type{R}) where {T,R} error("invalid eltype T=$T for raw eltype R=$R") end end + +# PyDenseArray + +ispy(::PyDenseArray) = true +Py(x::PyDenseArray) = x.py +Utils.ismutablearray(x::PyDenseArray{T,N,M}) where {T,N,M} = M + +function pydensearray_iscontiguous(strides, expected, size) + all(size[i] == 1 || strides[i] == expected[i] for i in eachindex(size)) +end + +# The size of the dense array if the data of x is contiguous, or nothing if it is not. +# Dimensions are reversed if the data is row-major. +function pydensearray_size(x::PyArray{T,N,M,L,T}) where {T,N,M,L} + if x.length == 0 + x.size + elseif pydensearray_iscontiguous(x.strides, + Utils.size_to_fstrides(sizeof(T), x.size), + x.size) + x.size + elseif pydensearray_iscontiguous(x.strides, + Utils.size_to_cstrides(sizeof(T), x.size), + x.size) + reverse(x.size) + else + nothing + end +end + +pydensearray_size(::PyArray) = nothing + +function PyDenseArray(x::PyArray{T,N,M,L,T}) where {T,N,M,L} + size = pydensearray_size(x) + if isnothing(size) + error("array data is not contiguous") + end + + PyDenseArray{T,N,M}(Val(:new), x.ptr, size, x.py, x.handle) +end + +function pydensearray_make( + ::Type{A}, + x::Py; + array::Bool = true, + buffer::Bool = true, + copy::Bool = true, +) where {A<:PyDenseArray} + r = pyarray_make(PyArray, x; array, buffer, copy) + if pyconvert_isunconverted(r) + return pyconvert_unconverted() + end + + p = pyconvert_result(PyArray, r) + if pydensearray_size(p) === nothing + return pyconvert_unconverted() + end + + d = PyDenseArray(p) + if d isa A + return pyconvert_return(d) + else + return pyconvert_unconverted() + end +end + +(::Type{A})( + x; + array::Bool = true, + buffer::Bool = true, + copy::Bool = true, +) where {A<:PyDenseArray} = @autopy x begin + r = pydensearray_make(A, x_; array, buffer, copy) + if pyconvert_isunconverted(r) + error("cannot convert this Python '$(pytype(x_).__name__)' to a '$A'") + else + return pyconvert_result(r)::A + end +end + +pyconvert_rule_densearray_nocopy(::Type{A}, x::Py) where {A<:PyDenseArray} = + pydensearray_make(A, x; copy = false) + +Base.size(x::PyDenseArray) = x.size +Base.IndexStyle(::Type{<:PyDenseArray}) = Base.IndexLinear() +Base.unsafe_convert(::Type{Ptr{T}}, x::PyDenseArray{T}) where {T} = x.ptr +Base.elsize(::Type{<:PyDenseArray{T}}) where {T} = sizeof(T) + +function Base.showarg(io::IO, x::PyDenseArray{T,N}, toplevel::Bool) where {T,N} + if !toplevel + print(io, "::") + end + + print(io, "PyDenseArray{") + show(io, T) + print(io, ", ", N, "}") +end + +@propagate_inbounds function Base.getindex(x::PyDenseArray, i::Int) + @boundscheck checkbounds(x, i) + unsafe_load(x.ptr, i) +end + +@propagate_inbounds function Base.setindex!(x::PyDenseArray{T,N,true}, v, i::Int) where {T,N} + @boundscheck checkbounds(x, i) + unsafe_store!(x.ptr, convert(T, v), i) + return x +end diff --git a/src/Wrap/Wrap.jl b/src/Wrap/Wrap.jl index d3b30ff2..c206bfae 100644 --- a/src/Wrap/Wrap.jl +++ b/src/Wrap/Wrap.jl @@ -14,7 +14,7 @@ using ..Convert using ..PyMacro import ..PythonCall: - PyArray, PyDict, PyIO, PyIterable, PyList, PyPandasDataFrame, PySet, PyTable + PyArray, PyDenseArray, PyDict, PyIO, PyIterable, PyList, PyPandasDataFrame, PySet, PyTable using Base: @propagate_inbounds using Tables: Tables @@ -81,6 +81,10 @@ function __init__() pyconvert_add_rule("", AbstractArray, pyconvert_rule_array, priority) pyconvert_add_rule("", AbstractArray, pyconvert_rule_array, priority) pyconvert_add_rule("", AbstractArray, pyconvert_rule_array, priority) + pyconvert_add_rule("", PyDenseArray, pyconvert_rule_densearray_nocopy, priority) + pyconvert_add_rule("", PyDenseArray, pyconvert_rule_densearray_nocopy, priority) + pyconvert_add_rule("", PyDenseArray, pyconvert_rule_densearray_nocopy, priority) + pyconvert_add_rule("", PyDenseArray, pyconvert_rule_densearray_nocopy, priority) end end diff --git a/test/Project.toml b/test/Project.toml index 3f088776..a7294ee4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -2,6 +2,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" PyCall = "438e738f-606a-5dbb-bf0a-cddfbfd45ab0" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" diff --git a/test/Wrap.jl b/test/Wrap.jl index 09d6ece7..b9fcb349 100644 --- a/test/Wrap.jl +++ b/test/Wrap.jl @@ -101,6 +101,114 @@ end end +@testitem "PyDenseArray" setup=[Setup] begin + using LinearAlgebra + x = pyimport("array").array("d", pylist(0:5)) + y = PyDenseArray(x) + + # Helper function to create a row-major Float64 array + function rowmajor(vals, shape) + arr = pyimport("array").array("d", pylist(vals)) + bytearr = pybuiltins.bytearray(arr.tobytes()) + + pybuiltins.memoryview(bytearr).cast("d", pylist(shape)) + end + + @testset "construct" begin + @test y isa PyDenseArray{Float64,1,true} + @test y isa StridedVector{Float64} + @test Py(y) === x + @test PyDenseArray{Float64,1,true}(x) isa PyDenseArray{Float64,1,true} + @test PyDenseArray(PyArray(x)) isa PyDenseArray{Float64,1,true} + @test PyDenseArray(pybytes(b"abc")) isa PyDenseArray{UInt8,1,false} + + @test pyconvert(PyDenseArray, x) isa PyDenseArray{Float64,1,true} + @test pyconvert(PyDenseArray{Float64,1}, x) isa PyDenseArray{Float64,1,true} + # Defaults are unchanged + @test pyconvert(Any, x) isa PyArray + @test pyconvert(DenseArray, x) isa Array + + @test_throws Exception PyDenseArray{Int}(x) + @test_throws Exception PyDenseArray{Float64,1,false}(x) + # Non-contiguous + strided = pybuiltins.memoryview(x)[pyslice(nothing, nothing, 2)] + @test_throws Exception PyDenseArray(strided) + @test_throws Exception pyconvert(PyDenseArray, strided) + + if Setup.devdeps + np = pyimport("numpy") + # Object arrays have eltype Py, which is not the buffer eltype + @test_throws Exception PyDenseArray(np.array(pylist([1, "a"]), dtype = np.object_)) + end + end + + @testset "shape" begin + # Row-major data is reversed + c = PyDenseArray(rowmajor(0:5, [2, 3])) + @test size(c) == (3, 2) + @test strides(c) == (1, 3) + @test c == transpose(PyArray(Py(c))) + + # Arrays with dimensions of size 1 are not + @test size(PyDenseArray(rowmajor(0:3, [1, 4]))) == (1, 4) + + # Nor is column-major data + if Setup.devdeps + np = pyimport("numpy") + f = PyDenseArray(np.asfortranarray(np.arange(6.0).reshape(2, 3))) + @test size(f) == (2, 3) + @test f == PyArray(Py(f)) + end + end + + @testset "indexing" begin + @test Base.IndexStyle(y) === Base.IndexLinear() + @test length(y) == 6 + @test pointer(y) == pointer(PyArray(x)) + @test pointer(y, 2) == pointer(y) + sizeof(Float64) # requires elsize() + @test y[2] == 1.0 + @test_throws BoundsError y[7] + + y[2] = 42 + @test pyeq(Bool, x[1], 42.0) + @test_throws Exception PyDenseArray(pybytes(b"abc"))[1] = 0x00 + end + + @testset "strided dispatch" begin + # dot() has a BLAS method for StridedVector{Float64} + @test which(dot, (typeof(y), typeof(y))) == + which(dot, (Vector{Float64}, Vector{Float64})) + @test which(dot, (typeof(y), typeof(y))) != + which(dot, (typeof(PyArray(x)), typeof(PyArray(x)))) + + a = PyDenseArray(rowmajor(0:5, [2, 3])) # 3×2 + b = PyDenseArray(rowmajor(0:5, [3, 2])) # 2×3 + @test mul!(zeros(3, 3), a, b) ≈ Matrix(a) * Matrix(b) + @test view(a, :, 1:2) isa StridedArray + @test copy(a) isa Matrix{Float64} + end + + @testset "serialize" begin + using Serialization: serialize, deserialize + arrays = Any[x] + + if Setup.devdeps + np = pyimport("numpy") + push!(arrays, np.arange(6.0).reshape(2, 3)) + end + + for a in arrays + c = PyDenseArray(a) + io = IOBuffer() + serialize(io, c) + seekstart(io) + c2 = deserialize(io) + @test typeof(c2) == typeof(c) + @test c2 == c + end + end +end + @testitem "PyDict" begin x = pydict(["foo" => 12]) y = PyDict(x)