diff --git a/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py b/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py index 8ee04f5b58..9fe19cd984 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py @@ -8,6 +8,7 @@ from typing import Any, Dict, List, Optional, Sequence +import numpy as np import pandas as pd from nemo_retriever.common.io.image_store import inline_image_b64 @@ -17,6 +18,11 @@ _CONTENT_COLUMNS = ("table", "chart", "infographic") +def _is_content_collection(value: Any) -> bool: + """Return whether a value is a supported extracted-content collection.""" + return isinstance(value, list) or (isinstance(value, np.ndarray) and value.ndim == 1 and value.dtype == object) + + def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequence[str]) -> str: """Combine page text with OCR content text for embedding.""" parts = [] @@ -25,7 +31,7 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ parts.append(base.strip()) for col in content_columns: content_list = row.get(col) - if isinstance(content_list, list): + if _is_content_collection(content_list): for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -43,7 +49,7 @@ def _deep_copy_row(row_dict: Dict[str, Any]) -> Dict[str, Any]: out: Dict[str, Any] = {} for key, value in row_dict.items(): - if isinstance(value, (dict, list)): + if isinstance(value, (dict, list)) or _is_content_collection(value): out[key] = copy.deepcopy(value) else: out[key] = value @@ -108,7 +114,7 @@ def explode_content_to_rows( for column in content_columns: content_list = row_dict.get(column) - if not isinstance(content_list, list): + if not _is_content_collection(content_list): continue for item in content_list: if not isinstance(item, dict): diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/records.py b/nemo_retriever/src/nemo_retriever/common/vdb/records.py index c492ede1c9..3bc10caeb5 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/records.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/records.py @@ -204,6 +204,19 @@ def _dict_or_empty(value: Any) -> dict[str, Any]: return dict(value) if isinstance(value, dict) else {} +def _bbox_from_graph_row(row: dict[str, Any]) -> list[Any] | None: + """Return a JSON-safe bbox without testing array truthiness.""" + for key in ("_bbox_xyxy_norm", "bbox_xyxy_norm"): + value = row.get(key) + if hasattr(value, "tolist"): + value = value.tolist() + elif isinstance(value, tuple): + value = list(value) + if isinstance(value, list) and value: + return value + return None + + def _is_image_backed_row(row: dict[str, Any]) -> bool: """Return whether a post-embed graph row retains its image or stored URI.""" return bool( @@ -265,8 +278,8 @@ def _client_record_from_graph_row(row: dict[str, Any], *, require_embedding: boo stored_image_uri = _first_str(row.get("_stored_image_uri"), row.get("stored_image_uri")) if stored_image_uri: content_metadata.setdefault("stored_image_uri", stored_image_uri) - bbox = row.get("_bbox_xyxy_norm") or row.get("bbox_xyxy_norm") - if bbox: + bbox = _bbox_from_graph_row(row) + if bbox is not None: content_metadata.setdefault("bbox_xyxy_norm", bbox) for key in ( diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index ca4f946c35..0f3b0a1c1f 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -37,6 +37,100 @@ _DEFAULT_GPU_OPERATOR_NUM_GPUS = OCR_GPUS_PER_ACTOR +def _contains_null_arrow_child(data_type: Any) -> bool: + """Return whether a nested Arrow type contains an inferred null child.""" + import pyarrow as pa + + if pa.types.is_null(data_type): + return True + if pa.types.is_struct(data_type): + return any(_contains_null_arrow_child(field.type) for field in data_type) + if pa.types.is_list(data_type) or pa.types.is_large_list(data_type) or pa.types.is_fixed_size_list(data_type): + return _contains_null_arrow_child(data_type.value_type) + if pa.types.is_map(data_type): + return _contains_null_arrow_child(data_type.key_type) or _contains_null_arrow_child(data_type.item_type) + return False + + +def _compact_vulnerable_arrow_columns(table: Any) -> Any: + """Reset offsets before Ray converts nested null children to pandas.""" + import pyarrow as pa + import pyarrow.compute as pc + + if not isinstance(table, pa.Table) or table.num_rows == 0: + return table + + indices = None + compacted = table + for index, field in enumerate(table.schema): + column = table.column(index) + if not _contains_null_arrow_child(field.type) or not any(chunk.offset for chunk in column.chunks): + continue + if indices is None: + indices = pa.array(range(table.num_rows), type=pa.int64()) + compacted = compacted.set_column(index, field, pc.take(column, indices)) + return compacted + + +def _normalize_pickled_object_columns(table: Any, frame: pd.DataFrame) -> pd.DataFrame: + """Convert Ray's pickled-object extension columns to plain pandas objects.""" + import pyarrow as pa + + if not isinstance(table, pa.Table): + return frame + + for index, field in enumerate(table.schema): + if getattr(field.type, "extension_name", None) != "ray.data.arrow_pickled_object": + continue + frame[field.name] = pd.Series(table.column(index).to_pylist(), index=frame.index, dtype=object) + return frame + + +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. + """ + if isinstance(table, pd.DataFrame): + return table + + from ray.data.block import BlockAccessor + + table = _compact_vulnerable_arrow_columns(table) + frame = BlockAccessor.for_block(table).to_pandas() + return _normalize_pickled_object_columns(table, frame) + + +def call_pandas_function_on_arrow( + table: Any, + *, + fn: Any, + fn_kwargs: dict[str, Any] | None = None, +) -> Any: + """Invoke a pandas batch function through the safe Arrow boundary.""" + return fn(arrow_table_to_pandas(table), **(fn_kwargs or {})) + + +class _ArrowPandasOperatorAdapter: + """Convert valid Arrow batches to pandas before invoking an NRL operator.""" + + def __init__(self, operator_class: type, operator_kwargs: dict[str, Any]) -> None: + self._operator = operator_class(**operator_kwargs) + + def __call__(self, table: Any) -> Any: + return self._operator(arrow_table_to_pandas(table)) + + +def _make_arrow_pandas_operator_adapter(operator_class: type) -> type[_ArrowPandasOperatorAdapter]: + """Keep the wrapped operator recognizable in Ray plans and worker logs.""" + adapter_name = f"{operator_class.__name__}ArrowPandasAdapter" + return type(adapter_name, (_ArrowPandasOperatorAdapter,), {}) + + def _concurrency_target(concurrency: Any) -> int: """Return the largest actor-pool size that resource planning can permit.""" if isinstance(concurrency, tuple): @@ -478,17 +572,27 @@ def build_dataset(self, data: Any, **kwargs: Any) -> Any: elif target_num_rows_per_block is not None and int(target_num_rows_per_block) > 0: ds = ds.repartition(target_num_rows_per_block=int(target_num_rows_per_block)) - # Pass the operator class directly to map_batches with - # fn_constructor_kwargs for deferred construction on workers. - # AbstractOperator.__call__ delegates to run(), so each stage - # executes the full preprocess -> process -> postprocess chain. + map_operator_class = node.operator_class + map_batch_format = batch_format + constructor_kwargs = node.operator_kwargs + if batch_format == "pandas": + # Ray's Arrow-backed pandas conversion can preserve unsafe + # offsets for sliced structs with inferred null children. + # Compact the valid Arrow batch before that conversion. + map_operator_class = _make_arrow_pandas_operator_adapter(node.operator_class) + map_batch_format = "pyarrow" + constructor_kwargs = { + "operator_class": node.operator_class, + "operator_kwargs": node.operator_kwargs, + } + ds = ds.map_batches( - node.operator_class, + map_operator_class, batch_size=batch_size, - batch_format=batch_format, + batch_format=map_batch_format, num_cpus=num_cpus, num_gpus=num_gpus, - fn_constructor_kwargs=node.operator_kwargs, + fn_constructor_kwargs=constructor_kwargs, **overrides, ) diff --git a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py index e769f3ddc4..fe2bb36f69 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py @@ -12,7 +12,7 @@ from typing import Any, Callable from nemo_retriever.graph import InprocessExecutor, RayDataExecutor -from nemo_retriever.graph.executor import preflight_executors +from nemo_retriever.graph.executor import call_pandas_function_on_arrow, preflight_executors from nemo_retriever.graph.ingestor_runtime import ( batch_tuning_to_node_overrides, build_graph, @@ -372,9 +372,12 @@ def normalize_ray_branch_datasets(branch_datasets: list[Any]) -> list[Any]: stable_columns = tuple(columns) return [ dataset.map_batches( - ensure_pandas_columns, - batch_format="pandas", - fn_kwargs={"columns": stable_columns}, + call_pandas_function_on_arrow, + batch_format="pyarrow", + fn_kwargs={ + "fn": ensure_pandas_columns, + "fn_kwargs": {"columns": stable_columns}, + }, ) for dataset in branch_datasets ] diff --git a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py index 08cc11d8cd..44e4a769bf 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py @@ -33,6 +33,7 @@ from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Self, Sequence, Tuple, Union from nemo_retriever.graph import InprocessExecutor, RayDataExecutor +from nemo_retriever.graph.executor import arrow_table_to_pandas, call_pandas_function_on_arrow from nemo_retriever.ingestor.branch_extraction import ExtractionBranchExecutor, merge_node_overrides from nemo_retriever.graph.ingestor_runtime import ( batch_tuning_to_node_overrides, @@ -1184,7 +1185,7 @@ def _stage_error_records(cls, batch: Any, *, columns: Iterable[str] | None = Non requested_columns = list(columns) if columns is not None else None if callable(iter_batches): - batches = iter_batches(batch_format="pandas") + batches = (arrow_table_to_pandas(batch_df) for batch_df in iter_batches(batch_format="pyarrow")) else: batches = (batch,) @@ -1366,7 +1367,11 @@ def get_error_rows(self, dataset: Any = None) -> Any: raise RuntimeError("No Ray Dataset available to inspect for errors.") if isinstance(target, pd.DataFrame): return self.extract_error_rows(target) - return target.map_batches(self.extract_error_rows, batch_format="pandas") + return target.map_batches( + call_pandas_function_on_arrow, + batch_format="pyarrow", + fn_kwargs={"fn": self.extract_error_rows}, + ) def get_dataset(self) -> Any: return self._rd_dataset diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/pdf/split.py b/nemo_retriever/src/nemo_retriever/operators/extract/pdf/split.py index e3ba0a8839..4e52525e56 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/pdf/split.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/pdf/split.py @@ -16,6 +16,7 @@ from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator from nemo_retriever.graph.designer import designer_component +from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.operators.operator_archetype import ArchetypeOperator try: @@ -225,4 +226,8 @@ def split_pdf(pdf_ds: Any, params: PdfSplitParams | None = None) -> Any: raise ImportError("split_pdf() requires Ray Data (`ray`).") from e # Note: returning a Dataset here creates the new dataset representing pages. - return pdf_ds.map_batches(PDFSplitActor(split_params=params), batch_format="pandas") + return pdf_ds.map_batches( + call_pandas_function_on_arrow, + batch_format="pyarrow", + fn_kwargs={"fn": PDFSplitActor(split_params=params)}, + ) diff --git a/nemo_retriever/src/nemo_retriever/operators/vdb.py b/nemo_retriever/src/nemo_retriever/operators/vdb.py index 35bc6002ee..35df3359a7 100644 --- a/nemo_retriever/src/nemo_retriever/operators/vdb.py +++ b/nemo_retriever/src/nemo_retriever/operators/vdb.py @@ -117,7 +117,7 @@ def __init__( ) -> None: merged = dict(vdb_kwargs or {}) clean_kwargs, sidecar = split_sidecar_from_vdb_kwargs(merged) - super().__init__(vdb=vdb, vdb_op=vdb_op, vdb_kwargs=clean_kwargs) + super().__init__(vdb=vdb, vdb_op=vdb_op, vdb_kwargs=merged) self._vdb_kwargs = clean_kwargs self._sidecar_spec = sidecar self._sidecar_lookup: dict[str, dict[str, Any]] | None = None diff --git a/nemo_retriever/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py new file mode 100644 index 0000000000..4f4131b413 --- /dev/null +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -0,0 +1,123 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Regression tests for Ray's Arrow-to-pandas operator boundary.""" + +from functools import partial +from typing import Any + +import numpy as np +import pandas as pd +import pyarrow as pa +from ray.data.block import BlockAccessor +from ray.data import DataContext + +from nemo_retriever.graph.executor import _ArrowPandasOperatorAdapter +from nemo_retriever.common.modality.content_transforms import collapse_content_to_page_rows, explode_content_to_rows +from nemo_retriever.operators.abstract_operator import AbstractOperator +from nemo_retriever.operators.graph_ops.custom_operator import UDFOperator + + +class _PassthroughOperator(AbstractOperator): + def preprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + def process(self, data: Any, **kwargs: Any) -> Any: + return data + + def postprocess(self, data: Any, **kwargs: Any) -> Any: + return data + + +def test_adapter_compacts_sliced_nested_arrow_columns() -> None: + table = pa.Table.from_pylist( + [ + { + "metadata": { + "has_text": True, + "source_path": "document.pdf", + "error": None, + }, + "text": f"page {page_number}", + } + for page_number in range(4) + ] + ) + + result = _ArrowPandasOperatorAdapter(_PassthroughOperator, {})(table.slice(2, 1)) + roundtripped = pa.Table.from_pandas(result, preserve_index=False) + + roundtripped.validate(full=True) + assert isinstance(result.dtypes["text"], pd.ArrowDtype) + + +def test_adapter_preserves_ray_pandas_conversion_policy() -> None: + context = DataContext.get_current() + original = context.enable_arrow_backed_pandas_conversion + context.enable_arrow_backed_pandas_conversion = False + try: + result = _ArrowPandasOperatorAdapter(_PassthroughOperator, {})(pa.table({"value": [1, 2]})) + finally: + context.enable_arrow_backed_pandas_conversion = original + + assert not isinstance(result.dtypes["value"], pd.ArrowDtype) + + +def test_adapter_normalizes_pickled_object_columns_for_pandas_row_operations() -> None: + table = BlockAccessor.batch_to_block( + pd.DataFrame( + { + "text": ["first", "second", "third"], + "tables": [np.array([], dtype=object) for _ in range(3)], + } + ) + ) + + result = _ArrowPandasOperatorAdapter(_PassthroughOperator, {})(table) + + assert result["tables"].dtype == object + assert all(isinstance(value, np.ndarray) and value.size == 0 for value in result["tables"]) + assert result.apply(lambda row: row["text"], axis=1).tolist() == ["first", "second", "third"] + + +def test_adapter_preserves_numeric_numpy_arrays() -> None: + table = BlockAccessor.batch_to_block(pd.DataFrame({"embedding": [np.array([0.1, 0.2])]})) + + result = _ArrowPandasOperatorAdapter(_PassthroughOperator, {})(table) + + assert isinstance(result.iloc[0]["embedding"], np.ndarray) + np.testing.assert_array_equal(result.iloc[0]["embedding"], np.array([0.1, 0.2])) + + +def test_adapter_preserves_multimodal_content_semantics() -> None: + table = BlockAccessor.batch_to_block( + pd.DataFrame( + { + "text": ["page text"], + "table": [np.array([{"text": "table text"}], dtype=object)], + "chart": [np.array([{"text": "chart text"}], dtype=object)], + } + ) + ) + + page_result = _ArrowPandasOperatorAdapter( + UDFOperator, + { + "fn": partial(collapse_content_to_page_rows, modality="text"), + "name": "CollapseContentToPageRows", + }, + )(table) + element_result = _ArrowPandasOperatorAdapter( + UDFOperator, + { + "fn": partial(explode_content_to_rows, modality="text"), + "name": "ExplodeContentToRows", + }, + )(table) + + assert page_result["text"].tolist() == ["page text\n\ntable text\n\nchart text"] + assert element_result["text"].tolist() == ["page text", "table text", "chart text"] + assert element_result["_content_type"].tolist() == ["text", "table", "chart"] + pa.Table.from_pandas(page_result, preserve_index=False).validate(full=True) + pa.Table.from_pandas(element_result, preserve_index=False).validate(full=True) diff --git a/nemo_retriever/tests/test_ingest_interface.py b/nemo_retriever/tests/test_ingest_interface.py index c99b6d6a77..647b190b4d 100644 --- a/nemo_retriever/tests/test_ingest_interface.py +++ b/nemo_retriever/tests/test_ingest_interface.py @@ -2,9 +2,12 @@ from pathlib import Path from types import SimpleNamespace +import numpy as np import pandas as pd +import pyarrow as pa import pytest from PIL import Image +from ray.data.block import BlockAccessor import nemo_retriever from nemo_retriever.graph.ingestor_runtime import build_graph @@ -743,21 +746,21 @@ class RayLikeDataset: columns = ["page_elements_v3", "metadata"] def iter_batches(self, *, batch_format: str): - assert batch_format == "pandas" - yield pd.DataFrame( - { - "page_elements_v3": [ - { + assert batch_format == "pyarrow" + yield pa.Table.from_pylist( + [ + { + "page_elements_v3": { "timing": None, "error": { "stage": "remote_inference", "type": "ConnectionError", "message": "connection refused", }, - } - ], - "metadata": [{"source": "test.pdf"}], - } + }, + "metadata": {"source": "test.pdf"}, + } + ] ) ingestor = GraphIngestor(run_mode="batch").extract( @@ -873,12 +876,12 @@ class RayLikeDataset: def __getitem__(self, key: str): raise AssertionError(f"expected map_batches path, got pandas access for {key}") - def map_batches(self, fn, *, batch_format: str): - assert batch_format == "pandas" - batch = pd.DataFrame( - { - "page_elements_v3": [ - { + def map_batches(self, fn, *, batch_format: str, fn_kwargs: dict[str, object]): + assert batch_format == "pyarrow" + batch = pa.Table.from_pylist( + [ + { + "page_elements_v3": { "timing": None, "error": { "stage": "remote_inference", @@ -886,14 +889,39 @@ def map_batches(self, fn, *, batch_format: str): "message": "connection refused", }, }, - {"timing": None, "error": None}, - ], - "text": ["first page", "second page"], - } + "text": "first page", + }, + {"page_elements_v3": {"timing": None, "error": None}, "text": "second page"}, + ] ) - return fn(batch) + return fn(batch, **fn_kwargs) errors = GraphIngestor(run_mode="batch").get_error_rows(RayLikeDataset()) assert len(errors) == 1 assert errors.iloc[0]["text"] == "first page" + + +def test_stage_error_records_normalizes_arrow_object_arrays_before_row_iteration() -> None: + table = BlockAccessor.batch_to_block( + pd.DataFrame( + { + "text": ["first page", "second page"], + "tables": [np.array([], dtype=object), np.array([], dtype=object)], + } + ) + ) + + class RayLikeDataset: + columns = ["text", "tables"] + + def iter_batches(self, *, batch_format: str): + if batch_format == "pandas": + yield BlockAccessor.for_block(table).to_pandas() + else: + assert batch_format == "pyarrow" + yield table + + records = GraphIngestor._stage_error_records(RayLikeDataset(), columns=[]) + + assert records == [] diff --git a/nemo_retriever/tests/test_ingest_manifest.py b/nemo_retriever/tests/test_ingest_manifest.py index 708576aeb9..822824cf4d 100644 --- a/nemo_retriever/tests/test_ingest_manifest.py +++ b/nemo_retriever/tests/test_ingest_manifest.py @@ -4,7 +4,9 @@ from typing import Any import pandas as pd +import pyarrow as pa import pytest +from ray.data.block import BlockAccessor from nemo_retriever.graph import Graph from nemo_retriever.operators.abstract_operator import AbstractOperator @@ -432,7 +434,8 @@ def schema(self) -> Any: return SimpleNamespace(names=self.columns) def map_batches(self, *_args: Any, **kwargs: Any) -> "_FakeDataset": - self.normalized_columns = kwargs["fn_kwargs"]["columns"] + assert kwargs["batch_format"] == "pyarrow" + self.normalized_columns = kwargs["fn_kwargs"]["fn_kwargs"]["columns"] return self def union(self, other: "_FakeDataset") -> "_FakeDataset": @@ -453,6 +456,26 @@ def map_batches(self, *_args: Any, **_kwargs: Any) -> "_LazySchemaDataset": return self +class _SlicedArrowDataset: + def __init__(self, table: pa.Table) -> None: + self.table = table + + def schema(self, *, fetch_if_missing: bool = True) -> Any: + assert fetch_if_missing is False + return self.table.schema + + def map_batches(self, fn: Any, *, batch_format: str, fn_kwargs: dict[str, Any]) -> Any: + if batch_format == "pandas": + batch = BlockAccessor.for_block(self.table).to_pandas() + else: + assert batch_format == "pyarrow" + batch = self.table + result = fn(batch, **fn_kwargs) + roundtripped = BlockAccessor.batch_to_block(result) + roundtripped.validate(full=True) + return result + + def test_ray_schema_normalization_does_not_trigger_lazy_schema_fetch() -> None: datasets = [_LazySchemaDataset(), _LazySchemaDataset()] @@ -462,6 +485,22 @@ def test_ray_schema_normalization_does_not_trigger_lazy_schema_fetch() -> None: assert all(not dataset.map_batches_called for dataset in datasets) +def test_ray_schema_normalization_compacts_sliced_nested_arrow_batches() -> None: + table = pa.Table.from_pylist( + [ + { + "path": f"document-{index}.pdf", + "metadata": {"source_path": f"document-{index}.pdf", "error": None}, + } + for index in range(4) + ] + ).slice(2, 2) + + result = normalize_ray_branch_datasets([_SlicedArrowDataset(table)])[0] + + assert result["path"].tolist() == ["document-2.pdf", "document-3.pdf"] + + def test_batch_branch_execution_uses_dataset_union(monkeypatch, tmp_path) -> None: pdf = tmp_path / "manual.pdf" image = tmp_path / "scan.png" diff --git a/nemo_retriever/tests/test_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index 5066adb3f0..cec96cd989 100644 --- a/nemo_retriever/tests/test_multimodal_embed.py +++ b/nemo_retriever/tests/test_multimodal_embed.py @@ -11,6 +11,7 @@ import sys from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd import pytest @@ -190,6 +191,28 @@ def test_text_mode_tags_modality(self): assert list(result["_embed_modality"]) == ["text", "text"] assert "_image_b64" not in result.columns + def test_arrow_backed_structured_arrays_expand_into_element_rows(self): + """Ray Arrow-backed list cells expand like their Python-list equivalents.""" + df = pd.DataFrame( + { + "text": ["page text"], + "table": [np.array([{"text": "table text"}], dtype=object)], + "chart": [np.array([{"text": "chart text"}], dtype=object)], + } + ) + + result = explode_content_to_rows(df) + + assert result["text"].tolist() == ["page text", "table text", "chart text"] + assert result["_content_type"].tolist() == ["text", "table", "chart"] + assert result.iloc[0]["table"] is not result.iloc[1]["table"] + + @pytest.mark.parametrize("value", [np.array(1, dtype=object), np.ones((2, 2))]) + def test_non_collection_arrays_are_not_expanded(self, value): + result = explode_content_to_rows(pd.DataFrame({"text": ["page text"], "table": [value]})) + + assert result["text"].tolist() == ["page text"] + @patch("nemo_retriever.common.modality.content_transforms._crop_b64_image_by_norm_bbox") def test_text_image_carries_image(self, mock_crop): """text_image mode copies page image to _image_b64, crops for structured content.""" @@ -246,6 +269,20 @@ def test_text_concatenation(self): assert result["text"].iloc[0] == "Hello world\n\ntable data\n\nchart data" assert result["_embed_modality"].iloc[0] == "text" + def test_arrow_backed_structured_arrays_are_collapsed_into_page_text(self): + """Ray Arrow-backed list cells contribute their text to the page row.""" + df = pd.DataFrame( + { + "text": ["page text"], + "table": [np.array([{"text": "table text"}], dtype=object)], + "chart": [np.array([{"text": "chart text"}], dtype=object)], + } + ) + + result = collapse_content_to_page_rows(df) + + assert result["text"].tolist() == ["page text\n\ntable text\n\nchart text"] + def test_full_page_image_used(self): """In image modalities, _image_b64 is the full page image (no cropping).""" df = pd.DataFrame( diff --git a/nemo_retriever/tests/test_nv_ingest_vdb_operator.py b/nemo_retriever/tests/test_nv_ingest_vdb_operator.py index d39fddd418..0245815a43 100644 --- a/nemo_retriever/tests/test_nv_ingest_vdb_operator.py +++ b/nemo_retriever/tests/test_nv_ingest_vdb_operator.py @@ -581,6 +581,24 @@ def test_put_operator_merges_sidecar_metadata_into_records_before_put() -> None: assert merged_content_meta["page_number"] == 7 +def test_ingest_operator_preserves_sidecar_kwargs_for_graph_reconstruction() -> None: + meta_df = pd.DataFrame({"source_id": ["/tmp/doc-a.pdf"], "category": ["legal"]}) + operator = IngestVdbOperator( + vdb=FakeVDB(), + vdb_kwargs={ + "meta_dataframe": meta_df, + "meta_source_field": "source_id", + "meta_fields": ["category"], + }, + ) + + reconstructed_kwargs = operator.get_constructor_kwargs()["vdb_kwargs"] + + assert reconstructed_kwargs["meta_dataframe"] is meta_df + assert reconstructed_kwargs["meta_source_field"] == "source_id" + assert reconstructed_kwargs["meta_fields"] == ["category"] + + def test_ingest_operator_preserves_canonical_batches_for_collection_write() -> None: vdb = FakeVDB() operator = IngestVdbOperator(vdb=vdb) diff --git a/nemo_retriever/tests/test_vdb_records.py b/nemo_retriever/tests/test_vdb_records.py index 7497960924..bd68725c24 100644 --- a/nemo_retriever/tests/test_vdb_records.py +++ b/nemo_retriever/tests/test_vdb_records.py @@ -4,6 +4,7 @@ from __future__ import annotations +import numpy as np import pytest from pydantic import ValidationError @@ -176,6 +177,21 @@ def test_graph_record_conversion_preserves_service_provenance() -> None: } +def test_graph_record_conversion_normalizes_arrow_backed_bbox_array() -> None: + records = to_client_vdb_records( + [ + { + "text": "table content", + "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, + "_content_type": "table", + "_bbox_xyxy_norm": np.array([0.1, 0.2, 0.8, 0.9]), + } + ] + ) + + assert records[0][0]["metadata"]["content_metadata"]["bbox_xyxy_norm"] == [0.1, 0.2, 0.8, 0.9] + + def test_narrow_lancedb_hit_promotes_canonical_multimodal_metadata() -> None: hit = _normalize_one( {