Late Materialization (Query-Then-Fetch) for pure Sort queries

[expani]

Late Materialization (Query-Then-Fetch) for top-N queries

Implements QTF for Sort + Limit queries: shards return only the sort columns plus a row-id helper, the coordinator picks the top-K rows globally, then fetches the remaining columns from the originating shards by row id and stitches them back into sort order.

Why

For a query like "select a few columns from a large index, ordered by one column, limit K", the existing path materializes every projected column on every matching row cross every shard before the global sort. With QTF, shards only ship the sort columns and row ids, the coordinator picks K winners, and only those K rows drive a second per-shard fetch for the remaining columns. This drops shard-to-coordinator network I/O and avoids per-row materialization on rows that were never going to win.

DAG shape

Four stages (three if the cluster has a single shard, where the rewriter short-circuits):

Stage 0: shard scan (returns sort columns + row id, per-shard ordinal stamped at the boundary)
Stage 1: coord reduce (Sort + Limit picks K winners)
Stage 2: late materialization (drain → group rows by shard → scatter fetch → stitch back into order)
Stage 3: coord reduce (post-LM derived projection / filter / aggregate)

---

Coordinator

Planner

• OpenSearchLateMaterializationRewriter — new rule that detects a Sort+Limit over a viable scan, inserts the LM marker, and short-circuits in single-shard case.
• OpenSearchLateMaterialization — new marker RelNode that records which columns are needed in the second-pass fetch.
• OpenSearchTableScan — keeps its overridden row type through Calcite normalization so the helper columns survive.
• PlannerImpl — adds Calcite's standard sort-project-transpose and project-merge rules so the planner output reaches the canonical shape QTF detection expects.
• RelNodeUtils — small helpers used by the rewriter when threading helper columns through the row type.
• ArrowCalciteTypes — new bidirectional Arrow ↔ Calcite type converter used when the LM stage builds its output schema.
• DAGBuilder — adds the cut around the LM stage and adjusts the Stage 1 input-scan placeholder so the schema includes the helper columns.
• FragmentConversionDriver — emits a schema-only Substrait Read for the LM stage (it has no source table to lower) and records the post-decoration schema for any stage where the wire schema differs from the producer's natural one.
• StagePlan — carries the post-decoration schema so the consuming reduce sink registers against the schema that actually crosses the wire.
• InputSinkDecorator — new SPI that lets a reduce stage wrap its input sink. Today the only implementation stamps the per-shard ordinal column.
• Stage, StageExecutionType — new LM stage type with the wiring needed to provide an exchange sink and an instruction handler factory.

Scheduler / Core

• LateMaterializationStageExecution — new stage orchestrator that runs the four QTF phases. It drains the upstream reduce output to recover row ids and per-shard ordinals, partitions rows by shard, fans out a fetch transport call per shard, and lets the stitcher write each response back into the K-row output in sort order.
• LateMaterializationStageExecutionFactory — builds the LM stage execution and resolves the descendant shard-stage that supplies the per-shard target lookup.
• Stitcher — new per-stage helper that owns the K-row output and copies cells from each shard's response into the correct row position. Emits one batch when all shards complete.
• OrdinalAppendingSink — new decorator that stamps a per-shard ordinal column onto each batch before it lands in the reduce sink. Used to identify which shard a winning row came from.
• VectorUtils.appendConstantInt — new zero-copy helper for adding the constant ordinal column. Works around an Arrow 18.x change that rejects the obvious "append at end" path.
• QueryContext — new side-table mapping a stage's resolved per-shard targets so a downstream stage can map an ordinal back to a node and shard.
• ReduceStageExecution — consults the optional input-sink decorator; multi-input dispatch is now keyed by the logical child-stage count rather than a marker
  interface.
• ReduceStageExecutionFactory — when registering child inputs, prefers the post-decoration schema over the producer's natural one.
• FetchByRowIdsAction, FetchByRowIdsRequest — new transport action and request type for the second-pass fetch.

DataFusion backend (coordinator-relevant pieces)

• DataFusionFragmentConvertor — new method that builds a schema-only Substrait Plan directly (bypassing isthmus, since the LM stage has no source table to lower).
• DataFusionInstructionHandlerFactory — factory methods now take a flag that requests row-id emission on the shard.
• DatafusionReduceSink — adds a schema-mismatch check on incoming batches and tolerates Timestamp precision/timezone divergence so the schema declared on the coordinator round-trips through the data-node parquet reader.

---

Data node

Core

• AnalyticsSearchService — new handler that, given a sorted list of row ids and a column projection, produces an Arrow stream of those exact rows in ascending row-id order.
• AnalyticsSearchTransportService — registers the new fetch action with the streaming transport and exposes the coordinator-side dispatch helper.
• ReaderContext / ReaderContextStore — keep-alive plumbing so the reader handle survives across the query → fetch round trip.

DataFusion backend (Java)

• ShardScanInstructionHandler — when row-id emission is requested, switches to the indexed execution path that can produce row ids.
• ShardScanWithDelegationHandler — same gating for the delegation path.
• NativeBridge — Java FFM signature for the indexed-session-creation call now carries the row-id flag.

DataFusion backend (Rust)

#21653

Row-id emission infrastructure:

• project_row_id_analyzer.rs, project_row_id_optimizer.rs, shard_table_provider.rs — new modules that wire row-id emission into the vanilla (non-indexed) execution path. The analyzer detects when a query references the row-id helper column and the optimizer rewrites the physical plan so the table provider produces row ids alongside the data columns.
• indexed_table/row_id_injection.rs — sibling injection point for the indexed execution path; row ids come from the position within the row group rather than from the table provider.
• api.rs — plan inspection now reports both whether the query is indexed and whether it requests row ids, so dispatch routes to the right executor.
• datafusion_query_config.rs — new enum that captures the strategy decision for the non-indexed-but-row-id-needed case.
• indexed_executor.rs — reads the row-id-emission flag from session config.
• New end-to-end and benchmark tests covering both emission paths and the strategy enum's branches.

Core-side opt-in:

• ffm.rs — the indexed-session FFM signature now carries a byte for the row-id-emission flag set on the Java side.
• session_context.rs — stores the flag on the indexed-execution config so the executor can read it without re-parsing the plan.

---

Cross-cutting (Framework SPI)

• ShardScanInstructionNode, ShardScanWithDelegationInstructionNode — new field that carries the row-id-emission request, with serde.
• FragmentInstructionHandlerFactory — factory methods accept the new flag.
• FragmentConvertor — new default-throwing schema-only-read method that backends opt into.
• ExchangeSink — new default method that lets sinks distinguish a per-source ordinal when one is available.
• FieldStorageInfo — exposes per-column storage hints so the late-materialization stage can build its fetch projection (i.e., decide which columns to ask the data
  node for in the second-pass fetch).

---

Tests

• New end-to-end IT that asserts QTF fires and that the post-LM stage runs derived expressions (e.g., UPPER(url)) over the stitched output, with per-row value
  assertions.
• New IT that verifies Substrait conversion at each stage.
• New planner-shape goldens for the rewrite.
• DAG-shape and fragment-conversion unit tests updated to cover the LM stage type and the post-decoration schema plumbing.

Known limitations / follow-ups

• Single-batch emission from the stitcher — the full stitched output is emitted in one feed when all shards complete, so the post-LM stage uses a buffered sink.
  Streaming the post-LM stage (per-shard sub-batch emission, dropping the position-sortedness invariant) is deferred until a real workload demands it.
• Type coverage in the Arrow ↔ Calcite converter — Timestamp currently hardcodes millisecond precision; Date, Time, smallint, tinyint, decimal, and timestamp-with-local-timezone are unmapped. Broaden before exposing QTF beyond keyword and date columns.

