From d70c4f34648c09f46ac15630b8c9a6f6068f408f Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Tue, 18 Aug 2026 14:45:24 +0000 Subject: [PATCH 1/3] fix: preserve native Ray blocks during materialization --- .../src/nemo_retriever/graph/executor.py | 41 +++++++++++++++---- .../tests/test_executor_arrow_pandas.py | 41 ++++++++++++++++++- 2 files changed, 72 insertions(+), 10 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index d23f8b61f..0d1e2f0b1 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/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() + for name in columns: + normalized[name] = pd.Series(frame[name].to_numpy().tolist(), index=frame.index, dtype=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,28 @@ 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")] + import ray + + frames = [] + for ref_bundle in dataset.iter_internal_ref_bundles(): + blocks = ray.get([entry.ref for entry in ref_bundle.blocks]) + frames.extend(arrow_table_to_pandas(block) for block in blocks) if frames: return pd.concat(frames, ignore_index=True) diff --git a/nemo_retriever/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py index dd3590cbf..1567a33a7 100644 --- a/nemo_retriever/tests/test_executor_arrow_pandas.py +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -5,11 +5,13 @@ """Regression tests for Ray's Arrow-to-pandas operator boundary.""" from functools import partial +from types import SimpleNamespace from typing import Any import numpy as np import pandas as pd import pyarrow as pa +import ray from ray.data.block import BlockAccessor from ray.data import DataContext from ray.data.extensions import TensorArray @@ -60,7 +62,7 @@ def test_adapter_compacts_sliced_nested_arrow_columns() -> None: assert isinstance(result.dtypes["text"], pd.ArrowDtype) -def test_dataset_materialization_returns_row_safe_pandas_dataframe() -> None: +def test_dataset_materialization_returns_row_safe_pandas_dataframe(monkeypatch) -> None: table = pa.Table.from_pylist( [ { @@ -76,15 +78,52 @@ def iter_batches(self, *, batch_format: str): assert batch_format == "pyarrow" yield table.slice(1, 1) + def iter_internal_ref_bundles(self): + yield SimpleNamespace(blocks=(SimpleNamespace(ref=table.slice(1, 1)),)) + def schema(self): return table.schema + monkeypatch.setattr(ray, "get", lambda refs: refs) result = ray_dataset_to_pandas(_Dataset()) assert [row.text for row in result.itertuples(index=False)] == ["page 1"] assert result.to_dict("records") == [{"metadata": {"error": None, "timing": None}, "text": "page 1"}] +def test_dataset_materialization_preserves_object_tensor_pandas_blocks(monkeypatch) -> None: + frame = pd.DataFrame( + { + "table": pd.Series( + TensorArray( + [ + np.array([{"text": "table text"}], dtype=object), + np.array([], dtype=object), + ] + ) + ) + } + ) + + class _Dataset: + def iter_batches(self, *, batch_format: str): + assert batch_format == "pyarrow" + yield BlockAccessor.for_block(frame).to_arrow() + + def iter_internal_ref_bundles(self): + yield SimpleNamespace(blocks=(SimpleNamespace(ref=frame),)) + + def schema(self): + return pa.schema([pa.field("table", pa.list_(pa.struct([pa.field("text", pa.string())])))]) + + monkeypatch.setattr(ray, "get", lambda refs: refs) + result = ray_dataset_to_pandas(_Dataset()) + + assert result["table"].dtype == object + assert result["table"].iloc[0].tolist() == [{"text": "table text"}] + assert result["table"].iloc[1].tolist() == [] + + def test_adapter_preserves_ray_pandas_conversion_policy() -> None: context = DataContext.get_current() original = context.enable_arrow_backed_pandas_conversion From da724755aa1251fcc0d07414e4c5938b75e88879 Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Tue, 18 Aug 2026 15:11:18 +0000 Subject: [PATCH 2/3] test: isolate recursive glob expansion --- nemo_retriever/tests/test_pipeline_graph.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/nemo_retriever/tests/test_pipeline_graph.py b/nemo_retriever/tests/test_pipeline_graph.py index 9517339cc..da2ec821c 100644 --- a/nemo_retriever/tests/test_pipeline_graph.py +++ b/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")]) From 8d5838df8b26d018d91d6f9f0f15bb7f3aea2749 Mon Sep 17 00:00:00 2001 From: jioffe502 Date: Tue, 18 Aug 2026 16:06:28 +0000 Subject: [PATCH 3/3] refactor: use native Ray batch iteration --- .../src/nemo_retriever/graph/executor.py | 11 +--- .../tests/test_executor_arrow_pandas.py | 65 ++++++++++--------- 2 files changed, 37 insertions(+), 39 deletions(-) diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index 0d1e2f0b1..257872334 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -100,9 +100,9 @@ def _normalize_object_tensor_columns(frame: pd.DataFrame) -> pd.DataFrame: if not columns: return frame - normalized = frame.copy() + normalized = frame.copy(deep=False) for name in columns: - normalized[name] = pd.Series(frame[name].to_numpy().tolist(), index=frame.index, dtype=object) + normalized[name] = frame[name].astype(object) return normalized @@ -146,12 +146,7 @@ def ray_dataset_to_pandas(dataset: ray.data.Dataset) -> pd.DataFrame: pandas.DataFrame Row-safe DataFrame containing all rows from ``dataset``. """ - import ray - - frames = [] - for ref_bundle in dataset.iter_internal_ref_bundles(): - blocks = ray.get([entry.ref for entry in ref_bundle.blocks]) - frames.extend(arrow_table_to_pandas(block) for block in blocks) + 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) diff --git a/nemo_retriever/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py index 1567a33a7..b7045ac8a 100644 --- a/nemo_retriever/tests/test_executor_arrow_pandas.py +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -5,13 +5,12 @@ """Regression tests for Ray's Arrow-to-pandas operator boundary.""" from functools import partial -from types import SimpleNamespace from typing import Any import numpy as np import pandas as pd import pyarrow as pa -import ray +import pytest from ray.data.block import BlockAccessor from ray.data import DataContext from ray.data.extensions import TensorArray @@ -62,7 +61,7 @@ def test_adapter_compacts_sliced_nested_arrow_columns() -> None: assert isinstance(result.dtypes["text"], pd.ArrowDtype) -def test_dataset_materialization_returns_row_safe_pandas_dataframe(monkeypatch) -> None: +def test_dataset_materialization_returns_row_safe_pandas_dataframe() -> None: table = pa.Table.from_pylist( [ { @@ -74,54 +73,58 @@ def test_dataset_materialization_returns_row_safe_pandas_dataframe(monkeypatch) ) 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 iter_internal_ref_bundles(self): - yield SimpleNamespace(blocks=(SimpleNamespace(ref=table.slice(1, 1)),)) - def schema(self): return table.schema - monkeypatch.setattr(ray, "get", lambda refs: refs) result = ray_dataset_to_pandas(_Dataset()) assert [row.text for row in result.itertuples(index=False)] == ["page 1"] assert result.to_dict("records") == [{"metadata": {"error": None, "timing": None}, "text": "page 1"}] -def test_dataset_materialization_preserves_object_tensor_pandas_blocks(monkeypatch) -> None: - frame = pd.DataFrame( - { - "table": pd.Series( - TensorArray( - [ - np.array([{"text": "table text"}], dtype=object), - np.array([], dtype=object), - ] - ) - ) - } - ) +@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: str): - assert batch_format == "pyarrow" - yield BlockAccessor.for_block(frame).to_arrow() - - def iter_internal_ref_bundles(self): - yield SimpleNamespace(blocks=(SimpleNamespace(ref=frame),)) + 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())])))]) - monkeypatch.setattr(ray, "get", lambda refs: refs) result = ray_dataset_to_pandas(_Dataset()) assert result["table"].dtype == object - assert result["table"].iloc[0].tolist() == [{"text": "table text"}] - assert result["table"].iloc[1].tolist() == [] + 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: