[Data] Fix `TensorArray` to Arrow tensor conversion by alexeykudinkin · Pull Request #59449 · ray-project/ray

alexeykudinkin · 2025-12-15T20:46:23Z

Description

This change addresses a long-standing problem when Pandas tensors holding null values couldn't be converted into Arrow ones.

More details are captured in #59445.

Following changes are made to address that:

Fixed _is_ndarray_variable_shaped_tensor
Numpy tensors with dtype='o' are cycled t/h Pyarrow to be converted into proper ndarrays
Path raveling and formatting are unified b/w fixed-shape and var-shaped tensors

Related issues

Addresses #59445

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

…ct ndarrays; Abstracted `AVSTA._ravel_tensors` Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

gemini-code-assist

Code Review

This pull request addresses an issue with converting TensorArray to Arrow tensors, particularly for data originating from pandas with nullable types. The core changes involve adding a specific handling path for np.object_ dtype arrays in ArrowTensorArray._from_numpy, refactoring tensor raveling logic into a reusable _ravel_tensors method, and updating an API call to use a keyword argument for clarity. The changes are logical and well-targeted. My review includes suggestions to address a TODO for ensuring correctness, a minor performance optimization, and a recommendation to add assertions to the new test case to make it more robust.

python/ray/air/util/tensor_extensions/arrow.py

python/ray/data/tests/test_tensor_extension.py

python/ray/air/util/tensor_extensions/arrow.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

This reverts commit b031dfd. Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

srinathk10

LGTM. Tests need fixing

python/ray/data/tests/test_tensor_extension.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaning up useless permutations Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin · 2025-12-17T02:11:33Z

/gemini review

gemini-code-assist

Code Review

This pull request effectively addresses a long-standing issue with converting Pandas tensors containing null values into Arrow tensors. The core of the fix, which involves cycling object-dtyped NumPy arrays through Pyarrow to handle nulls correctly, is sound and well-implemented in the new _ensure_scalar_ndarray function.

The accompanying changes are also positive:

The bug fix in _is_ndarray_variable_shaped_tensor is correct and crucial.
Refactoring the tensor raveling logic into _ravel_tensors improves code clarity and maintainability.
The API is improved by making column_name a keyword-only argument in ArrowTensorArray.from_numpy.
The new tests in test_tensor_extension.py provide good coverage for the fix.

I've left a few minor comments regarding typos in comments and test code.

One thing to note is the removal of tensor_format parametrization from several tests. While this simplifies the test suite, it seems to reduce test coverage for the v1 tensor format, as the tests will now only run against the default v2 format. This might be intentional, but it would be good to confirm.

python/ray/air/util/tensor_extensions/arrow.py

python/ray/data/tests/test_tensor_extension.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Alexey Kudinkin <alexey.kudinkin@gmail.com>

matthewdeng

STAMP

python/ray/air/util/tensor_extensions/arrow.py

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

## Description This change addresses a long-standing problem when Pandas tensors holding null values couldn't be converted into Arrow ones. More details are captured in ray-project#59445. Following changes are made to address that: - Fixed `_is_ndarray_variable_shaped_tensor` - Numpy tensors with `dtype='o'` are cycled t/h Pyarrow to be converted into proper ndarrays - Path raveling and formatting are unified b/w fixed-shape and var-shaped tensors ## Related issues Addresses ray-project#59445 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Alexey Kudinkin <alexey.kudinkin@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

## Description This change addresses a long-standing problem when Pandas tensors holding null values couldn't be converted into Arrow ones. More details are captured in ray-project#59445. Following changes are made to address that: - Fixed `_is_ndarray_variable_shaped_tensor` - Numpy tensors with `dtype='o'` are cycled t/h Pyarrow to be converted into proper ndarrays - Path raveling and formatting are unified b/w fixed-shape and var-shaped tensors ## Related issues Addresses ray-project#59445 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Alexey Kudinkin <alexey.kudinkin@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

alexeykudinkin added 5 commits December 15, 2025 12:42

