From c53718400d8f2e559832284d351faabcf0a4e6c3 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 12 Aug 2026 14:14:37 +0000 Subject: [PATCH 1/5] fix: compact sliced nested Arrow batches Signed-off-by: Jacob Ioffe --- .../src/nemo_retriever/graph/executor.py | 79 +++++++++++++++++-- .../tests/test_executor_arrow_pandas.py | 59 ++++++++++++++ 2 files changed, 131 insertions(+), 7 deletions(-) create mode 100644 nemo_retriever/tests/test_executor_arrow_pandas.py diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index 1d55233a35..9a922a3a61 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -36,6 +36,61 @@ _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 + + +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: + from ray.data.block import BlockAccessor + + table = _compact_vulnerable_arrow_columns(table) + frame = BlockAccessor.for_block(table).to_pandas() + return self._operator(frame) + + +def _named_arrow_pandas_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,), {}) + + class AbstractExecutor(ABC): """Base class for pipeline executors. @@ -352,17 +407,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 = _named_arrow_pandas_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/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py new file mode 100644 index 0000000000..c36afd4af6 --- /dev/null +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -0,0 +1,59 @@ +# 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 typing import Any + +import pandas as pd +import pyarrow as pa +from ray.data import DataContext + +from nemo_retriever.graph.executor import _ArrowPandasOperatorAdapter +from nemo_retriever.operators.abstract_operator import AbstractOperator + + +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) From 9653185465403a0203186e03a3ce5487e8e97dae Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 12 Aug 2026 15:03:44 +0000 Subject: [PATCH 2/5] fix: preserve batch object columns and sidecar metadata Normalize Ray pickled-object extension columns before pandas row operations and retain sidecar settings when graph actors are reconstructed. Signed-off-by: Jacob Ioffe --- .../src/nemo_retriever/graph/executor.py | 15 +++++++++++++++ .../src/nemo_retriever/operators/vdb.py | 2 +- .../tests/test_executor_arrow_pandas.py | 18 ++++++++++++++++++ .../tests/test_nv_ingest_vdb_operator.py | 18 ++++++++++++++++++ 4 files changed, 52 insertions(+), 1 deletion(-) diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index 9a922a3a61..12620ee55d 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -71,6 +71,20 @@ def _compact_vulnerable_arrow_columns(table: Any) -> Any: 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 + + class _ArrowPandasOperatorAdapter: """Convert valid Arrow batches to pandas before invoking an NRL operator.""" @@ -82,6 +96,7 @@ def __call__(self, table: Any) -> Any: table = _compact_vulnerable_arrow_columns(table) frame = BlockAccessor.for_block(table).to_pandas() + frame = _normalize_pickled_object_columns(table, frame) return self._operator(frame) 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 index c36afd4af6..3e45694a7d 100644 --- a/nemo_retriever/tests/test_executor_arrow_pandas.py +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -6,8 +6,10 @@ 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 @@ -57,3 +59,19 @@ def test_adapter_preserves_ray_pandas_conversion_policy() -> None: 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 result.apply(lambda row: row["text"], axis=1).tolist() == ["first", "second", "third"] 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) From dc6a44ce31dd2ed48b45d52071e0d78070116d39 Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 12 Aug 2026 18:31:26 +0000 Subject: [PATCH 3/5] fix: preserve Arrow-backed multimodal content Signed-off-by: Jacob Ioffe --- .../common/modality/content_transforms.py | 6 ++- .../src/nemo_retriever/common/vdb/records.py | 17 ++++++++- .../tests/test_executor_arrow_pandas.py | 38 +++++++++++++++++++ nemo_retriever/tests/test_multimodal_embed.py | 30 +++++++++++++++ nemo_retriever/tests/test_vdb_records.py | 16 ++++++++ 5 files changed, 103 insertions(+), 4 deletions(-) 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..6e3d5e74d2 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 @@ -15,6 +16,7 @@ from nemo_retriever.common.params.models import IMAGE_MODALITIES _CONTENT_COLUMNS = ("table", "chart", "infographic") +_CONTENT_COLLECTION_TYPES = (list, np.ndarray) def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequence[str]) -> str: @@ -25,7 +27,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 isinstance(content_list, _CONTENT_COLLECTION_TYPES): for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -108,7 +110,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 isinstance(content_list, _CONTENT_COLLECTION_TYPES): 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/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py index 3e45694a7d..a82a623673 100644 --- a/nemo_retriever/tests/test_executor_arrow_pandas.py +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -4,6 +4,7 @@ """Regression tests for Ray's Arrow-to-pandas operator boundary.""" +from functools import partial from typing import Any import numpy as np @@ -13,7 +14,9 @@ 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): @@ -75,3 +78,38 @@ def test_adapter_normalizes_pickled_object_columns_for_pandas_row_operations() - assert result["tables"].dtype == object assert result.apply(lambda row: row["text"], axis=1).tolist() == ["first", "second", "third"] + + +def test_adapter_preserves_page_and_element_content_semantics_for_sliced_pdf_rows() -> None: + table = pa.Table.from_pylist( + [ + { + "text": f"page {page_number}", + "table": [{"text": f"table {page_number}"}], + "chart": [{"text": f"chart {page_number}"}], + "metadata": {"source_path": "document.pdf", "error": None}, + } + for page_number in range(3) + ] + ).slice(2, 1) + + 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 2\n\ntable 2\n\nchart 2"] + assert element_result["text"].tolist() == ["page 2", "table 2", "chart 2"] + 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_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index 5066adb3f0..170f76a6f6 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,21 @@ 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"] + @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 +262,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_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( { From 32484e965c20a424ab31e3549941fb28732df02a Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 12 Aug 2026 19:23:56 +0000 Subject: [PATCH 4/5] fix: normalize all Ray Arrow pandas boundaries Signed-off-by: Jacob Ioffe --- .../common/api/internal/mutate/deduplicate.py | 8 ++- .../common/modality/collections.py | 41 +++++++++++ .../common/modality/content_transforms.py | 17 +++-- .../common/modality/ocr/shared.py | 11 +-- .../common/modality/page_elements/shared.py | 2 +- .../common/modality/pipeline/content.py | 15 ++-- .../common/modality/table_and_chart.py | 6 +- .../common/vdb/lancedb_collections.py | 6 +- .../common/vdb/lancedb_schema.py | 33 +++++---- .../src/nemo_retriever/common/vdb/records.py | 5 +- .../src/nemo_retriever/graph/executor.py | 40 ++++++++--- .../harness/portal/static/views/designer.jsx | 19 ++++-- .../ingestor/branch_extraction.py | 10 ++- .../nemo_retriever/ingestor/graph_ingestor.py | 9 ++- .../src/nemo_retriever/operators/dedup.py | 17 ++--- .../operators/embed/text_embed.py | 2 +- .../operators/extract/caption/caption.py | 13 ++-- .../extract/page_elements/gpu_actor.py | 7 +- .../operators/extract/pdf/split.py | 7 +- .../src/nemo_retriever/operators/rerank.py | 50 ++++++++------ .../tools/benchmark/audio_extract_actor.py | 11 +-- .../tools/benchmark/extract_actor.py | 6 +- .../tools/benchmark/ocr_actor.py | 14 ++-- .../tools/benchmark/page_elements_actor.py | 10 ++- .../tools/benchmark/split_actor.py | 6 +- nemo_retriever/tests/test_caption.py | 16 +++++ nemo_retriever/tests/test_dedup.py | 36 ++++++++++ nemo_retriever/tests/test_ingest_interface.py | 68 +++++++++++++------ nemo_retriever/tests/test_ingest_manifest.py | 41 ++++++++++- .../tests/test_lancedb_collections.py | 10 +++ .../tests/test_lancedb_row_metadata.py | 22 ++++++ .../tests/test_nemotron_rerank_vl_v2_hf.py | 29 ++++++++ nemo_retriever/tests/test_vdb_records.py | 15 ++++ 33 files changed, 466 insertions(+), 136 deletions(-) create mode 100644 nemo_retriever/src/nemo_retriever/common/modality/collections.py diff --git a/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py b/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py index 3d5890aac6..973671e231 100644 --- a/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py +++ b/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py @@ -9,6 +9,8 @@ import pandas as pd +from nemo_retriever.common.modality.collections import bbox_coordinates + from nemo_retriever.common.api.internal.enums.common import ContentTypeEnum from nemo_retriever.common.api.internal.schemas.mutate.mutate_image_dedup_schema import ImageDedupSchema @@ -112,7 +114,8 @@ def _get_image_bbox_info(row: pd.Series) -> Optional[Dict[str, Any]]: bbox = image_metadata.get("image_location") max_dims = image_metadata.get("image_location_max_dimensions", (0, 0)) - if bbox is None or not isinstance(bbox, (tuple, list)) or len(bbox) < 4: + bbox = bbox_coordinates(bbox) + if bbox is None: return None # Normalize bbox by max dimensions @@ -153,7 +156,8 @@ def _get_structured_bbox_info(row: pd.Series) -> Optional[Dict[str, Any]]: bbox = table_metadata.get("table_location") max_dims = table_metadata.get("table_location_max_dimensions", (0, 0)) - if bbox is None or not isinstance(bbox, (tuple, list)) or len(bbox) < 4: + bbox = bbox_coordinates(bbox) + if bbox is None: return None # Normalize bbox by max dimensions diff --git a/nemo_retriever/src/nemo_retriever/common/modality/collections.py b/nemo_retriever/src/nemo_retriever/common/modality/collections.py new file mode 100644 index 0000000000..3c959ce574 --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/common/modality/collections.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Collection adapters for multimodal values crossing Ray's pandas boundary.""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + + +def multimodal_collection(value: Any) -> list[Any] | None: + """Return a mutable list view of a structured multimodal collection. + + Ray may preserve list-valued pandas object cells as one-dimensional NumPy + object arrays. Restrict ndarray support to one dimension so image tensors + and other genuine multidimensional model inputs are never reinterpreted as + collections of extracted objects. + """ + if isinstance(value, list): + return value + if isinstance(value, tuple): + return list(value) + if isinstance(value, np.ndarray) and value.ndim == 1: + return list(value) + return None + + +def bbox_coordinates(value: Any) -> list[Any] | None: + """Return four-or-more bbox coordinates without array truth testing.""" + if isinstance(value, np.ndarray): + if value.ndim != 1: + return None + value = value.tolist() + elif isinstance(value, tuple): + value = list(value) + if isinstance(value, list) and len(value) >= 4: + return value + return None 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 6e3d5e74d2..8c32fcfb15 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py @@ -8,15 +8,14 @@ from typing import Any, Dict, List, Optional, Sequence -import numpy as np import pandas as pd +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.io.image_store import inline_image_b64 from nemo_retriever.operators.extract.ocr.ocr import _crop_b64_image_by_norm_bbox from nemo_retriever.common.params.models import IMAGE_MODALITIES _CONTENT_COLUMNS = ("table", "chart", "infographic") -_CONTENT_COLLECTION_TYPES = (list, np.ndarray) def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequence[str]) -> str: @@ -26,8 +25,8 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ if isinstance(base, str) and base.strip(): parts.append(base.strip()) for col in content_columns: - content_list = row.get(col) - if isinstance(content_list, _CONTENT_COLLECTION_TYPES): + content_list = multimodal_collection(row.get(col)) + if content_list is not None: for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -45,7 +44,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, tuple)) or multimodal_collection(value) is not None: out[key] = copy.deepcopy(value) else: out[key] = value @@ -109,8 +108,8 @@ def explode_content_to_rows( exploded_any = True for column in content_columns: - content_list = row_dict.get(column) - if not isinstance(content_list, _CONTENT_COLLECTION_TYPES): + content_list = multimodal_collection(row_dict.get(column)) + if content_list is None: continue for item in content_list: if not isinstance(item, dict): @@ -129,8 +128,8 @@ def explode_content_to_rows( if item_b64: content_row["_image_b64"] = item_b64 elif page_image_b64: - bbox = item.get("bbox_xyxy_norm") - if bbox and len(bbox) == 4: + bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) + if bbox is not None: cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox) content_row["_image_b64"] = cropped_b64 else: diff --git a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py index 81f2e0951b..242b24bb0b 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py @@ -25,6 +25,7 @@ import numpy as np import pandas as pd +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.params import RemoteRetryParams from nemo_retriever.models.nim.nim import NIMClient, invoke_image_inference_batches from nemo_retriever.common.modality.table_and_chart import join_table_structure_and_ocr_output @@ -175,7 +176,8 @@ def _clamp_int(v: float, lo: int, hi: int) -> int: continue bbox = det.get("bbox_xyxy_norm") - if not isinstance(bbox, (list, tuple)) or len(bbox) != 4: + bbox = bbox_coordinates(bbox) + if bbox is None or len(bbox) != 4: continue try: @@ -620,7 +622,8 @@ def _find_ts_detections_for_bbox( if not isinstance(region, dict): continue region_bbox = region.get("bbox_xyxy_norm") - if not isinstance(region_bbox, (list, tuple)) or len(region_bbox) != 4: + region_bbox = bbox_coordinates(region_bbox) + if region_bbox is None or len(region_bbox) != 4: continue if not _bboxes_close(table_bbox, region_bbox): continue @@ -698,9 +701,7 @@ def _prepare_ocr_rows( page_elements = getattr(row, "page_elements_v3", None) detections: List[Dict[str, Any]] = [] if isinstance(page_elements, dict): - detections = page_elements.get("detections") or [] - if not isinstance(detections, list): - detections = [] + detections = multimodal_collection(page_elements.get("detections")) or [] page_image = getattr(row, "page_image", None) or {} page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None diff --git a/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py index 856e0844db..e73b13927d 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py @@ -461,7 +461,7 @@ def detect_page_elements_v3( Run Nemotron Page Elements v3 on a pandas batch. Input: - - `pages_df`: pandas.DataFrame (typical Ray Data `batch_format="pandas"`) + - `pages_df`: pandas.DataFrame (provided by NRL's safe Arrow adapter in Ray Data pipelines) Must contain an image base64 source either in `image_b64` or one of `images`/`tables`/`charts`/`infographics` (each as list[{"image_b64": ...}]). diff --git a/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py b/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py index e8d84c0db4..6f447211d1 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py @@ -10,6 +10,7 @@ import pandas as pd +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.io.image_store import inline_image_b64 from nemo_retriever.operators.extract.ocr.ocr import _crop_b64_image_by_norm_bbox from nemo_retriever.common.params.models import IMAGE_MODALITIES @@ -23,8 +24,8 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ if isinstance(base, str) and base.strip(): parts.append(base.strip()) for col in content_columns: - content_list = row.get(col) - if isinstance(content_list, list): + content_list = multimodal_collection(row.get(col)) + if content_list is not None: for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -41,7 +42,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, tuple)) or multimodal_collection(value) is not None: out[key] = copy.deepcopy(value) else: out[key] = value @@ -105,8 +106,8 @@ def explode_content_to_rows( exploded_any = True for column in content_columns: - content_list = row_dict.get(column) - if not isinstance(content_list, list): + content_list = multimodal_collection(row_dict.get(column)) + if content_list is None: continue for item in content_list: if not isinstance(item, dict): @@ -124,8 +125,8 @@ def explode_content_to_rows( if item_b64: content_row["_image_b64"] = item_b64 elif page_image_b64: - bbox = item.get("bbox_xyxy_norm") - if bbox and len(bbox) == 4: + bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) + if bbox is not None: cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox) content_row["_image_b64"] = cropped_b64 else: diff --git a/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py b/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py index 0ba6d73197..0b82bd15ac 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py @@ -15,6 +15,8 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple # noqa: F401 import numpy as np + +from nemo_retriever.common.modality.collections import bbox_coordinates import pandas as pd logger = logging.getLogger(__name__) @@ -396,8 +398,8 @@ def _structure_dets_to_class_boxes( grouped: Dict[str, list] = {} for d in dets: name = d.get("label_name", "") - bbox = d.get("bbox_xyxy_norm") - if not bbox or len(bbox) != 4: + bbox = bbox_coordinates(d.get("bbox_xyxy_norm")) + if bbox is None or len(bbox) != 4: continue x1, y1, x2, y2 = float(bbox[0]) * W, float(bbox[1]) * H, float(bbox[2]) * W, float(bbox[3]) * H grouped.setdefault(name, []).append([x1, y1, x2, y2]) diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py index ddfa0ec533..712da211b3 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py @@ -128,7 +128,9 @@ def _decode_cursor( def _json_string(value: Any) -> str: if isinstance(value, str): return value - return json.dumps(value or {}, ensure_ascii=False, separators=(",", ":"), default=str) + if hasattr(value, "tolist"): + value = value.tolist() + return json.dumps({} if value is None else value, ensure_ascii=False, separators=(",", ":"), default=str) def _content_text(record: dict[str, Any], metadata: dict[str, Any]) -> str: @@ -227,7 +229,7 @@ def _collection_rows( "metadata": _json_string(content_metadata), "stored_image_uri": stored_image_uri, "content_type": content_type, - "bbox_xyxy_norm": _json_string(bbox) if bbox else "", + "bbox_xyxy_norm": _json_string(bbox) if bbox is not None else "", "chunk_id": hashlib.sha256( f"{context.document_id}\0{context.document_version}\0{row_index}".encode() ).hexdigest(), diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py index d684add889..6a8deab636 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py @@ -10,9 +10,18 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.vdb.records import normalize_content_type +def _embedding_list(value: Any) -> Optional[List[float]]: + if hasattr(value, "tolist"): + value = value.tolist() + elif isinstance(value, tuple): + value = list(value) + return value if isinstance(value, list) and value else None + + def extract_embedding_from_row( row: Any, *, @@ -22,15 +31,15 @@ def extract_embedding_from_row( """Extract an embedding vector from a row.""" meta = getattr(row, "metadata", None) if isinstance(meta, dict): - emb = meta.get("embedding") - if isinstance(emb, list) and emb: - return emb # type: ignore[return-value] + emb = _embedding_list(meta.get("embedding")) + if emb is not None: + return emb payload = getattr(row, embedding_column, None) if isinstance(payload, dict): - emb = payload.get(embedding_key) - if isinstance(emb, list) and emb: - return emb # type: ignore[return-value] + emb = _embedding_list(payload.get(embedding_key)) + if emb is not None: + return emb return None @@ -114,8 +123,8 @@ def _build_detection_metadata(row: Any) -> Dict[str, Any]: } for ocr_column in ("table", "chart", "infographic"): - entries = getattr(row, ocr_column, None) - if isinstance(entries, list): + entries = multimodal_collection(getattr(row, ocr_column, None)) + if entries is not None: out[f"ocr_{ocr_column}_detections"] = int(len(entries)) return out @@ -186,8 +195,8 @@ def build_lancedb_row( content_type = getattr(row, "_content_type", None) row_out["content_type"] = str(content_type) if content_type else "" - bbox = getattr(row, "_bbox_xyxy_norm", None) - row_out["bbox_xyxy_norm"] = json.dumps(bbox) if bbox else "" + bbox = bbox_coordinates(getattr(row, "_bbox_xyxy_norm", None)) + row_out["bbox_xyxy_norm"] = json.dumps(bbox) if bbox is not None else "" return row_out @@ -259,8 +268,8 @@ def lancedb_schema(vector_dim: int = 2048, *, collection_managed: bool = False) def infer_vector_dim(rows: List[Dict[str, Any]]) -> int: """Return the embedding dimension from the first row that has a vector.""" for row in rows: - vector = row.get("vector") - if isinstance(vector, list) and vector: + vector = _embedding_list(row.get("vector")) + if vector is not None: return len(vector) return 0 diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/records.py b/nemo_retriever/src/nemo_retriever/common/vdb/records.py index 3bc10caeb5..ac87ebf5b1 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/records.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/records.py @@ -13,6 +13,7 @@ from pydantic import ValidationError +from nemo_retriever.common.modality.collections import multimodal_collection from nemo_retriever.common.schemas.collections import QueryHit _CONTENT_TYPE_ALIASES: dict[str, str] = { @@ -195,8 +196,8 @@ def _add_detection_metadata( content_metadata.setdefault("page_elements_v3_counts_by_label", normalized_counts) for content_type in ("table", "chart", "infographic"): - detections = row.get(content_type) - if isinstance(detections, list): + detections = multimodal_collection(row.get(content_type)) + if detections is not None: content_metadata.setdefault(f"ocr_{content_type}_detections", len(detections)) diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index 12620ee55d..b4ef0eb1b6 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -85,6 +85,35 @@ def _normalize_pickled_object_columns(table: Any, frame: pd.DataFrame) -> pd.Dat 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.""" @@ -92,15 +121,10 @@ def __init__(self, operator_class: type, operator_kwargs: dict[str, Any]) -> Non self._operator = operator_class(**operator_kwargs) def __call__(self, table: Any) -> Any: - from ray.data.block import BlockAccessor - - table = _compact_vulnerable_arrow_columns(table) - frame = BlockAccessor.for_block(table).to_pandas() - frame = _normalize_pickled_object_columns(table, frame) - return self._operator(frame) + return self._operator(arrow_table_to_pandas(table)) -def _named_arrow_pandas_adapter(operator_class: type) -> type[_ArrowPandasOperatorAdapter]: +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,), {}) @@ -429,7 +453,7 @@ def build_dataset(self, data: Any, **kwargs: Any) -> Any: # 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 = _named_arrow_pandas_adapter(node.operator_class) + map_operator_class = make_arrow_pandas_operator_adapter(node.operator_class) map_batch_format = "pyarrow" constructor_kwargs = { "operator_class": node.operator_class, diff --git a/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx b/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx index 2b9e0e9652..99c7743aac 100644 --- a/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx +++ b/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx @@ -105,6 +105,7 @@ function _generateRayDataCode(nodes, edges) { const imports = new Set(); imports.add('import ray'); imports.add('import ray.data'); + imports.add('from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter'); chain.forEach(n => { if (n.operator && n.operator.type !== 'ray_data_source') { imports.add(n.operator.import_path); @@ -161,10 +162,13 @@ function _generateRayDataCode(nodes, edges) { lines.push(`# Stage: ${op.display_name || op.class_name}`); lines.push(`ds = ds.map_batches(`); - lines.push(` ${op.class_name},`); - lines.push(` fn_constructor_kwargs=${kwargsStr},`); + lines.push(` make_arrow_pandas_operator_adapter(${op.class_name}),`); + lines.push(` fn_constructor_kwargs={`); + lines.push(` "operator_class": ${op.class_name},`); + lines.push(` "operator_kwargs": ${kwargsStr},`); + lines.push(` },`); lines.push(` batch_size=1,`); - lines.push(` batch_format="pandas",`); + lines.push(` batch_format="pyarrow",`); if (isGpu) { lines.push(` num_cpus=0,`); lines.push(` num_gpus=1,`); @@ -191,10 +195,13 @@ function _generateRayDataCode(nodes, edges) { lines.push(`# Sink: ${op.display_name || op.class_name}`); lines.push(`ds = ds.map_batches(`); - lines.push(` ${op.class_name},`); - lines.push(` fn_constructor_kwargs=${kwargsStr},`); + lines.push(` make_arrow_pandas_operator_adapter(${op.class_name}),`); + lines.push(` fn_constructor_kwargs={`); + lines.push(` "operator_class": ${op.class_name},`); + lines.push(` "operator_kwargs": ${kwargsStr},`); + lines.push(` },`); lines.push(` batch_size=1,`); - lines.push(` batch_format="pandas",`); + lines.push(` batch_format="pyarrow",`); if (sinkConcurrency !== undefined && sinkConcurrency !== '') { lines.push(` concurrency=${parseInt(sinkConcurrency, 10) || 1},`); } diff --git a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py index 8f414049b8..95a84632c8 100644 --- a/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py +++ b/nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py @@ -12,6 +12,7 @@ from typing import Any, Callable from nemo_retriever.graph import InprocessExecutor, RayDataExecutor +from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.graph.ingestor_runtime import batch_tuning_to_node_overrides, build_graph, build_post_extract_graph from nemo_retriever.ingestor.manifest import ( ExtractionBranchPlan, @@ -341,9 +342,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 e25286b99f..e82bb6a324 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, build_graph from nemo_retriever.ingestor.manifest import ( @@ -1171,7 +1172,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,) @@ -1353,7 +1354,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/dedup.py b/nemo_retriever/src/nemo_retriever/operators/dedup.py index 4ecb1b3580..a4a83f4cb4 100644 --- a/nemo_retriever/src/nemo_retriever/operators/dedup.py +++ b/nemo_retriever/src/nemo_retriever/operators/dedup.py @@ -10,6 +10,7 @@ import pandas as pd from nemo_retriever.common.api.internal.mutate.deduplicate import calculate_iou +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.params import DedupParams _STRUCTURED_COLUMNS = ("table", "chart", "infographic") @@ -19,14 +20,14 @@ def _collect_structured_bboxes(row: pd.Series) -> List[Tuple[float, ...]]: """Gather all bounding boxes from tables, charts, and infographics columns.""" bboxes: List[Tuple[float, ...]] = [] for col in _STRUCTURED_COLUMNS: - items = row.get(col) - if not isinstance(items, list): + items = multimodal_collection(row.get(col)) + if items is None: continue for item in items: if not isinstance(item, dict): continue - bbox = item.get("bbox_xyxy_norm") - if bbox and len(bbox) >= 4: + bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) + if bbox is not None: bboxes.append(tuple(bbox[:4])) return bboxes @@ -55,8 +56,8 @@ def dedup_images( return batch_df for row_idx, row in batch_df.iterrows(): - images = row.get("images") - if not isinstance(images, list) or not images: + images = multimodal_collection(row.get("images")) + if not images: continue filtered = list(images) @@ -88,8 +89,8 @@ def dedup_images( if not isinstance(item, dict): surviving.append(item) continue - img_bbox = item.get("bbox_xyxy_norm") - if not img_bbox or len(img_bbox) < 4: + img_bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) + if img_bbox is None: surviving.append(item) continue img_bbox_t = tuple(img_bbox[:4]) diff --git a/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py b/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py index 49151edc51..e529aaebfe 100644 --- a/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py +++ b/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py @@ -80,7 +80,7 @@ def embed_text_1b_v2( Embed a batch of text rows using the local `LlamaNemotronEmbed1BV2Embedder`. Input: - - `batch_df`: pandas.DataFrame (Ray Data `batch_format="pandas"` compatible) + - `batch_df`: pandas.DataFrame (provided by NRL's safe Arrow adapter in Ray Data pipelines) - `text_column`: preferred column name to read text from (defaults to `"text"`) Output: diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py b/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py index a7000af694..56d259a224 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py @@ -17,6 +17,7 @@ merge_request_extras, resolve_caption_model_name, ) +from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator from nemo_retriever.graph.designer import designer_component @@ -372,8 +373,8 @@ def caption_images( for row_idx, row in batch_df.iterrows(): # Unstructured images. if has_images: - images = row.get("images") - if isinstance(images, list): + images = multimodal_collection(row.get("images")) + if images is not None: for item_idx, item in enumerate(images): if not isinstance(item, dict): continue @@ -385,8 +386,8 @@ def caption_images( # Infographics — crop from page image. if has_infographics: - infographics = row.get("infographic") - if isinstance(infographics, list): + infographics = multimodal_collection(row.get("infographic")) + if infographics is not None: page_image = row.get("page_image") page_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None if page_b64: @@ -395,8 +396,8 @@ def caption_images( continue if item.get("caption"): continue # already captioned - bbox = item.get("bbox_xyxy_norm") - if not bbox or len(bbox) < 4: + bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) + if bbox is None: continue cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_b64, bbox_xyxy_norm=bbox) if cropped_b64 and _image_meets_min_size(cropped_b64): diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py b/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py index a3f88e4e3d..d720ab64e4 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py @@ -18,8 +18,11 @@ class PageElementDetectionActor(AbstractOperator, GPUOperator): """ Ray-friendly callable that initializes Nemotron Page Elements v3 once. - Use with Ray Data: - ds = ds.map_batches(PageElementDetectionActor, fn_constructor_kwargs={...}, batch_format="pandas") + Ray Data callers should use + :func:`nemo_retriever.graph.executor.make_arrow_pandas_operator_adapter` + with ``batch_format="pyarrow"`` so nested Arrow batches are normalized + before this pandas operator runs. ``RayDataExecutor`` applies that adapter + automatically. """ def __init__(self, **detect_kwargs: Any) -> None: 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/rerank.py b/nemo_retriever/src/nemo_retriever/operators/rerank.py index 48cce39cdd..25ff1b86b7 100644 --- a/nemo_retriever/src/nemo_retriever/operators/rerank.py +++ b/nemo_retriever/src/nemo_retriever/operators/rerank.py @@ -32,19 +32,23 @@ Ray Data actor usage:: import ray + from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter ds = ds.map_batches( - NemotronRerankActor, + make_arrow_pandas_operator_adapter(NemotronRerankActor), batch_size=64, - batch_format="pandas", + batch_format="pyarrow", num_gpus=1, compute=ray.data.ActorPoolStrategy(size=4), fn_constructor_kwargs={ - "model_name": "nvidia/llama-nemotron-rerank-1b-v2", - "query_column": "query", - "text_column": "text", - "score_column": "rerank_score", - "max_length": 512, - "batch_size": 32, + "operator_class": NemotronRerankActor, + "operator_kwargs": { + "model_name": "nvidia/llama-nemotron-rerank-1b-v2", + "query_column": "query", + "text_column": "text", + "score_column": "rerank_score", + "max_length": 512, + "batch_size": 32, + }, }, ) """ @@ -58,6 +62,7 @@ from typing import Any, Dict, List, Optional import pandas as pd +from nemo_retriever.common.modality.collections import bbox_coordinates from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.graph.designer import designer_component from nemo_retriever.operators.cpu_operator import CPUOperator @@ -324,11 +329,12 @@ def rerank_hits( render_cache[cache_key] = render_page_image_b64(h["path"], int(h["page_number"])) page_b64 = render_cache[cache_key] # Crop to element bbox if available (tables/charts). - bbox_str = h.get("bbox_xyxy_norm", "") - if page_b64 and bbox_str: + bbox_value = h.get("bbox_xyxy_norm") + if page_b64 and bbox_value is not None: try: - bbox = json.loads(bbox_str) - if isinstance(bbox, list) and len(bbox) == 4: + bbox = json.loads(bbox_value) if isinstance(bbox_value, str) else bbox_value + bbox = bbox_coordinates(bbox) + if bbox is not None and len(bbox) == 4: cropped, _ = _crop_b64_image_by_norm_bbox(page_b64, bbox_xyxy_norm=bbox) images_b64.append(cropped) continue @@ -414,19 +420,23 @@ class NemotronRerankGPUActor(AbstractOperator, GPUOperator): Usage with Ray Data:: import ray + from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter ds = ds.map_batches( - NemotronRerankActor, + make_arrow_pandas_operator_adapter(NemotronRerankActor), batch_size=64, - batch_format="pandas", + batch_format="pyarrow", num_gpus=1, compute=ray.data.ActorPoolStrategy(size=4), fn_constructor_kwargs={ - "model_name": "nvidia/llama-nemotron-rerank-1b-v2", - "query_column": "query", - "text_column": "text", - "score_column": "rerank_score", - "max_length": 512, - "batch_size": 32, + "operator_class": NemotronRerankActor, + "operator_kwargs": { + "model_name": "nvidia/llama-nemotron-rerank-1b-v2", + "query_column": "query", + "text_column": "text", + "score_column": "rerank_score", + "max_length": 512, + "batch_size": 32, + }, }, ) diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py index 4bdb16a92c..776cba2cf2 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py @@ -23,6 +23,7 @@ from nemo_retriever.common.modality.audio.media_interface import is_media_available from nemo_retriever.common.modality.audio.media_interface import media_dependency_error_message from nemo_retriever.common.params import AudioChunkParams +from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -91,20 +92,22 @@ def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: asr_actor = ASRActor(params=asr_params_from_env()) ds = ds.map_batches( - chunk_actor, + call_pandas_function_on_arrow, batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), + fn_kwargs={"fn": chunk_actor}, ) ds = ds.map_batches( - asr_actor, + call_pandas_function_on_arrow, batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=1, num_gpus=0.25 if not mock_asr else 0, compute=rd.TaskPoolStrategy(size=int(worker_count)), + fn_kwargs={"fn": asr_actor}, ) return ds diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py index cfbb7bb6e9..d22bb1880c 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py @@ -11,6 +11,7 @@ import typer from nemo_retriever.operators.extract.pdf.extract import PDFExtractionActor +from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -52,12 +53,13 @@ def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: dpi=int(dpi), ) return ds.map_batches( - actor, + call_pandas_function_on_arrow, batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), + fn_kwargs={"fn": actor}, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py index 8234ba850e..5f047e8849 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py @@ -11,6 +11,7 @@ import typer from nemo_retriever.operators.extract.ocr.gpu_ocr import OCRActor +from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -54,16 +55,19 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - OCRActor, + make_arrow_pandas_operator_adapter(OCRActor), batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=float(num_cpus), num_gpus=float(num_gpus), compute=rd.ActorPoolStrategy(size=int(worker_count)), fn_constructor_kwargs={ - "extract_tables": bool(extract_tables), - "extract_charts": bool(extract_charts), - "extract_infographics": bool(extract_infographics), + "operator_class": OCRActor, + "operator_kwargs": { + "extract_tables": bool(extract_tables), + "extract_charts": bool(extract_charts), + "extract_infographics": bool(extract_infographics), + }, }, ) diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py index 9cdad6b0e9..706db3fb67 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py @@ -11,6 +11,7 @@ import typer from nemo_retriever.operators.extract.page_elements.gpu_actor import PageElementDetectionActor +from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -45,13 +46,16 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - PageElementDetectionActor, + make_arrow_pandas_operator_adapter(PageElementDetectionActor), batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=float(num_cpus), num_gpus=float(num_gpus), compute=rd.ActorPoolStrategy(size=int(worker_count)), - fn_constructor_kwargs={"inference_batch_size": int(inference_batch_size)}, + fn_constructor_kwargs={ + "operator_class": PageElementDetectionActor, + "operator_kwargs": {"inference_batch_size": int(inference_batch_size)}, + }, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py index ac461be1e8..c44cca5995 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py @@ -11,6 +11,7 @@ import typer from nemo_retriever.operators.extract.pdf.split import PDFSplitActor +from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -41,12 +42,13 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - PDFSplitActor(), + call_pandas_function_on_arrow, batch_size=int(batch_size), - batch_format="pandas", + batch_format="pyarrow", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), + fn_kwargs={"fn": PDFSplitActor()}, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/tests/test_caption.py b/nemo_retriever/tests/test_caption.py index acc27e03f6..1f75a7cc87 100644 --- a/nemo_retriever/tests/test_caption.py +++ b/nemo_retriever/tests/test_caption.py @@ -8,6 +8,7 @@ import io from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd import pytest @@ -41,6 +42,21 @@ def test_caption_images_writes_back(): assert result.iloc[0]["images"][1]["text"] == "cap2" +def test_caption_images_writes_back_to_arrow_object_array(): + from nemo_retriever.operators.extract.caption.caption import caption_images + + source = _make_page_df(num_images=1).iloc[0].to_dict() + source["images"] = np.array(source["images"], dtype=object) + df = pd.DataFrame([source]) + mock_model = MagicMock() + mock_model.caption_batch.return_value = ["captioned"] + + result = caption_images(df, model=mock_model) + + mock_model.caption_batch.assert_called_once() + assert result.iloc[0]["images"][0]["text"] == "captioned" + + def test_caption_images_skips_already_captioned(): from nemo_retriever.operators.extract.caption.caption import caption_images diff --git a/nemo_retriever/tests/test_dedup.py b/nemo_retriever/tests/test_dedup.py index 0f29ffb684..55e08f8038 100644 --- a/nemo_retriever/tests/test_dedup.py +++ b/nemo_retriever/tests/test_dedup.py @@ -4,6 +4,7 @@ """Unit tests for nemo_retriever.operators.dedup.""" +import numpy as np import pandas as pd from nemo_retriever.common.api.internal.mutate.deduplicate import calculate_iou @@ -65,6 +66,22 @@ def test_keeps_unique_images(self): result = dedup_images(df, content_hash=True, bbox_iou=False) assert len(result.iloc[0]["images"]) == 2 + def test_removes_duplicate_images_from_arrow_object_array(self): + df = _make_df( + [ + { + "images": np.array( + [_make_image("AAA"), _make_image("BBB"), _make_image("AAA")], + dtype=object, + ), + } + ] + ) + + result = dedup_images(df, content_hash=True, bbox_iou=False) + + assert [image["image_b64"] for image in result.iloc[0]["images"]] == ["AAA", "BBB"] + # ── bbox IoU dedup ───────────────────────────────────────────────────────── @@ -126,6 +143,25 @@ def test_image_without_bbox_is_kept(self): result = dedup_images(df, content_hash=False, bbox_iou=True) assert len(result.iloc[0]["images"]) == 1 + def test_drops_arrow_object_array_image_overlapping_with_structured_content(self): + df = _make_df( + [ + { + "images": np.array([_make_image("X", bbox=[0.0, 0.0, 1.0, 1.0])], dtype=object), + "table": np.array( + [{"bbox_xyxy_norm": [0.0, 0.0, 1.0, 1.0], "text": "t"}], + dtype=object, + ), + "chart": np.array([], dtype=object), + "infographic": np.array([], dtype=object), + } + ] + ) + + result = dedup_images(df, content_hash=False, bbox_iou=True, iou_threshold=0.45) + + assert result.iloc[0]["images"] == [] + # ── edge cases ───────────────────────────────────────────────────────────── 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 f754792bdc..40e66b44b3 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_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py index 3606cbaf32..e3b59d4d1f 100644 --- a/nemo_retriever/tests/test_lancedb_collections.py +++ b/nemo_retriever/tests/test_lancedb_collections.py @@ -15,6 +15,7 @@ from dataclasses import replace import lancedb +import numpy as np import pytest import nemo_retriever.common.vdb.lancedb_collections as collections_module @@ -163,6 +164,15 @@ def test_collection_row_conversion_preserves_identity_and_provenance(): assert row["chunk_id"] == hashlib.sha256(b"document-a\x00v1\x000").hexdigest() +def test_collection_row_conversion_normalizes_arrow_backed_bbox_array(): + records = _records() + records[0][0]["metadata"]["content_metadata"]["bbox_xyxy_norm"] = np.array([0.1, 0.2, 0.8, 0.9]) + + rows = _collection_rows(records, context=_context()) + + assert json.loads(rows[0]["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] + + @pytest.mark.parametrize( ("raw_type", "expected_type"), [("table_caption", "table"), ("chart_caption", "chart")], diff --git a/nemo_retriever/tests/test_lancedb_row_metadata.py b/nemo_retriever/tests/test_lancedb_row_metadata.py index d2123f1d2e..0b29303ef4 100644 --- a/nemo_retriever/tests/test_lancedb_row_metadata.py +++ b/nemo_retriever/tests/test_lancedb_row_metadata.py @@ -2,6 +2,8 @@ import json from types import SimpleNamespace +import numpy as np + from nemo_retriever.common.vdb.lancedb_bulk import _build_lancedb_rows_from_df from nemo_retriever.common.vdb.lancedb_schema import build_lancedb_row @@ -22,6 +24,26 @@ def test_build_lancedb_row_persists_normalized_content_type() -> None: assert metadata["_content_type"] == "table" +def test_build_lancedb_row_normalizes_arrow_backed_arrays() -> None: + row = SimpleNamespace( + path="/tmp/doc_arrays.pdf", + page_number=2, + metadata={"embedding": np.array([0.1, 0.2], dtype=np.float32)}, + text="table text", + table=np.array([{"text": "first"}, {"text": "second"}], dtype=object), + chart=np.array([], dtype=object), + infographic=np.array([], dtype=object), + _bbox_xyxy_norm=np.array([0.1, 0.2, 0.8, 0.9]), + ) + + row_out = build_lancedb_row(row) + + assert row_out is not None + assert row_out["vector"] == [np.float32(0.1), np.float32(0.2)] + assert json.loads(row_out["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] + assert json.loads(row_out["metadata"])["ocr_table_detections"] == 2 + + def test_build_lancedb_rows_from_df_persists_normalized_content_type() -> None: rows = [ { diff --git a/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py b/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py index 38f379e470..302a6a44c7 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py +++ b/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py @@ -14,6 +14,7 @@ from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd import pytest @@ -265,6 +266,34 @@ def test_rerank_hits_graceful_without_uri(self): assert call_kwargs.kwargs["images_b64"] == [None, None] assert len(out) == 2 + def test_rerank_hits_crops_arrow_backed_bbox_array(self): + from nemo_retriever.operators.rerank import rerank_hits + + model = MagicMock() + model.model_name = "nvidia/llama-nemotron-rerank-vl-1b-v2" + model.score.return_value = [1.0] + hits = [ + { + "text": "table", + "path": "/tmp/document.pdf", + "page_number": 1, + "bbox_xyxy_norm": np.array([0.1, 0.2, 0.8, 0.9]), + } + ] + + with ( + patch("nemo_retriever.common.io.image_store.render_page_image_b64", return_value="page") as render, + patch( + "nemo_retriever.operators.extract.ocr.ocr._crop_b64_image_by_norm_bbox", + return_value=("crop", None), + ) as crop, + ): + rerank_hits("q", hits, model=model, modality="text_image") + + render.assert_called_once_with("/tmp/document.pdf", 1) + crop.assert_called_once_with("page", bbox_xyxy_norm=[0.1, 0.2, 0.8, 0.9]) + assert model.score.call_args.kwargs["images_b64"] == ["crop"] + # --------------------------------------------------------------------------- # _rerank_via_endpoint with images diff --git a/nemo_retriever/tests/test_vdb_records.py b/nemo_retriever/tests/test_vdb_records.py index bd68725c24..5b0220f3fc 100644 --- a/nemo_retriever/tests/test_vdb_records.py +++ b/nemo_retriever/tests/test_vdb_records.py @@ -192,6 +192,21 @@ def test_graph_record_conversion_normalizes_arrow_backed_bbox_array() -> None: assert records[0][0]["metadata"]["content_metadata"]["bbox_xyxy_norm"] == [0.1, 0.2, 0.8, 0.9] +def test_graph_record_conversion_counts_arrow_backed_detection_array() -> None: + records = to_client_vdb_records( + [ + { + "text": "table content", + "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, + "_content_type": "table", + "table": np.array([{"content": "a"}, {"content": "b"}], dtype=object), + } + ] + ) + + assert records[0][0]["metadata"]["content_metadata"]["ocr_table_detections"] == 2 + + def test_narrow_lancedb_hit_promotes_canonical_multimodal_metadata() -> None: hit = _normalize_one( { From 9cbe3d21a1df0aec26b410254a988f80bc69eefa Mon Sep 17 00:00:00 2001 From: Jacob Ioffe Date: Wed, 12 Aug 2026 20:18:34 +0000 Subject: [PATCH 5/5] fix: narrow Arrow normalization to Ray boundaries Signed-off-by: Jacob Ioffe --- .../common/api/internal/mutate/deduplicate.py | 8 +-- .../common/modality/collections.py | 41 --------------- .../common/modality/content_transforms.py | 21 +++++--- .../common/modality/ocr/shared.py | 11 ++-- .../common/modality/page_elements/shared.py | 2 +- .../common/modality/pipeline/content.py | 15 +++--- .../common/modality/table_and_chart.py | 6 +-- .../common/vdb/lancedb_collections.py | 6 +-- .../common/vdb/lancedb_schema.py | 33 +++++------- .../src/nemo_retriever/common/vdb/records.py | 5 +- .../src/nemo_retriever/graph/executor.py | 4 +- .../harness/portal/static/views/designer.jsx | 19 +++---- .../src/nemo_retriever/operators/dedup.py | 17 +++---- .../operators/embed/text_embed.py | 2 +- .../operators/extract/caption/caption.py | 13 +++-- .../extract/page_elements/gpu_actor.py | 7 +-- .../src/nemo_retriever/operators/rerank.py | 50 ++++++++----------- .../tools/benchmark/audio_extract_actor.py | 11 ++-- .../tools/benchmark/extract_actor.py | 6 +-- .../tools/benchmark/ocr_actor.py | 14 ++---- .../tools/benchmark/page_elements_actor.py | 10 ++-- .../tools/benchmark/split_actor.py | 6 +-- nemo_retriever/tests/test_caption.py | 16 ------ nemo_retriever/tests/test_dedup.py | 36 ------------- .../tests/test_executor_arrow_pandas.py | 32 +++++++----- .../tests/test_lancedb_collections.py | 10 ---- .../tests/test_lancedb_row_metadata.py | 22 -------- nemo_retriever/tests/test_multimodal_embed.py | 7 +++ .../tests/test_nemotron_rerank_vl_v2_hf.py | 29 ----------- nemo_retriever/tests/test_vdb_records.py | 15 ------ 30 files changed, 134 insertions(+), 340 deletions(-) delete mode 100644 nemo_retriever/src/nemo_retriever/common/modality/collections.py diff --git a/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py b/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py index 973671e231..3d5890aac6 100644 --- a/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py +++ b/nemo_retriever/src/nemo_retriever/common/api/internal/mutate/deduplicate.py @@ -9,8 +9,6 @@ import pandas as pd -from nemo_retriever.common.modality.collections import bbox_coordinates - from nemo_retriever.common.api.internal.enums.common import ContentTypeEnum from nemo_retriever.common.api.internal.schemas.mutate.mutate_image_dedup_schema import ImageDedupSchema @@ -114,8 +112,7 @@ def _get_image_bbox_info(row: pd.Series) -> Optional[Dict[str, Any]]: bbox = image_metadata.get("image_location") max_dims = image_metadata.get("image_location_max_dimensions", (0, 0)) - bbox = bbox_coordinates(bbox) - if bbox is None: + if bbox is None or not isinstance(bbox, (tuple, list)) or len(bbox) < 4: return None # Normalize bbox by max dimensions @@ -156,8 +153,7 @@ def _get_structured_bbox_info(row: pd.Series) -> Optional[Dict[str, Any]]: bbox = table_metadata.get("table_location") max_dims = table_metadata.get("table_location_max_dimensions", (0, 0)) - bbox = bbox_coordinates(bbox) - if bbox is None: + if bbox is None or not isinstance(bbox, (tuple, list)) or len(bbox) < 4: return None # Normalize bbox by max dimensions diff --git a/nemo_retriever/src/nemo_retriever/common/modality/collections.py b/nemo_retriever/src/nemo_retriever/common/modality/collections.py deleted file mode 100644 index 3c959ce574..0000000000 --- a/nemo_retriever/src/nemo_retriever/common/modality/collections.py +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Collection adapters for multimodal values crossing Ray's pandas boundary.""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - - -def multimodal_collection(value: Any) -> list[Any] | None: - """Return a mutable list view of a structured multimodal collection. - - Ray may preserve list-valued pandas object cells as one-dimensional NumPy - object arrays. Restrict ndarray support to one dimension so image tensors - and other genuine multidimensional model inputs are never reinterpreted as - collections of extracted objects. - """ - if isinstance(value, list): - return value - if isinstance(value, tuple): - return list(value) - if isinstance(value, np.ndarray) and value.ndim == 1: - return list(value) - return None - - -def bbox_coordinates(value: Any) -> list[Any] | None: - """Return four-or-more bbox coordinates without array truth testing.""" - if isinstance(value, np.ndarray): - if value.ndim != 1: - return None - value = value.tolist() - elif isinstance(value, tuple): - value = list(value) - if isinstance(value, list) and len(value) >= 4: - return value - return None 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 8c32fcfb15..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,9 +8,9 @@ from typing import Any, Dict, List, Optional, Sequence +import numpy as np import pandas as pd -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.io.image_store import inline_image_b64 from nemo_retriever.operators.extract.ocr.ocr import _crop_b64_image_by_norm_bbox from nemo_retriever.common.params.models import IMAGE_MODALITIES @@ -18,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,8 +30,8 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ if isinstance(base, str) and base.strip(): parts.append(base.strip()) for col in content_columns: - content_list = multimodal_collection(row.get(col)) - if content_list is not None: + content_list = row.get(col) + if _is_content_collection(content_list): for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -44,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, tuple)) or multimodal_collection(value) is not None: + if isinstance(value, (dict, list)) or _is_content_collection(value): out[key] = copy.deepcopy(value) else: out[key] = value @@ -108,8 +113,8 @@ def explode_content_to_rows( exploded_any = True for column in content_columns: - content_list = multimodal_collection(row_dict.get(column)) - if content_list is None: + content_list = row_dict.get(column) + if not _is_content_collection(content_list): continue for item in content_list: if not isinstance(item, dict): @@ -128,8 +133,8 @@ def explode_content_to_rows( if item_b64: content_row["_image_b64"] = item_b64 elif page_image_b64: - bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) - if bbox is not None: + bbox = item.get("bbox_xyxy_norm") + if bbox and len(bbox) == 4: cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox) content_row["_image_b64"] = cropped_b64 else: diff --git a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py index 242b24bb0b..81f2e0951b 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/ocr/shared.py @@ -25,7 +25,6 @@ import numpy as np import pandas as pd -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.params import RemoteRetryParams from nemo_retriever.models.nim.nim import NIMClient, invoke_image_inference_batches from nemo_retriever.common.modality.table_and_chart import join_table_structure_and_ocr_output @@ -176,8 +175,7 @@ def _clamp_int(v: float, lo: int, hi: int) -> int: continue bbox = det.get("bbox_xyxy_norm") - bbox = bbox_coordinates(bbox) - if bbox is None or len(bbox) != 4: + if not isinstance(bbox, (list, tuple)) or len(bbox) != 4: continue try: @@ -622,8 +620,7 @@ def _find_ts_detections_for_bbox( if not isinstance(region, dict): continue region_bbox = region.get("bbox_xyxy_norm") - region_bbox = bbox_coordinates(region_bbox) - if region_bbox is None or len(region_bbox) != 4: + if not isinstance(region_bbox, (list, tuple)) or len(region_bbox) != 4: continue if not _bboxes_close(table_bbox, region_bbox): continue @@ -701,7 +698,9 @@ def _prepare_ocr_rows( page_elements = getattr(row, "page_elements_v3", None) detections: List[Dict[str, Any]] = [] if isinstance(page_elements, dict): - detections = multimodal_collection(page_elements.get("detections")) or [] + detections = page_elements.get("detections") or [] + if not isinstance(detections, list): + detections = [] page_image = getattr(row, "page_image", None) or {} page_image_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None diff --git a/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py b/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py index e73b13927d..856e0844db 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/page_elements/shared.py @@ -461,7 +461,7 @@ def detect_page_elements_v3( Run Nemotron Page Elements v3 on a pandas batch. Input: - - `pages_df`: pandas.DataFrame (provided by NRL's safe Arrow adapter in Ray Data pipelines) + - `pages_df`: pandas.DataFrame (typical Ray Data `batch_format="pandas"`) Must contain an image base64 source either in `image_b64` or one of `images`/`tables`/`charts`/`infographics` (each as list[{"image_b64": ...}]). diff --git a/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py b/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py index 6f447211d1..e8d84c0db4 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/pipeline/content.py @@ -10,7 +10,6 @@ import pandas as pd -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.io.image_store import inline_image_b64 from nemo_retriever.operators.extract.ocr.ocr import _crop_b64_image_by_norm_bbox from nemo_retriever.common.params.models import IMAGE_MODALITIES @@ -24,8 +23,8 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ if isinstance(base, str) and base.strip(): parts.append(base.strip()) for col in content_columns: - content_list = multimodal_collection(row.get(col)) - if content_list is not None: + content_list = row.get(col) + if isinstance(content_list, list): for item in content_list: if isinstance(item, dict): text = item.get("text", "") @@ -42,7 +41,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, tuple)) or multimodal_collection(value) is not None: + if isinstance(value, (dict, list)): out[key] = copy.deepcopy(value) else: out[key] = value @@ -106,8 +105,8 @@ def explode_content_to_rows( exploded_any = True for column in content_columns: - content_list = multimodal_collection(row_dict.get(column)) - if content_list is None: + content_list = row_dict.get(column) + if not isinstance(content_list, list): continue for item in content_list: if not isinstance(item, dict): @@ -125,8 +124,8 @@ def explode_content_to_rows( if item_b64: content_row["_image_b64"] = item_b64 elif page_image_b64: - bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) - if bbox is not None: + bbox = item.get("bbox_xyxy_norm") + if bbox and len(bbox) == 4: cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_image_b64, bbox_xyxy_norm=bbox) content_row["_image_b64"] = cropped_b64 else: diff --git a/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py b/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py index 0b82bd15ac..0ba6d73197 100644 --- a/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py +++ b/nemo_retriever/src/nemo_retriever/common/modality/table_and_chart.py @@ -15,8 +15,6 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple # noqa: F401 import numpy as np - -from nemo_retriever.common.modality.collections import bbox_coordinates import pandas as pd logger = logging.getLogger(__name__) @@ -398,8 +396,8 @@ def _structure_dets_to_class_boxes( grouped: Dict[str, list] = {} for d in dets: name = d.get("label_name", "") - bbox = bbox_coordinates(d.get("bbox_xyxy_norm")) - if bbox is None or len(bbox) != 4: + bbox = d.get("bbox_xyxy_norm") + if not bbox or len(bbox) != 4: continue x1, y1, x2, y2 = float(bbox[0]) * W, float(bbox[1]) * H, float(bbox[2]) * W, float(bbox[3]) * H grouped.setdefault(name, []).append([x1, y1, x2, y2]) diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py index 712da211b3..ddfa0ec533 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py @@ -128,9 +128,7 @@ def _decode_cursor( def _json_string(value: Any) -> str: if isinstance(value, str): return value - if hasattr(value, "tolist"): - value = value.tolist() - return json.dumps({} if value is None else value, ensure_ascii=False, separators=(",", ":"), default=str) + return json.dumps(value or {}, ensure_ascii=False, separators=(",", ":"), default=str) def _content_text(record: dict[str, Any], metadata: dict[str, Any]) -> str: @@ -229,7 +227,7 @@ def _collection_rows( "metadata": _json_string(content_metadata), "stored_image_uri": stored_image_uri, "content_type": content_type, - "bbox_xyxy_norm": _json_string(bbox) if bbox is not None else "", + "bbox_xyxy_norm": _json_string(bbox) if bbox else "", "chunk_id": hashlib.sha256( f"{context.document_id}\0{context.document_version}\0{row_index}".encode() ).hexdigest(), diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py index 6a8deab636..d684add889 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_schema.py @@ -10,18 +10,9 @@ from pathlib import Path from typing import Any, Dict, List, Optional, Tuple -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.vdb.records import normalize_content_type -def _embedding_list(value: Any) -> Optional[List[float]]: - if hasattr(value, "tolist"): - value = value.tolist() - elif isinstance(value, tuple): - value = list(value) - return value if isinstance(value, list) and value else None - - def extract_embedding_from_row( row: Any, *, @@ -31,15 +22,15 @@ def extract_embedding_from_row( """Extract an embedding vector from a row.""" meta = getattr(row, "metadata", None) if isinstance(meta, dict): - emb = _embedding_list(meta.get("embedding")) - if emb is not None: - return emb + emb = meta.get("embedding") + if isinstance(emb, list) and emb: + return emb # type: ignore[return-value] payload = getattr(row, embedding_column, None) if isinstance(payload, dict): - emb = _embedding_list(payload.get(embedding_key)) - if emb is not None: - return emb + emb = payload.get(embedding_key) + if isinstance(emb, list) and emb: + return emb # type: ignore[return-value] return None @@ -123,8 +114,8 @@ def _build_detection_metadata(row: Any) -> Dict[str, Any]: } for ocr_column in ("table", "chart", "infographic"): - entries = multimodal_collection(getattr(row, ocr_column, None)) - if entries is not None: + entries = getattr(row, ocr_column, None) + if isinstance(entries, list): out[f"ocr_{ocr_column}_detections"] = int(len(entries)) return out @@ -195,8 +186,8 @@ def build_lancedb_row( content_type = getattr(row, "_content_type", None) row_out["content_type"] = str(content_type) if content_type else "" - bbox = bbox_coordinates(getattr(row, "_bbox_xyxy_norm", None)) - row_out["bbox_xyxy_norm"] = json.dumps(bbox) if bbox is not None else "" + bbox = getattr(row, "_bbox_xyxy_norm", None) + row_out["bbox_xyxy_norm"] = json.dumps(bbox) if bbox else "" return row_out @@ -268,8 +259,8 @@ def lancedb_schema(vector_dim: int = 2048, *, collection_managed: bool = False) def infer_vector_dim(rows: List[Dict[str, Any]]) -> int: """Return the embedding dimension from the first row that has a vector.""" for row in rows: - vector = _embedding_list(row.get("vector")) - if vector is not None: + vector = row.get("vector") + if isinstance(vector, list) and vector: return len(vector) return 0 diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/records.py b/nemo_retriever/src/nemo_retriever/common/vdb/records.py index ac87ebf5b1..3bc10caeb5 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/records.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/records.py @@ -13,7 +13,6 @@ from pydantic import ValidationError -from nemo_retriever.common.modality.collections import multimodal_collection from nemo_retriever.common.schemas.collections import QueryHit _CONTENT_TYPE_ALIASES: dict[str, str] = { @@ -196,8 +195,8 @@ def _add_detection_metadata( content_metadata.setdefault("page_elements_v3_counts_by_label", normalized_counts) for content_type in ("table", "chart", "infographic"): - detections = multimodal_collection(row.get(content_type)) - if detections is not None: + detections = row.get(content_type) + if isinstance(detections, list): content_metadata.setdefault(f"ocr_{content_type}_detections", len(detections)) diff --git a/nemo_retriever/src/nemo_retriever/graph/executor.py b/nemo_retriever/src/nemo_retriever/graph/executor.py index b4ef0eb1b6..a10ad5f10c 100644 --- a/nemo_retriever/src/nemo_retriever/graph/executor.py +++ b/nemo_retriever/src/nemo_retriever/graph/executor.py @@ -124,7 +124,7 @@ 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]: +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,), {}) @@ -453,7 +453,7 @@ def build_dataset(self, data: Any, **kwargs: Any) -> Any: # 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_operator_class = _make_arrow_pandas_operator_adapter(node.operator_class) map_batch_format = "pyarrow" constructor_kwargs = { "operator_class": node.operator_class, diff --git a/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx b/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx index 99c7743aac..2b9e0e9652 100644 --- a/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx +++ b/nemo_retriever/src/nemo_retriever/harness/portal/static/views/designer.jsx @@ -105,7 +105,6 @@ function _generateRayDataCode(nodes, edges) { const imports = new Set(); imports.add('import ray'); imports.add('import ray.data'); - imports.add('from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter'); chain.forEach(n => { if (n.operator && n.operator.type !== 'ray_data_source') { imports.add(n.operator.import_path); @@ -162,13 +161,10 @@ function _generateRayDataCode(nodes, edges) { lines.push(`# Stage: ${op.display_name || op.class_name}`); lines.push(`ds = ds.map_batches(`); - lines.push(` make_arrow_pandas_operator_adapter(${op.class_name}),`); - lines.push(` fn_constructor_kwargs={`); - lines.push(` "operator_class": ${op.class_name},`); - lines.push(` "operator_kwargs": ${kwargsStr},`); - lines.push(` },`); + lines.push(` ${op.class_name},`); + lines.push(` fn_constructor_kwargs=${kwargsStr},`); lines.push(` batch_size=1,`); - lines.push(` batch_format="pyarrow",`); + lines.push(` batch_format="pandas",`); if (isGpu) { lines.push(` num_cpus=0,`); lines.push(` num_gpus=1,`); @@ -195,13 +191,10 @@ function _generateRayDataCode(nodes, edges) { lines.push(`# Sink: ${op.display_name || op.class_name}`); lines.push(`ds = ds.map_batches(`); - lines.push(` make_arrow_pandas_operator_adapter(${op.class_name}),`); - lines.push(` fn_constructor_kwargs={`); - lines.push(` "operator_class": ${op.class_name},`); - lines.push(` "operator_kwargs": ${kwargsStr},`); - lines.push(` },`); + lines.push(` ${op.class_name},`); + lines.push(` fn_constructor_kwargs=${kwargsStr},`); lines.push(` batch_size=1,`); - lines.push(` batch_format="pyarrow",`); + lines.push(` batch_format="pandas",`); if (sinkConcurrency !== undefined && sinkConcurrency !== '') { lines.push(` concurrency=${parseInt(sinkConcurrency, 10) || 1},`); } diff --git a/nemo_retriever/src/nemo_retriever/operators/dedup.py b/nemo_retriever/src/nemo_retriever/operators/dedup.py index a4a83f4cb4..4ecb1b3580 100644 --- a/nemo_retriever/src/nemo_retriever/operators/dedup.py +++ b/nemo_retriever/src/nemo_retriever/operators/dedup.py @@ -10,7 +10,6 @@ import pandas as pd from nemo_retriever.common.api.internal.mutate.deduplicate import calculate_iou -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.common.params import DedupParams _STRUCTURED_COLUMNS = ("table", "chart", "infographic") @@ -20,14 +19,14 @@ def _collect_structured_bboxes(row: pd.Series) -> List[Tuple[float, ...]]: """Gather all bounding boxes from tables, charts, and infographics columns.""" bboxes: List[Tuple[float, ...]] = [] for col in _STRUCTURED_COLUMNS: - items = multimodal_collection(row.get(col)) - if items is None: + items = row.get(col) + if not isinstance(items, list): continue for item in items: if not isinstance(item, dict): continue - bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) - if bbox is not None: + bbox = item.get("bbox_xyxy_norm") + if bbox and len(bbox) >= 4: bboxes.append(tuple(bbox[:4])) return bboxes @@ -56,8 +55,8 @@ def dedup_images( return batch_df for row_idx, row in batch_df.iterrows(): - images = multimodal_collection(row.get("images")) - if not images: + images = row.get("images") + if not isinstance(images, list) or not images: continue filtered = list(images) @@ -89,8 +88,8 @@ def dedup_images( if not isinstance(item, dict): surviving.append(item) continue - img_bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) - if img_bbox is None: + img_bbox = item.get("bbox_xyxy_norm") + if not img_bbox or len(img_bbox) < 4: surviving.append(item) continue img_bbox_t = tuple(img_bbox[:4]) diff --git a/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py b/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py index e529aaebfe..49151edc51 100644 --- a/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py +++ b/nemo_retriever/src/nemo_retriever/operators/embed/text_embed.py @@ -80,7 +80,7 @@ def embed_text_1b_v2( Embed a batch of text rows using the local `LlamaNemotronEmbed1BV2Embedder`. Input: - - `batch_df`: pandas.DataFrame (provided by NRL's safe Arrow adapter in Ray Data pipelines) + - `batch_df`: pandas.DataFrame (Ray Data `batch_format="pandas"` compatible) - `text_column`: preferred column name to read text from (defaults to `"text"`) Output: diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py b/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py index 56d259a224..a7000af694 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/caption/caption.py @@ -17,7 +17,6 @@ merge_request_extras, resolve_caption_model_name, ) -from nemo_retriever.common.modality.collections import bbox_coordinates, multimodal_collection from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.operators.cpu_operator import CPUOperator from nemo_retriever.graph.designer import designer_component @@ -373,8 +372,8 @@ def caption_images( for row_idx, row in batch_df.iterrows(): # Unstructured images. if has_images: - images = multimodal_collection(row.get("images")) - if images is not None: + images = row.get("images") + if isinstance(images, list): for item_idx, item in enumerate(images): if not isinstance(item, dict): continue @@ -386,8 +385,8 @@ def caption_images( # Infographics — crop from page image. if has_infographics: - infographics = multimodal_collection(row.get("infographic")) - if infographics is not None: + infographics = row.get("infographic") + if isinstance(infographics, list): page_image = row.get("page_image") page_b64 = page_image.get("image_b64") if isinstance(page_image, dict) else None if page_b64: @@ -396,8 +395,8 @@ def caption_images( continue if item.get("caption"): continue # already captioned - bbox = bbox_coordinates(item.get("bbox_xyxy_norm")) - if bbox is None: + bbox = item.get("bbox_xyxy_norm") + if not bbox or len(bbox) < 4: continue cropped_b64, _ = _crop_b64_image_by_norm_bbox(page_b64, bbox_xyxy_norm=bbox) if cropped_b64 and _image_meets_min_size(cropped_b64): diff --git a/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py b/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py index d720ab64e4..a3f88e4e3d 100644 --- a/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py +++ b/nemo_retriever/src/nemo_retriever/operators/extract/page_elements/gpu_actor.py @@ -18,11 +18,8 @@ class PageElementDetectionActor(AbstractOperator, GPUOperator): """ Ray-friendly callable that initializes Nemotron Page Elements v3 once. - Ray Data callers should use - :func:`nemo_retriever.graph.executor.make_arrow_pandas_operator_adapter` - with ``batch_format="pyarrow"`` so nested Arrow batches are normalized - before this pandas operator runs. ``RayDataExecutor`` applies that adapter - automatically. + Use with Ray Data: + ds = ds.map_batches(PageElementDetectionActor, fn_constructor_kwargs={...}, batch_format="pandas") """ def __init__(self, **detect_kwargs: Any) -> None: diff --git a/nemo_retriever/src/nemo_retriever/operators/rerank.py b/nemo_retriever/src/nemo_retriever/operators/rerank.py index 25ff1b86b7..48cce39cdd 100644 --- a/nemo_retriever/src/nemo_retriever/operators/rerank.py +++ b/nemo_retriever/src/nemo_retriever/operators/rerank.py @@ -32,23 +32,19 @@ Ray Data actor usage:: import ray - from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter ds = ds.map_batches( - make_arrow_pandas_operator_adapter(NemotronRerankActor), + NemotronRerankActor, batch_size=64, - batch_format="pyarrow", + batch_format="pandas", num_gpus=1, compute=ray.data.ActorPoolStrategy(size=4), fn_constructor_kwargs={ - "operator_class": NemotronRerankActor, - "operator_kwargs": { - "model_name": "nvidia/llama-nemotron-rerank-1b-v2", - "query_column": "query", - "text_column": "text", - "score_column": "rerank_score", - "max_length": 512, - "batch_size": 32, - }, + "model_name": "nvidia/llama-nemotron-rerank-1b-v2", + "query_column": "query", + "text_column": "text", + "score_column": "rerank_score", + "max_length": 512, + "batch_size": 32, }, ) """ @@ -62,7 +58,6 @@ from typing import Any, Dict, List, Optional import pandas as pd -from nemo_retriever.common.modality.collections import bbox_coordinates from nemo_retriever.operators.abstract_operator import AbstractOperator from nemo_retriever.graph.designer import designer_component from nemo_retriever.operators.cpu_operator import CPUOperator @@ -329,12 +324,11 @@ def rerank_hits( render_cache[cache_key] = render_page_image_b64(h["path"], int(h["page_number"])) page_b64 = render_cache[cache_key] # Crop to element bbox if available (tables/charts). - bbox_value = h.get("bbox_xyxy_norm") - if page_b64 and bbox_value is not None: + bbox_str = h.get("bbox_xyxy_norm", "") + if page_b64 and bbox_str: try: - bbox = json.loads(bbox_value) if isinstance(bbox_value, str) else bbox_value - bbox = bbox_coordinates(bbox) - if bbox is not None and len(bbox) == 4: + bbox = json.loads(bbox_str) + if isinstance(bbox, list) and len(bbox) == 4: cropped, _ = _crop_b64_image_by_norm_bbox(page_b64, bbox_xyxy_norm=bbox) images_b64.append(cropped) continue @@ -420,23 +414,19 @@ class NemotronRerankGPUActor(AbstractOperator, GPUOperator): Usage with Ray Data:: import ray - from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter ds = ds.map_batches( - make_arrow_pandas_operator_adapter(NemotronRerankActor), + NemotronRerankActor, batch_size=64, - batch_format="pyarrow", + batch_format="pandas", num_gpus=1, compute=ray.data.ActorPoolStrategy(size=4), fn_constructor_kwargs={ - "operator_class": NemotronRerankActor, - "operator_kwargs": { - "model_name": "nvidia/llama-nemotron-rerank-1b-v2", - "query_column": "query", - "text_column": "text", - "score_column": "rerank_score", - "max_length": 512, - "batch_size": 32, - }, + "model_name": "nvidia/llama-nemotron-rerank-1b-v2", + "query_column": "query", + "text_column": "text", + "score_column": "rerank_score", + "max_length": 512, + "batch_size": 32, }, ) diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py index 776cba2cf2..4bdb16a92c 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/audio_extract_actor.py @@ -23,7 +23,6 @@ from nemo_retriever.common.modality.audio.media_interface import is_media_available from nemo_retriever.common.modality.audio.media_interface import media_dependency_error_message from nemo_retriever.common.params import AudioChunkParams -from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -92,22 +91,20 @@ def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: asr_actor = ASRActor(params=asr_params_from_env()) ds = ds.map_batches( - call_pandas_function_on_arrow, + chunk_actor, batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), - fn_kwargs={"fn": chunk_actor}, ) ds = ds.map_batches( - call_pandas_function_on_arrow, + asr_actor, batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=1, num_gpus=0.25 if not mock_asr else 0, compute=rd.TaskPoolStrategy(size=int(worker_count)), - fn_kwargs={"fn": asr_actor}, ) return ds diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py index d22bb1880c..cfbb7bb6e9 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/extract_actor.py @@ -11,7 +11,6 @@ import typer from nemo_retriever.operators.extract.pdf.extract import PDFExtractionActor -from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -53,13 +52,12 @@ def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: dpi=int(dpi), ) return ds.map_batches( - call_pandas_function_on_arrow, + actor, batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), - fn_kwargs={"fn": actor}, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py index 5f047e8849..8234ba850e 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/ocr_actor.py @@ -11,7 +11,6 @@ import typer from nemo_retriever.operators.extract.ocr.gpu_ocr import OCRActor -from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -55,19 +54,16 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - make_arrow_pandas_operator_adapter(OCRActor), + OCRActor, batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=float(num_cpus), num_gpus=float(num_gpus), compute=rd.ActorPoolStrategy(size=int(worker_count)), fn_constructor_kwargs={ - "operator_class": OCRActor, - "operator_kwargs": { - "extract_tables": bool(extract_tables), - "extract_charts": bool(extract_charts), - "extract_infographics": bool(extract_infographics), - }, + "extract_tables": bool(extract_tables), + "extract_charts": bool(extract_charts), + "extract_infographics": bool(extract_infographics), }, ) diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py index 706db3fb67..9cdad6b0e9 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/page_elements_actor.py @@ -11,7 +11,6 @@ import typer from nemo_retriever.operators.extract.page_elements.gpu_actor import PageElementDetectionActor -from nemo_retriever.graph.executor import make_arrow_pandas_operator_adapter from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -46,16 +45,13 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - make_arrow_pandas_operator_adapter(PageElementDetectionActor), + PageElementDetectionActor, batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=float(num_cpus), num_gpus=float(num_gpus), compute=rd.ActorPoolStrategy(size=int(worker_count)), - fn_constructor_kwargs={ - "operator_class": PageElementDetectionActor, - "operator_kwargs": {"inference_batch_size": int(inference_batch_size)}, - }, + fn_constructor_kwargs={"inference_batch_size": int(inference_batch_size)}, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py b/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py index c44cca5995..ac461be1e8 100644 --- a/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py +++ b/nemo_retriever/src/nemo_retriever/tools/benchmark/split_actor.py @@ -11,7 +11,6 @@ import typer from nemo_retriever.operators.extract.pdf.split import PDFSplitActor -from nemo_retriever.graph.executor import call_pandas_function_on_arrow from nemo_retriever.tools.benchmark.common import ( benchmark_sweep, @@ -42,13 +41,12 @@ def run( def _map(ds: rd.Dataset, worker_count: int, batch_size: int) -> rd.Dataset: return ds.map_batches( - call_pandas_function_on_arrow, + PDFSplitActor(), batch_size=int(batch_size), - batch_format="pyarrow", + batch_format="pandas", num_cpus=1, num_gpus=0, compute=rd.TaskPoolStrategy(size=int(worker_count)), - fn_kwargs={"fn": PDFSplitActor()}, ) best, results = benchmark_sweep( diff --git a/nemo_retriever/tests/test_caption.py b/nemo_retriever/tests/test_caption.py index 1f75a7cc87..acc27e03f6 100644 --- a/nemo_retriever/tests/test_caption.py +++ b/nemo_retriever/tests/test_caption.py @@ -8,7 +8,6 @@ import io from unittest.mock import MagicMock, patch -import numpy as np import pandas as pd import pytest @@ -42,21 +41,6 @@ def test_caption_images_writes_back(): assert result.iloc[0]["images"][1]["text"] == "cap2" -def test_caption_images_writes_back_to_arrow_object_array(): - from nemo_retriever.operators.extract.caption.caption import caption_images - - source = _make_page_df(num_images=1).iloc[0].to_dict() - source["images"] = np.array(source["images"], dtype=object) - df = pd.DataFrame([source]) - mock_model = MagicMock() - mock_model.caption_batch.return_value = ["captioned"] - - result = caption_images(df, model=mock_model) - - mock_model.caption_batch.assert_called_once() - assert result.iloc[0]["images"][0]["text"] == "captioned" - - def test_caption_images_skips_already_captioned(): from nemo_retriever.operators.extract.caption.caption import caption_images diff --git a/nemo_retriever/tests/test_dedup.py b/nemo_retriever/tests/test_dedup.py index 55e08f8038..0f29ffb684 100644 --- a/nemo_retriever/tests/test_dedup.py +++ b/nemo_retriever/tests/test_dedup.py @@ -4,7 +4,6 @@ """Unit tests for nemo_retriever.operators.dedup.""" -import numpy as np import pandas as pd from nemo_retriever.common.api.internal.mutate.deduplicate import calculate_iou @@ -66,22 +65,6 @@ def test_keeps_unique_images(self): result = dedup_images(df, content_hash=True, bbox_iou=False) assert len(result.iloc[0]["images"]) == 2 - def test_removes_duplicate_images_from_arrow_object_array(self): - df = _make_df( - [ - { - "images": np.array( - [_make_image("AAA"), _make_image("BBB"), _make_image("AAA")], - dtype=object, - ), - } - ] - ) - - result = dedup_images(df, content_hash=True, bbox_iou=False) - - assert [image["image_b64"] for image in result.iloc[0]["images"]] == ["AAA", "BBB"] - # ── bbox IoU dedup ───────────────────────────────────────────────────────── @@ -143,25 +126,6 @@ def test_image_without_bbox_is_kept(self): result = dedup_images(df, content_hash=False, bbox_iou=True) assert len(result.iloc[0]["images"]) == 1 - def test_drops_arrow_object_array_image_overlapping_with_structured_content(self): - df = _make_df( - [ - { - "images": np.array([_make_image("X", bbox=[0.0, 0.0, 1.0, 1.0])], dtype=object), - "table": np.array( - [{"bbox_xyxy_norm": [0.0, 0.0, 1.0, 1.0], "text": "t"}], - dtype=object, - ), - "chart": np.array([], dtype=object), - "infographic": np.array([], dtype=object), - } - ] - ) - - result = dedup_images(df, content_hash=False, bbox_iou=True, iou_threshold=0.45) - - assert result.iloc[0]["images"] == [] - # ── edge cases ───────────────────────────────────────────────────────────── diff --git a/nemo_retriever/tests/test_executor_arrow_pandas.py b/nemo_retriever/tests/test_executor_arrow_pandas.py index a82a623673..4f4131b413 100644 --- a/nemo_retriever/tests/test_executor_arrow_pandas.py +++ b/nemo_retriever/tests/test_executor_arrow_pandas.py @@ -77,21 +77,29 @@ def test_adapter_normalizes_pickled_object_columns_for_pandas_row_operations() - 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_page_and_element_content_semantics_for_sliced_pdf_rows() -> None: - table = pa.Table.from_pylist( - [ +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": f"page {page_number}", - "table": [{"text": f"table {page_number}"}], - "chart": [{"text": f"chart {page_number}"}], - "metadata": {"source_path": "document.pdf", "error": None}, + "text": ["page text"], + "table": [np.array([{"text": "table text"}], dtype=object)], + "chart": [np.array([{"text": "chart text"}], dtype=object)], } - for page_number in range(3) - ] - ).slice(2, 1) + ) + ) page_result = _ArrowPandasOperatorAdapter( UDFOperator, @@ -108,8 +116,8 @@ def test_adapter_preserves_page_and_element_content_semantics_for_sliced_pdf_row }, )(table) - assert page_result["text"].tolist() == ["page 2\n\ntable 2\n\nchart 2"] - assert element_result["text"].tolist() == ["page 2", "table 2", "chart 2"] + 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_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py index e3b59d4d1f..3606cbaf32 100644 --- a/nemo_retriever/tests/test_lancedb_collections.py +++ b/nemo_retriever/tests/test_lancedb_collections.py @@ -15,7 +15,6 @@ from dataclasses import replace import lancedb -import numpy as np import pytest import nemo_retriever.common.vdb.lancedb_collections as collections_module @@ -164,15 +163,6 @@ def test_collection_row_conversion_preserves_identity_and_provenance(): assert row["chunk_id"] == hashlib.sha256(b"document-a\x00v1\x000").hexdigest() -def test_collection_row_conversion_normalizes_arrow_backed_bbox_array(): - records = _records() - records[0][0]["metadata"]["content_metadata"]["bbox_xyxy_norm"] = np.array([0.1, 0.2, 0.8, 0.9]) - - rows = _collection_rows(records, context=_context()) - - assert json.loads(rows[0]["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] - - @pytest.mark.parametrize( ("raw_type", "expected_type"), [("table_caption", "table"), ("chart_caption", "chart")], diff --git a/nemo_retriever/tests/test_lancedb_row_metadata.py b/nemo_retriever/tests/test_lancedb_row_metadata.py index 0b29303ef4..d2123f1d2e 100644 --- a/nemo_retriever/tests/test_lancedb_row_metadata.py +++ b/nemo_retriever/tests/test_lancedb_row_metadata.py @@ -2,8 +2,6 @@ import json from types import SimpleNamespace -import numpy as np - from nemo_retriever.common.vdb.lancedb_bulk import _build_lancedb_rows_from_df from nemo_retriever.common.vdb.lancedb_schema import build_lancedb_row @@ -24,26 +22,6 @@ def test_build_lancedb_row_persists_normalized_content_type() -> None: assert metadata["_content_type"] == "table" -def test_build_lancedb_row_normalizes_arrow_backed_arrays() -> None: - row = SimpleNamespace( - path="/tmp/doc_arrays.pdf", - page_number=2, - metadata={"embedding": np.array([0.1, 0.2], dtype=np.float32)}, - text="table text", - table=np.array([{"text": "first"}, {"text": "second"}], dtype=object), - chart=np.array([], dtype=object), - infographic=np.array([], dtype=object), - _bbox_xyxy_norm=np.array([0.1, 0.2, 0.8, 0.9]), - ) - - row_out = build_lancedb_row(row) - - assert row_out is not None - assert row_out["vector"] == [np.float32(0.1), np.float32(0.2)] - assert json.loads(row_out["bbox_xyxy_norm"]) == [0.1, 0.2, 0.8, 0.9] - assert json.loads(row_out["metadata"])["ocr_table_detections"] == 2 - - def test_build_lancedb_rows_from_df_persists_normalized_content_type() -> None: rows = [ { diff --git a/nemo_retriever/tests/test_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index 170f76a6f6..cec96cd989 100644 --- a/nemo_retriever/tests/test_multimodal_embed.py +++ b/nemo_retriever/tests/test_multimodal_embed.py @@ -205,6 +205,13 @@ def test_arrow_backed_structured_arrays_expand_into_element_rows(self): 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): diff --git a/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py b/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py index 302a6a44c7..38f379e470 100644 --- a/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py +++ b/nemo_retriever/tests/test_nemotron_rerank_vl_v2_hf.py @@ -14,7 +14,6 @@ from unittest.mock import MagicMock, patch -import numpy as np import pandas as pd import pytest @@ -266,34 +265,6 @@ def test_rerank_hits_graceful_without_uri(self): assert call_kwargs.kwargs["images_b64"] == [None, None] assert len(out) == 2 - def test_rerank_hits_crops_arrow_backed_bbox_array(self): - from nemo_retriever.operators.rerank import rerank_hits - - model = MagicMock() - model.model_name = "nvidia/llama-nemotron-rerank-vl-1b-v2" - model.score.return_value = [1.0] - hits = [ - { - "text": "table", - "path": "/tmp/document.pdf", - "page_number": 1, - "bbox_xyxy_norm": np.array([0.1, 0.2, 0.8, 0.9]), - } - ] - - with ( - patch("nemo_retriever.common.io.image_store.render_page_image_b64", return_value="page") as render, - patch( - "nemo_retriever.operators.extract.ocr.ocr._crop_b64_image_by_norm_bbox", - return_value=("crop", None), - ) as crop, - ): - rerank_hits("q", hits, model=model, modality="text_image") - - render.assert_called_once_with("/tmp/document.pdf", 1) - crop.assert_called_once_with("page", bbox_xyxy_norm=[0.1, 0.2, 0.8, 0.9]) - assert model.score.call_args.kwargs["images_b64"] == ["crop"] - # --------------------------------------------------------------------------- # _rerank_via_endpoint with images diff --git a/nemo_retriever/tests/test_vdb_records.py b/nemo_retriever/tests/test_vdb_records.py index 5b0220f3fc..bd68725c24 100644 --- a/nemo_retriever/tests/test_vdb_records.py +++ b/nemo_retriever/tests/test_vdb_records.py @@ -192,21 +192,6 @@ def test_graph_record_conversion_normalizes_arrow_backed_bbox_array() -> None: assert records[0][0]["metadata"]["content_metadata"]["bbox_xyxy_norm"] == [0.1, 0.2, 0.8, 0.9] -def test_graph_record_conversion_counts_arrow_backed_detection_array() -> None: - records = to_client_vdb_records( - [ - { - "text": "table content", - "text_embeddings_1b_v2": {"embedding": [0.1, 0.2]}, - "_content_type": "table", - "table": np.array([{"content": "a"}, {"content": "b"}], dtype=object), - } - ] - ) - - assert records[0][0]["metadata"]["content_metadata"]["ocr_table_detections"] == 2 - - def test_narrow_lancedb_hit_promotes_canonical_multimodal_metadata() -> None: hit = _normalize_one( {