diff --git a/python/pyarrow/lib.pxd b/python/pyarrow/lib.pxd index 48350212c2076..b1187a77c2a6e 100644 --- a/python/pyarrow/lib.pxd +++ b/python/pyarrow/lib.pxd @@ -295,6 +295,8 @@ cdef class Tensor(_Weakrefable): cdef readonly: DataType type + bytes _ssize_t_shape + bytes _ssize_t_strides cdef void init(self, const shared_ptr[CTensor]& sp_tensor) diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi index 1afce7f4a1015..6fb4fc99d7cbc 100644 --- a/python/pyarrow/tensor.pxi +++ b/python/pyarrow/tensor.pxi @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. +# Avoid name clash with `pa.struct` function +import struct as _struct + cdef class Tensor(_Weakrefable): """ @@ -40,6 +43,14 @@ cdef class Tensor(_Weakrefable): self.sp_tensor = sp_tensor self.tp = sp_tensor.get() self.type = pyarrow_wrap_data_type(self.tp.type()) + self._ssize_t_shape = self._make_shape_or_strides_buffer(self.shape) + self._ssize_t_strides = self._make_shape_or_strides_buffer(self.strides) + + def _make_shape_or_strides_buffer(self, values): + """ + Make a bytes object holding an array of `values` cast to `Py_ssize_t`. + """ + return _struct.pack(f"{len(values)}n", *values) def __repr__(self): return """ @@ -282,10 +293,8 @@ strides: {0.strides}""".format(self) buffer.readonly = 0 else: buffer.readonly = 1 - # NOTE: This assumes Py_ssize_t == int64_t, and that the shape - # and strides arrays lifetime is tied to the tensor's - buffer.shape = &self.tp.shape()[0] - buffer.strides = &self.tp.strides()[0] + buffer.shape = cp.PyBytes_AsString(self._ssize_t_shape) + buffer.strides = cp.PyBytes_AsString(self._ssize_t_strides) buffer.suboffsets = NULL