GH-40348: [Python] Add Python wrapper for VariableShapeTensor by rok · Pull Request #40354 · apache/arrow

rok · 2024-03-04T16:43:54Z

Rationale for this change

Once C++ implementation of VariableShapeTensor is complete (#38007) we want a Python wrapper for VariableShapeTensor.

What changes are included in this PR?

This adds a Python wrapper for VariableShapeTensor.

Are these changes tested?

This adds appropriate Python tests.

Are there any user-facing changes?

Yes, a new extension and array type are exposed.

GitHub Issue: [Python] Add Python wrapper for VariableShapeTensor #40348

AlenkaF

Small nits =) Thank you for working on this, takes lots of work!

python/pyarrow/array.pxi

AlenkaF · 2024-04-06T18:04:40Z

The failures are connected.

rok · 2026-02-27T16:30:14Z

cpp/src/arrow/extension/variable_shape_tensor.cc

+  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);



I'm thinking thinking about this still. It's probably a bug or a different interpretation at least.

rok · 2026-03-17T20:55:35Z

@AlenkaF could you do a quick pass on this? I'm a little unsure about the *row_splits APIs, here's more about the idea.

AlenkaF · 2026-03-18T07:49:23Z

@rok sure! Will have a look later today.

AlenkaF · 2026-03-19T10:10:54Z

I did a very quick look at the code now. I would suggest splitting this PR into two as there is quite a lot of lines added and there might be even something that can be moved out.

Moving the *row_splits API related code into a separate issue and PR would be my suggestion. Otherwise, from first look, they sound useful. Might also be good to have a chat on the issue if it would be in fact useful for somebody?

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

rok · 2026-04-14T18:15:46Z

Moving the *row_splits API related code into a separate issue and PR would be my suggestion. Otherwise, from first look, they sound useful. Might also be good to have a chat on the issue if it would be in fact useful for somebody?

Agreed, I moved these into a separate PR and rebased here. I think this is ready for a more final review.

AlenkaF

Thank you for adding this @rok! Huge amount of work was needed here (and in the C++ implementation).

I did a first check but still need to go through most of the Cython code and all the tests.
As for the C++ permutation change - maybe that also fits better into a separate issue&PR?

AlenkaF · 2026-04-15T07:45:55Z

python/pyarrow/includes/libarrow.pxd


        shared_ptr[CArray] storage()

+cdef extern from "arrow/extension/variable_shape_tensor.h" namespace "arrow::extension" nogil:


Could this be moved little lower, after fixed shape tensor ref?

AlenkaF · 2026-04-15T07:49:10Z

python/pyarrow/includes/libarrow.pxd

 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:


Is there a reason nogil is added here and used in the variable version also?

AlenkaF · 2026-04-15T07:55:52Z

python/pyarrow/includes/libarrow.pxd

            const shared_ptr[CBuffer] null_bitmap,
        )

+        @staticmethod


I am not sure where this is needed?

AlenkaF · 2026-04-15T07:56:36Z

python/pyarrow/includes/libarrow.pxd

        void set_chunksize(int64_t chunksize)

    cdef cppclass CTensor" arrow::Tensor":
