GH-38007: [Python] Add VariableShapeTensor Python bindings

Add PyArrow bindings for the VariableShapeTensor extension type,
including VariableShapeTensorType, VariableShapeTensorArray, and
VariableShapeTensorScalar with support for converting to/from
NumPy tensors.

Author: rok

docs/source/python/api/arrays.rst

@@ -101,6 +101,7 @@ may expose data type-specific methods or properties.
    JsonArray
    UuidArray
    Bool8Array
+   VariableShapeTensorArray
 
 .. _api.scalar:
 

python/pyarrow/__init__.py

@@ -165,6 +165,7 @@ def print_entry(label, value):
                          dictionary,
                          run_end_encoded,
                          bool8, fixed_shape_tensor, json_, opaque, uuid,
+                         variable_shape_tensor,
                          field,
                          type_for_alias,
                          DataType, DictionaryType, StructType,
@@ -178,6 +179,7 @@ def print_entry(label, value):
                          RunEndEncodedType, Bool8Type, FixedShapeTensorType,
                          JsonType, OpaqueType, UuidType,
                          UnknownExtensionType,
+                         VariableShapeTensorType,
                          register_extension_type, unregister_extension_type,
                          DictionaryMemo,
                          KeyValueMetadata,
@@ -214,6 +216,7 @@ def print_entry(label, value):
                          StructArray, ExtensionArray,
                          RunEndEncodedArray, Bool8Array, FixedShapeTensorArray,
                          JsonArray, OpaqueArray, UuidArray,
+                         VariableShapeTensorArray,
                          scalar, NA, _NULL as NULL, Scalar,
                          NullScalar, BooleanScalar,
                          Int8Scalar, Int16Scalar, Int32Scalar, Int64Scalar,

python/pyarrow/array.pxi

@@ -4621,7 +4621,7 @@ cdef class FixedShapeTensorArray(ExtensionArray):
         and the rest of the dimensions will match the permuted shape of the fixed
         shape tensor.
 
-        The conversion is zero-copy.
+        The conversion is zero-copy if data is primitive numeric and without nulls.
 
         Returns
         -------
@@ -4874,6 +4874,137 @@ cdef class Bool8Array(ExtensionArray):
         return Bool8Array.from_storage(storage_arr)
 
 
