Add SPI for doc_values based collection

[harshavamsi]

Description

This PR adds an execution path from Lucene doc values to DataFusion:

1. Read Lucene doc values into Arrow batches.
2. Pass batches through Java FFM callbacks.
3. Expose the batches as a Rust DataFusion TableProvider.
4. Execute DataFusion plans against the Java-backed source.

Changes

• Adds the ArrowBatchSource, ArrowBatchSourceFactory, and ArrowBatchSourcePlan SPI contracts.
• Adds a Lucene DocValuesBatchSource and factory backed by real Lucene indexes.
• Adds Java callback registration and native bridge integration.
• Adds a Rust DataFusion table provider for Java-backed Arrow batches.
• Supports scalar columns:
  • LONG
  • TIMESTAMP
  • KEYWORD
  • BOOLEAN
  • FLOAT
  • DOUBLE
  • BINARY
  • IP
• Supports nullable multi-valued numeric, Boolean, floating-point, keyword, and IP columns as Arrow lists.
• Extends InputColumn serialization with multiValued while preserving the existing two-argument constructor.
• Uses safe Arrow buffer writes without broad native-access permissions.

Related Issues

Resolves #[Issue number to be closed when this PR is merged]

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Code Analyzer ❗

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

⛔ Hard block: Issues at Medium severity or above will block this PR from merging.

Path	Line	Severity	Description
sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1	1	high	Dependency version change: jackson-core bumped from 2.22.1 to 2.22.2 with a new SHA1 checksum. Artifact authenticity cannot be confirmed without verifying against the upstream Maven Central release.
sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1	1	high	Dependency version change: jackson-databind bumped from 2.22.1 to 2.22.2 with a new SHA1 checksum. Artifact authenticity cannot be confirmed without verifying against the upstream Maven Central release.
sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1	1	high	Dependency version change: jackson-datatype-jdk8 bumped from 2.22.1 to 2.22.2 with a new SHA1 checksum. Artifact authenticity cannot be confirmed without verifying against the upstream Maven Central release.
sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.2.jar.sha1	1	high	Dependency version change: jackson-dataformat-yaml bumped from 2.22.1 to 2.22.2 with a new SHA1 checksum. Artifact authenticity cannot be confirmed without verifying against the upstream Maven Central release.

The table above displays the top 10 most important findings.

Total: 4 | Critical: 0 | High: 4 | Medium: 0 | Low: 0

---

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.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 413f9f7)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 Multiple PR themes Sub-PR theme: Bump Jackson dependencies from 2.22.1 to 2.22.2 Relevant files: sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1 sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1 sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1 sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.2.jar.sha1 Sub-PR theme: Add Lucene DocValues batch source and factory Relevant files: sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSourceFactory.java sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/LuceneDocValuesBatchSourceTests.java Sub-PR theme: Add Rust DataFusion Arrow batch source table provider Relevant files: sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/mod.rs sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/table_provider.rs sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs

⚡ Recommended focus areas for review Possible Issue decodeSorted packs the ord+1 into the upper 44 bits and the row index into the lower 20 bits ((ord+1) << 20 | i). This limits the batch size to 2^20 = 1,048,576 rows and the number of unique ords to fit in 44 bits, but more importantly, BATCH_SIZE is 65,536 which fits, yet the row-index mask & 0xFFFFF (20 bits) also only allows rows up to 1,048,575. If BATCH_SIZE is ever raised above 1M or the ord exceeds 2^44, the packed value silently corrupts row assignments. Consider using two arrays or clearly documenting/enforcing the bit-packing limits. private void decodeSorted(SortedColumn column, FieldVector vector, int size) throws IOException { SortedDocValues values = column.values(); for (int i = 0; i < size; i++) { ordScratch[i] = values.advanceExact(docs[i]) ? values.ordValue() : -1; } allocateBytes(vector, size); for (int i = 0; i < size; i++) { ordRowScratch[i] = ((long) (ordScratch[i] + 1) << 20) | i; } Arrays.sort(ordRowScratch, 0, size); int lastOrd = Integer.MIN_VALUE; BytesRef term = null; for (int i = 0; i < size; i++) { int row = (int) (ordRowScratch[i] & 0xFFFFF); int ord = (int) (ordRowScratch[i] >>> 20) - 1; if (ord < 0) { vector.setNull(row); } else { if (ord != lastOrd) { term = values.lookupOrd(ord); lastOrd = ord; } setBytes(vector, row, term); } } vector.setValueCount(size); Possible Issue nextBatch's status contract collides EMPTY_BATCH (-3) with a legitimate zero-row batch case: root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(). However, EMPTY_BATCH (-3) is not handled distinctly on the Rust side vs 0 (EOF) in a way that preserves the exported ArrowArray — import_batch is called with expected_rows = 0 for EMPTY_BATCH but 0 status is treated as EOF (returns Ok(None)) without importing the exported array/schema. For a zero-row batch, Java exports arrays but the Rust side (for status 0) never releases them; only EMPTY_BATCH path imports. Ensure Java never returns 0 for an exported batch (only for true EOF where nothing was exported) — this appears to be handled, but verify a null nextBatch() never runs the export code and worth documenting more explicitly. VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } try { ArrowArray array = ArrowArray.wrap(arrayPointer.address()); ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); } finally { root.close(); } } Concurrency SourceEntry.exportNextBatch is synchronized on the entry, but cancel() acquires a separate cancellationLock and calls source.cancel() concurrently. The ArrowBatchSource.cancel() javadoc states cancel can run concurrently with nextBatch(). However, SourceEntry.close() is synchronized(this) and calls cancel() (which takes a different lock) and then source.close(). If close() runs while exportNextBatch is executing on another thread, they contend on the entry monitor — but a caller invoking binding.release() from a callback that itself synchronizes on the Binding may block waiting on the entry lock held by nextBatch, defeating cooperative cancellation semantics. Verify release/close paths cannot dead-block a long-running pull. private static final class SourceEntry { private final ArrowBatchSource source; private final Object cancellationLock = new Object(); private final AtomicBoolean cancellationRequested = new AtomicBoolean(); private boolean closed; private SourceEntry(ArrowBatchSource source) { this.source = source; } private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { if (closed) { throw new IllegalStateException("Arrow batch source is closed"); } VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } try { ArrowArray array = ArrowArray.wrap(arrayPointer.address()); ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); } finally { root.close(); } } private void cancel() { synchronized (cancellationLock) { if (cancellationRequested.compareAndSet(false, true)) { source.cancel(); } } } private synchronized void close() { cancel(); if (closed == false) { closed = true; source.close(); } } } Possible Issue In next_batch, when status is 0 (EOF), the local array and schema FFI structs go out of scope with their default empty() contents — this is fine only if Java did not populate them. Ensure Java writes nothing to arrayPointer/schemaPointer when returning 0. Currently SourceEntry.exportNextBatch only exports when root is non-null, so this holds, but the invariant is worth an explicit assertion or comment because a spurious export at EOF would leak the child buffers (from_ffi is not called and the release callback in the FFI structs would still fire, but only if populated — otherwise nothing is released and any allocated Java-side memory leaks). pub fn next_batch(&self) -> Result<Option<RecordBatch>, DataFusionError> { let callback = next_callback()?; let mut array = FFI_ArrowArray::empty(); let mut schema = FFI_ArrowSchema::empty(); let mut error = [0u8; ERROR_CAPACITY]; let status = unsafe { callback( self.binding_id, self.source_key, &mut array, &mut schema, error.as_mut_ptr(), error.len() as i64, ) }; match status { 0 => Ok(None), CANCELLED => Err(DataFusionError::Execution("query cancelled".into())), ERROR => Err(callback_error( &error, "Java Arrow batch source failed".into(), )), EMPTY_BATCH => import_batch(array, schema, 0), rows if rows > 0 => import_batch(array, schema, rows as usize), other => Err(DataFusionError::Execution(format!( "unexpected Arrow batch callback status {other}" ))), } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 413f9f7

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Only cancel native query when started In the setup failure path, NativeBridge.cancelQuery(taskId) is invoked even when the native execution never started (e.g. when registerArrowBatchSourceProvider or executeLocalPlan itself threw). Guard the cancellation so it runs only after output was successfully created, otherwise cancellation is dispatched for a task that has no active native query, which can produce spurious errors or interfere with unrelated tasks sharing the ID space. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [85-95] -if (taskId != 0L) { - closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); -} if (output != null) { closeAfterFailure(throwable, output::close); + if (taskId != 0L) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); + } } if (session != null) { closeAfterFailure(throwable, session::close); } closeAfterFailure(throwable, registration::close); throw throwable; Suggestion importance[1-10]: 6 __ Why: Valid observation that calling NativeBridge.cancelQuery(taskId) when the native query was never established (e.g., if registerArrowBatchSourceProvider failed) could produce spurious effects. The suggested reorder is a reasonable defensive improvement, though impact is moderate.	Low
General	Avoid blocking cancellation via synchronization exportNextBatch is synchronized on the SourceEntry, and close() is also synchronized. A cooperative cancel() invoked while exportNextBatch is running acquires the same monitor via close() transitively (through cancel -> source.cancel() under cancellationLock, which is fine), but calls to close() from the Registration/tearDown path will block until nextBatch returns — defeating cooperative cancellation. Consider not synchronizing exportNextBatch on the entry (single-consumer contract) or using a non-blocking cancellation path so close() can proceed. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [392-408] -private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { +private long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { if (closed) { throw new IllegalStateException("Arrow batch source is closed"); } VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } try { ArrowArray array = ArrowArray.wrap(arrayPointer.address()); ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); } finally { root.close(); } } Suggestion importance[1-10]: 5 __ Why: Points to a potentially valid concern where synchronized close() on SourceEntry could block behind exportNextBatch, undermining cooperative cancellation. However, the improved_code is essentially identical to existing_code (only comment/wording), lowering the concrete value of the suggestion.	Low
	Guard packed key width invariant The packed key uses only 20 bits for the row index and the remaining 44 bits for the ordinal, but BATCH_SIZE is 65_536 which requires 17 bits for row and works today. However, if BATCH_SIZE is ever increased beyond 2^20 the row index will silently collide with the ordinal bits, corrupting decoded output. Add an assertion or derive the shift from BATCH_SIZE to make this invariant explicit and safe against future changes. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java [265-275] private void decodeSorted(SortedColumn column, FieldVector vector, int size) throws IOException { + assert BATCH_SIZE <= (1 << 20) : "row index must fit in 20 bits"; SortedDocValues values = column.values(); for (int i = 0; i < size; i++) { ordScratch[i] = values.advanceExact(docs[i]) ? values.ordValue() : -1; } allocateBytes(vector, size); for (int i = 0; i < size; i++) { ordRowScratch[i] = ((long) (ordScratch[i] + 1) << 20) | i; } Suggestion importance[1-10]: 5 __ Why: Reasonable defensive assertion for a fragile bit-packing invariant tied to BATCH_SIZE. Currently correct, but future-proofing has moderate value.	Low
	Use idiomatic boolean negation swap returns the previous value; when the previous value is false the handle has not yet been cancelled, so the condition should be !self.cancelled.swap(...). The current == false comparison is correct in Rust but non-idiomatic and easy to misread; more importantly ensure the intent (invoke callback exactly once) is preserved. Consider using if !self.cancelled.swap(true, Ordering::AcqRel) for clarity. sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs [145-151] pub fn cancel(&self) { - if self.cancelled.swap(true, Ordering::AcqRel) == false { + if !self.cancelled.swap(true, Ordering::AcqRel) { if let Some(callback) = cancel_callback() { unsafe { callback(self.binding_id, self.source_key) }; } } } Suggestion importance[1-10]: 2 __ Why: Purely stylistic idiomatic change; existing code is functionally correct. Minimal impact.	Low

---

Previous suggestions

Suggestions up to commit baa9a18

Category	Suggestion	Impact
Possible issue	Fix resource teardown ordering on setup failure The failure cleanup order calls NativeBridge.cancelQuery(taskId) before closing the native output stream and session. However, cancelQuery may cause native code to still be processing while session::close frees the LocalSession (which owns the registered batch source provider). Close output and session first (they own native resources depending on the binding) and close registration last so Java callbacks remain servable while native teardown drains. Also, closing registration before native cancellation completes can cause native calls into a closed binding. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [85-94] +if (output != null) { + closeAfterFailure(throwable, output::close); +} if (taskId != 0L) { closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); -} -if (output != null) { - closeAfterFailure(throwable, output::close); } if (session != null) { closeAfterFailure(throwable, session::close); } closeAfterFailure(throwable, registration::close); Suggestion importance[1-10]: 5 __ Why: The suggestion raises a plausible concern about teardown ordering between cancelQuery, output::close, and session::close, but the reasoning is speculative and the current order may well be intentional. Moderate impact if the concern is real.	Low
General	Validate projection length range explicitly Math.toIntExact throws ArithmeticException (not caught cleanly with a helpful message) for values that overflow int, and never returns a negative value from a non-negative input. If native passes a negative projectionLength, toIntExact yields a negative int and the check works, but if it passes a huge unsigned-like value cast from Rust's i64, the ArithmeticException propagates to the catch-all and gets stringified. Validate the range explicitly with a clear message before converting. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [104-107] -int length = Math.toIntExact(projectionLength); -if (length < 0) { - throw new IllegalArgumentException("projection length must be non-negative"); +if (projectionLength < 0L || projectionLength > Integer.MAX_VALUE) { + throw new IllegalArgumentException("projection length out of range: " + projectionLength); } +int length = (int) projectionLength; Suggestion importance[1-10]: 3 __ Why: Minor improvement in error message clarity when projection length overflows int; Math.toIntExact already throws which is caught by the surrounding handler, so functional impact is small.	Low
	Use idiomatic boolean negation The expression self.cancelled.swap(true, Ordering::AcqRel) == false is Clippy-lint-worthy and, more importantly, swap returns the previous value: if it was already true, we skip; if it was false, we proceed. The current logic is correct, but idiomatic style using ! improves readability and avoids a subtle inversion bug if refactored. Consider using if !self.cancelled.swap(true, Ordering::AcqRel). sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs [146-151] pub fn cancel(&self) { - if self.cancelled.swap(true, Ordering::AcqRel) == false { + if !self.cancelled.swap(true, Ordering::AcqRel) { if let Some(callback) = cancel_callback() { unsafe { callback(self.binding_id, self.source_key) }; } } } Suggestion importance[1-10]: 2 __ Why: Purely a stylistic Clippy-lint improvement; the current code is functionally correct. Very low impact.	Low
	Explicitly drop in-flight batch on cancellation When cancellation was requested while a batch was in-flight, discarding a successful Ok(Ok(Some(batch))) result leaks the imported Arrow buffers (their FFI-imported memory stays charged to the Java allocator until the source drops). Explicitly drop the batch or forward it before signaling cancellation, or ensure the pending task's returned RecordBatch is dropped here (which Rust does when result goes out of scope, but the concern is the batch's imported buffers must be released — verify FFI buffers are released on drop). Consider forcing an explicit drop and adding a comment. sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/table_provider.rs [269-273] Ok(result) => { self.pending = None; if self.cancellation_requested { + // Drop any imported batch to release FFI buffers before cancelling. + drop(result); return self.cancelled(); } Suggestion importance[1-10]: 2 __ Why: The suggested explicit drop(result) is redundant since Rust already drops result when it goes out of scope, and FFI buffer release happens on drop regardless. The suggestion adds no functional benefit.	Low

Suggestions up to commit e6640fd

Category	Suggestion	Impact
Possible issue	Fix native-resource teardown order The failure-cleanup order closes output before session, but output is a stream produced by executeLocalPlan on that session's context; closing the session first (or in the wrong order relative to the native stream) commonly triggers use-after-free in native code. Standard teardown is stream first, then session, then registration — which matches what OwnedResultStream.close already does. Also, calling cancelQuery before closing the stream can race with the stream's own poll; canceling after the stream closes is safer and mirrors the success path. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [81-96] } catch (RuntimeException | Error throwable) { if (shardTask != null) { shardTask.clearCancellationListener(); - } - if (taskId != 0L) { - closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); } if (output != null) { closeAfterFailure(throwable, output::close); } if (session != null) { closeAfterFailure(throwable, session::close); } + if (taskId != 0L) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); + } closeAfterFailure(throwable, registration::close); throw throwable; } Suggestion importance[1-10]: 6 __ Why: The existing code already closes output before session (matching the suggestion's stream-first order), so part of the reasoning is inaccurate. However, the point about moving cancelQuery after stream close to avoid races with the native poll is potentially valid and worth considering.	Low
	Add version guards to Writeable payload ArrowBatchSourcePlan implements Writeable and is intended to travel on the wire, but the read/write pair has no version guard. If the record's layout (fields, encodings, or ColumnKind enum) ever changes, mixed-version clusters will corrupt the stream. Even for the initial introduction, add a Version guard on both writeTo and the StreamInput constructor so peers before(V) never see this payload (or gate it with a symmetric onOrAfter(V) check on both sides). sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourcePlan.java [44-46] public ArrowBatchSourcePlan(StreamInput input) throws IOException { + // TODO: guard with in.getVersion().onOrAfter(V_INTRODUCED) once the target version constant is defined this(input.readString(), input.readByteArray(), input.readList(ArrowBatchSourcePlan::readInputColumn)); } Suggestion importance[1-10]: 4 __ Why: Version guards for Writeable payloads can be relevant for BWC in mixed-version clusters, but the suggestion only proposes a TODO comment without a concrete change, and this is in a sandbox module where BWC guarantees may not yet apply.	Low
General	Validate projection length before narrowing Math.toIntExact throws ArithmeticException for values that don't fit in an int (including negatives like large longs that wrap), so the following length < 0 check is unreachable while overflow surfaces as an unclear exception. Validate the raw projectionLength bounds before narrowing, so oversized/negative inputs produce a clear error. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [104-107] -int length = Math.toIntExact(projectionLength); -if (length < 0) { - throw new IllegalArgumentException("projection length must be non-negative"); +if (projectionLength < 0L || projectionLength > Integer.MAX_VALUE) { + throw new IllegalArgumentException("projection length out of range: " + projectionLength); } +int length = (int) projectionLength; Suggestion importance[1-10]: 5 __ Why: The observation is correct: Math.toIntExact throws ArithmeticException making the length < 0 check unreachable and producing less clear errors. The fix is a minor robustness/clarity improvement.	Low
	Ensure paired cancel/release callback availability Drop calls self.cancel() and then unconditionally calls the release callback. Because cancel uses swap(true), a previous external cancel() call sets cancelled to true and the Drop-time cancel becomes a no-op — but the release callback still fires, which is correct. However, if create returned Err you never construct the handle, so release is not called; that's fine. The real hazard is that cancel_callback() may be None (callbacks not installed) while release_callback() is Some, or vice-versa, leaving Java state inconsistent. Validate both callbacks are present at construction (they are registered atomically) or assert on their pairing to avoid partial teardown. sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs [145-151] pub fn cancel(&self) { if self.cancelled.swap(true, Ordering::AcqRel) == false { if let Some(callback) = cancel_callback() { unsafe { callback(self.binding_id, self.source_key) }; + } else { + // Callbacks were unregistered after creation; nothing to cancel on the Java side. } } } Suggestion importance[1-10]: 2 __ Why: The improved_code only adds a comment on the else branch and doesn't actually implement the suggested pairing validation. Since callbacks are registered together atomically, the concern is largely theoretical.	Low

Suggestions up to commit e6640fd

Category	Suggestion	Impact
General	Validate projection length before narrowing cast Math.toIntExact(projectionLength) on a negative long returns a negative int without throwing (it only throws on overflow), so the subsequent length < 0 check runs on a negative value — but then (long) length * Integer.BYTES proceeds if the caller passes a projection length whose int cast is negative. More importantly, toIntExact throws ArithmeticException on overflow, which is caught by the broad Throwable handler and reported as an opaque error; validate before casting so the error message is actionable. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [104-107] - int length = Math.toIntExact(projectionLength); - if (length < 0) { - throw new IllegalArgumentException("projection length must be non-negative"); + if (projectionLength < 0 || projectionLength > Integer.MAX_VALUE) { + throw new IllegalArgumentException("projection length out of range: " + projectionLength); } + int length = (int) projectionLength; Suggestion importance[1-10]: 6 __ Why: Valid observation: Math.toIntExact throws on overflow (caught by broad handler with opaque message), and the length < 0 check is unreachable since toIntExact on negative long returns negative but only overflow throws. Explicit validation provides clearer errors.	Low
	Check cancellation before exporting batch exportNextBatch checks closed at entry but not after source.nextBatch() returns. If close() runs concurrently (it also synchronizes on this, so this is serialized, but cancel() may have been called and then close() runs after nextBatch returns), the exported batch could still be handed to native code after the source was closed. Additionally, check the cancellationRequested flag before exporting to avoid exporting a stale batch after cancellation. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [392-399] private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { if (closed) { throw new IllegalStateException("Arrow batch source is closed"); + } + if (cancellationRequested.get()) { + throw new TaskCancelledException("Arrow batch source cancelled"); } VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } Suggestion importance[1-10]: 5 __ Why: Adding an early cancellation check could reduce wasted work, but exportNextBatch is already synchronized and the native side handles CANCELLED status. The improvement is minor and defensive.	Low
	Reorder native resource cleanup on failure path On the failure path, NativeBridge.cancelQuery(taskId) is invoked before output and session are closed, but taskId is derived from task.getId() and is not tied to the native session's lifetime. Cancellation is instead ordered before native resources are freed, and if executeLocalPlan failed after registerArrowBatchSourceProvider, cancelling by taskId may target the wrong/no query. Reorder to close native resources first (output, session) and only invoke cancelQuery when the plan actually started, to avoid spurious cancellations and clarify ownership. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [85-95] - if (taskId != 0L) { - closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); - } if (output != null) { closeAfterFailure(throwable, output::close); } if (session != null) { closeAfterFailure(throwable, session::close); } + if (taskId != 0L && output != null) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); + } closeAfterFailure(throwable, registration::close); throw throwable; Suggestion importance[1-10]: 4 __ Why: The reordering argument is plausible but not clearly correct—cancelling before closing native output/session may actually be intentional to unblock the stream. The suggestion's rationale about taskId targeting the wrong query is speculative.	Low
	Document release invariant for source handle Drop invokes release_source unconditionally, but if ArrowBatchSourceHandle::create failed after Java allocated a sourceKey, that key would leak because the Rust handle was never constructed. That path is safe here (Java returns -1 on failure and doesn't allocate a key), but if create ever succeeds and something later panics inside the constructor logic, no drop runs. Consider documenting or asserting that create only returns Ok after source_key is set, and ensure Java never returns a positive key on the error path. sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs [145-161] pub fn cancel(&self) { if self.cancelled.swap(true, Ordering::AcqRel) == false { if let Some(callback) = cancel_callback() { unsafe { callback(self.binding_id, self.source_key) }; } } } } impl Drop for ArrowBatchSourceHandle { fn drop(&mut self) { + // Invariant: source_key is only set when Java successfully allocated it; release is safe. self.cancel(); if let Some(callback) = release_callback() { unsafe { callback(self.binding_id, self.source_key) }; } } } Suggestion importance[1-10]: 2 __ Why: The suggestion only adds a comment documenting an existing invariant and acknowledges the current code is safe. Minimal impact.	Low

Suggestions up to commit 4cb89e2

Category	Suggestion	Impact
Possible issue	Avoid double-free after Arrow C Data export Data.exportVectorSchemaRoot transfers ownership of Arrow buffers via C Data Interface; the release callback embedded in FFI_ArrowArray is responsible for freeing. Calling root.close() in the finally block immediately after export may double-release the underlying buffers, because export typically retains references and expects the importer to invoke the release callback. Verify Arrow Java's export semantics — typically the caller should NOT close root after exportVectorSchemaRoot since ownership transfers to the C Data consumer. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [392-408] private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { if (closed) { throw new IllegalStateException("Arrow batch source is closed"); } VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } - try { - ArrowArray array = ArrowArray.wrap(arrayPointer.address()); - ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); - Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); - return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); - } finally { - root.close(); - } + ArrowArray array = ArrowArray.wrap(arrayPointer.address()); + ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); + Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); + // Ownership of buffers transfers to the C Data importer via the release callback. + return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); } Suggestion importance[1-10]: 6 __ Why: Potentially valid concern about Arrow C Data ownership semantics; however, exportVectorSchemaRoot in Arrow Java typically retains buffers via the release callback, so closing the root on the Java side is standard and does not double-free. The suggestion is speculative but points to a subtle area worth verification.	Low
General	Assert batch size fits the packed row bits The 20-bit shift limits row indices to values < 2^20 (1,048,576) and ordinals to < 2^43 - 1. BATCH_SIZE is 65,536 which fits in 20 bits, but sorted-set ordinal counts can exceed 2^20 in large segments — wait, this is packing row | ord, and row is bounded by BATCH_SIZE so 20 bits (1M) is sufficient. However, the ordinal is stored in bits 20..63 (44 bits), which is fine. The & 0xFFFFF mask correctly extracts the row. This works, but the mask magic-number is fragile: if BATCH_SIZE ever grows past 2^20, extraction will silently corrupt row values. Consider deriving the mask from BATCH_SIZE or asserting the invariant. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java [270-280] allocateBytes(vector, size); +assert BATCH_SIZE <= (1 << 20) : "row index must fit in 20 bits"; for (int i = 0; i < size; i++) { ordRowScratch[i] = ((long) (ordScratch[i] + 1) << 20) | i; } Arrays.sort(ordRowScratch, 0, size); int lastOrd = Integer.MIN_VALUE; BytesRef term = null; for (int i = 0; i < size; i++) { int row = (int) (ordRowScratch[i] & 0xFFFFF); int ord = (int) (ordRowScratch[i] >>> 20) - 1; Suggestion importance[1-10]: 4 __ Why: Adding an assertion for the packed-bit invariant is a reasonable maintainability improvement to prevent silent corruption if BATCH_SIZE grows, but it's a minor safeguard.	Low
	Guard against zero-address projection pointer Math.toIntExact on projectionLength will only throw for values outside int range, but the subsequent length < 0 check is unreachable for negative long inputs that fit an int — actually negative values do get caught, but very large values that overflow to negative after cast do not, since toIntExact throws first. However, a projectionPointer allocated via MemorySegment.ofArray(int[]) in tests has an exact byteSize, but a native-supplied raw address will have byteSize()==0, so asSlice is never taken. Consider always using reinterpret for pointers whose byte size is unknown, or explicitly checking projectionPointer.address() != 0 to avoid dereferencing NULL. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [104-121] int length = Math.toIntExact(projectionLength); if (length < 0) { throw new IllegalArgumentException("projection length must be non-negative"); } int[] projection; if (length == 0) { projection = new int[0]; } else { - if (projectionPointer == null || projectionPointer.equals(MemorySegment.NULL)) { + if (projectionPointer == null + || projectionPointer.equals(MemorySegment.NULL) + || projectionPointer.address() == 0L) { throw new IllegalArgumentException("projection pointer is null"); } long byteLength = (long) length * Integer.BYTES; MemorySegment projectionView = projectionPointer.byteSize() >= byteLength ? projectionPointer.asSlice(0L, byteLength) : projectionPointer.reinterpret(byteLength); projection = projectionView.toArray(ValueLayout.JAVA_INT); } return lease.binding.open(projection); Suggestion importance[1-10]: 3 __ Why: Minor defensive hardening; MemorySegment.NULL check likely already covers the zero-address case, so the additional address() == 0L check is marginal.	Low
	Reorder failure cleanup to close native first On setup failure, session::close is invoked while registration::close still holds the callback binding. The native session teardown may attempt to release Arrow batch sources through the Java upcalls, but ordering here closes the session before the registration, which is fine — however, calling NativeBridge.cancelQuery and output.close() before session.close() may race with the native side still holding session references. Consider closing output first, then session, then invoking cancelQuery only after session teardown, to ensure native cleanup completes before Java-side binding is unregistered. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [85-94] -if (taskId != 0L) { - closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); -} if (output != null) { closeAfterFailure(throwable, output::close); } if (session != null) { closeAfterFailure(throwable, session::close); } +if (taskId != 0L) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); +} closeAfterFailure(throwable, registration::close); Suggestion importance[1-10]: 3 __ Why: Speculative reordering; the suggestion admits the original ordering "is fine" and provides no concrete evidence of a race, making the impact unclear.	Low

Suggestions up to commit 1ef5c2f

Category	Suggestion	Impact
Possible issue	Fix packed ordinal overflow risk The packed ordinal/row encoding uses only 20 bits for the row index (& 0xFFFFF masks to 20 bits, max 1,048,575), but BATCH_SIZE is 65,536 which fits. However, only 20 bits for the ordinal is unsafe: term dictionaries commonly exceed 1M unique values, and (ordScratch[i] + 1) << 20 can overflow into the sign bit or lose high bits, producing wrong term lookups. Widen the ord field (e.g., pack row in the low 20 bits and ord in the high 44 bits) or use two separate parallel arrays. sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java [270-280] allocateBytes(vector, size); for (int i = 0; i < size; i++) { - ordRowScratch[i] = ((long) (ordScratch[i] + 1) << 20) | i; + ordRowScratch[i] = (((long) ordScratch[i] + 1L) << 20) | i; } +// NOTE: pack ord in high 44 bits, row in low 20 bits; ord range must fit in 44 bits. Arrays.sort(ordRowScratch, 0, size); int lastOrd = Integer.MIN_VALUE; BytesRef term = null; for (int i = 0; i < size; i++) { - int row = (int) (ordRowScratch[i] & 0xFFFFF); - int ord = (int) (ordRowScratch[i] >>> 20) - 1; + int row = (int) (ordRowScratch[i] & 0xFFFFFL); + int ord = (int) ((ordRowScratch[i] >>> 20) - 1L); Suggestion importance[1-10]: 8 __ Why: Legitimate correctness bug: packing (ord + 1) << 20 in a long with only 20 bits reserved for row is fine for row, but ordinals >= 2^43 would overflow. More critically, 20 bits for row (max 1,048,575) fits BATCH_SIZE but the ordinal shift uses only ~44 bits which is generally safe; however the analysis about potential issues with large term dictionaries is valid and worth addressing.	Medium
	Fix failure cleanup ordering and guard The failure cleanup closes output (native stream) before session, but the native stream is created from the session and may reference session-owned resources; closing the session first can also leave the stream orphaned. More importantly, output is only assigned after executeLocalPlan succeeds, so on failure of registerArrowBatchSourceProvider the session is closed but the just-registered provider binding is torn down via registration.close() afterward — reversing the natural stack order. Close output first, then session, then registration, and only cancel the query when output was actually created (the current taskId != 0 check always passes when task != null). sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java [81-96] } catch (RuntimeException | Error throwable) { if (shardTask != null) { shardTask.clearCancellationListener(); } - if (taskId != 0L) { - closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); - } if (output != null) { + if (taskId != 0L) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); + } closeAfterFailure(throwable, output::close); } if (session != null) { closeAfterFailure(throwable, session::close); } closeAfterFailure(throwable, registration::close); throw throwable; } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that output is only assigned after executeLocalPlan, so cancelQuery may be called before any native stream exists. However, the existing code already handles all failure paths safely with suppressed exceptions, so the impact is moderate ordering/hygiene improvement.	Low
