[jioffe502]

Description

Problem. Every embedded row carries its vector twice: once as metadata["embedding"] and once inside the text_embeddings_1b_v2 struct. Both copies travel through the Ray object store and into the driver, which doubles vector memory per row.

Change. The embed stage writes one Arrow list<float32> column, embedding (EmbedParams.output_column), plus a string column embedding_error. Every writer goes through assign_embedding_columns (models/inference/main_text_embed.py). Readers (common/vdb/records.py, operators/vdb.py, ingestor/results.py, and the LanceDB write and query paths) read only the configured column. Replaces the stacked #2781 (closed). It doesn't depend on #2778.

Scope (breaking, no compatibility shim).

• Removed graph columns: text_embeddings_1b_v2*, metadata["embedding"] on graph rows, _contains_embeddings.
• Removed public config fields: embedding_key, embedding_dim_column, has_embedding_column, embed_output_column.
• Default change: output_column now defaults to embedding instead of text_embeddings_1b_v2.

Unchanged.

• Vector normalization still happens in the embedder (local _l2_normalize in float32, or vLLM pooling).
• The LanceDB vector: list<float32>[dim] schema.
• The client-record metadata["embedding"] handed to VDB backends.
• The NIM request.

Measured (NRB bo767, 1x H100 NVL, stock config, n=1 per arm; details in the PR comment). main 8c95ad7c → this PR:

• Throughput: 24.00 → 26.75 pages/s.
• Recall: R@5 0.8819 / R@10 0.9233 in both arms.
• nDCG@10: 0.7752 → 0.7746.
• Rows stored: 83,024 in both arms.

The same change measured as #2781 (stacked on #2778, bo767 n=3): driver PSS peak 45.1 → 38.7 GiB, and 24.98 → 28.20 pages/s.

Tests. pytest tests --ignore=tests/integration: 3790 passed, 185 skipped.

Port notes vs #2781 (conflict resolutions against current main)