+cdef class VariableShapeTensorArray(ExtensionArray):
+    """
+    Concrete class for variable shape tensor extension arrays.
+
+    Examples
+    --------
+    Define the extension type for tensor array
+
+    >>> import pyarrow as pa
+    >>> tensor_type = pa.variable_shape_tensor(pa.float64(), 2)
+
+    Create an extension array
+
+    >>> shapes = pa.array([[2, 3], [1, 2]], pa.list_(pa.int32(), 2))
+    >>> values = pa.array([[1, 2, 3, 4, 5, 6], [7, 8]], pa.list_(pa.float64()))
+    >>> arr = pa.StructArray.from_arrays([values, shapes], names=["data", "shape"])
+    >>> pa.ExtensionArray.from_storage(tensor_type, arr)
+    <pyarrow.lib.VariableShapeTensorArray object at ...>
+    -- is_valid: all not null
+    -- child 0 type: list<item: double>
+      [
+        [
+          1,
+          2,
+          3,
+          4,
+          5,
+          6
+        ],
+        [
+          7,
+          8
+        ]
+      ]
+    -- child 1 type: fixed_size_list<item: int32>[2]
+      [
+        [
+          2,
+          3
+        ],
+        [
+          1,
+          2
+        ]
+      ]
+    """
+
+    @staticmethod
+    def from_numpy_ndarray(obj):
+        """
+        Convert a list of numpy.ndarrays to a variable shape tensor extension array.
+        The length of the input list will become the length of the variable shape tensor array.
+
+        Parameters
+        ----------
+        obj : list of numpy.ndarray
+
+        Examples
+        --------
+        >>> import pyarrow as pa
+        >>> import numpy as np
+
+        >>> ndarray_list = [
+        ...         np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32),
+        ...         np.array([[7, 8]], dtype=np.float32),
+        ...     ]
+        >>> arr = pa.VariableShapeTensorArray.from_numpy_ndarray(ndarray_list)
+        >>> assert len(ndarray_list) == len(arr)
+        >>> arr.type
+        VariableShapeTensorType(extension<arrow.variable_shape_tensor[value_type=float, ndim=2, permutation=[0,1]]>)
+        >>> arr
+        <pyarrow.lib.VariableShapeTensorArray object at ...>
+        -- is_valid: all not null
+        -- child 0 type: list<item: float>
+          [
+            [
+              1,
+              2,
+              3,
+              4,
+              5,
+              6
+            ],
+            [
+              7,
+              8
+            ]
+          ]
+        -- child 1 type: fixed_size_list<item: int32>[2]
+          [
+            [
+              2,
+              3
+            ],
+            [
+              1,
+              2
+            ]
+          ]
+        """
+        assert isinstance(obj, list), 'obj must be a list of numpy arrays'
+        numpy_type = obj[0].dtype
+        arrow_type = from_numpy_dtype(numpy_type)
+        ndim = obj[0].ndim
+        permutations = [(-np.array(o.strides)).argsort(kind="stable") for o in obj]
+        permutation = permutations[0]
+        shapes = [np.take(o.shape, permutation) for o in obj]
+
+        if not all([o.dtype == numpy_type for o in obj]):
+            raise TypeError('All numpy arrays must have matching dtype.')
+
+        if not all([o.ndim == ndim for o in obj]):
+            raise ValueError('All numpy arrays must have matching ndim.')
+
+        if not all([np.array_equal(p, permutation) for p in permutations]):
+            raise ValueError('All numpy arrays must have matching permutation.')
+
+        for shape in shapes:
+            if len(shape) < 2:
+                raise ValueError(
+                    "Cannot convert 1D array or scalar to fixed shape tensor array")
+            if np.prod(shape) == 0:
+                raise ValueError("Expected a non-empty ndarray")
+
+        values = array([np.ravel(o, order="K") for o in obj], list_(arrow_type))
+        shapes = array(shapes, list_(int32(), list_size=ndim))
+        struct_arr = StructArray.from_arrays([values, shapes], names=["data", "shape"])
+
+        return ExtensionArray.from_storage(variable_shape_tensor(arrow_type, ndim, permutation=permutation), struct_arr)
+
+
 cdef dict _array_classes = {
     _Type_NA: NullArray,
     _Type_BOOL: BooleanArray,

python/pyarrow/includes/libarrow.pxd

@@ -908,6 +908,14 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
             const shared_ptr[CBuffer] null_bitmap,
         )
 
+        @staticmethod
+        CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"(
+            shared_ptr[CDataType],
+            const shared_ptr[CArray]& offsets,
+            const shared_ptr[CArray]& keys,
+            const shared_ptr[CArray]& items,
+            CMemoryPool* pool)
+
         shared_ptr[CArray] keys()
         shared_ptr[CArray] items()
         CMapType* map_type()
@@ -1184,6 +1192,11 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         void set_chunksize(int64_t chunksize)
 
     cdef cppclass CTensor" arrow::Tensor":
+        CTensor(const shared_ptr[CDataType]& type,
+                const shared_ptr[CBuffer]& data,
+                const vector[int64_t]& shape,
+                const vector[int64_t]& strides,
+                const vector[c_string]& dim_names)
         shared_ptr[CDataType] type()
         shared_ptr[CBuffer] data()
 
@@ -3014,6 +3027,24 @@ cdef extern from "arrow/extension_type.h" namespace "arrow":
 
         shared_ptr[CArray] storage()
 
+cdef extern from "arrow/extension/variable_shape_tensor.h" namespace "arrow::extension" nogil:
+    cdef cppclass CVariableShapeTensorType \
+            " arrow::extension::VariableShapeTensorType"(CExtensionType):
+
+        CResult[shared_ptr[CTensor]] MakeTensor(const shared_ptr[CExtensionScalar]& scalar) const
+
+        @staticmethod
+        CResult[shared_ptr[CDataType]] Make(const shared_ptr[CDataType]& value_type,
+                                            const int32_t ndim,
+                                            const vector[int64_t] permutation,
+                                            const vector[c_string] dim_names,
+                                            const vector[optional[int64_t]] uniform_shape)
+
+        const shared_ptr[CDataType] value_type()
+        const int32_t ndim()
+        const vector[int64_t] permutation()
+        const vector[c_string] dim_names()
+        const vector[optional[int64_t]] uniform_shape()
 
 cdef extern from "arrow/extension/json.h" namespace "arrow::extension" nogil:
     cdef cppclass CJsonType" arrow::extension::JsonExtensionType"(CExtensionType):