+        CTensor(const shared_ptr[CDataType]& type,


Here also, I am not sure if we actually need this?

AlenkaF · 2026-04-15T08:17:54Z

python/pyarrow/array.pxi

+        if isinstance(obj, np.ndarray):
+            raise TypeError("obj must be a sequence of numpy arrays")


I think the first check falls into the second one and could be omitted?

AlenkaF · 2026-04-15T08:22:06Z

python/pyarrow/array.pxi

+        array_ndim = arrays[0].ndim
+        arrow_type = from_numpy_dtype(base_dtype)
+
+        if value_type is not None and value_type != arrow_type:


Do we need to accept value_type if the input is not empty as it has to be same as the base type anyways? Not sure what are other options, but feels a bit funny as we are giving this option but it then has to be specific.

rok requested a review from westonpace as a code owner March 4, 2024 16:43

github-actions bot added Component: C++ Component: Python Component: Documentation labels Mar 4, 2024

rok mentioned this pull request Mar 4, 2024

GH-38007: [C++] Add VariableShapeTensor implementation #38008

Merged

github-actions bot added the awaiting committer review Awaiting committer review label Mar 4, 2024

AlenkaF reviewed Mar 12, 2024

View reviewed changes

python/pyarrow/array.pxi Outdated Show resolved Hide resolved

python/pyarrow/array.pxi Outdated Show resolved Hide resolved

rok force-pushed the 38007_python branch 8 times, most recently from b6b63a6 to a372fe6 Compare April 6, 2024 10:24

github-actions bot added awaiting changes Awaiting changes and removed awaiting committer review Awaiting committer review labels Apr 6, 2024

rok force-pushed the 38007_python branch from 78d5501 to ff0af17 Compare April 6, 2024 19:00

github-actions bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Apr 6, 2024

rok force-pushed the 38007_python branch from ff0af17 to 3edefd2 Compare April 6, 2024 19:18

rok mentioned this pull request May 8, 2024

GH-35747: Add FixedSizeListArray.to_numpy_ndarray and FixedSizeListArray.from_numpy_ndarray #35864

Closed

rok force-pushed the 38007_python branch 4 times, most recently from fab047e to f58c973 Compare June 7, 2024 00:28

rok force-pushed the 38007_python branch 2 times, most recently from 5197532 to 55cc40b Compare August 26, 2024 17:55

github-actions bot added the Status: stale-warning Issues and PRs flagged as stale which are due to be closed if no indication otherwise label Nov 18, 2025

github-actions bot added the Component: C++ label Feb 24, 2026

rok force-pushed the 38007_python branch 2 times, most recently from 4429dd2 to b4cb6cd Compare February 24, 2026 23:58

github-actions bot removed the Component: C++ label Feb 24, 2026

rok force-pushed the 38007_python branch from aa7d6b0 to 059429a Compare February 25, 2026 10:05

github-actions bot added the Component: C++ label Feb 25, 2026

rok force-pushed the 38007_python branch from 28f369b to dee2f30 Compare February 26, 2026 15:35

rok mentioned this pull request Feb 26, 2026

Use pyarrow Tensor dtype huggingface/datasets#5272

Open

rok force-pushed the 38007_python branch from 5db5916 to 9cb9a46 Compare February 26, 2026 21:11

rok mentioned this pull request Feb 27, 2026

Allow saving multi-dimensional ndarray with dynamic shapes huggingface/datasets#7738

Open

rok commented Feb 27, 2026

View reviewed changes

github-actions bot added awaiting changes Awaiting changes awaiting change review Awaiting change review and removed awaiting change review Awaiting change review awaiting changes Awaiting changes labels Feb 27, 2026

rok mentioned this pull request Feb 27, 2026

[Python] Consider renaming FixedShapeTensorArray.to_numpy_ndarray to FixedShapeTensorArray.to_numpy #39991

Open

apache deleted a comment from github-actions bot Feb 27, 2026

rok requested review from AlenkaF and jorisvandenbossche February 27, 2026 20:13

rok force-pushed the 38007_python branch 3 times, most recently from 0df8139 to d5a548d Compare March 3, 2026 17:45

rok force-pushed the 38007_python branch from d5a548d to 79d059d Compare March 14, 2026 07:07

apacheGH-38007: [Python] Add VariableShapeTensor Python bindings

a2fbf8c

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

rok force-pushed the 38007_python branch from 79d059d to a2fbf8c Compare April 14, 2026 18:06

AlenkaF reviewed Apr 15, 2026

View reviewed changes


		shared_ptr[CArray] storage()

		cdef extern from "arrow/extension/variable_shape_tensor.h" namespace "arrow::extension" nogil:

		if isinstance(obj, np.ndarray):
		raise TypeError("obj must be a sequence of numpy arrays")

Conversation

rok commented Mar 4, 2024 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

AlenkaF left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

AlenkaF commented Apr 6, 2024

Uh oh!

rok Feb 27, 2026

Choose a reason for hiding this comment

Uh oh!

rok commented Mar 17, 2026

Uh oh!

AlenkaF commented Mar 18, 2026

Uh oh!

AlenkaF commented Mar 19, 2026

Uh oh!

rok commented Apr 14, 2026

Uh oh!

AlenkaF left a comment

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

rok commented Mar 4, 2024 •

edited by github-actions bot

Loading