Add Java Bindings for Hybrid Scan Parquet Reader - #22456
Conversation
* Add `HybridScanReader`, a Java binding for
`cudf::io::parquet::experimental::hybrid_scan_reader`, targeting
highly-selective filter expressions over Parquet files.
* Expose the full two-phase pipeline: filter columns are materialized
first and an AST filter is applied to produce a boolean row mask;
payload columns are then read only for the rows that survive,
avoiding unnecessary I/O and decompression.
* Support both single-shot and chunked/streaming materialization. In
the chunked filter path, the C++ wrapper holds the row mask column
across chunk calls and mutates it in place; `takeFilterRowMask()`
transfers ownership to Java without copying.
* Add `ColumnNameReference`, a new AST node that references columns by
name rather than index. This is required because filter expressions
are constructed before the Parquet schema is resolved; the reader
resolves names to positions internally.
* Introduce a new `@Experimental` annotation and apply it to the
entire new surface area to signal that these APIs may change as the
upstream C++ experimental API stabilizes.
* Use `MemoryCleaner` for native resource management, consistent with
the rest of cudf-java; leaked readers are logged with their handle
address.
* Fix `pom.xml` resource configuration to allow the JDT language
server to correctly import the Java project in VS Code / Cursor.
The LICENSE resource was declared with `${basedir}/..` as its
source directory, which caused the JDT LS to fail to recognize
the project entirely. Replace it with an explicit `copy-license`
plugin execution that stages only the LICENSE file under
`target/license-resources/` before the resource directory is read.
* Add `HybridScanReaderTest` covering the full pipeline end-to-end:
single-shot and chunked filter + payload materialization, row-group
pruning via stats and dictionary pages, page-index-seeded row masks,
pass construction, and leak detection.
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds an experimental HybridScanReader JNI-backed Parquet hybrid-scan API, supporting page-index and dictionary-based row-group pruning, byte-range planning, single-shot and chunked materialization. Also adds AST column-name references, value types, native JNI helpers, build changes, and comprehensive tests. ChangesHybridScanReader Feature Implementation
🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
java/src/main/native/src/jni_compiled_expr.hpp (1)
50-55: ⚡ Quick winAdd a doxygen comment for the new helper.
add_column_name_reffollows the same public helper pattern as neighboring methods; documenting it keeps generated native docs consistent.Suggested diff
+ /** + * Add a column-name reference node and return a reference to the stored node. + */ cudf::ast::column_name_reference& add_column_name_ref( std::unique_ptr<cudf::ast::column_name_reference> ref_ptr) {As per coding guidelines, "Use doxygen for documentation generation and as a documentation linter on C++/CUDA code".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/native/src/jni_compiled_expr.hpp` around lines 50 - 55, Add a Doxygen comment block above the helper function add_column_name_ref describing its purpose, parameters and return value: explain that it appends a std::unique_ptr<cudf::ast::column_name_reference> to the expressions container (transferring ownership), document the ref_ptr parameter and that the function returns a reference to the stored cudf::ast::column_name_reference, and note any lifetime/ownership expectations regarding the expressions vector; keep the style consistent with neighboring helper methods' comments.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@java/src/main/java/ai/rapids/cudf/MemoryCleaner.java`:
- Around line 373-375: The register method currently adds a CleanerWeakReference
for HybridScanReader with the RMM-blocker flag set to false; change the call in
register(HybridScanReader reader, Cleaner cleaner) so the CleanerWeakReference
is constructed with the RMM-blocker boolean set to true (i.e., replace the final
argument false with true) so leaked HybridScanReader instances are treated as
RMM blockers and will be processed by cleanAllRmmBlockers(); update any related
comments if present in MemoryCleaner and ensure the call remains
all.put(cleaner.id, new CleanerWeakReference(reader, cleaner, collected, true)).
In `@java/src/main/native/src/HybridScanReaderJniInternal.cpp`:
- Around line 155-168: The make_device_spans function must validate j_addrs and
j_lens lengths and values before constructing device_span objects: check that
addrs.size() == lens.size(), verify each lens[i] is non-negative and reasonably
bounded before casting to size_t, and if any check fails release/cancel
native_jlongArray resources (addrs.cancel(), lens.cancel()) and throw a JNI
exception instead of proceeding; ensure these checks occur before the loop that
calls out.emplace_back(reinterpret_cast<uint8_t const*>(addrs[i]),
static_cast<size_t>(lens[i])) so you never read past the j_lens array or convert
a negative length into a huge size_t.
- Around line 109-112: The HybridScanReader currently only copies
filter.getNativeHandle() into native code (using filter_handle → compiled_expr)
and does not keep a Java-side reference to the CompiledExpression, which allows
its cleaner to free the AST and leaves parquet_reader_options (in
hybrid_scan_reader_wrapper) with a dangling pointer; fix this by retaining
ownership of the CompiledExpression: either add a CompiledExpression field to
the Java HybridScanReader and set it to the provided filter (so the Java object
retains a strong reference until close), or transfer ownership into the native
wrapper by storing a persistent/native-owning handle in
hybrid_scan_reader_wrapper and ensuring it is freed when the reader is closed;
ensure the code paths that call set_filter(filter_expr->get_top_expression())
use the retained/owned compiled_expr and that the retained object is
closed/freed together with HybridScanReader (constructor and close/destructor
paths).
In `@java/src/main/native/src/HybridScanReaderJniMaterialize.cpp`:
- Around line 173-174: The native JNI code currently casts jlong parameters like
chunk_read_limit and pass_read_limit directly to size_t in
wrapper->reader->setup_chunking_for_filter_columns (and similarly in the
payload/all-columns setup calls and construct_row_group_passes), which converts
negative inputs into huge unsigned values; add an explicit checked conversion:
validate each incoming jlong (e.g., chunk_read_limit, pass_read_limit, and the
analogous payload/all-columns limits) for negative values before casting, and if
any are negative throw a proper JNI exception (e.g., IllegalArgumentException
via env->ThrowNew) or return with an error so the native reader is never passed
a size_t converted from negative jlongs.
In `@java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java`:
- Around line 1141-1158: withFilter currently creates newFilter/newReader then
closes the old this.filter/this.reader directly; if closing the old filter
throws, the newly created resources leak. Fix by surrounding the old-resource
close calls with their own try/catch so that any exception thrown while closing
the old resources will first close newReader and newFilter (if non-null) before
rethrowing; only assign this.reader and this.filter to the new instances after
the old resources are successfully closed. Refer to the withFilter method and
the symbols newReader, newFilter, this.filter.close(), this.reader.close() when
making the change.
---
Nitpick comments:
In `@java/src/main/native/src/jni_compiled_expr.hpp`:
- Around line 50-55: Add a Doxygen comment block above the helper function
add_column_name_ref describing its purpose, parameters and return value: explain
that it appends a std::unique_ptr<cudf::ast::column_name_reference> to the
expressions container (transferring ownership), document the ref_ptr parameter
and that the function returns a reference to the stored
cudf::ast::column_name_reference, and note any lifetime/ownership expectations
regarding the expressions vector; keep the style consistent with neighboring
helper methods' comments.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f94f6010-0ac1-4e71-b546-3c7bc30776b7
📒 Files selected for processing (21)
.gitignorejava/pom.xmljava/src/main/java/ai/rapids/cudf/ByteRange.javajava/src/main/java/ai/rapids/cudf/Experimental.javajava/src/main/java/ai/rapids/cudf/HybridScanReader.javajava/src/main/java/ai/rapids/cudf/MemoryCleaner.javajava/src/main/java/ai/rapids/cudf/ParquetWriterOptions.javajava/src/main/java/ai/rapids/cudf/SecondaryFilterRanges.javajava/src/main/java/ai/rapids/cudf/UseDataPageMask.javajava/src/main/java/ai/rapids/cudf/ast/AstExpression.javajava/src/main/java/ai/rapids/cudf/ast/ColumnNameReference.javajava/src/main/native/CMakeLists.txtjava/src/main/native/include/hybrid_scan_jni_internal.hppjava/src/main/native/src/CompiledExpression.cppjava/src/main/native/src/HybridScanReaderJni.cppjava/src/main/native/src/HybridScanReaderJniInternal.cppjava/src/main/native/src/HybridScanReaderJniMaterialize.cppjava/src/main/native/src/jni_compiled_expr.hppjava/src/test/java/ai/rapids/cudf/HybridScanReaderTest.javajava/src/test/java/ai/rapids/cudf/ast/ColumnNameReferenceTest.javajava/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
java/src/main/native/src/HybridScanReaderJniMaterialize.cpp (1)
22-28: ⚡ Quick winAdd doxygen blocks for the new JNI entrypoints.
This file introduces a sizable exported native surface, but the entrypoints are only described with regular comments. Please add doxygen on the
JNIEXPORTfunctions so ownership, preconditions, and failure behavior stay documented and linted.As per coding guidelines, "Use doxygen for documentation generation and as a documentation linter on C++/CUDA code".
Also applies to: 134-138, 213-216, 351-352
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/main/native/src/HybridScanReaderJniMaterialize.cpp` around lines 22 - 28, Add proper Doxygen blocks to every JNIEXPORT function in this file—starting with Java_ai_rapids_cudf_HybridScanReader_materializeFilterColumnsWithKind and the other JNIEXPORT entrypoints later in the file—documenting: a one-line purpose, parameter descriptions (including JNIEnv*, jobject, and any jlong/jint arrays), ownership semantics for returned handles (who frees the jlongArray contents and any GPU/host resources), preconditions (e.g., env and inputs non-null, valid reader state), and failure behavior (what is returned or which Java exception is thrown on error). Ensure each Doxygen block uses `@param`, `@return`, `@throws/`@note for error cases, and a brief `@thread_safety` or `@precondition` note so the exported JNI surface is fully specified for linting.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@java/src/main/native/src/HybridScanReaderJniMaterialize.cpp`:
- Around line 61-68: convert_table_for_return(...) and release_as_jlong(...)
currently hand native ownership to Java before JNI calls complete, risking leaks
if NewLongArray or SetLongArrayRegion fail; change the flow in
HybridScanReaderJniMaterialize.cpp so you retain ownership of table_handles
(cudf::jni::native_jlongArray) and row_mask_col (do not call release_as_jlong
yet) until after env->NewLongArray and both env->SetLongArrayRegion calls
succeed, then perform the releases/convert-to-jlongs and set them into the Java
array; alternatively wrap the released handles with RAII/finalizers that free on
any JNI exception path so convert_table_for_return, release_as_jlong,
NewLongArray, SetLongArrayRegion, table_handles and row_mask_col are all cleaned
up on failure.
In `@java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java`:
- Around line 243-250: The test
testSecondaryFiltersByteRangesEmptyForHighCardinalityInts currently creates the
OpenReader via OpenReader.pageIndex(tmp).withFilter(...) but does not enable
page-index metadata, causing secondaryFiltersByteRanges(...) to be empty for the
wrong reason; fix by constructing the reader with page-index metadata enabled
(call withPageIndex() on the OpenReader chain before invoking
secondaryFiltersByteRanges) so the page index is actually loaded and the
assertion exercises the high-cardinality path (also apply the same change to the
analogous test around lines 334-347 that uses OpenReader.pageIndex and
secondaryFiltersByteRanges).
---
Nitpick comments:
In `@java/src/main/native/src/HybridScanReaderJniMaterialize.cpp`:
- Around line 22-28: Add proper Doxygen blocks to every JNIEXPORT function in
this file—starting with
Java_ai_rapids_cudf_HybridScanReader_materializeFilterColumnsWithKind and the
other JNIEXPORT entrypoints later in the file—documenting: a one-line purpose,
parameter descriptions (including JNIEnv*, jobject, and any jlong/jint arrays),
ownership semantics for returned handles (who frees the jlongArray contents and
any GPU/host resources), preconditions (e.g., env and inputs non-null, valid
reader state), and failure behavior (what is returned or which Java exception is
thrown on error). Ensure each Doxygen block uses `@param`, `@return`, `@throws/`@note
for error cases, and a brief `@thread_safety` or `@precondition` note so the
exported JNI surface is fully specified for linting.
🪄 Autofix (Beta)
❌ Autofix failed (check again to retry)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 51f92943-b236-497e-8432-79ea03508920
📒 Files selected for processing (7)
java/src/main/java/ai/rapids/cudf/HybridScanReader.javajava/src/main/java/ai/rapids/cudf/MemoryCleaner.javajava/src/main/native/include/hybrid_scan_jni_internal.hppjava/src/main/native/src/HybridScanReaderJniInternal.cppjava/src/main/native/src/HybridScanReaderJniMaterialize.cppjava/src/main/native/src/jni_compiled_expr.hppjava/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java
🚧 Files skipped from review as they are similar to previous changes (4)
- java/src/main/native/src/HybridScanReaderJniInternal.cpp
- java/src/main/native/src/jni_compiled_expr.hpp
- java/src/main/native/include/hybrid_scan_jni_internal.hpp
- java/src/main/java/ai/rapids/cudf/HybridScanReader.java
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. An unexpected error occurred while generating fixes: Not Found - https://docs.github.com/rest/git/refs#get-a-reference |
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java (1)
1020-1040: ⚡ Quick winAdd negative
passReadLimitcoverage for payload/all chunking.The matrix only exercises the swapped-limit case for
setupChunkingForFilterColumns. A regression insetupChunkingForPayloadColumns(..., passReadLimit)orsetupChunkingForAllColumns(..., passReadLimit)would slip through even though this test block claims both arguments are covered across the chunking entrypoints.Suggested test additions
return Stream.of( invocation("setupChunkingForFilterColumns", r -> r.setupChunkingForFilterColumns( -1L, 0L, new int[]{0}, UseDataPageMask.NO, HybridScanReader.RowMaskKind.ALL_TRUE, new DeviceMemoryBuffer[0])), invocation("setupChunkingForPayloadColumns", r -> { try (ColumnVector mask = ColumnVector.fromBooleans(true)) { r.setupChunkingForPayloadColumns(-1L, 0L, new int[]{0}, mask, UseDataPageMask.NO, new DeviceMemoryBuffer[0]); } }), invocation("setupChunkingForAllColumns", r -> r.setupChunkingForAllColumns(-1L, 0L, new int[]{0}, new DeviceMemoryBuffer[0])), invocation("constructRowGroupPasses", r -> r.constructRowGroupPasses(new int[]{0}, -1L)), + invocation("setupChunkingForPayloadColumnsPassLimit", r -> { + try (ColumnVector mask = ColumnVector.fromBooleans(true)) { + r.setupChunkingForPayloadColumns(0L, -1L, new int[]{0}, mask, + UseDataPageMask.NO, new DeviceMemoryBuffer[0]); + } + }), + invocation("setupChunkingForAllColumnsPassLimit", r -> + r.setupChunkingForAllColumns(0L, -1L, new int[]{0}, new DeviceMemoryBuffer[0])), invocation("setupChunkingForFilterColumnsPassLimit", r -> r.setupChunkingForFilterColumns(0L, -1L, new int[]{0}, UseDataPageMask.NO, HybridScanReader.RowMaskKind.ALL_TRUE, new DeviceMemoryBuffer[0])) );🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java` around lines 1020 - 1040, Add test invocations that exercise negative passReadLimit for the payload/all chunking entrypoints: add invocations similar to the existing swapped-case but calling setupChunkingForPayloadColumns(r -> { try (ColumnVector mask = ColumnVector.fromBooleans(true)) { r.setupChunkingForPayloadColumns(0L, -1L, new int[]{0}, mask, UseDataPageMask.NO, new DeviceMemoryBuffer[0]); } }) and setupChunkingForAllColumns(r -> r.setupChunkingForAllColumns(0L, -1L, new int[]{0}, new DeviceMemoryBuffer[0])); this ensures setupChunkingForPayloadColumns and setupChunkingForAllColumns validate passReadLimit the same way as setupChunkingForFilterColumns.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@java/src/main/native/src/HybridScanReaderJni.cpp`:
- Around line 29-265: Add Doxygen comment blocks for the new JNI exports so the
native API is documented and passes lint: add a brief file-level description and
for each exported function
(Java_ai_rapids_cudf_HybridScanReader_createFromFooter,
Java_ai_rapids_cudf_HybridScanReader_destroy,
Java_ai_rapids_cudf_HybridScanReader_pageIndexByteRange,
Java_ai_rapids_cudf_HybridScanReader_setupPageIndex,
Java_ai_rapids_cudf_HybridScanReader_allRowGroups,
Java_ai_rapids_cudf_HybridScanReader_totalRowsInRowGroups,
Java_ai_rapids_cudf_HybridScanReader_filterRowGroupsWithStats,
Java_ai_rapids_cudf_HybridScanReader_secondaryFiltersByteRanges,
Java_ai_rapids_cudf_HybridScanReader_filterRowGroupsWithDictionaryPages,
Java_ai_rapids_cudf_HybridScanReader_filterColumnChunksByteRanges,
Java_ai_rapids_cudf_HybridScanReader_payloadColumnChunksByteRanges,
Java_ai_rapids_cudf_HybridScanReader_allColumnChunksByteRanges) add a Doxygen
block immediately above the JNIEXPORT that describes the function purpose, lists
parameters with `@param` (env, jclass, handle, addresses/lengths, arrays, etc.),
specifies the return value with `@return` (including nullptr/0 on error), and
notes exceptions/error behavior (e.g., JNI_NULL_CHECK/JNI_CATCH semantics); keep
descriptions short and consistent with project style so the file-level and
per-function documentation satisfy the repository's C++/CUDA doxygen linting.
In `@java/src/main/native/src/HybridScanReaderJniMaterialize.cpp`:
- Around line 27-380: Add Doxygen comment blocks for each JNI-exported function
(e.g., Java_ai_rapids_cudf_HybridScanReader_materializeFilterColumnsWithKind,
Java_ai_rapids_cudf_HybridScanReader_materializePayloadColumns,
Java_ai_rapids_cudf_HybridScanReader_materializeAllColumns,
Java_ai_rapids_cudf_HybridScanReader_setupChunkingForFilterColumnsWithKind,
Java_ai_rapids_cudf_HybridScanReader_materializeFilterColumnsChunk,
Java_ai_rapids_cudf_HybridScanReader_takeFilterRowMask,
Java_ai_rapids_cudf_HybridScanReader_setupChunkingForPayloadColumns,
Java_ai_rapids_cudf_HybridScanReader_materializePayloadColumnsChunk,
Java_ai_rapids_cudf_HybridScanReader_setupChunkingForAllColumns,
Java_ai_rapids_cudf_HybridScanReader_materializeAllColumnsChunk,
Java_ai_rapids_cudf_HybridScanReader_hasNextTableChunk,
Java_ai_rapids_cudf_HybridScanReader_constructRowGroupPasses) describing
purpose, parameters, return value, and exceptions/throw behavior; place each
Doxygen block immediately above the corresponding JNIEXPORT function signature,
follow the repository's Doxygen style (brief description, `@param` for each
JNIEnv*, jclass, handle and other parameters, `@return` for
jlong/jlongArray/jobjectArray/jboolean/void, and note error conditions for
JNI_NULL_CHECK/JNI_THROW_NEW), and ensure formatting satisfies the project's
linter rules.
In `@java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java`:
- Around line 1163-1170: The catch block currently closes reader, footer, and
file one-after-another so a failure in reader.close() can mask the original
setup throwable; change the cleanup to close each partially-initialized resource
(reader, footer, file) in separate try/catch blocks and when a close() call
itself throws, add that exception as a suppressed exception to the original
throwable before rethrowing. Locate the setup around extractFooter(...), the
HybridScanReader(...) constructor, and optsForColumns(cols) where footer,
reader, and file are declared/initialized and implement per-resource guarded
closes (try { resource.close(); } catch(Throwable closeEx) {
original.addSuppressed(closeEx); }) so all resources are attempted to be closed
without losing the original error.
- Around line 1371-1385: The current loop allocates a DeviceMemoryBuffer dev for
each ByteRange but only assigns it into out[i] after the copy succeeds, so if
allocation/slice/copy throws the newly allocated dev leaks; fix by ensuring dev
is closed on all failure paths — either wrap the allocation+copy in an inner
try/finally (or try-with-resources) that sets a success flag and only transfers
ownership to out[i] when successful, or in the outer catch iterate and close the
current dev if it was allocated but not stored; reference the loop variables
ranges, dev (DeviceMemoryBuffer), out and the surrounding try/catch in
HybridScanReaderTest so the allocated buffer is always closed on exceptions.
---
Nitpick comments:
In `@java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java`:
- Around line 1020-1040: Add test invocations that exercise negative
passReadLimit for the payload/all chunking entrypoints: add invocations similar
to the existing swapped-case but calling setupChunkingForPayloadColumns(r -> {
try (ColumnVector mask = ColumnVector.fromBooleans(true)) {
r.setupChunkingForPayloadColumns(0L, -1L, new int[]{0}, mask,
UseDataPageMask.NO, new DeviceMemoryBuffer[0]); } }) and
setupChunkingForAllColumns(r -> r.setupChunkingForAllColumns(0L, -1L, new
int[]{0}, new DeviceMemoryBuffer[0])); this ensures
setupChunkingForPayloadColumns and setupChunkingForAllColumns validate
passReadLimit the same way as setupChunkingForFilterColumns.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4bfe3be2-356c-4347-b576-39241068647f
📒 Files selected for processing (4)
java/src/main/java/ai/rapids/cudf/HybridScanReader.javajava/src/main/native/src/HybridScanReaderJni.cppjava/src/main/native/src/HybridScanReaderJniMaterialize.cppjava/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java
🚧 Files skipped from review as they are similar to previous changes (1)
- java/src/main/java/ai/rapids/cudf/HybridScanReader.java
| jlongArray j_addrs, | ||
| jlongArray j_lens, | ||
| jboolean use_data_page_mask, | ||
| jboolean all_true) |
There was a problem hiding this comment.
I don't think we need this boolean. We can decide based on the use_data_page_mask. If it's set to YES, then do reader->build_row_mask_with_page_index_stats, otherwise reader->build_all_true_row_mask
There was a problem hiding this comment.
The separation of the use_data_page_mask flag and the row mask initialization mirrors cuDF C++ to a large extent. Using use_data_page_mask to drive the row mask initialization logic implies that users cannot set use_data_page_mask=false and all_true=false, where page stats prunning can occur without the overhead of creating data page masks.
There was a problem hiding this comment.
Right, I think let me layout what makes sense and then you can choose the best path forward.
- If
build_all_true_mask:use_data_page_maskmust always be NO for filter columns,YES/NOfor payload columns - If
build_row_mask_from_page_stats:use_data_page_maskshould be NO for filter columns (no need) but can be YES/NO for payload columns.
Side note: If try { build_row_mask_from_page_stats } fails saying there's no page index available, you can safely fallback to build_all_true + hardcoded use_page_mask::NO for both columns
There was a problem hiding this comment.
A new resolve_data_page_mask_mode function has been added to reflect this logic.
| std::unique_ptr<cudf::column> row_mask_col; | ||
| if (all_true) { | ||
| row_mask_col = wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr); | ||
| } else { | ||
| row_mask_col = wrapper->reader->build_row_mask_with_page_index_stats( | ||
| holder.span(), wrapper->options, stream, mr); | ||
| } |
There was a problem hiding this comment.
Let's prefer initialization to assignment.
| std::unique_ptr<cudf::column> row_mask_col; | |
| if (all_true) { | |
| row_mask_col = wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr); | |
| } else { | |
| row_mask_col = wrapper->reader->build_row_mask_with_page_index_stats( | |
| holder.span(), wrapper->options, stream, mr); | |
| } | |
| std::unique_ptr<cudf::column> row_mask_col = (all_true)? | |
| row_mask_col = wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr) | |
| : row_mask_col = wrapper->reader->build_row_mask_with_page_index_stats( | |
| holder.span(), wrapper->options, stream, mr); |
|
Hello @pmattione-nvidia, I understand you were working on a C++ implementation with cuDF hybrid scan. Would you please work with @paul-aiyedun to help him land these new Java bindings? |
There was a problem hiding this comment.
Pull request overview
Adds an experimental Java/JNI surface area for the cuDF Parquet hybrid scan reader, enabling a two-phase “filter columns first, payload columns second” pipeline (including chunked/streaming materialization), plus AST support for column references by name and end-to-end Java tests.
Changes:
- Introduces
HybridScanReaderJava API (+ supportingByteRange,SecondaryFilterRanges,UseDataPageMask,@Experimental) and corresponding JNI/C++ wrappers. - Extends the Java AST with
ColumnNameReferenceand wires it through the native compiled-expression deserializer. - Adds comprehensive JUnit coverage for reader metadata, pruning, materialization (single-shot + chunked), and lifecycle/argument validation; adjusts Maven resource handling and gitignore entries.
Reviewed changes
Copilot reviewed 20 out of 21 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| java/src/test/java/ai/rapids/cudf/HybridScanReaderTest.java | End-to-end and behavioral tests for the new HybridScanReader pipeline. |
| java/src/test/java/ai/rapids/cudf/ast/CompiledExpressionTest.java | Adds tests covering ColumnNameReference behavior and type ID expectations. |
| java/src/test/java/ai/rapids/cudf/ast/ColumnNameReferenceTest.java | Unit tests for ColumnNameReference serialization/validation. |
| java/src/main/native/src/jni_compiled_expr.hpp | Adds native compiled-expression storage for column_name_reference. |
| java/src/main/native/src/HybridScanReaderJniMaterialize.cpp | JNI entrypoints for single-shot + chunked materialization, row-mask ownership transfer. |
| java/src/main/native/src/HybridScanReaderJniInternal.cpp | Shared JNI helpers: option building, spans, range packing, and validation helpers. |
| java/src/main/native/src/HybridScanReaderJni.cpp | JNI entrypoints for reader creation/destruction, metadata, pruning, and byte-range discovery. |
| java/src/main/native/src/CompiledExpression.cpp | Extends serialized AST decoding to support COLUMN_NAME_REFERENCE (type id 5). |
| java/src/main/native/include/hybrid_scan_jni_internal.hpp | Declares shared JNI helper interfaces and wrapper struct. |
| java/src/main/native/CMakeLists.txt | Adds new HybridScan JNI sources to the native build. |
| java/src/main/java/ai/rapids/cudf/UseDataPageMask.java | Java mirror of use_data_page_mask for controlling page-mask computation. |
| java/src/main/java/ai/rapids/cudf/SecondaryFilterRanges.java | Java holder for bloom/dictionary byte-range results from secondary filtering. |
| java/src/main/java/ai/rapids/cudf/ParquetWriterOptions.java | Adds StatisticsFrequency.COLUMN needed to emit page index in test fixtures. |
| java/src/main/java/ai/rapids/cudf/MemoryCleaner.java | Registers HybridScanReader as an RMM-blocking native resource. |
| java/src/main/java/ai/rapids/cudf/HybridScanReader.java | New public experimental Java binding for hybrid scan reader API. |
| java/src/main/java/ai/rapids/cudf/Experimental.java | New annotation marking experimental public types. |
| java/src/main/java/ai/rapids/cudf/ByteRange.java | Immutable offset+size type used for IO byte range plumbing. |
| java/src/main/java/ai/rapids/cudf/ast/ColumnNameReference.java | New AST node for referencing columns by name (serialized into native plan). |
| java/src/main/java/ai/rapids/cudf/ast/AstExpression.java | Adds COLUMN_NAME_REFERENCE to serialized AST ExpressionType enum. |
| java/pom.xml | Changes LICENSE resource staging to improve JDT LS project import reliability. |
| .gitignore | Ignores local Maven repo and generated java/bin/ artifacts. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| public static void register(HybridScanReader reader, Cleaner cleaner) { | ||
| // RMM blocker: wrapper can hold device memory (chunked_filter_row_mask). | ||
| all.put(cleaner.id, new CleanerWeakReference(reader, cleaner, collected, true)); | ||
| } |
| NULL_LITERAL(1), | ||
| COLUMN_REFERENCE(2), | ||
| UNARY_EXPRESSION(3), | ||
| BINARY_EXPRESSION(4); | ||
| BINARY_EXPRESSION(4), | ||
| COLUMN_NAME_REFERENCE(5); |
| /** | ||
| * Enumeration of the AST expression types that can appear in the serialized data. | ||
| * NOTE: This must be kept in sync with the NodeType enumeration in AstNode.java! | ||
| */ |
| private static long[] bufferAddrs(DeviceMemoryBuffer[] buffers) { | ||
| if (buffers == null) { | ||
| return new long[0]; | ||
| } | ||
| long[] addrs = new long[buffers.length]; | ||
| for (int i = 0; i < buffers.length; i++) { | ||
| addrs[i] = buffers[i].getAddress(); | ||
| } | ||
| return addrs; | ||
| } | ||
|
|
||
| private static long[] bufferLens(DeviceMemoryBuffer[] buffers) { | ||
| if (buffers == null) { | ||
| return new long[0]; | ||
| } | ||
| long[] lens = new long[buffers.length]; | ||
| for (int i = 0; i < buffers.length; i++) { | ||
| lens[i] = buffers[i].getLength(); | ||
| } | ||
| return lens; | ||
| } |
| planned_copy_result result; | ||
| if (ranges.empty()) { return result; } | ||
|
|
||
| if (are_ranges_contiguous(ranges)) { |
There was a problem hiding this comment.
This is all-or-nothing. It might be worth doing something more robust here. We could assume that the ranges are in order but coalesce together each set of ranges that are adjacent to each other.
There was a problem hiding this comment.
This function is no longer used in the current PR revision, so I have removed it.
Yes, I added a new |
| * @brief Build a parquet_reader_options from the supplied JNI args. The footer is provided | ||
| * separately because the hybrid_scan_reader does not consume it via the options. | ||
| */ | ||
| cudf::io::parquet_reader_options build_options(JNIEnv* env, |
There was a problem hiding this comment.
Q: Does this not already exist as it's also used by the regular parquet reader?
There was a problem hiding this comment.
There is currently no shared build_options C++ function as the different parquet_reader_options building use cases take in different parameters.
| * | ||
| * <p>The APIs in this file are experimental and subject to change. | ||
| */ | ||
| @Experimental |
There was a problem hiding this comment.
Should this be experimental since the cudf::io::test::byte_range_info isn't experimental
There was a problem hiding this comment.
ByteRange is annotated as experimental, since it is only used by HybridScanReader, which is experimental. We can remove the experimental annotation later, if any non-experimental class requires this class.
| */ | ||
| public HybridScanReader(HostMemoryBuffer footerBuffer, | ||
| ParquetOptions opts, | ||
| CompiledExpression filter) { |
There was a problem hiding this comment.
I think we should not take in the filter here and allow the user to resetColumnSelection and provide a new filter (as needed) to all filter_** APIs. Doing this means one can't change their filter at will.
There was a problem hiding this comment.
This is okay to keep as is for now in interest of keeping things moving.
There was a problem hiding this comment.
Updated to use a new setFilter method instead of the constructor.
| requireNonNullRowGroups(rowGroupIndices); | ||
| long[] addrs = bufferAddrs(columnChunkData); | ||
| long[] lens = bufferLens(columnChunkData); | ||
| boolean allTrue = (kind == RowMaskKind.ALL_TRUE); |
There was a problem hiding this comment.
Can we do the following instead:
allTrue = (kind == RowMaskKind.ALL_TRUE) or (UseDataPageMask == NO)
mode = (allTrue)? UseDataPageMask::NO : mode
There's no point of initializing the data page mask with page index stats if we aren't going to use it.
Similarly, if the mask is AllTrue, there's no point in (UseDataPageMask==YES); though there's early exit in libcudf for this.
There was a problem hiding this comment.
Based on all the RowMaskKind and UseDataPageMask related comments, I have simplified the APIs to use a usePageLevelPruning flag instead. This should avoid flag combinations that can cause performance issues.
| private static native long[] payloadColumnChunksByteRanges(long handle, int[] rowGroupIndices); | ||
| private static native long[] allColumnChunksByteRanges(long handle, int[] rowGroupIndices); | ||
|
|
||
| // Single-shot materialize |
There was a problem hiding this comment.
This is two-step materialize no?
There was a problem hiding this comment.
Comment updated.
| extern "C" { | ||
|
|
||
| // ---------------------------------------------------------------------- | ||
| // Single-shot materialize |
There was a problem hiding this comment.
| // Single-shot materialize | |
| // Two-step materialize |
| // When the row mask is all-true (filter path with RowMaskKind::ALL_TRUE), the computed data page | ||
| // mask is also all-true and bit-identical to passing NO, so short-circuit to skip the compute. | ||
| // Payload paths pass row_mask_all_true = false (the default). | ||
| exp_pq::use_data_page_mask resolve_data_page_mask_mode(bool use_data_page_mask, |
| auto row_mask_col = all_true | ||
| ? wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr) | ||
| : wrapper->reader->build_row_mask_with_page_index_stats( |
There was a problem hiding this comment.
Need to take care of if input mode == use_data_page_mask::NO, then build_all_true_row_mask anyway
There was a problem hiding this comment.
I believe the new use_page_level_pruning based API should address this.
| auto row_mask_col = all_true | ||
| ? wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr) | ||
| : wrapper->reader->build_row_mask_with_page_index_stats( | ||
| holder.span(), wrapper->options, stream, mr); | ||
| auto mode = resolve_data_page_mask_mode(use_data_page_mask, all_true); |
|
Looks like we might be missing some hybrid scan APIs like these: cudf/cpp/include/cudf/io/experimental/hybrid_scan.hpp Lines 338 to 358 in 08f48c2 |
| * <p>The filter and payload materialization paths accept a boolean that toggles | ||
| * page-level pruning: skips decode of pages the filter (or row mask) proves empty, in | ||
| * exchange for a per-page stats scan and a carried row-mask column. Enable when the | ||
| * workload prunes many pages; on the filter path this requires prior | ||
| * {@link #setupPageIndex(HostMemoryBuffer)}. |
There was a problem hiding this comment.
This can only be done if page index is available. We should fall back to allTrue+use_data_page_mask::NO when page index isn't present.
There was a problem hiding this comment.
Calling setupPageIndex with an invalid page index buffer currently throws an exception, and setting usePageLevelPruning=true without first calling setupPageIndex also throws an exception. Users can detect absence of a page index by calling the pageIndexByteRange (and choose to skip the setupPageIndex call). I think silently falling back in this case would be misleading. A caller who explicitly requested page-level pruning and silently did not get it would likely not notice a resulting performance regression.
There was a problem hiding this comment.
I guess it depends on the intent of the Java side reader. I am okay if no silent fallback is okay with the java codeowners
| * | ||
| * @param filter the new filter expression, or {@code null} to clear | ||
| */ | ||
| public void setFilter(CompiledExpression filter) { |
mhaseeb123
left a comment
There was a problem hiding this comment.
Perhaps we can cache if the file has page_index and then override (fallback to all true) the usePageLevelPruning based on that?
| auto row_mask_col = use_page_level_pruning | ||
| ? wrapper->reader->build_row_mask_with_page_index_stats( | ||
| holder.span(), wrapper->options, stream, mr) | ||
| : wrapper->reader->build_all_true_row_mask(holder.span(), stream, mr); |
There was a problem hiding this comment.
Perhaps here we could take in account the availability of page index and only call build_row_mask_with_page_index_stats if so. Perhaps log a warning instead of throwing
There was a problem hiding this comment.
I believe my response here also applies to this comment.
These APIs exist in the current |
I thought |
| * <p>The filter and payload materialization paths accept a boolean that toggles | ||
| * page-level pruning: skips decode of pages the filter (or row mask) proves empty, in | ||
| * exchange for a per-page stats scan and a carried row-mask column. Enable when the | ||
| * workload prunes many pages; on the filter path this requires prior | ||
| * {@link #setupPageIndex(HostMemoryBuffer)}. |
There was a problem hiding this comment.
I guess it depends on the intent of the Java side reader. I am okay if no silent fallback is okay with the java codeowners
|
/merge |
|
/ok to test 984dd44 |
|
/ok to test 8f7ba95 |
Description
Add
HybridScanReader, a Java binding forcudf::io::parquet::experimental::hybrid_scan_reader, targeting highly-selective filter expressions over Parquet files.Expose the full two-phase pipeline: filter columns are materialized first and an AST filter is applied to produce a boolean row mask; payload columns are then read only for the rows that survive, avoiding unnecessary I/O and decompression.
Support both single-shot and chunked/streaming materialization. In the chunked filter path, the C++ wrapper holds the row mask column across chunk calls and mutates it in place;
takeFilterRowMask()transfers ownership to Java without copying.Add
ColumnNameReference, a new AST node that references columns by name rather than index. This is required because filter expressions are constructed before the Parquet schema is resolved; the reader resolves names to positions internally.Introduce a new
@Experimentalannotation and apply it to the entire new surface area to signal that these APIs may change as the upstream C++ experimental API stabilizes.Use
MemoryCleanerfor native resource management, consistent with the rest of cudf-java; leaked readers are logged with their handle address.Fix
pom.xmlresource configuration to allow the JDT language server to correctly import the Java project in VS Code / Cursor. The LICENSE resource was declared with${basedir}/..as its source directory, which caused the JDT LS to fail to recognize the project entirely. Replace it with an explicitcopy-licenseplugin execution that stages only the LICENSE file undertarget/license-resources/before the resource directory is read.Add
HybridScanReaderTestcovering the full pipeline end-to-end: single-shot and chunked filter + payload materialization, row-group pruning via stats and dictionary pages, page-index-seeded row masks, pass construction, and leak detection.Related issue: #22271
Checklist