@@ -3034,7 +3065,7 @@ cdef extern from "arrow/extension/uuid.h" namespace "arrow::extension" nogil:
 
 cdef extern from "arrow/extension/fixed_shape_tensor.h" namespace "arrow::extension" nogil:
     cdef cppclass CFixedShapeTensorType \
-            " arrow::extension::FixedShapeTensorType"(CExtensionType):
+            " arrow::extension::FixedShapeTensorType"(CExtensionType) nogil:
 
         CResult[shared_ptr[CTensor]] MakeTensor(const shared_ptr[CExtensionScalar]& scalar) const
 

python/pyarrow/lib.pxd

@@ -195,6 +195,11 @@ cdef class ExtensionType(BaseExtensionType):
         const CPyExtensionType* cpy_ext_type
 
 
+cdef class VariableShapeTensorType(BaseExtensionType):
+    cdef:
+        const CVariableShapeTensorType* tensor_ext_type
+
+
 cdef class FixedShapeTensorType(BaseExtensionType):
     cdef:
         const CFixedShapeTensorType* tensor_ext_type

python/pyarrow/public-api.pxi

@@ -131,6 +131,8 @@ cdef api object pyarrow_wrap_data_type(
             out = Bool8Type.__new__(Bool8Type)
         elif extension_name == b"arrow.fixed_shape_tensor":
             out = FixedShapeTensorType.__new__(FixedShapeTensorType)
+        elif extension_name == b"arrow.variable_shape_tensor":
+            out = VariableShapeTensorType.__new__(VariableShapeTensorType)
         elif extension_name == b"arrow.opaque":
             out = OpaqueType.__new__(OpaqueType)
         elif extension_name == b"arrow.uuid":

python/pyarrow/scalar.pxi

@@ -1484,7 +1484,7 @@ cdef class FixedShapeTensorScalar(ExtensionScalar):
 
         The resulting ndarray's shape matches the permuted shape of the
         fixed shape tensor scalar.
-        The conversion is zero-copy.
+        The conversion is zero-copy if data is primitive numeric and without nulls.
 
         Returns
         -------
@@ -1539,6 +1539,43 @@ cdef class Bool8Scalar(ExtensionScalar):
         py_val = super().as_py()
         return None if py_val is None else py_val != 0
 
+
+cdef class VariableShapeTensorScalar(ExtensionScalar):
+    """
+    Concrete class for variable shape tensor extension scalar.
+    """
+
+    def to_numpy_ndarray(self):
+        """
+        Convert variable shape tensor extension scalar to a numpy array.
+
+        The conversion is zero-copy if data is primitive numeric and without nulls.
+
+        Returns
+        -------
+        numpy.ndarray
+        """
+        return self.to_tensor().to_numpy()
+
+    def to_tensor(self):
+        """
+        Convert variable shape tensor extension scalar to a pyarrow.Tensor.
+
+        Returns
+        -------
+        tensor : pyarrow.Tensor
+        """
+        cdef:
+            CVariableShapeTensorType* c_type = static_pointer_cast[CVariableShapeTensorType, CDataType](
+                self.wrapped.get().type).get()
+            shared_ptr[CExtensionScalar] scalar = static_pointer_cast[CExtensionScalar, CScalar](self.wrapped)
+            shared_ptr[CTensor] ctensor
+
+        with nogil:
+            ctensor = GetResultValue(c_type.MakeTensor(scalar))
+        return pyarrow_wrap_tensor(ctensor)
+
+
 cdef dict _scalar_classes = {
     _Type_BOOL: BooleanScalar,
     _Type_UINT8: UInt8Scalar,