• runtime.py, operators/vdb.py, graph/ingestor_runtime.py, graph/retriever.py: current main logic (query splitting/grouping from Fix long-query retrieval failures and incorrect result grouping #2788, _embedding_params, _query_text_column) is kept as is. embedding_column is passed through.
• Test fakes added to main after the stack (test_split_query_retrieval.py, test_nemotron35_embedder.py, test_graph_embedding_identity.py) now write the typed embedding column.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

Link to Devin session: https://nvidia-cloud.devinenterprise.com/sessions/07b4aa3f6ffd42ddba417a3c9a38a168
Open in Devin Desktop: https://nvidia-cloud.devinenterprise.com/desktop/session/07b4aa3f6ffd42ddba417a3c9a38a168?variant=devin
Requested by: @jioffe502

---

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access. I'll skip comments containing "(aside)".

• Disable automatic comment, CI, and merge conflict monitoring

[devin-ai-integration]

Note

Newer findings are available below. Devin Review posted a newer report on this PR, in addition to the findings presented here.

Devin Review found 2 potential issues.

[devin-ai-integration]

🟡 Compact results hide embedding failures

When assign_embedding_columns records a failed embedding, compact service results omit its error. compact_result_record reads only the old error fields, so clients see no failure.

Learn more

The embedding stage now represents failures in a separate string column. Service transport projects a DataFrame into compact records with compact_result_record. That projection discards most columns and looks only at top-level error or metadata['error'], so the newly emitted error cannot reach compact clients. This affects the compact result schema even when the legacy schema carries the error column.

Example: A row has embedding=None and embedding_error='RuntimeError: embedding batch failed'. Compact transport returns text and provenance without any error, leaving the client unable to distinguish a failed request from absent embeddable content.

Recommended fix: Make compact_result_record accept the configured embedding column, read its matching error column, and expose a compact error entry when populated. Test the default and custom output column with result_schema='compact'.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

This is not a regression from this PR. compact_result_record's error handling is identical to main: it reads only the top-level error / metadata["error"]. On main, the embedding error lived inside the text_embeddings_1b_v2 payload, which compact mode also never read. So compact clients didn't see embedding failures before this change either. Surfacing {output_column}_error in compact results would be a behavior change to the compact schema, so I'm leaving it for a separate PR rather than widening this one.

[devin-ai-integration]

This was already the case on main, so this PR doesn't introduce it. On main the embed stage puts its error inside the text_embeddings_1b_v2 payload (_error_payload → {"embedding": None, "error": {...}}). compact_result_record on main only reads row["error"] / metadata["error"] (ingestor/results.py:428, identical to this branch), so embedding failures were never shown in compact results there either. This PR only moves the error into the embedding_error column. The VDB path still fails closed on it (records.py:395). Showing embedding_error in compact results would be a behaviour change for service clients, so I'd do it in a follow-up rather than widen this PR.

[devin-ai-integration]

Devin Review found 1 new potential issue.

[devin-ai-integration]

🟡 Arrow scalar embeddings block VDB uploads

When a row contains a PyArrow list scalar, embedding_vector returns None instead of its vector. to_client_vdb_records then rejects the valid row as missing an embedding.

Learn more

A graph row can carry an embedding as a PyArrow list scalar or a pandas extension array. The conversion helper previously converted these values into lists, but now recognizes only NumPy arrays and tuples. to_client_vdb_records interprets a rejected embedding as absent and raises a VDB upload error for searchable content. The same helper feeds build_lancedb_row, which silently skips rejected embeddings.

Example: For {'text': 'example', 'embedding': pa.scalar([0.5, 0.25], type=pa.list_(pa.float32()))}, the vector is rejected instead of stored as [0.5, 0.25].

Recommended fix: Restore conversion of supported PyArrow scalars with as_py() and pandas extension arrays with to_numpy() or tolist() before the nonempty-list check. Retain NumPy and tuple support.

---

Was this helpful? React with 👍 or 👎 to provide feedback.

[devin-ai-integration]

I couldn't find a path that reaches this. The embed stage stores the vector as a pandas ArrowExtensionArray (list<float32>). I checked pandas' row access on that column: .iloc, to_dict(orient="records") and itertuples return a Python list, and iterrows and apply(axis=1) return a numpy.ndarray. embedding_vector already handles all of those. A pyarrow.ListScalar only shows up if you index a raw pa.Array/pa.Table directly, and the record builders don't do that. I'm leaving the code unchanged. If we want to be defensive anyway, it's a one-line as_py() branch, and that's the maintainer's call.

[devin-ai-integration]

Main vs this PR, bo767, 1x H100 NVL (NRB, stock config 26.08-core-single-gpu-library-main-nemotron3-embed, n=1 per arm, ./nrb results diff: comparable)

	main 8c95ad7c (Oct 20 run, same box and config hash)	this PR 22be3dfc	Δ
pages/s	24.00	26.75	+11.4%
R@5 / R@10	0.8819 / 0.9233	0.8819 / 0.9233	0
nDCG@10	0.7752	0.7746	−0.0006
rows stored	83,024	83,024	identical
failed embeds	0	0	–
driver PSS peak	n/a	38.3 GiB	–
host RAM peak	n/a	77.5 GiB	–

The control is a reused run from two days earlier, and the throughput difference comes from one run per arm, so treat it as indicative. The control run recorded no memory telemetry. For memory, see the supplementary numbers below (driver PSS 45.1 → 38.7 GiB on the same H100).

Supplementary: the stacked #2781 version (same change, measured on main bf9205d1 vs 7895d46e, which is main + #2778 + #2781, so it also includes #2778's NumPy row conversion; NRB, bo767 n=3 and ViDoRe English n=2 per arm, alternating):

box	dataset	pages/s main → PR	peak driver Pss	recall@5 / nDCG@10
H100 NVL	bo767	24.98 → 28.20 (+12.9%)	45.1 → 38.7 GiB	0.8648 / 0.7767 → 0.8644 / 0.7770
H100 NVL	ViDoRe (5)	+1.9% to +3.9% on 4 workloads, −0.9% on 1	~flat	unchanged (±0.001)
L40S (fixed per-stage GPU limits)	bo767	18.19 → 18.26 (+0.4%)	46.3 → 40.4 GiB	0.8671 / 0.7773 → 0.8665 / 0.7772
L40S (fixed per-stage GPU limits)	ViDoRe (5)	5.27 → 5.38 (+2.1%)	~flat	unchanged (±0.003)

Rows processed and stored matched on every workload. Full details are in #2781.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge with its explicitly intended breaking changes.

Summary

The embedding stage now writes one typed embedding column and a matching embedding_error column instead of carrying duplicate vectors.

• Embedding rows carry one typed vector and a matching error column.
• VDB uploads read vectors from the configured embedding column.
• Query retrieval reads vectors from the configured embedding column.
• LanceDB writes take vectors from the embedding column.

No actionable new issues were found. Tests were not run during this review.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Local model or remote endpoint] --> B[assign_embedding_columns]
  B --> C[Configured vector column]
  B --> D[Matching error column]
  C --> E[Upload and retrieval readers]
  C --> F[Result transport]
  D --> G[Ingestion error reporting]
  E --> H[Canonical records and LanceDB vectors]

Loading

_{Reviews (4) · Last reviewed commit: "Merge origin/main into devin/1792516687-..." · Reviewed by Greptile}

[jperez999]

[P2] Account for this additional breaking surface

This deletes the importable nemo_retriever.operators.graph_ops.lancedb_write_operator module and LanceDBWriteOperator class, but the PR description’s explicit breaking-change list only covers embedding columns and configuration fields. The absence of in-repo imports does not protect downstream callers from an import failure. To keep this standalone representation change focused, could we retain and update this operator (or leave a compatibility shim)? If its removal is intentional, please at least identify the module/class removal in the PR description and release notes, with the rationale for including it here.