General	Allow cancel to preempt in-flight pulls exportNextBatch is synchronized, so a concurrent cancel() on the same SourceEntry (which also acquires the monitor via its own synchronized close()) cannot preempt an in-progress pull. Since the class documents cancel() as being able to run concurrently with nextBatch(), guard the pull with a non-monitor mechanism (e.g., an AtomicBoolean inFlight or the existing cancellationLock) so a cancel request does not have to wait for nextBatch to complete before invoking source.cancel(). sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [392-408] -private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { - if (closed) { - throw new IllegalStateException("Arrow batch source is closed"); +private long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { + synchronized (cancellationLock) { + if (closed) { + throw new IllegalStateException("Arrow batch source is closed"); + } } VectorSchemaRoot root = source.nextBatch(); if (root == null) { return 0L; } try { ArrowArray array = ArrowArray.wrap(arrayPointer.address()); ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); } finally { root.close(); } } Suggestion importance[1-10]: 6 __ Why: Valid concurrency concern: exportNextBatch being synchronized combined with close() being synchronized means cancel() invoked via close() blocks until the pull completes, contradicting the documented concurrent cancellation contract. The proposed fix is reasonable though partial.	Low
	Avoid truncating UTF-8 mid-character When capacity is 1, length becomes 0 and the null terminator is written at index 0, which is fine, but when bytes.length fits within capacity - 1, length truncates otherwise valid UTF-8 mid-character on the boundary. Consider trimming to the nearest UTF-8 code-point boundary before writing so downstream native readers do not decode a malformed trailing byte sequence. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java [198-201] -int length = Math.toIntExact(Math.min(bytes.length, capacity - 1L)); +int max = Math.toIntExact(Math.min(bytes.length, capacity - 1L)); +int length = max; +while (length > 0 && (bytes[length] & 0xC0) == 0x80) { + length--; +} MemorySegment view = output.reinterpret(capacity); MemorySegment.copy(bytes, 0, view, ValueLayout.JAVA_BYTE, 0L, length); view.set(ValueLayout.JAVA_BYTE, length, (byte) 0); Suggestion importance[1-10]: 3 __ Why: Minor cosmetic improvement for error messages only; a malformed trailing byte in a diagnostic error string has negligible impact.	Low

[github-actions]

Persistent review updated to latest commit 1ef5c2f

[github-actions]

Persistent review updated to latest commit 1ef5c2f

[github-actions]

✅ Gradle check result for 1ef5c2f: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.48%. Comparing base (150dc57) to head (413f9f7).

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22844      +/-   ##
============================================
- Coverage     71.54%   71.48%   -0.06%     
+ Complexity    77274    77234      -40     
============================================
  Files          6170     6170              
  Lines        359771   359771              
  Branches      52478    52478              
============================================
- Hits         257399   257185     -214     
- Misses        81854    82110     +256     
+ Partials      20518    20476      -42     

☔ View full report in Codecov by Harness.
📢 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 4cb89e2

[github-actions]

✅ Gradle check result for 4cb89e2: SUCCESS

[github-actions]

Persistent review updated to latest commit e6640fd

[github-actions]

Persistent review updated to latest commit e6640fd

[github-actions]

Persistent review updated to latest commit baa9a18

[github-actions]

✅ Gradle check result for baa9a18: SUCCESS

[github-actions]

Persistent review updated to latest commit 413f9f7

[github-actions]

✅ Gradle check result for 413f9f7: SUCCESS