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

cpp/src/arrow/extension/tensor_extension_array_test.cc

@@ -897,6 +897,29 @@ TEST_F(TestVariableShapeTensorType, ComputeStrides) {
   ASSERT_TRUE(tensor->Equals(*t));
 }
 
+TEST_F(TestVariableShapeTensorType, ComputeStridesWithNonTrivialPermutation) {
+  auto permuted_ext_type = internal::checked_pointer_cast<VariableShapeTensorType>(
+      variable_shape_tensor(value_type_, ndim_, {1, 0, 2}, dim_names_, uniform_shape_));
+
+  auto shape = ArrayFromJSON(shape_type_, "[[2,3,1]]");
+  auto data = ArrayFromJSON(data_type_, "[[1,1,2,3,4,5]]");
+  std::vector<std::shared_ptr<Field>> fields = {field("data", data_type_),
+                                                field("shape", shape_type_)};
+  ASSERT_OK_AND_ASSIGN(auto storage_arr, StructArray::Make({data, shape}, fields));
+  auto ext_arr = ExtensionType::WrapArray(permuted_ext_type, storage_arr);
+  auto ext_array = std::static_pointer_cast<VariableShapeTensorArray>(ext_arr);
+
+  ASSERT_OK_AND_ASSIGN(auto scalar, ext_array->GetScalar(0));
+  ASSERT_OK_AND_ASSIGN(auto tensor,
+                       permuted_ext_type->MakeTensor(
+                           internal::checked_pointer_cast<ExtensionScalar>(scalar)));
+
+  ASSERT_EQ(tensor->shape(), (std::vector<int64_t>{3, 2, 1}));
+  ASSERT_EQ(tensor->strides(), (std::vector<int64_t>{sizeof(int64_t), sizeof(int64_t) * 2,
+                                                     sizeof(int64_t)}));
+  ASSERT_EQ(tensor->dim_names(), (std::vector<std::string>{"y", "x", "z"}));
+}
+
 TEST_F(TestVariableShapeTensorType, ToString) {
   auto exact_ext_type =
       internal::checked_pointer_cast<VariableShapeTensorType>(ext_type_);

cpp/src/arrow/extension/variable_shape_tensor.cc

@@ -262,9 +262,9 @@ Result<std::shared_ptr<Tensor>> VariableShapeTensorType::MakeTensor(
     internal::Permute<std::string>(permutation, &dim_names);
   }
 
+  internal::Permute<int64_t>(permutation, &shape);
   ARROW_ASSIGN_OR_RAISE(
       auto strides, internal::ComputeStrides(ext_type.value_type(), shape, permutation));
-  internal::Permute<int64_t>(permutation, &shape);
 
   ARROW_ASSIGN_OR_RAISE(const auto buffer,
                         internal::SliceTensorBuffer(*data_array, value_type, shape));

docs/source/python/api/arrays.rst

@@ -101,6 +101,7 @@ may expose data type-specific methods or properties.
    JsonArray
    UuidArray
    Bool8Array
+   VariableShapeTensorArray
 
 .. _api.scalar:
 
@@ -165,6 +166,7 @@ classes may expose data type-specific methods or properties.
    UnionScalar
    ExtensionScalar
    FixedShapeTensorScalar
+   VariableShapeTensorScalar
    OpaqueScalar
    JsonScalar
    UuidScalar

docs/source/python/api/datatypes.rst

@@ -68,6 +68,7 @@ These should be used to create Arrow data types and schemas.
    dictionary
    run_end_encoded
    fixed_shape_tensor
+   variable_shape_tensor
    union
    dense_union
    sparse_union
@@ -142,6 +143,7 @@ implemented by PyArrow.
    :toctree: ../generated/
 
    FixedShapeTensorType
+   VariableShapeTensorType
    OpaqueType
    JsonType
    UuidType

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,
@@ -231,7 +234,8 @@ def print_entry(label, value):
                          FixedSizeBinaryScalar, DictionaryScalar,
                          MapScalar, StructScalar, UnionScalar,
                          RunEndEncodedScalar, Bool8Scalar, ExtensionScalar,
-                         FixedShapeTensorScalar, JsonScalar, OpaqueScalar, UuidScalar)
+                         FixedShapeTensorScalar, JsonScalar, OpaqueScalar, UuidScalar,
+                         VariableShapeTensorScalar)
 
 # Buffers, allocation
 from pyarrow.lib import (DeviceAllocationType, Device, MemoryManager,