Added _ensure_scalar_ndarray to cycle t/h Pyarrow array for np.obje…

8339991

…ct ndarrays; Abstracted `AVSTA._ravel_tensors` Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Ravel object ndarrays in ATA.from_numpy

abddebd

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

0c143d0

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added test

97f088b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

64dbf86

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin requested review from a team as code owners December 15, 2025 20:46

alexeykudinkin added the go add ONLY when ready to merge, run all tests label Dec 15, 2025

alexeykudinkin linked an issue Dec 15, 2025 that may be closed by this pull request

[Data] Failure when converting to Arrow from TensorArray #59445

Closed

gemini-code-assist bot reviewed Dec 15, 2025

View reviewed changes

python/ray/air/util/tensor_extensions/arrow.py Outdated Show resolved Hide resolved

python/ray/air/util/tensor_extensions/arrow.py Outdated Show resolved Hide resolved

python/ray/data/tests/test_tensor_extension.py Outdated Show resolved Hide resolved

cursor bot reviewed Dec 15, 2025

View reviewed changes

python/ray/air/util/tensor_extensions/arrow.py Outdated Show resolved Hide resolved

alexeykudinkin added 5 commits December 15, 2025 16:48

Fixed _is_ndarray_variable_shaped_tensor

f98f7b8

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

XXX

b031dfd

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Revert "XXX"

dcd383c

This reverts commit b031dfd. Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

8060a80

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

2facbb1

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin changed the title ~~[WIP][Data] Fix TensorArray to Arrow tensor conversion~~ [Data] Fix TensorArray to Arrow tensor conversion Dec 16, 2025

ray-gardener bot added the data Ray Data-related issues label Dec 16, 2025

srinathk10 approved these changes Dec 16, 2025

View reviewed changes

python/ray/data/tests/test_tensor_extension.py Outdated Show resolved Hide resolved

alexeykudinkin added 7 commits December 16, 2025 17:23

Fixed tests

d53b856

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed _ensure_scalar_ndarray to properly reconstruct tensor's shape

154c7c4

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Simplified handling of FS tensors

71ab42a

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

f942f85

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Using multi-dim tensor shapes;

55c9ea6

Cleaning up useless permutations Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

ce72c90

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

44787a0

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

gemini-code-assist bot reviewed Dec 17, 2025

View reviewed changes

python/ray/air/util/tensor_extensions/arrow.py Outdated Show resolved Hide resolved

python/ray/data/tests/test_tensor_extension.py Outdated Show resolved Hide resolved

python/ray/data/tests/test_tensor_extension.py Outdated Show resolved Hide resolved

Tidying up

0a82f34

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Alexey Kudinkin <alexey.kudinkin@gmail.com>

matthewdeng approved these changes Dec 17, 2025

View reviewed changes

alexeykudinkin enabled auto-merge (squash) December 17, 2025 02:45

cursor bot reviewed Dec 17, 2025

View reviewed changes

python/ray/air/util/tensor_extensions/arrow.py Outdated Show resolved Hide resolved

alexeykudinkin added 2 commits December 17, 2025 11:25

Restored proper tensor raveling (that was reverted by Claude)

6012d85

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

b50134c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

github-actions bot disabled auto-merge December 17, 2025 19:37

alexeykudinkin enabled auto-merge (squash) December 17, 2025 19:37

alexeykudinkin merged commit 2015205 into ray-project:master Dec 17, 2025
7 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] Fix `TensorArray` to Arrow tensor conversion#59449

[Data] Fix `TensorArray` to Arrow tensor conversion#59449
alexeykudinkin merged 20 commits intoray-project:masterfrom
alexeykudinkin:ak/pd-tnsr-fix

alexeykudinkin commented Dec 15, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

srinathk10 left a comment

Uh oh!

Uh oh!

alexeykudinkin commented Dec 17, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

matthewdeng left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

alexeykudinkin commented Dec 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related issues

Additional information

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

srinathk10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alexeykudinkin commented Dec 17, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

matthewdeng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

alexeykudinkin commented Dec 15, 2025 •

edited

Loading