[github-actions]

PR Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit 5c0ccdb.

'Diff too large, requires skip by maintainers after manual review'

---

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

---

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

[mch2]

thanks @expani still going through this but pasting initial batch

[mch2]

this is a huge feature thanks @expani and @alchemist51 for taking this on

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit d28655a)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

📝 TODO sections sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [101]: Don't Delete until fixed. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [102]: We need to create Rewriting for Distributed Query Execution as a separate PlannerPhase. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [103]: One categorization that applies is Correctness v/s Performance. So, a RewritePhase -> LateMatRewriter. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [104]: Late Materialization is a rewrite done for performance. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [105]: Rewrite for TopK Approximation is done for Correctness as a response to an User ExecutionHint in request.

🔀 No multiple PR themes

⚡ Recommended focus areas for review Possible Issue In rebuildBelowProject, the method fills outputRemap for kept columns but leaves dropped columns at -1. If the anchor's collation references a dropped column, rebuildAnchor will throw an exception at line 492-495. However, the rewriter's detection phase should have already validated that all sort columns are in belowAnchorPhysicalFields (line 279-282), so this scenario should be impossible. If it can occur, the detection logic has a gap; if it cannot, the exception message at line 493-495 is misleading (it suggests a rewriter bug rather than an invariant violation). Consider adding an assertion in detection that every anchor collation field is in belowAnchorPhysicalFields, or clarifying the exception message to indicate this is an internal invariant failure. private static BelowProjectRebuild rebuildBelowProject( OpenSearchProject p, RelNode newChild, int[] scanIdxRemap, RelDataTypeFactory typeFactory ) { // Below-Project is passthrough; its output→scan map was captured during analyzeBelow, // but we recompute here from p's projects since each project is a RexInputRef. int[] origOutToScan = new int[p.getProjects().size()]; for (int i = 0; i < p.getProjects().size(); i++) { origOutToScan[i] = ((RexInputRef) p.getProjects().get(i)).getIndex(); } List<RexNode> newProjects = new ArrayList<>(); List<String> newNames = new ArrayList<>(); int[] outputRemap = new int[origOutToScan.length]; Arrays.fill(outputRemap, -1); for (int origOut = 0; origOut < origOutToScan.length; origOut++) { int newScanIdx = scanIdxRemap[origOutToScan[origOut]]; if (newScanIdx < 0) continue; // source dropped (now fetched, not in narrowed Scan) RelDataTypeField field = newChild.getRowType().getFieldList().get(newScanIdx); outputRemap[origOut] = newProjects.size(); newProjects.add(new RexInputRef(newScanIdx, field.getType())); newNames.add(p.getRowType().getFieldList().get(origOut).getName()); } // Pass through ___row_id (always last in narrowed Scan rowType). int rowIdIdx = newChild.getRowType().getFieldCount() - 1; RelDataTypeField rowIdField = newChild.getRowType().getFieldList().get(rowIdIdx); newProjects.add(new RexInputRef(rowIdIdx, rowIdField.getType())); newNames.add(OpenSearchLateMaterialization.ROW_ID_FIELD); RelDataTypeFactory.Builder pb = typeFactory.builder(); for (int j = 0; j < newProjects.size(); j++) { pb.add(newNames.get(j), newProjects.get(j).getType()); } OpenSearchProject rebuilt = new OpenSearchProject( p.getCluster(), p.getTraitSet(), newChild, newProjects, pb.build(), p.getViableBackends() ); return new BelowProjectRebuild(rebuilt, outputRemap); } Possible Issue In drainAndGroupByUgsi, the method pre-sizes drainedRowIds and drainedUgsis arrays using childInputBuffer.getRowCount() (line 456-460). If getRowCount() returns a value exceeding Integer.MAX_VALUE, the cast at line 460 will silently overflow, resulting in a negative or incorrect array size, which will then cause an exception when the arrays are allocated. The check at line 457-459 throws an exception, but only after the cast has already happened. The cast should occur after the bounds check, or the check should operate on the long value directly before casting. private void drainAndGroupByUgsi() { long total = childInputBuffer.getRowCount(); if (total > Integer.MAX_VALUE) { throw new IllegalStateException("Drained row count exceeds Integer.MAX_VALUE: " + total); } int K = (int) total; drainedRowIds = new long[K]; drainedUgsis = new int[K]; Possible Issue In drainFetchByRowIds, the method acquires a ReaderContext at line 223 and releases it at line 239 or 266 on failure paths. However, if an exception is thrown between line 223 and the try-with-resources block starting at line 274, the context is never released. For example, if backend.fetchByRowIds at line 264 throws an exception before the FragmentResources is constructed, the readerContextStore.releaseContext call at line 270 executes, but if an exception occurs after line 223 and before line 256 (e.g., during the assertAscending call at line 254 in a dev build), the context leaks. The acquisition should be moved inside the try block, or a finally block should ensure release on all paths. ReaderContext readerContext = readerContextStore.acquireContext(request.getQueryId(), shard.shardId()); if (readerContext == null) { responseHandler.onFailure( new IllegalStateException( "No ReaderContext for queryId=" + request.getQueryId() + " on " + shard.shardId() + " — query phase missing or context expired" ) ); return; } assert assertFetchInvariants(readerContext, request.getQueryId()); AnalyticsSearchBackendPlugin backend = backends.get(request.getBackendId()); if (backend == null) { readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); responseHandler.onFailure( new IllegalStateException( "No backend registered for backendId=" + request.getBackendId() + " on " + shard.shardId() + "; available: " + backends.keySet() ) ); return; } // Caller contract: rowIds must already be sorted ascending (RowSelection invariant on // native side). Asserted here so violations are caught in dev builds before the FFM call. assert assertAscending(rowIds); BigIntVector rowIdVector = null; FragmentResources resources = null; try { rowIdVector = new BigIntVector(DocumentInput.ROW_ID_FIELD, allocator); rowIdVector.allocateNew(rowIds.length); for (int i = 0; i < rowIds.length; i++) { rowIdVector.set(i, rowIds[i]); } rowIdVector.setValueCount(rowIds.length); EngineResultStream stream = backend.fetchByRowIds(readerContext.getReader(), rowIdVector, columns, allocator); // FragmentResources keeps the rowIdVector alive until the stream drains — closing // it earlier would pull off-heap memory out from under the native FFM call. resources = new FragmentResources(readerContextStore, readerContext, null, stream, null, rowIdVector); } catch (Exception e) { if (rowIdVector != null) rowIdVector.close(); readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); responseHandler.onFailure(new RuntimeException("Failed to execute fetch-by-row-ids on " + shard.shardId(), e)); return; } try (FragmentResources ctx = resources) { Iterator<EngineResultBatch> it = ctx.stream().iterator(); while (it.hasNext()) { responseHandler.onBatch(it.next()); } responseHandler.onComplete(); } catch (Exception e) { responseHandler.onFailure(e); } } Possible Issue In acceptBatch, the method checks rowsCopiedSoFar + batchRows > positions.length at line 117-125, but does not validate that dstRow (computed at line 129 as positions[rowsCopiedSoFar + srcRow]) is within bounds of the output VSR (i.e., dstRow < totalRows). If a shard's response contains a corrupted or malicious positions array with values >= totalRows, the copyFromSafe call at line 140 will write out of bounds. This can occur if the data node is compromised or if there is a bug in the Phase B sort logic. Add a bounds check: if (dstRow < 0 || dstRow >= totalRows) throw new IllegalStateException(...) before the copy loop. for (int srcRow = 0; srcRow < batchRows; srcRow++) { int dstRow = positions[rowsCopiedSoFar + srcRow]; int outCol = 0; for (int srcCol = 0; srcCol < responseColCount; srcCol++) { if (srcCol == rowIdIdx) continue; if (outCol >= aboveColCount) { throw new IllegalStateException( "Response column count exceeds output schema: outCol=" + outCol + " aboveColCount=" + aboveColCount ); } FieldVector srcVec = batch.getVector(srcCol); FieldVector dstVec = output.getVector(outCol); dstVec.copyFromSafe(srcRow, dstRow, srcVec); outCol++; } } Possible Issue In cutAtLateMaterialization, the method picks the first viable backend from lm.getViableBackends() to resolve the ExchangeSinkProvider at line 188. The TODO at line 33-45 notes this is ambiguous when multiple backends are viable. If the first backend in the list does not actually support the reduce operation (e.g., due to a capability-resolution bug elsewhere), the sink provider may be incorrect, leading to a runtime failure when the stage tries to instantiate the sink. The method should either assert that exactly one backend is viable (throwing an exception if not), or defer this cut until after PlanForker has collapsed the viable list to a single resolved backend per alternative, as the TODO suggests. private static RelNode cutAtLateMaterialization( OpenSearchLateMaterialization lm, int[] counter, List<Stage> parentChildStages, CapabilityRegistry registry, ClusterService clusterService, IndexNameExpressionResolver indexNameExpressionResolver ) { // 1. Reduce child — Sort+Limit reduce above shard scans. Multi-shard QTF only. List<Stage> reduceChildren = new ArrayList<>(); RelNode reduceFragment = sever(lm.getInput(), counter, reduceChildren, registry, clusterService, indexNameExpressionResolver); if (reduceChildren.isEmpty()) { throw new IllegalStateException( "QTF rewriter fired but the wrapper's input has no ExchangeReducer below it — " + "single-shard collapse should have short-circuited the rewriter." ); } int reduceStageId = counter[0]++; List<String> reduceViable = CapabilityResolutionUtils.filterByReduceCapability(registry, lm.getViableBackends()); ExchangeSinkProvider reduceSinkProvider = registry.getBackend(reduceViable.getFirst()).getExchangeSinkProvider(); Stage reduceStage = new Stage( reduceStageId, reduceFragment, reduceChildren, /*exchangeInfo=*/ null, reduceSinkProvider, /*targetResolver=*/ null ); // Reducer feeds the LM stage. Stamp every shard's batches with their target.ordinal() // as ___ugsi BEFORE the backend's reduce sees them so the LM stage can group rows by // source shard for fan-out fetches. reduceStage.setInputSinkDecorator( (sink, allocator) -> new OrdinalAppendingSink(sink, allocator, OpenSearchLateMaterialization.UGSI_FIELD) );

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to d28655a

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Enforce row ID bounds in release The bounds check for gid is only active in debug builds. If a caller passes an out-of-range row ID in release mode, partition_point can return an invalid index, causing a panic or silent corruption. Promote the bounds validation to a runtime check that always executes. sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs [673-686] for &gid in &row_ids { - debug_assert!(gid >= 0, "fetch_by_row_ids: negative row id {}", gid); + if gid < 0 { + return Err(DataFusionError::Execution( + format!("fetch_by_row_ids: negative row id {}", gid) + )); + } let seg_idx = segments .partition_point(|s| s.global_base <= gid as u64) .saturating_sub(1); let seg = &segments[seg_idx]; - debug_assert!( - (gid as u64) >= seg.global_base && (gid as u64) < seg.global_base + seg.max_doc as u64, - ... - ); + if (gid as u64) < seg.global_base || (gid as u64) >= seg.global_base + seg.max_doc as u64 { + return Err(DataFusionError::Execution( + format!("fetch_by_row_ids: row id {} out of bounds for segment {} (base={}, max_doc={})", + gid, seg_idx, seg.global_base, seg.max_doc) + )); + } let local_pos = (gid as u64 - seg.global_base) as u32; per_segment.entry(seg_idx).or_default().insert(local_pos); } Suggestion importance[1-10]: 9 __ Why: The suggestion correctly identifies that bounds checks are only active in debug builds via debug_assert!. In release mode, invalid row IDs can cause panics or silent corruption. Promoting these to runtime checks is critical for production safety. This is a high-impact correctness fix.	High
	Fix resource release ordering Release the reader context before closing rowIdVector to prevent use-after-free. The rowIdVector buffer is passed to the native fetch call, and closing it before releasing the context could invalidate memory still referenced by the reader if another thread acquires the context immediately after release. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java [90-100] -first = closeQuietly(rowIdVector, first); // Release (not close) — the store's reaper closes after keepAlive, and the QTF // fetch phase may still need this reader before then. if (readerContext != null) { try { readerContextStore.releaseContext(readerContext.getQueryId(), readerContext.getShardId()); } catch (Exception e) { if (first == null) first = e; else first.addSuppressed(e); } } +first = closeQuietly(rowIdVector, first); Suggestion importance[1-10]: 8 __ Why: The suggestion identifies a subtle use-after-free risk: closing rowIdVector before releasing the readerContext could invalidate memory still referenced by the native fetch call if another thread acquires the context immediately after release. Reordering the operations eliminates this race condition and is a meaningful correctness improvement.	Medium
	Fail fast on schema mismatches The column mapping logic silently creates null arrays when a field is missing from the batch schema. This masks schema mismatches that should be hard errors. If row_id_idx is out of bounds or a required column is absent, the function produces a valid-looking batch with incorrect data instead of failing fast. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs [73-82] -pub fn inject_row_ids( - output: &RecordBatch, - ctx: &RowIdContext, - batch_len: usize, - current_mask: Option<&BooleanArray>, - mask_offset_before: usize, - row_id_idx: usize, - schema: &SchemaRef, -) -> Result<RecordBatch> { - let num_surviving = output.num_rows(); - ... - let columns: Vec<Arc<dyn Array>> = schema - .fields() - .iter() - .enumerate() - .map(|(i, field)| match i { - idx if idx == row_id_idx => Arc::clone(&row_id_array), - _ => batch_schema - .index_of(field.name()) - .map(|col_idx| Arc::clone(output.column(col_idx))) - .unwrap_or_else(|_| datafusion::arrow::array::new_null_array(field.data_type(), num_surviving)), - }) - .collect(); +if row_id_idx >= schema.fields().len() { + return Err(datafusion::common::DataFusionError::Internal( + format!("inject_row_ids: row_id_idx {} out of bounds for schema with {} fields", row_id_idx, schema.fields().len()) + )); +} +let columns: Vec<Arc<dyn Array>> = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| match i { + idx if idx == row_id_idx => Arc::clone(&row_id_array), + _ => batch_schema + .index_of(field.name()) + .map(|col_idx| Arc::clone(output.column(col_idx))) + .ok_or_else(|| datafusion::common::DataFusionError::Internal( + format!("inject_row_ids: field {} not found in batch schema", field.name()) + )) + }) + .collect::<Result<Vec<_>>>()?; Suggestion importance[1-10]: 8 __ Why: The suggestion identifies a real issue: unwrap_or_else silently creates null arrays when fields are missing, which masks schema mismatches. The improved code adds bounds checking for row_id_idx and converts the fallback to a hard error. This is a significant correctness improvement that prevents silent data corruption.	Medium
	Add array bounds validation Add bounds checking before accessing belowProjOutToScan[slot] to prevent potential ArrayIndexOutOfBoundsException. When a below-Project exists, verify that slot is within the array bounds before dereferencing. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [320-330] private static List<String> buildAnchorSlotToPhysicalField(BelowChain belowChain) { RelNode topBelowOp = belowChain.chain.isEmpty() ? belowChain.scan : belowChain.chain.get(0); int slotCount = topBelowOp.getRowType().getFieldCount(); List<RelDataTypeField> scanFields = belowChain.scan.getRowType().getFieldList(); List<String> out = new ArrayList<>(slotCount); for (int slot = 0; slot < slotCount; slot++) { - int scanIdx = (belowChain.belowProjOutToScan == null) ? slot : belowChain.belowProjOutToScan[slot]; + int scanIdx; + if (belowChain.belowProjOutToScan == null) { + scanIdx = slot; + } else { + if (slot >= belowChain.belowProjOutToScan.length) { + throw new IllegalStateException("slot " + slot + " exceeds belowProjOutToScan length " + belowChain.belowProjOutToScan.length); + } + scanIdx = belowChain.belowProjOutToScan[slot]; + } out.add(scanFields.get(scanIdx).getName()); } return out; } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential ArrayIndexOutOfBoundsException when belowProjOutToScan is non-null. Adding bounds checking before array access improves robustness, though the planner's internal invariants may already prevent this scenario in practice.	Medium
	Validate destination row bounds Add bounds checking for dstRow before calling copyFromSafe to prevent writing beyond the pre-allocated output buffer. Verify that dstRow < totalRows to catch position-array corruption or oversized responses early. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java [128-142] public void acceptBatch(VectorSchemaRoot batch, int[] positions, int rowsCopiedSoFar) { synchronized (outputLock) { ... for (int srcRow = 0; srcRow < batchRows; srcRow++) { int dstRow = positions[rowsCopiedSoFar + srcRow]; + if (dstRow < 0 || dstRow >= totalRows) { + throw new IllegalStateException( + "Invalid destination row position: dstRow=" + dstRow + " totalRows=" + totalRows + ); + } int outCol = 0; for (int srcCol = 0; srcCol < responseColCount; srcCol++) { if (srcCol == rowIdIdx) continue; if (outCol >= aboveColCount) { throw new IllegalStateException( "Response column count exceeds output schema: outCol=" + outCol + " aboveColCount=" + aboveColCount ); } FieldVector srcVec = batch.getVector(srcCol); FieldVector dstVec = output.getVector(outCol); dstVec.copyFromSafe(srcRow, dstRow, srcVec); outCol++; } } } } Suggestion importance[1-10]: 7 __ Why: Adding bounds checking for dstRow before copyFromSafe is a valid defensive measure. It catches position-array corruption or oversized responses early, improving robustness. The check is cheap and prevents potential buffer overruns.	Medium
	Prevent context overwrite on duplicate creation Check for existing context before creating a new one to prevent overwriting. Multi-shard QTF queries may call createContext multiple times with the same queryId but different shardId. However, if called twice with identical (queryId, shardId), the second call silently overwrites the first, leaking the original reader. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java [69-74] public ReaderContext createContext(String queryId, ShardId shardId, GatedCloseable<Reader> gatedReader) { + Key key = new Key(queryId, shardId); + ReaderContext existing = activeContexts.get(key); + if (existing != null) { + throw new IllegalStateException("Context already exists for queryId=" + queryId + " shardId=" + shardId); + } ReaderContext ctx = new ReaderContext(queryId, shardId, gatedReader, defaultKeepAliveMillis); ctx.markInUse(); - activeContexts.put(new Key(queryId, shardId), ctx); + activeContexts.put(key, ctx); return ctx; } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential issue where calling createContext twice with the same (queryId, shardId) would silently overwrite the first context, leaking the original reader. Adding a check prevents this resource leak, though the PR's Javadoc already notes the multi-shard keying design makes collisions unlikely in normal operation.	Medium
	Validate segment ordering before distribution The row ID distribution logic assumes segments are sorted by global_base but does not verify this invariant. If segments arrive unsorted, partition_point will produce incorrect indices, causing silent data corruption. Add an explicit runtime check before the loop to validate segment ordering. sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs [672-686] -pub async unsafe fn fetch_by_row_ids( - shard_view: &ShardView, - runtime: &DataFusionRuntime, - manager: &crate::runtime_manager::RuntimeManager, - row_ids: Vec<i64>, - columns: Vec<String>, -) -> Result<i64, DataFusionError> { - ... - debug_assert!(!segments.is_empty(), "fetch_by_row_ids: build_segments returned empty for non-empty shard view"); - let mut per_segment: HashMap<usize, RoaringBitmap> = HashMap::new(); - for &gid in &row_ids { - debug_assert!(gid >= 0, "fetch_by_row_ids: negative row id {}", gid); - let seg_idx = segments - .partition_point(|s| s.global_base <= gid as u64) - .saturating_sub(1); - let seg = &segments[seg_idx]; - debug_assert!( - (gid as u64) >= seg.global_base && (gid as u64) < seg.global_base + seg.max_doc as u64, - "fetch_by_row_ids: row id {} out of bounds for segment {} (base={}, max_doc={})", - gid, seg_idx, seg.global_base, seg.max_doc - ); +debug_assert!(!segments.is_empty(), "fetch_by_row_ids: build_segments returned empty for non-empty shard view"); +if !segments.windows(2).all(|w| w[0].global_base <= w[1].global_base) { + return Err(DataFusionError::Execution( + "fetch_by_row_ids: segments not sorted by global_base".into() + )); +} +let mut per_segment: HashMap<usize, RoaringBitmap> = HashMap::new(); +for &gid in &row_ids { + debug_assert!(gid >= 0, "fetch_by_row_ids: negative row id {}", gid); + let seg_idx = segments + .partition_point(|s| s.global_base <= gid as u64) + .saturating_sub(1); Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that partition_point assumes sorted segments. However, the code already has a debug_assert! at line 671 checking non-empty segments, and build_segments is documented to return sorted segments. Adding a runtime check adds safety but may be redundant if the contract is upheld. The suggestion is valid but the impact is moderate since the invariant is already enforced elsewhere.	Medium
	Fail on missing position map entries The function silently falls back to delivered_idx when rg_position() returns None, which can produce incorrect row IDs if the position map is incomplete. This fallback masks data corruption. Return an error instead of guessing the position. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs [132-140] #[inline] -fn position_to_global_id(delivered_idx: usize, pm: Option<&PositionMap>, base: u64) -> u64 { +fn position_to_global_id(delivered_idx: usize, pm: Option<&PositionMap>, base: u64) -> Result<u64, String> { let rg_pos = match pm { - Some(p) => p.rg_position(delivered_idx).unwrap_or(delivered_idx), + Some(p) => p.rg_position(delivered_idx) + .ok_or_else(|| format!("position_to_global_id: position {} not in map", delivered_idx))?, None => delivered_idx, }; - let id = base + rg_pos as u64; - debug_assert!(id >= base, "position_to_global_id: underflow base={} rg_pos={}", base, rg_pos); - id + let id = base.checked_add(rg_pos as u64) + .ok_or_else(|| format!("position_to_global_id: overflow base={} rg_pos={}", base, rg_pos))?; + Ok(id) } Suggestion importance[1-10]: 7 __ Why: The suggestion identifies that unwrap_or(delivered_idx) silently falls back when rg_position() returns None, which can produce incorrect row IDs. However, the improved code changes the function signature to return Result, which would require updating all call sites. The suggestion is valid but the implementation impact is broader than shown. The overflow check via checked_add is a good addition.	Medium
	Validate before casting signed to unsigned The max_doc variable is of type i32 (from offset), and casting it to u64 could silently convert negative values into large positive numbers. Verify that max_doc is non-negative before the cast to prevent potential overflow issues in cumulative_rows. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs [69-98] let mut cumulative_rows: u64 = 0; for (seg_ord, meta) in object_metas.iter().enumerate() { ... let max_doc = offset; + if max_doc < 0 { + return Err(format!("Invalid max_doc {} for segment {}", max_doc, seg_ord)); + } let global_base = cumulative_rows; cumulative_rows += max_doc as u64; Suggestion importance[1-10]: 7 __ Why: Valid concern about casting i32 to u64 without validation. Since offset accumulates num_rows values, negative values could indicate data corruption or logic errors. Adding validation would prevent silent overflow bugs in cumulative_rows calculation.	Medium
	Validate column indices before vector access Validate that rowIdIdx and ugsiIdx are non-negative before using them as vector indices. The current check happens after the indices are computed but before vector access; move the validation immediately after indexOf to fail fast if required columns are missing. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [466-479] private void drainAndGroupByUgsi() { long total = childInputBuffer.getRowCount(); if (total > Integer.MAX_VALUE) { throw new IllegalStateException("Drained row count exceeds Integer.MAX_VALUE: " + total); } int K = (int) total; drainedRowIds = new long[K]; drainedUgsis = new int[K]; int position = 0; for (VectorSchemaRoot vsr : childInputBuffer.readResult()) { int rowIdIdx = vsr.getSchema().getFields().indexOf(vsr.getSchema().findField(OpenSearchLateMaterialization.ROW_ID_FIELD)); int ugsiIdx = vsr.getSchema().getFields().indexOf(vsr.getSchema().findField(OpenSearchLateMaterialization.UGSI_FIELD)); - ... + if (rowIdIdx < 0 || ugsiIdx < 0) { + throw new IllegalStateException( + "LM stage drain expected [" + + OpenSearchLateMaterialization.ROW_ID_FIELD + + ", " + + OpenSearchLateMaterialization.UGSI_FIELD + + "] in batch schema; got " + + vsr.getSchema() + ); + } BigIntVector rowIds = (BigIntVector) vsr.getVector(rowIdIdx); IntVector ugsis = (IntVector) vsr.getVector(ugsiIdx); Suggestion importance[1-10]: 3 __ Why: The validation already exists at lines 468-477 in the existing_code snippet shown in the diff. The suggestion asks to move it "immediately after indexOf", but the current placement (before vector access) is already correct and fails fast. No meaningful improvement.	Low
General	Preserve original exception in failure path Suppress the original exception when wrapping send failures. Throwing RuntimeException(sendException) discards the original failure e, losing critical diagnostic context. Chain both exceptions so the root cause is preserved in logs and stack traces. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [156-162] @Override public void onFailure(Exception e) { try { channel.sendResponse(e); } catch (Exception sendException) { - throw new RuntimeException(sendException); + RuntimeException wrapped = new RuntimeException("Failed to send error response", sendException); + wrapped.addSuppressed(e); + throw wrapped; } } Suggestion importance[1-10]: 6 __ Why: The suggestion correctly notes that wrapping sendException in a RuntimeException discards the original failure e, losing diagnostic context. Chaining both exceptions via addSuppressed improves debuggability, though the impact is moderate since the original exception is already logged upstream in most failure paths.	Low
	Add null check for vector parameter Validate rowIdVector nullability before dereferencing. The method does not null-check rowIdVector before calling getDataBuffer(), which will throw NullPointerException if the vector is null. Add an explicit null check with a clear error message. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java [837-851] +if (rowIdVector == null) { + throw new IllegalArgumentException("rowIdVector must not be null"); +} long bufAddr = rowIdVector.getDataBuffer().memoryAddress(); int count = rowIdVector.getValueCount(); long streamPtr; if (bufAddr != 0 && count > 0) { streamPtr = NativeBridge.fetchByRowIds( dfReader.getReaderHandle().getPointer(), bufAddr, count, columns, dataFusionService.getNativeRuntime().get() ); } else { throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0"); } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that rowIdVector is not null-checked before dereferencing, which would throw NullPointerException if null. Adding an explicit null check with a clear error message improves robustness, though the method signature and call sites likely enforce non-null in practice, making this a defensive improvement rather than a critical fix.	Low
	Ensure resource cleanup in all paths Ensure readerContextStore.releaseContext is called in all failure paths after a successful acquireContext. The current code releases on backend-not-found but not on the earlier null-context check, creating an asymmetry. Add a try-finally block to guarantee release. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [223-251] private void drainFetchByRowIds( FetchByRowIdsRequest request, IndexShard shard, AnalyticsShardTask task, StreamingFragmentResponseHandler responseHandler ) { ... ReaderContext readerContext = readerContextStore.acquireContext(request.getQueryId(), shard.shardId()); if (readerContext == null) { responseHandler.onFailure( new IllegalStateException( "No ReaderContext for queryId=" + request.getQueryId() + " on " + shard.shardId() + " — query phase missing or context expired" ) ); return; } - assert assertFetchInvariants(readerContext, request.getQueryId()); - AnalyticsSearchBackendPlugin backend = backends.get(request.getBackendId()); - if (backend == null) { + try { + assert assertFetchInvariants(readerContext, request.getQueryId()); + AnalyticsSearchBackendPlugin backend = backends.get(request.getBackendId()); + if (backend == null) { + responseHandler.onFailure( + new IllegalStateException( + "No backend registered for backendId=" + + request.getBackendId() + + " on " + + shard.shardId() + + "; available: " + + backends.keySet() + ) + ); + return; + } + ... + } finally { readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + } Suggestion importance[1-10]: 4 __ Why: The suggestion identifies an asymmetry in resource cleanup: releaseContext is called on backend-not-found but not on null-context. However, the null-context path returns immediately without acquiring the context, so there's nothing to release. The try-finally pattern would be cleaner but doesn't fix a real leak here.	Low
	Simplify redundant pattern matching The None case and empty bitmap case both return Ok(None), making the pattern redundant. Simplify by handling None first, then checking if the unwrapped bitmap is empty. This reduces code duplication and improves readability. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs [73-77] let candidates = match universe_bitmap_from_page_ranges(&page_ranges, rg) { + None => return Ok(None), Some(bm) if bm.is_empty() => return Ok(None), Some(bm) => bm, - None => return Ok(None), }; Suggestion importance[1-10]: 4 __ Why: The suggestion correctly identifies redundant pattern matching where both None and empty bitmap return Ok(None). However, the proposed reordering doesn't significantly improve functionality or readability—both versions are equally clear. The impact is minimal.	Low

---

Previous suggestions

Suggestions up to commit d28655a

Category	Suggestion	Impact
General	Prevent reader context leak on failure Ensure readerContext is released in the catch block even when rowIdVector allocation fails before the FragmentResources wrapper is created. Currently if allocateNew throws, the reader context leaks because the release only happens inside the catch after rowIdVector is non-null. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [256-273] private void drainFetchByRowIds( FetchByRowIdsRequest request, IndexShard shard, AnalyticsShardTask task, StreamingFragmentResponseHandler responseHandler ) { ... BigIntVector rowIdVector = null; FragmentResources resources = null; try { rowIdVector = new BigIntVector(DocumentInput.ROW_ID_FIELD, allocator); rowIdVector.allocateNew(rowIds.length); for (int i = 0; i < rowIds.length; i++) { rowIdVector.set(i, rowIds[i]); } rowIdVector.setValueCount(rowIds.length); EngineResultStream stream = backend.fetchByRowIds(readerContext.getReader(), rowIdVector, columns, allocator); resources = new FragmentResources(readerContextStore, readerContext, null, stream, null, rowIdVector); } catch (Exception e) { - if (rowIdVector != null) rowIdVector.close(); - readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + try { + if (rowIdVector != null) rowIdVector.close(); + } finally { + readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + } responseHandler.onFailure(new RuntimeException("Failed to execute fetch-by-row-ids on " + shard.shardId(), e)); return; } ... } Suggestion importance[1-10]: 9 __ Why: Critical resource leak fix. If rowIdVector.allocateNew or the loop throws before FragmentResources is created, the readerContext is never released. The finally block ensures cleanup happens even when rowIdVector is null, preventing reader context leaks.	High
	Add null check for vector parameter Validate that rowIdVector is not null before dereferencing it. A null vector would cause a NullPointerException before reaching the buffer address check, making the error message misleading. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java [837-851] +if (rowIdVector == null) { + throw new IllegalArgumentException("rowIdVector must not be null"); +} long bufAddr = rowIdVector.getDataBuffer().memoryAddress(); int count = rowIdVector.getValueCount(); -long streamPtr; -if (bufAddr != 0 && count > 0) { - streamPtr = NativeBridge.fetchByRowIds( - dfReader.getReaderHandle().getPointer(), - bufAddr, - count, - columns, - dataFusionService.getNativeRuntime().get() - ); -} else { +if (bufAddr == 0 || count == 0) { throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0"); } +long streamPtr = NativeBridge.fetchByRowIds( + dfReader.getReaderHandle().getPointer(), + bufAddr, + count, + columns, + dataFusionService.getNativeRuntime().get() +); Suggestion importance[1-10]: 6 __ Why: Adding a null check for rowIdVector before dereferencing it prevents a misleading NullPointerException and provides a clearer error message. The suggestion improves error handling, though the impact is moderate.	Low
	Simplify pattern matching logic The pattern matching can be simplified by combining the empty and None cases. Both conditions return Ok(None), so they can be handled together to reduce code duplication and improve readability. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs [73-77] let candidates = match universe_bitmap_from_page_ranges(&page_ranges, rg) { - Some(bm) if bm.is_empty() => return Ok(None), - Some(bm) => bm, - None => return Ok(None), + Some(bm) if !bm.is_empty() => bm, + _ => return Ok(None), }; Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies that the pattern matching can be simplified by combining the empty and None cases. Both return Ok(None), so using a catch-all pattern _ improves readability and reduces duplication.	Low
	Remove redundant null check The readerContext null check is redundant because the constructor asserts both readerContextStore and readerContext are non-null. Remove the check to avoid confusion and align with the constructor's invariants. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java [93-100] first = closeQuietly(rowIdVector, first); // Release (not close) — the store's reaper closes after keepAlive, and the QTF // fetch phase may still need this reader before then. -if (readerContext != null) { - try { - readerContextStore.releaseContext(readerContext.getQueryId(), readerContext.getShardId()); - } catch (Exception e) { - if (first == null) first = e; - else first.addSuppressed(e); - } +try { + readerContextStore.releaseContext(readerContext.getQueryId(), readerContext.getShardId()); +} catch (Exception e) { + if (first == null) first = e; + else first.addSuppressed(e); } Suggestion importance[1-10]: 4 __ Why: The constructor asserts that readerContext is non-null, making the null check redundant. Removing it improves code clarity and aligns with the constructor's invariants, though the impact is minor.	Low
Possible issue	Validate parallel array lengths match Validate that rowIds and positions arrays have matching lengths before sorting. Mismatched array lengths would cause ArrayIndexOutOfBoundsException when accessing positions[order[i]] if positions is shorter than rowIds. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [536-551] static void sortParallelByRowIdAscending(long[] rowIds, int[] positions) { int n = rowIds.length; + if (n != positions.length) { + throw new IllegalArgumentException("rowIds length " + n + " != positions length " + positions.length); + } if (n <= 1) return; Integer[] order = new Integer[n]; for (int i = 0; i < n; i++) order[i] = i; Arrays.sort(order, (a, b) -> Long.compare(rowIds[a], rowIds[b])); long[] sortedRowIds = new long[n]; int[] sortedPositions = new int[n]; for (int i = 0; i < n; i++) { sortedRowIds[i] = rowIds[order[i]]; sortedPositions[i] = positions[order[i]]; } System.arraycopy(sortedRowIds, 0, rowIds, 0, n); System.arraycopy(sortedPositions, 0, positions, 0, n); } Suggestion importance[1-10]: 8 __ Why: Critical validation. The method assumes rowIds and positions are parallel arrays; mismatched lengths would cause ArrayIndexOutOfBoundsException at positions[order[i]]. The check prevents silent corruption.	Medium
	Validate destination row index bounds Validate that dstRow (derived from positions[rowsCopiedSoFar + srcRow]) is within bounds of the output VSR before calling copyFromSafe. An out-of-range position value would cause silent data corruption or buffer overrun in the Arrow vector. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java [108-142] public void acceptBatch(VectorSchemaRoot batch, int[] positions, int rowsCopiedSoFar) { synchronized (outputLock) { ... for (int srcRow = 0; srcRow < batchRows; srcRow++) { int dstRow = positions[rowsCopiedSoFar + srcRow]; + if (dstRow < 0 || dstRow >= totalRows) { + throw new IllegalStateException( + "Position out of range: dstRow=" + dstRow + " totalRows=" + totalRows + ); + } int outCol = 0; for (int srcCol = 0; srcCol < responseColCount; srcCol++) { if (srcCol == rowIdIdx) continue; if (outCol >= aboveColCount) { throw new IllegalStateException( "Response column count exceeds output schema: outCol=" + outCol + " aboveColCount=" + aboveColCount ); } FieldVector srcVec = batch.getVector(srcCol); FieldVector dstVec = output.getVector(outCol); dstVec.copyFromSafe(srcRow, dstRow, srcVec); outCol++; } } } } Suggestion importance[1-10]: 8 __ Why: Important safety check. The dstRow value comes from the positions array built during Phase B; if a position is out of range (due to a bug in grouping or a corrupted response), copyFromSafe would write to an invalid index. Validating bounds prevents silent data corruption.	Medium
	Prevent context overwrite resource leak Check for existing context before creating a new one. If a context with the same queryId and shardId already exists, the current implementation silently overwrites it without closing the previous reader, causing a resource leak. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java [69-74] public ReaderContext createContext(String queryId, ShardId shardId, GatedCloseable<Reader> gatedReader) { + Key key = new Key(queryId, shardId); + ReaderContext existing = activeContexts.get(key); + if (existing != null) { + throw new IllegalStateException("Context already exists for query=" + queryId + " shard=" + shardId); + } ReaderContext ctx = new ReaderContext(queryId, shardId, gatedReader, defaultKeepAliveMillis); ctx.markInUse(); - activeContexts.put(new Key(queryId, shardId), ctx); + activeContexts.put(key, ctx); return ctx; } Suggestion importance[1-10]: 8 __ Why: Silently overwriting an existing ReaderContext without closing the previous reader causes a resource leak. The suggestion correctly identifies this issue and proposes throwing an exception to prevent it, which is a critical fix for resource management.	Medium
	Prevent segment index underflow The partition_point call can return 0 when all segments have global_base > gid, causing saturating_sub(1) to wrap to 0 and access the wrong segment. This leads to incorrect segment selection and potential out-of-bounds row IDs. Add an explicit bounds check before saturating_sub to catch this case and return an error instead of silently proceeding with wrong data. sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs [674-681] -let seg_idx = segments - .partition_point(|s| s.global_base <= gid as u64) - .saturating_sub(1); +let partition_idx = segments.partition_point(|s| s.global_base <= gid as u64); +if partition_idx == 0 { + return Err(DataFusionError::Execution(format!( + "fetch_by_row_ids: row id {} is before first segment (base={})", + gid, segments[0].global_base + ))); +} +let seg_idx = partition_idx - 1; let seg = &segments[seg_idx]; debug_assert!( (gid as u64) >= seg.global_base && (gid as u64) < seg.global_base + seg.max_doc as u64, "fetch_by_row_ids: row id {} out of bounds for segment {} (base={}, max_doc={})", gid, seg_idx, seg.global_base, seg.max_doc ); Suggestion importance[1-10]: 8 __ Why: The partition_point call can return 0 when all segments have global_base > gid, causing saturating_sub(1) to wrap to 0 and access the wrong segment. This is a critical correctness issue that would silently produce incorrect results. The suggestion correctly identifies the bug and provides a proper fix with an explicit bounds check.	Medium
	Add array bounds validation Add bounds checking before accessing belowProjOutToScan[slot] to prevent potential ArrayIndexOutOfBoundsException. The slot index may exceed the array length if the projection mapping is inconsistent with the operator's field count. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java [320-330] private static List<String> buildAnchorSlotToPhysicalField(BelowChain belowChain) { RelNode topBelowOp = belowChain.chain.isEmpty() ? belowChain.scan : belowChain.chain.get(0); int slotCount = topBelowOp.getRowType().getFieldCount(); List<RelDataTypeField> scanFields = belowChain.scan.getRowType().getFieldList(); List<String> out = new ArrayList<>(slotCount); for (int slot = 0; slot < slotCount; slot++) { - int scanIdx = (belowChain.belowProjOutToScan == null) ? slot : belowChain.belowProjOutToScan[slot]; + int scanIdx; + if (belowChain.belowProjOutToScan == null) { + scanIdx = slot; + } else { + if (slot >= belowChain.belowProjOutToScan.length) { + throw new IllegalStateException("Slot index " + slot + " exceeds belowProjOutToScan length " + belowChain.belowProjOutToScan.length); + } + scanIdx = belowChain.belowProjOutToScan[slot]; + } out.add(scanFields.get(scanIdx).getName()); } return out; } Suggestion importance[1-10]: 7 __ Why: Valid defensive check. The belowProjOutToScan array is built during analyzeBelow and should match slotCount, but an explicit bounds check prevents potential ArrayIndexOutOfBoundsException if the projection mapping is inconsistent.	Medium
	Prevent thread crash on send failure Throwing a RuntimeException from onFailure can propagate unchecked exceptions up the call stack, potentially crashing the handler thread. Log the failure instead of rethrowing to prevent thread termination. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java [156-162] @Override public void onFailure(Exception e) { try { channel.sendResponse(e); } catch (Exception sendException) { - throw new RuntimeException(sendException); + logger.error("Failed to send error response for exception: {}", e.getMessage(), sendException); } } Suggestion importance[1-10]: 7 __ Why: Throwing a RuntimeException from onFailure can crash the handler thread. Logging the failure instead prevents thread termination and improves robustness, though the original code's intent to surface the error is also valid.	Medium
	Detect row ID overflow Casting u64 row IDs to i64 can silently overflow for IDs >= 2^63, producing negative values that corrupt query results. Since Arrow's Int64Array is signed, large shards with billions of rows will wrap around. Check for overflow before the cast and return an error if any ID exceeds i64::MAX. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs [54-68] let ids = compute_row_ids( &ctx.eval_mask, current_mask, mask_offset_before, batch_len, ctx.batch_offset, ctx.position_map.as_ref(), ctx.base, ); -Arc::new(Int64Array::from_iter_values(ids.into_iter().map(|id| id as i64))) +let i64_ids: Result<Vec<i64>, _> = ids.into_iter().map(|id| { + i64::try_from(id).map_err(|_| datafusion::common::DataFusionError::Execution( + format!("inject_row_ids: row ID {} exceeds i64::MAX", id) + )) +}).collect(); +Arc::new(Int64Array::from_iter_values(i64_ids?.into_iter())) Suggestion importance[1-10]: 7 __ Why: Casting u64 row IDs to i64 can overflow for IDs >= 2^63, producing negative values. While this is a real issue for very large shards, it's less likely to occur in practice than the segment index underflow. The suggestion provides a correct fix using try_from to detect overflow.	Medium
	Prevent row ID base overflow The base computation self.global_base + self.current_rg_first_row as u64 can overflow when global_base is near u64::MAX and current_rg_first_row is large. This produces a wrapped-around base that corrupts all row IDs in the batch. Use checked arithmetic to detect overflow and return an error instead of silently producing wrong IDs. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs [599-609] let row_id_ctx = if self.row_id_output_index.is_some() { + let base = self.global_base.checked_add(self.current_rg_first_row as u64) + .ok_or_else(|| datafusion::error::DataFusionError::Execution( + format!("IndexedStream: row ID base overflow (global_base={}, rg_first_row={})", + self.global_base, self.current_rg_first_row) + ))?; Some(super::row_id_injection::RowIdContext { batch_offset: self.batch_offset, position_map: self.current_position_map.as_ref().cloned(), - base: self.global_base + self.current_rg_first_row as u64, + base, eval_mask: eval_mask.clone(), }) } else { None }; Suggestion importance[1-10]: 7 __ Why: The base computation can overflow when global_base is near u64::MAX. This is a real correctness issue that would corrupt all row IDs in the batch. The suggestion provides a proper fix using checked_add, though the scenario is unlikely in practice given typical shard sizes.	Medium
	Prevent negative value corruption The max_doc value (i32) is cast to u64 without checking for negative values. If max_doc is negative, the cast will produce a large u64 value, corrupting cumulative_rows and causing incorrect global_base calculations for subsequent segments. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs [69-98] let mut cumulative_rows: u64 = 0; for (seg_ord, meta) in object_metas.iter().enumerate() { ... let global_base = cumulative_rows; - cumulative_rows += max_doc as u64; + cumulative_rows = cumulative_rows.saturating_add(max_doc.max(0) as u64); Suggestion importance[1-10]: 7 __ Why: Valid concern about casting max_doc (i32) to u64 without checking for negative values. While unlikely in practice (row counts should be non-negative), adding max(0) or validation prevents potential corruption of cumulative_rows and ensures robustness.	Medium
Security	Validate pointer alignment The FFM boundary reads row IDs from a raw pointer without verifying alignment. If Java passes a misaligned buffer (e.g., from a non-direct ByteBuffer or corrupted ArrowBuf offset), this triggers undefined behavior on platforms requiring aligned i64 access. Add an alignment check before dereferencing the pointer to catch this at the FFM boundary. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [299-303] +if row_ids_ptr as usize % std::mem::align_of::<i64>() != 0 { + return Err(format!( + "df_fetch_by_row_ids: row_ids_ptr 0x{:x} is misaligned (must be {}-byte aligned)", + row_ids_ptr, std::mem::align_of::<i64>() + )); +} let row_ids: Vec<i64> = slice::from_raw_parts( row_ids_ptr as *const i64, row_ids_count as usize, ).to_vec(); Suggestion importance[1-10]: 6 __ Why: The FFM boundary reads row IDs from a raw pointer without verifying alignment. While this is a valid concern for undefined behavior, the Java side (ArrowBuf) typically provides aligned buffers. The check adds safety but may be overly defensive for the expected usage pattern.	Low

Suggestions up to commit b6c1cc1

Category	Suggestion	Impact
Possible issue	Release reader context on failure The readerContext acquired at the start is not released if an exception occurs before FragmentResources is constructed. If rowIdVector.allocateNew or the loop throws, the context remains locked in the store, causing a resource leak. Add readerContextStore.releaseContext to the early-return path after readerContext acquisition to ensure cleanup on all failure branches. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [196-283] private void drainFetchByRowIds( FetchByRowIdsRequest request, IndexShard shard, AnalyticsShardTask task, StreamingFragmentResponseHandler responseHandler ) { ... ReaderContext readerContext = readerContextStore.acquireContext(request.getQueryId(), shard.shardId()); if (readerContext == null) { responseHandler.onFailure(...); return; } - ... + assert assertFetchInvariants(readerContext, request.getQueryId()); + AnalyticsSearchBackendPlugin backend = backends.get(request.getBackendId()); + if (backend == null) { + readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + responseHandler.onFailure(...); + return; + } + assert assertAscending(rowIds); BigIntVector rowIdVector = null; FragmentResources resources = null; try { rowIdVector = new BigIntVector(DocumentInput.ROW_ID_FIELD, allocator); rowIdVector.allocateNew(rowIds.length); for (int i = 0; i < rowIds.length; i++) { rowIdVector.set(i, rowIds[i]); } - ... + rowIdVector.setValueCount(rowIds.length); + EngineResultStream stream = backend.fetchByRowIds(readerContext.getReader(), rowIdVector, columns, allocator); + resources = new FragmentResources(readerContextStore, readerContext, null, stream, null, rowIdVector); } catch (Exception e) { if (rowIdVector != null) rowIdVector.close(); readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); - responseHandler.onFailure(...); + responseHandler.onFailure(new RuntimeException("Failed to execute fetch-by-row-ids on " + shard.shardId(), e)); return; } ... } Suggestion importance[1-10]: 9 __ Why: Critical resource leak fix. The readerContext acquired at line 223 is not released if an exception occurs before FragmentResources is constructed (lines 256-267). If rowIdVector.allocateNew or the loop throws, the context remains locked in the store, causing a resource leak. The existing code already releases on the backend-null path (line 239), but the try-catch block at lines 256-270 only closes rowIdVector and doesn't release the context. This is a real bug.	High
	Validate parallel array lengths match Validate that rowIds and positions arrays have matching lengths before sorting. Mismatched array lengths would cause ArrayIndexOutOfBoundsException when accessing positions[order[i]] if positions.length < rowIds.length. This validation prevents silent data corruption or runtime crashes during the parallel sort operation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [536-551] static void sortParallelByRowIdAscending(long[] rowIds, int[] positions) { int n = rowIds.length; + if (n != positions.length) { + throw new IllegalArgumentException("Array length mismatch: rowIds.length=" + n + " != positions.length=" + positions.length); + } if (n <= 1) return; Integer[] order = new Integer[n]; for (int i = 0; i < n; i++) order[i] = i; Arrays.sort(order, (a, b) -> Long.compare(rowIds[a], rowIds[b])); long[] sortedRowIds = new long[n]; int[] sortedPositions = new int[n]; for (int i = 0; i < n; i++) { sortedRowIds[i] = rowIds[order[i]]; sortedPositions[i] = positions[order[i]]; } System.arraycopy(sortedRowIds, 0, rowIds, 0, n); System.arraycopy(sortedPositions, 0, positions, 0, n); } Suggestion importance[1-10]: 8 __ Why: Critical validation. Mismatched array lengths would cause ArrayIndexOutOfBoundsException when accessing positions[order[i]] if positions.length < rowIds.length. This check prevents silent data corruption or runtime crashes during the parallel sort operation.	Medium
	Validate destination row index bounds Add bounds checking for dstRow before calling copyFromSafe to prevent writing beyond the pre-allocated output buffer. If positions[rowsCopiedSoFar + srcRow] yields a value >= totalRows, the copy will corrupt memory or throw an Arrow buffer overflow exception. Validate that dstRow < totalRows to catch position-array corruption early. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java [108-145] public void acceptBatch(VectorSchemaRoot batch, int[] positions, int rowsCopiedSoFar) { synchronized (outputLock) { ... for (int srcRow = 0; srcRow < batchRows; srcRow++) { int dstRow = positions[rowsCopiedSoFar + srcRow]; + if (dstRow < 0 || dstRow >= totalRows) { + throw new IllegalStateException("Position out of bounds: dstRow=" + dstRow + " totalRows=" + totalRows); + } int outCol = 0; for (int srcCol = 0; srcCol < responseColCount; srcCol++) { if (srcCol == rowIdIdx) continue; if (outCol >= aboveColCount) { throw new IllegalStateException(...); } FieldVector srcVec = batch.getVector(srcCol); FieldVector dstVec = output.getVector(outCol); dstVec.copyFromSafe(srcRow, dstRow, srcVec); outCol++; } } } } Suggestion importance[1-10]: 8 __ Why: Important bounds check. If positions[rowsCopiedSoFar + srcRow] yields a value >= totalRows, the copyFromSafe call will corrupt memory or throw an Arrow buffer overflow exception. Validating dstRow < totalRows catches position-array corruption early and prevents buffer overruns.	Medium
	Fix race between close and markInUse Race condition between closed check and compareAndSet. Another thread could close the context after the closed check but before compareAndSet succeeds, allowing a closed context to be marked in-use. Move the closed check inside a synchronized block or use atomic state transitions to prevent this race. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContext.java [66-70] -public boolean markInUse() { +public synchronized boolean markInUse() { if (closed) return false; + if (!inUse.compareAndSet(false, true)) return false; lastAccessTime.set(System.currentTimeMillis()); - return inUse.compareAndSet(false, true); + return true; } Suggestion importance[1-10]: 8 __ Why: The suggestion identifies a legitimate race condition where a context could be marked in-use after being closed. The proposed synchronized fix prevents this race, though it introduces a broader lock. This is a correctness issue with potentially significant impact on concurrent access patterns.	Medium
	Prevent silent overflow in row ID cast The cast from u64 to i64 can silently overflow for row IDs exceeding i64::MAX. This would produce negative row IDs, corrupting the fetch phase. Add an overflow check before the cast to fail fast with a clear error message instead of silently producing invalid data. sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs [54-68] let ids = compute_row_ids( &ctx.eval_mask, current_mask, mask_offset_before, batch_len, ctx.batch_offset, ctx.position_map.as_ref(), ctx.base, ); -Arc::new(Int64Array::from_iter_values(ids.into_iter().map(|id| id as i64))) +let checked_ids: Result<Vec<i64>, _> = ids.into_iter().map(|id| { + i64::try_from(id).map_err(|_| { + datafusion::common::DataFusionError::Execution(format!( + "inject_row_ids: row ID {} exceeds i64::MAX", id + )) + }) +}).collect(); +Arc::new(Int64Array::from_iter_values(checked_ids?.into_iter())) Suggestion importance[1-10]: 8 __ Why: The suggestion identifies a real correctness issue where casting u64 to i64 can overflow for large row IDs, producing negative values. The proposed checked conversion prevents silent data corruption and provides clear error messages. This is a critical c...

[github-actions]

✅ Gradle check result for 4fb7d7b: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.49%. Comparing base (6105940) to head (d28655a).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21836      +/-   ##
============================================
+ Coverage     73.37%   73.49%   +0.12%     
- Complexity    75448    75555     +107     
============================================
  Files          6034     6033       -1     
  Lines        342504   342572      +68     
  Branches      49259    49276      +17     
============================================
+ Hits         251310   251776     +466     
+ Misses        71175    70789     -386     
+ Partials      20019    20007      -12     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit b6c1cc1

[github-actions]

✅ Gradle check result for b6c1cc1: SUCCESS

[github-actions]

Persistent review updated to latest commit d28655a

[github-actions]

❌ Gradle check result for d28655a: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit d28655a

[github-actions]

✅ Gradle check result for d28655a: SUCCESS