fix: repair Ray Arrow-to-pandas batch ingestion

nemo_retriever/src/nemo_retriever/common/modality/content_transforms.py

@@ -8,6 +8,7 @@
 
 from typing import Any, Dict, List, Optional, Sequence
 
+import numpy as np
 import pandas as pd
 
 from nemo_retriever.common.io.image_store import inline_image_b64
@@ -17,6 +18,11 @@
 _CONTENT_COLUMNS = ("table", "chart", "infographic")
 
 
+def _is_content_collection(value: Any) -> bool:
+    """Return whether a value is a supported extracted-content collection."""
+    return isinstance(value, list) or (isinstance(value, np.ndarray) and value.ndim == 1 and value.dtype == object)
+
+
 def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequence[str]) -> str:
     """Combine page text with OCR content text for embedding."""
     parts = []
@@ -25,7 +31,7 @@ def _combine_text_with_content(row: Any, text_column: str, content_columns: Sequ
         parts.append(base.strip())
     for col in content_columns:
         content_list = row.get(col)
-        if isinstance(content_list, list):
+        if _is_content_collection(content_list):
             for item in content_list:
                 if isinstance(item, dict):
                     text = item.get("text", "")
@@ -43,7 +49,7 @@ def _deep_copy_row(row_dict: Dict[str, Any]) -> Dict[str, Any]:
 
     out: Dict[str, Any] = {}
     for key, value in row_dict.items():
-        if isinstance(value, (dict, list)):
+        if isinstance(value, (dict, list)) or _is_content_collection(value):
             out[key] = copy.deepcopy(value)
         else:
             out[key] = value
@@ -108,7 +114,7 @@ def explode_content_to_rows(
 
         for column in content_columns:
             content_list = row_dict.get(column)
-            if not isinstance(content_list, list):
+            if not _is_content_collection(content_list):
                 continue
             for item in content_list:
                 if not isinstance(item, dict):

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 (

nemo_retriever/src/nemo_retriever/graph/executor.py

@@ -37,6 +37,100 @@
 _DEFAULT_GPU_OPERATOR_NUM_GPUS = OCR_GPUS_PER_ACTOR
 
 
+def _contains_null_arrow_child(data_type: Any) -> bool:
+    """Return whether a nested Arrow type contains an inferred null child."""
+    import pyarrow as pa
+
+    if pa.types.is_null(data_type):
+        return True
+    if pa.types.is_struct(data_type):
+        return any(_contains_null_arrow_child(field.type) for field in data_type)
+    if pa.types.is_list(data_type) or pa.types.is_large_list(data_type) or pa.types.is_fixed_size_list(data_type):
+        return _contains_null_arrow_child(data_type.value_type)
+    if pa.types.is_map(data_type):
+        return _contains_null_arrow_child(data_type.key_type) or _contains_null_arrow_child(data_type.item_type)
+    return False
+
+
+def _compact_vulnerable_arrow_columns(table: Any) -> Any:
+    """Reset offsets before Ray converts nested null children to pandas."""
+    import pyarrow as pa
+    import pyarrow.compute as pc
+
+    if not isinstance(table, pa.Table) or table.num_rows == 0:
+        return table
+
+    indices = None
+    compacted = table
+    for index, field in enumerate(table.schema):
+        column = table.column(index)
+        if not _contains_null_arrow_child(field.type) or not any(chunk.offset for chunk in column.chunks):
+            continue
+        if indices is None:
+            indices = pa.array(range(table.num_rows), type=pa.int64())
+        compacted = compacted.set_column(index, field, pc.take(column, indices))
+    return compacted
+
+
+def _normalize_pickled_object_columns(table: Any, frame: pd.DataFrame) -> pd.DataFrame:
+    """Convert Ray's pickled-object extension columns to plain pandas objects."""
+    import pyarrow as pa
+
+    if not isinstance(table, pa.Table):
+        return frame
+
+    for index, field in enumerate(table.schema):
+        if getattr(field.type, "extension_name", None) != "ray.data.arrow_pickled_object":
+            continue
+        frame[field.name] = pd.Series(table.column(index).to_pylist(), index=frame.index, dtype=object)
+    return frame
+
+
+def arrow_table_to_pandas(table: Any) -> pd.DataFrame:
+    """Convert a Ray Arrow batch to a row-safe pandas DataFrame.
+
+    Ray 2.56+ preserves Arrow-backed pandas dtypes. Before conversion, sliced
+    nested columns with inferred null children must be compacted. Ray's
+    pickled-object extension columns also need to be materialized as ordinary
+    object columns so pandas row operations do not interpret their payloads as
+    malformed extension arrays.
+    """
+    if isinstance(table, pd.DataFrame):
+        return table
+
+    from ray.data.block import BlockAccessor
+
+    table = _compact_vulnerable_arrow_columns(table)
+    frame = BlockAccessor.for_block(table).to_pandas()
+    return _normalize_pickled_object_columns(table, frame)
+
+
+def call_pandas_function_on_arrow(
+    table: Any,
+    *,
+    fn: Any,
+    fn_kwargs: dict[str, Any] | None = None,
+) -> Any:
+    """Invoke a pandas batch function through the safe Arrow boundary."""
+    return fn(arrow_table_to_pandas(table), **(fn_kwargs or {}))
+
+
+class _ArrowPandasOperatorAdapter:
+    """Convert valid Arrow batches to pandas before invoking an NRL operator."""
+
+    def __init__(self, operator_class: type, operator_kwargs: dict[str, Any]) -> None:
+        self._operator = operator_class(**operator_kwargs)
+
+    def __call__(self, table: Any) -> Any:
+        return self._operator(arrow_table_to_pandas(table))
+
+
+def _make_arrow_pandas_operator_adapter(operator_class: type) -> type[_ArrowPandasOperatorAdapter]:
+    """Keep the wrapped operator recognizable in Ray plans and worker logs."""
+    adapter_name = f"{operator_class.__name__}ArrowPandasAdapter"
+    return type(adapter_name, (_ArrowPandasOperatorAdapter,), {})
+
+
 def _concurrency_target(concurrency: Any) -> int:
     """Return the largest actor-pool size that resource planning can permit."""
     if isinstance(concurrency, tuple):
@@ -478,17 +572,27 @@ def build_dataset(self, data: Any, **kwargs: Any) -> Any:
             elif target_num_rows_per_block is not None and int(target_num_rows_per_block) > 0:
                 ds = ds.repartition(target_num_rows_per_block=int(target_num_rows_per_block))
 
-            # Pass the operator class directly to map_batches with
-            # fn_constructor_kwargs for deferred construction on workers.
-            # AbstractOperator.__call__ delegates to run(), so each stage
-            # executes the full preprocess -> process -> postprocess chain.
+            map_operator_class = node.operator_class
+            map_batch_format = batch_format
+            constructor_kwargs = node.operator_kwargs
+            if batch_format == "pandas":
+                # Ray's Arrow-backed pandas conversion can preserve unsafe
+                # offsets for sliced structs with inferred null children.
+                # Compact the valid Arrow batch before that conversion.
+                map_operator_class = _make_arrow_pandas_operator_adapter(node.operator_class)
+                map_batch_format = "pyarrow"
+                constructor_kwargs = {
+                    "operator_class": node.operator_class,
+                    "operator_kwargs": node.operator_kwargs,
+                }
+
             ds = ds.map_batches(
-                node.operator_class,
+                map_operator_class,
                 batch_size=batch_size,
-                batch_format=batch_format,
+                batch_format=map_batch_format,
                 num_cpus=num_cpus,
                 num_gpus=num_gpus,
-                fn_constructor_kwargs=node.operator_kwargs,
+                fn_constructor_kwargs=constructor_kwargs,
                 **overrides,
             )
 

nemo_retriever/src/nemo_retriever/ingestor/branch_extraction.py

@@ -12,7 +12,7 @@
 from typing import Any, Callable
 
 from nemo_retriever.graph import InprocessExecutor, RayDataExecutor
-from nemo_retriever.graph.executor import preflight_executors
+from nemo_retriever.graph.executor import call_pandas_function_on_arrow, preflight_executors
 from nemo_retriever.graph.ingestor_runtime import (
     batch_tuning_to_node_overrides,
     build_graph,
@@ -372,9 +372,12 @@ def normalize_ray_branch_datasets(branch_datasets: list[Any]) -> list[Any]:
     stable_columns = tuple(columns)
     return [
         dataset.map_batches(
-            ensure_pandas_columns,
-            batch_format="pandas",
-            fn_kwargs={"columns": stable_columns},
+            call_pandas_function_on_arrow,
+            batch_format="pyarrow",
+            fn_kwargs={
+                "fn": ensure_pandas_columns,
+                "fn_kwargs": {"columns": stable_columns},
+            },
         )
         for dataset in branch_datasets
     ]

nemo_retriever/src/nemo_retriever/ingestor/graph_ingestor.py

@@ -33,6 +33,7 @@
 from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Self, Sequence, Tuple, Union
 
 from nemo_retriever.graph import InprocessExecutor, RayDataExecutor
+from nemo_retriever.graph.executor import arrow_table_to_pandas, call_pandas_function_on_arrow
 from nemo_retriever.ingestor.branch_extraction import ExtractionBranchExecutor, merge_node_overrides
 from nemo_retriever.graph.ingestor_runtime import (
     batch_tuning_to_node_overrides,
@@ -1184,7 +1185,7 @@ def _stage_error_records(cls, batch: Any, *, columns: Iterable[str] | None = Non
         requested_columns = list(columns) if columns is not None else None
 
         if callable(iter_batches):
-            batches = iter_batches(batch_format="pandas")
+            batches = (arrow_table_to_pandas(batch_df) for batch_df in iter_batches(batch_format="pyarrow"))
         else:
             batches = (batch,)
 
@@ -1366,7 +1367,11 @@ def get_error_rows(self, dataset: Any = None) -> Any:
             raise RuntimeError("No Ray Dataset available to inspect for errors.")
         if isinstance(target, pd.DataFrame):
             return self.extract_error_rows(target)
-        return target.map_batches(self.extract_error_rows, batch_format="pandas")
+        return target.map_batches(
+            call_pandas_function_on_arrow,
+            batch_format="pyarrow",
+            fn_kwargs={"fn": self.extract_error_rows},
+        )
 
     def get_dataset(self) -> Any:
         return self._rd_dataset

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)},
+    )

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