Implement PyDenseArray

This is a minimal array type to allow dispatching to things like BLAS functions
that require `DenseArray`.

Author: JamesWrigley

docs/src/pythoncall-reference.md

@@ -185,6 +185,7 @@ PySet
 PyDict
 PyIterable
 PyArray
+PyDenseArray
 PyIO
 PyTable
 PyPandasDataFrame

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

src/API/exports.jl

@@ -113,6 +113,7 @@ export pyconvert_unconverted
 
 # Wrap
 export PyArray
+export PyDenseArray
 export PyDict
 export PyIO
 export PyIterable

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

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

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

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("<arrayinterface>", AbstractArray, pyconvert_rule_array, priority)
     pyconvert_add_rule("<array>", AbstractArray, pyconvert_rule_array, priority)
     pyconvert_add_rule("<buffer>", AbstractArray, pyconvert_rule_array, priority)
+    pyconvert_add_rule("<arraystruct>", PyDenseArray, pyconvert_rule_densearray_nocopy, priority)
+    pyconvert_add_rule("<arrayinterface>", PyDenseArray, pyconvert_rule_densearray_nocopy, priority)
+    pyconvert_add_rule("<array>", PyDenseArray, pyconvert_rule_densearray_nocopy, priority)
+    pyconvert_add_rule("<buffer>", PyDenseArray, pyconvert_rule_densearray_nocopy, priority)
 end
 
 end

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"

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)