fix: preserve native Ray blocks during pandas materialization

nemo_retriever/src/nemo_retriever/graph/executor.py

@@ -90,17 +90,33 @@ def _normalize_pickled_object_columns(table: Any, frame: pd.DataFrame) -> pd.Dat
     return frame
 
 
+def _normalize_object_tensor_columns(frame: pd.DataFrame) -> pd.DataFrame:
+    """Convert object-backed Ray tensor columns to ordinary pandas objects."""
+    from ray.data.extensions import TensorDtype
+
+    columns = [
+        name for name, dtype in frame.dtypes.items() if isinstance(dtype, TensorDtype) and dtype.element_dtype.hasobject
+    ]
+    if not columns:
+        return frame
+
+    normalized = frame.copy(deep=False)
+    for name in columns:
+        normalized[name] = frame[name].astype(object)
+    return normalized
+
+
 def arrow_table_to_pandas(table: Any) -> pd.DataFrame:
     """Convert a Ray Arrow batch to a row-safe pandas DataFrame.
 
     Ray 2.56+ preserves Arrow-backed pandas dtypes. Before conversion, sliced
     nested columns with inferred null children must be compacted. Ray's
-    pickled-object extension columns also need to be materialized as ordinary
-    object columns so pandas row operations do not interpret their payloads as
-    malformed extension arrays.
+    pickled-object and object-backed tensor extension columns also need to be
+    materialized as ordinary object columns so pandas row operations do not
+    interpret their payloads as malformed extension arrays.
     """
     if isinstance(table, pd.DataFrame):
-        return table
+        return _normalize_object_tensor_columns(table)
 
     from ray.data.block import BlockAccessor
 
@@ -114,21 +130,23 @@ def ray_dataset_to_pandas(dataset: ray.data.Dataset) -> pd.DataFrame:
 
     Ray 2.56+ enables Arrow-backed pandas conversion by default. Calling
     ``Dataset.to_pandas()`` directly can therefore expose sliced nested Arrow
-    columns whose child offsets are invalid for pandas row access. Convert
-    each Arrow block through :func:`arrow_table_to_pandas` before concatenating
-    so the public SDK result is safe to consume with standard pandas APIs.
+    columns whose child offsets are invalid for pandas row access. Forcing a
+    pandas block to Arrow can also fail for object-backed tensor columns. Read
+    each block in its native format and convert it through
+    :func:`arrow_table_to_pandas` before concatenating so the public SDK result
+    is safe to consume with standard pandas APIs.
 
     Parameters
     ----------
     dataset
-        Ray dataset to materialize as Arrow batches.
+        Ray dataset to materialize in its native block formats.
 
     Returns
     -------
     pandas.DataFrame
         Row-safe DataFrame containing all rows from ``dataset``.
     """
-    frames = [arrow_table_to_pandas(batch) for batch in dataset.iter_batches(batch_format="pyarrow")]
+    frames = [arrow_table_to_pandas(block) for block in dataset.iter_batches(batch_format=None, batch_size=None)]
     if frames:
         return pd.concat(frames, ignore_index=True)
 

nemo_retriever/tests/test_executor_arrow_pandas.py

@@ -10,6 +10,7 @@
 import numpy as np
 import pandas as pd
 import pyarrow as pa
+import pytest
 from ray.data.block import BlockAccessor
 from ray.data import DataContext
 from ray.data.extensions import TensorArray
@@ -72,8 +73,9 @@ def test_dataset_materialization_returns_row_safe_pandas_dataframe() -> None:
     )
 
     class _Dataset:
-        def iter_batches(self, *, batch_format: str):
-            assert batch_format == "pyarrow"
+        def iter_batches(self, *, batch_format, batch_size):
+            assert batch_format is None
+            assert batch_size is None
             yield table.slice(1, 1)
 
         def schema(self):
@@ -85,6 +87,46 @@ def schema(self):
     assert result.to_dict("records") == [{"metadata": {"error": None, "timing": None}, "text": "page 1"}]
 
 
+@pytest.mark.parametrize(
+    ("values", "expected"),
+    [
+        pytest.param(
+            [
+                np.array([{"text": "first table"}], dtype=object),
+                np.array([{"text": "second table"}], dtype=object),
+            ],
+            [[{"text": "first table"}], [{"text": "second table"}]],
+            id="fixed-shape",
+        ),
+        pytest.param(
+            [
+                np.array([{"text": "table text"}], dtype=object),
+                np.array([], dtype=object),
+            ],
+            [[{"text": "table text"}], []],
+            id="ragged",
+        ),
+    ],
+)
+def test_dataset_materialization_normalizes_object_tensor_pandas_blocks(values, expected) -> None:
+    frame = pd.DataFrame({"table": pd.Series(TensorArray(values))})
+
+    class _Dataset:
+        def iter_batches(self, *, batch_format, batch_size):
+            assert batch_format is None
+            assert batch_size is None
+            yield frame
+
+        def schema(self):
+            return pa.schema([pa.field("table", pa.list_(pa.struct([pa.field("text", pa.string())])))])
+
+    result = ray_dataset_to_pandas(_Dataset())
+
+    assert result["table"].dtype == object
+    assert all(isinstance(value, np.ndarray) for value in result["table"])
+    assert [value.tolist() for value in result["table"]] == expected
+
+
 def test_adapter_preserves_ray_pandas_conversion_policy() -> None:
     context = DataContext.get_current()
     original = context.enable_arrow_backed_pandas_conversion

nemo_retriever/tests/test_pipeline_graph.py

@@ -1379,13 +1379,9 @@ def test_ingest_expands_recursive_glob_patterns(self, tmp_path, monkeypatch):
         pdf_path.write_bytes(b"pdf")
 
         class _FakeDataset:
-            def iter_batches(self, *, batch_format):
-                assert batch_format == "pyarrow"
-                return iter([])
-
-            def schema(self):
-                return SimpleNamespace(names=[])
+            pass
 
+        fake_dataset = _FakeDataset()
         captured: dict[str, object] = {}
 
         class _FakeDataContext:
@@ -1399,7 +1395,11 @@ def get_current(cls):
         def _fake_read_binary_files(paths, include_paths=True):
             captured["paths"] = list(paths)
             captured["include_paths"] = include_paths
-            return _FakeDataset()
+            return fake_dataset
+
+        def _fake_ray_dataset_to_pandas(dataset):
+            assert dataset is fake_dataset
+            return pd.DataFrame()
 
         fake_ray_data = SimpleNamespace(
             Dataset=_FakeDataset,
@@ -1415,6 +1415,7 @@ def _fake_read_binary_files(paths, include_paths=True):
             lambda ray: SimpleNamespace(available_gpu_count=lambda: 0),
         )
         monkeypatch.setattr("nemo_retriever.graph.executor.resolve_graph", lambda graph, cluster: graph)
+        monkeypatch.setattr("nemo_retriever.graph.executor.ray_dataset_to_pandas", _fake_ray_dataset_to_pandas)
 
         executor = RayDataExecutor(Graph())
         result = executor.ingest([str(tmp_path / "**" / "*.pdf")])