Skip to content

fix: return errors instead of panicking on malformed JNI arguments - #7186

Merged
viirya merged 1 commit into
apache:mainfrom
viirya:jni-arg-errors
Oct 24, 2026
Merged

viirya merged 1 commit into
apache:mainfrom
viirya:jni-arg-errors

Conversation

@viirya

@viirya viirya commented Oct 23, 2026 •

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #7183.

Rationale for this change

Some JNI entry points convert their arguments with unwrap() / expect(), so a malformed argument panics. try_unwrap_or_throw turns that panic into a CometNativeException, but the exception carries only the panic message, such as called `Result::unwrap()` on an `Err` value: ..., without saying which argument was bad. These conversions should return regular errors, the same way the other arguments of these entry points already do.

What changes are included in this PR?

All changes are in native/core/src/execution/jni_api.rs. Each conversion now returns a CometError::Internal through ?, so the exception class stays CometNativeException and only the message changes:

  • convert_datatype_arrays, used by columnarToRowInit, writeSortedFileNative and writeSortedPartitionsNative, now decodes each type with a new deserialize_arrow_datatype helper. When decoding fails, the error says which column it was, for example Cannot deserialize the data type of column 2: ....
  • The file_path and compression_codec arguments of writeSortedFileNative, and the name arguments of traceBegin, traceEnd and logMemoryUsage, are now read with a new string_arg helper. When a string cannot be read, for example because it is null, the error names the argument and the entry point, for example Cannot read argument 'event' of traceBegin: ....
  • get_execution_context now returns an error for a null handle (Native execution context handle is null). Its callers are executePlan, getShufflePartitionOffsets and setShufflePartitionPusher. setShufflePartitionPusher already rejected a zero handle with a NullPointerException before reaching this check, and it still does.

A data type that decodes but is invalid, such as an unknown type id, can still panic in serde::to_arrow_datatype. That function is shared with the planner and has many callers, so making it fallible is left out of this PR.

How are these changes tested?

New Rust unit tests in jni_api.rs, none of which need a JVM:

  • deserialize_arrow_datatype_decodes_a_serialized_type
  • deserialize_arrow_datatype_rejects_malformed_bytes
  • get_execution_context_rejects_a_null_handle

The two error tests also check that the error is CometError::Internal, which maps to CometNativeException.

The string conversions need a JVM to call, so no new tests cover them. They are covered by the existing tests of CometNativeColumnarToRowSuite and CometShuffleSuite.

Several JNI entry points converted their arguments with unwrap() or
expect(), so a malformed argument surfaced as a CometNativeException
carrying only the panic message. Return a CometError::Internal that
names the bad argument instead:

- convert_datatype_arrays reports which column's data type failed to
  deserialize.
- writeSortedFileNative, traceBegin, traceEnd and logMemoryUsage report
  which string argument could not be read.
- get_execution_context rejects a null handle.

The exception class stays CometNativeException.
@github-actions github-actions Bot added the bug Something isn't working label Oct 23, 2026

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed full head 35e2e2bf6461003c41de2dfa2951e32c8f0ec76b. No introduced P1/P2 issues found within this review.

The PR is non-draft. The full PR diff against supplied base 69111f68a5ff62904342910c445d191f53615aff, through merge-base 785fdbcb4cbf9057dfcfae3bd9dd45a77bb505bb, contains one commit changing native/core/src/execution/jni_api.rs. The base-only memory-management commit is not a PR-introduced reversal. There are no stacked prerequisites. The supplied snapshot and live discussion checks contained no reviews, issue comments, inline comments, or review threads.

Routed skills: review-comet-pr, review-comet-ffi-pr, and review-comet-shuffle-pr. Read AGENTS.md and the relevant contributor documentation.

Summary

  • Prior state and problem: Malformed serialized types, unreadable strings, and null execution handles reached unwrap() or expect(). The JNI boundary caught these panics, but diagnostics lacked argument context and invoked panic logging.
  • Design approach: Replace those panic paths with fallible helpers and propagate CometError::Internal through the existing JNI exception boundary.
  • Correctness: Checked every changed caller, exception mapping, and early-return cleanup. Valid inputs retain the same conversions. JNI array guards still drop on errors, and decoding fails before the affected writer or converter performs its work. No expression evaluation, partitioning, or row-format semantics change.
  • Compatibility analysis: JNI signatures and protobuf encoding remain unchanged. Both old panic errors and new internal errors map to CometNativeException. Checked Spark's ShuffleWriteProcessor sources at configured versions 3.4.3, 3.5.9, 4.0.4, 4.1.3, and 4.2.0: each stops the writer on an exception and rethrows it without depending on these diagnostic messages. The affected file on latest release branch branch-1.2, at 889856dc5af821de1f7c6ca2e9538bd663431d04, is byte-identical to the PR merge-base, so the release-relative changes are exactly this intended diagnostic change.
  • Key design decisions: Retaining the established exception class avoids changing JVM failure classification. setShufflePartitionPusher preserves its earlier explicit NullPointerException check. Semantic validation inside to_arrow_datatype remains unchanged, including its existing panic cases.
  • Implementation sketch: deserialize_arrow_datatype adds the column index to decode errors. string_arg identifies the export and argument. get_execution_context returns a result, and all three callers propagate it. Three Rust tests cover successful decoding, malformed bytes, and a null handle.
  • Performance: Error-message formatting occurs only on failure. Successful calls retain the existing decoding, string conversion, and pointer check without additional batch copies, buffering, or allocations for diagnostics. Avoiding unwinding and panic-backtrace logging improves the malformed-input path structurally, but no throughput claim was benchmarked.
  • Design: The change is localized and preserves existing ownership, resource lifecycle, and shuffle contracts. The current FFI and shuffle documentation remains accurate.
  • Abstraction & complexity: The two small conversion helpers centralize repeated diagnostics without introducing a broader framework. Returning CometResult from the datatype-array helper makes error propagation straightforward.
  • Behavioral changes worth calling out: The covered malformed arguments now produce contextual error messages through ordinary error propagation instead of panic handling. Valid query results and shuffle output remain unchanged.
  • Suggested improvements: No corrective change meeting the P1/P2 reporting bar was identified.

Exact-head CI: title checks, labeling, CodeQL, and Analyze Actions passed. Comet CI remained in Preflight. No failed check was reported, but there was no completed Rust build/test or Spark-suite verdict.

Validation: git diff --check passed. Two disposable harness tests passed using the unchanged production string and handle helpers, jni 0.22.4, and OpenJDK 21 with -Xcheck:jni. They exercised null-string diagnostics for all five call sites, empty/Unicode/embedded-NUL strings, null handles, and valid pointer identity. The harness used minimal error/context scaffolding and does not replace Comet integration tests. The focused repository test could not start because the offline registry lacked locked dependency aws-config 1.13.0. Consequently, the added repository tests and end-to-end Comet/Spark suites were not run locally. The project checkout remains unchanged.

@viirya
viirya added this pull request to the merge queue Oct 24, 2026
Merged via the queue into apache:main with commit 8263260 Oct 24, 2026
45 checks passed
@viirya
viirya deleted the jni-arg-errors branch October 24, 2026 02:15
@viirya

viirya commented Oct 24, 2026

Copy link
Copy Markdown
Member Author

Thanks @sunchao

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Malformed JNI arguments panic instead of returning an error

2 participants