Skip to content

fix: preserve native Ray blocks during pandas materialization - #2539

Merged
edknv merged 3 commits into
NVIDIA:mainfrom
jioffe502:jioffe502/fix-6628603-ray-dataset-to-pandas
Aug 18, 2026
Merged

fix: preserve native Ray blocks during pandas materialization#2539
edknv merged 3 commits into
NVIDIA:mainfrom
jioffe502:jioffe502/fix-6628603-ray-dataset-to-pandas

Conversation

@jioffe502

Copy link
Copy Markdown
Collaborator

Summary

Root cause

ray_dataset_to_pandas() forced every final dataset block through iter_batches(batch_format="pyarrow"). After text splitting, structured content can live in a native pandas block as TensorDtype(object). Converting that block back to Arrow calls pa.from_numpy_dtype(object) and raises ArrowNotImplementedError: Unsupported numpy type 17.

The split had completed successfully; the crash occurred only while materializing the final SDK result. This fix keeps native block ownership at that boundary and avoids both the new pandas-to-Arrow failure and the sliced-Arrow row-access failure addressed by #2523.

Fixes NVBug 6628603.

Validation

  • exact multimodal_test.pdf batch reproduction: extract returns 3 rows, split returns 9 rows, and row iteration succeeds
  • reproduced and fixed on Ray 2.57.0
  • verified the regression and fix on the locked Ray 2.56.1 environment
  • uv run pytest -q tests/test_executor_arrow_pandas.py: 17 passed
  • test_batch_inline_text_matches_text_file: passed
  • repository pre-commit checks: passed

This is intentionally a boundary-local fix. It does not mutate Ray DataContext, add an operator-specific exception, or change the sticky pandas-output policy.

@jioffe502
jioffe502 marked this pull request as ready for review August 18, 2026 14:41
@jioffe502
jioffe502 requested review from a team as code owners August 18, 2026 14:41
@jioffe502
jioffe502 requested a review from edknv August 18, 2026 14:41
@greptile-apps

greptile-apps Bot commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR updates Ray Dataset materialization to preserve each block’s native Arrow or pandas representation while retaining row-safe Arrow conversion.

  • Iterates native Ray blocks without forcing pandas blocks through Arrow.
  • Normalizes object-backed Ray tensor columns before final pandas concatenation.
  • Adds regression coverage for fixed-shape and ragged object tensor columns and updates executor tests for the native-block boundary.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
nemo_retriever/src/nemo_retriever/graph/executor.py Preserves native Ray block formats during final materialization and converts object-backed tensor extension columns to ordinary pandas object columns.
nemo_retriever/tests/test_executor_arrow_pandas.py Adds focused regression coverage for Arrow row safety and fixed-shape and ragged object-backed tensor materialization.
nemo_retriever/tests/test_pipeline_graph.py Updates pipeline test doubles to isolate final dataset materialization through the shared conversion helper.

Reviews (3): Last reviewed commit: "refactor: use native Ray batch iteration" | Re-trigger Greptile

@edknv
edknv merged commit 57a6718 into NVIDIA:main Aug 18, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants