GH-38007: [C++] Add VariableShapeTensor implementation by rok · Pull Request #38008 · apache/arrow

rok · 2023-10-04T09:55:28Z

Rationale for this change

We want to add VariableShapeTensor extension type definition for arrays containing tensors with variable shapes.

What changes are included in this PR?

This adds a C++ implementation.

Are these changes tested?

Yes.

Are there any user-facing changes?

This adds a new extension type C++.

Closes: [C++][Python] Add VariableShapeTensor implementation #38007
GitHub Issue: [C++][Python] Add VariableShapeTensor implementation #38007

AlenkaF · 2023-11-16T15:10:36Z

Hi @rok! It would be great to have this in 15.0.0 (added the milestone for visibility). Do you think you have enough bandwidth to work on it at the moment?

rok · 2023-11-16T15:42:49Z

@AlenkaF yes I'd very much like to continue working on this, I was waiting for the release before chasing reviewers :D
I'll rebase tonight and would very much appreciate reviews afterwards.

python/pyarrow/scalar.pxi

pitrou

Some comments. I haven't looked at the tests too closely.

cpp/src/arrow/compute/kernels/CMakeLists.txt

cpp/src/arrow/extension/tensor_internal.h

pitrou · 2026-02-17T14:23:01Z

cpp/src/arrow/extension/tensor_internal.h

-  const auto size = static_cast<int64_t>(permutation.size());
-  std::vector<uint8_t> dim_seen(size, 0);
+ARROW_EXPORT
+Status IsPermutationValid(const std::vector<int64_t>& permutation);


These helper functions all could take std::span<const int64_t> as inputs, though not important either.

Done. I hope std::span won't cause CRAN issues.

pitrou · 2026-02-17T14:39:06Z

cpp/src/arrow/extension/tensor_internal.cc

+                      std::vector<int64_t>* strides) {
+  auto fixed_width_type = internal::checked_pointer_cast<FixedWidthType>(value_type);
+  if (permutation.empty()) {
+    return internal::ComputeRowMajorStrides(*fixed_width_type.get(), shape, strides);


This one is doing the exact same thing except no permutation, please reconcile the two functions instead of duplicating code.

I think I addressed this.

pitrou · 2026-02-17T14:39:40Z

cpp/src/arrow/extension/variable_shape_tensor.cc

+    return false;
+  }
+
+  auto is_permutation_trivial = [](const std::vector<int64_t>& permutation) {


Again, please don't duplicate code like this, factor out into common helpers.

Factored out.

pitrou · 2026-02-17T14:41:25Z

cpp/src/arrow/extension/variable_shape_tensor.cc

+    for (const auto& x : document["permutation"].GetArray()) {
+      permutation.emplace_back(x.GetInt64());
+    }


What happens if "permutation" is not an array of ints?

Added some checks.

python/pyarrow/tests/test_extension_type.py

pitrou · 2026-02-17T14:45:53Z

python/pyarrow/tests/test_extension_type.py

    assert result.to_tensor().shape == (1, 3, 2, 2)
    assert result.to_tensor().strides == (12 * bw, 1 * bw, 6 * bw, 2 * bw)

+    tensor_type = pa.fixed_shape_tensor(arrow_type, [2, 2, 3], permutation=[2, 1, 0])


I suppose it's unrelated but it's testing something useful anyway?

I forget what this was about. I think it might just be additional test of the same path, but called from fixed_shape_tensor_here.

pitrou · 2026-02-17T14:49:04Z

cpp/src/arrow/extension/variable_shape_tensor.cc

+  const auto start_position = data_array->offset() * byte_width;
+  const auto size = std::accumulate(shape.begin(), shape.end(), static_cast<int64_t>(1),
+                                    std::multiplies<>());
+  ARROW_CHECK_EQ(size * byte_width, data_array->length() * byte_width);


This can be evidently simplified.

Indeed! Done.

rok · 2026-02-21T00:27:22Z

Thanks for the review @pitrou! I think I addressed all comments and some more things I found. Could you please do another pass so we can merge this while it's fresh in our minds?

github-actions bot added Component: C++ Component: Python Component: Documentation awaiting committer review Awaiting committer review labels Oct 4, 2023

This was referenced Oct 4, 2023

GH-37988: [Format] Add VariableShapeTensor canonical extension type definition #37992

Closed

GH-24868: [C++] Add a Tensor logical value type with varying dimensions, implemented using ExtensionType #37166

Merged

rok mentioned this pull request Oct 11, 2023

feat(python) Add variable shape image type lance-format/lance#1387

Open

rok force-pushed the 38007 branch 2 times, most recently from eec16ec to 6287402 Compare October 11, 2023 15:47

rok requested a review from pitrou October 12, 2023 11:14

rok force-pushed the 38007 branch from 81e1994 to 7e3feb6 Compare October 12, 2023 13:48

rok assigned rok and unassigned rok Oct 12, 2023

rok requested a review from jorisvandenbossche October 17, 2023 13:36

rok force-pushed the 38007 branch 2 times, most recently from dff5441 to 9c3b464 Compare October 29, 2023 05:12

AlenkaF added this to the 15.0.0 milestone Nov 16, 2023

rok force-pushed the 38007 branch from 076da28 to 4a491a1 Compare November 16, 2023 16:46

rok commented Nov 23, 2023

View reviewed changes

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

github-actions bot added awaiting changes Awaiting changes and removed awaiting committer review Awaiting committer review labels Nov 23, 2023

pitrou reviewed Nov 27, 2023

View reviewed changes

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

github-actions bot added awaiting change review Awaiting change review awaiting changes Awaiting changes and removed awaiting changes Awaiting changes awaiting change review Awaiting change review labels Nov 28, 2023

rok added 5 commits February 12, 2026 23:13

Refactoring ComputeStrides

19bb6b6

Move RoundtripBatch to gtest_util.cc

b888bd7

Post rebase changes

ed4ff5a

Post rebase changes

7daa31a

post rebase fixes

440ce55

rok force-pushed the 38007 branch from add3403 to 440ce55 Compare February 12, 2026 22:13

pitrou requested changes Feb 17, 2026

View reviewed changes

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

review feedback

81f3d02

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

span

19f90e3

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

std::span

5754832

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

some changes

0e7027c

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

rok added 2 commits February 21, 2026 00:02

permutation equivalency check

ac57d2b

lint

1ca46dc

rok force-pushed the 38007 branch from 1a4b067 to f522f59 Compare February 20, 2026 23:44

some typos, style changes, etc

8eecdc5

rok force-pushed the 38007 branch from f522f59 to 8eecdc5 Compare February 20, 2026 23:49

lint

775195a

Comments

Conversation

rok commented Oct 4, 2023 • edited 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 commented Nov 16, 2023

Uh oh!

rok commented Nov 16, 2023

Uh oh!

Uh oh!

Uh oh!

pitrou left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rok commented Feb 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

rok commented Oct 4, 2023 •

edited

Loading