Add a protocol buffer decode kernel - #4107
Conversation
There was a problem hiding this comment.
Pull request overview
This PR adds a GPU-accelerated protocol buffer decoder with intentionally limited features, focusing on simple scalar field types. The implementation provides a JNI interface for decoding binary protobuf messages into cuDF STRUCT columns.
Key changes:
- Implements GPU kernels for decoding protobuf varint, fixed32/64, and length-delimited (string) fields
- Adds JNI bindings between Java and CUDA implementation
- Provides basic test coverage for INT64 and STRING field types
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 14 comments.
Show a summary per file
| File | Description |
|---|---|
| src/main/java/com/nvidia/spark/rapids/jni/ProtobufSimple.java | Java API providing decodeToStruct() method with parameter validation |
| src/test/java/com/nvidia/spark/rapids/jni/ProtobufSimpleTest.java | Basic test case covering varint (INT64) and string decoding with missing fields and null messages |
| src/main/cpp/src/protobuf_simple.hpp | C++ API declaration with documentation of supported types |
| src/main/cpp/src/protobuf_simple.cu | CUDA implementation with three specialized kernels for varint, fixed-width, and string extraction |
| src/main/cpp/src/ProtobufSimpleJni.cpp | JNI bridge translating Java arrays to C++ vectors and invoking decode logic |
| src/main/cpp/CMakeLists.txt | Build configuration adding new source files to compilation targets |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@greptile full review |
Greptile SummaryThis PR adds a GPU-accelerated protocol buffer decode kernel using CUDA, implementing a multi-pass algorithm (count → scan → extract → build) that supports nested STRUCT, LIST, STRING, scalar types, proto2 closed enum semantics with deferred error reporting, and singular nested message merging via
Confidence Score: 3/5Largely safe but two P1 logic issues need author confirmation before merge: premature deferred enum error under last-wins semantics, and a potentially dead drop_unknown_repeated_enum_values function. The PR is a large, well-structured addition with a clear design. Two P1 findings were identified: (1) the deferred enum error is set for every invalid occurrence during the scan pass, meaning an intermediate invalid occurrence that is later overridden by a valid last-wins value still leaves the deferred error flag set, which may cause spurious errors in strict mode; (2) drop_unknown_repeated_enum_values is defined but has no call site in the PR, suggesting it is dead code or a missing call. Two P2 findings cover a stack-size/depth-guard decoupling in skip_group and incomplete fragment validation coverage. The overall architecture is sound and the Java-side schema validation improvements are correct. Files Needing Attention: src/main/cpp/src/protobuf/protobuf_kernels.cu (deferred enum error timing), src/main/cpp/src/protobuf/protobuf_builders.cu (drop_unknown_repeated_enum_values dead code) Important Files Changed
Sequence DiagramsequenceDiagram
participant Java as Protobuf.java (JNI)
participant Entry as protobuf.cu
participant Scan as scan_all_field_occurrences_kernel
participant Val as validate_message_fragments_kernel
participant Scalar as extract_scalar_batched_kernel
participant Build as protobuf_builders.cu
Java->>Entry: decode(data, schema, options)
Entry->>Entry: allocate deferred_enum_error, field_scan_view
Entry->>Scan: count pass - count occurrences per field per row
Scan-->>Entry: h_occurrence_counts
Entry->>Entry: prefix-sum to offsets, allocate occurrence buffers
Entry->>Scan: scan pass - fill field_scan_view with locations
Scan-->>Entry: field locations, wire type errors, deferred enum errors
Entry->>Entry: check h_error and h_deferred_enum_error
Entry->>Scalar: extract_scalar_batched_kernel for all scalar_kinds
Scalar-->>Entry: typed scalar buffers
Entry->>Val: validate_message_fragments_kernel for singular nested messages
Val-->>Entry: merged fragment validation results
Entry->>Build: build_protobuf_field_values_column_shared
Build->>Build: build_nested_struct_column recursive for STRUCT fields
Build->>Build: build_merged_singular_struct_column via cub DeviceMemcpy Batched
Build-->>Entry: cudf column result
Entry-->>Java: column handle
|
There was a problem hiding this comment.
Additional Comments (6)
-
src/main/cpp/src/protobuf_simple.cu, line 89-91 (link)logic: potential overflow:
len64can be up to 2^64-1, but casting tointon line 90 can overflow iflen64 > INT_MAX -
src/main/cpp/src/protobuf_simple.cu, line 323-324 (link)logic: potential overflow:
len64can be larger thanINT_MAX, but casting tointon line 324 will overflow -
src/main/cpp/src/protobuf_simple.cu, line 375-376 (link)logic: race condition: multiple threads write to
*error_flagwithout atomics, causing undefined behavior when multiple threads encounter errors simultaneouslyThen in kernels, use
atomicOr(error_flag, 1)instead of*error_flag = 1 -
src/main/cpp/src/protobuf_simple.cu, line 398-407 (link)logic: protobuf uses zigzag encoding for signed integers (sint32/sint64), but varint decoding here treats them as unsigned - decoding negative values will produce incorrect results. Are you only supporting unsigned int32/int64, or should zigzag decoding be implemented for signed types?
-
src/main/cpp/src/protobuf_simple.cu, line 240 (link)syntax: type punning through
reinterpret_castof incompatible pointer types is undefined behavior in C++ -
src/main/cpp/src/protobuf_simple.cu, line 248 (link)syntax: type punning through
reinterpret_castof incompatible pointer types is undefined behavior in C++
6 files reviewed, 6 comments
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
@greptile full review |
There was a problem hiding this comment.
Additional Comments (2)
-
src/main/cpp/src/protobuf_simple.cu, line 186 (link)logic: zigzag decode uses signed right shift which is correct for signed types, but applied to unsigned
v -
src/main/cpp/src/ProtobufSimpleJni.cpp, line 55 (link)logic:
encodingsis constructed fromn_type_scalesbut thenout_typesis also constructed usingn_type_scales[i]as the scale parameter, which would be wrong for non-decimal types where this represents encoding
6 files reviewed, 2 comments
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 13 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 10 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
…ven/spark-rapids-jni into protocol_buffer_jni_dev
|
@greptileai full review |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
NOTE: release/26.02 has been created from main. Please retarget your PR to release/26.02 if it should be included in the release. |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
…en (#4644) Part 3 of #4107. Implements **JNI-3b.1** and **JNI-3b.2** from the split plan in NVIDIA/cudf-spark#14069. ## Summary Adds nested message field scan utilities and the first `build_nested_struct_column` path, wiring top-level `STRUCT` fields through the orchestrator for scalar numeric/bool children. After this PR, the following decode correctly at the top level: - `STRUCT { int32 / int64 / uint32 / uint64 / float / double / bool ... }` Nested STRING / LIST<UINT8> / recursive STRUCT children and repeated-in-nested fields still produce typed null columns; those land in 3b.3 – 3b.6. ## Components **Kernels** (`protobuf_kernels.cu`): - `scan_nested_message_fields_kernel` — one thread per parent row, records last-one-wins locations for direct singleton children of a nested message. Repeated children are intentionally left to the dedicated count/scan path in 3b.5/3b.6. - `extract_strided_locations_kernel` — pulls one field's per-row locations out of the 2D nested-locations array, replacing the previous D2H + CPU loop + H2D pattern. **Builders** (`protobuf_builders.cu`): - `build_nested_struct_column` — scans child field locations on device, then delegates each numeric/bool child to `extract_typed_column` with a `nested_location_provider`. Child types not yet supported in nested context are filled with schema-aware typed null columns so the output schema stays well-formed. **Orchestrator** (`protobuf.cu`): - After repeated-field processing, calls `build_nested_struct_column` for each top-level STRUCT field, using `launch_extract_strided_locations` to derive per-row parent locations from the existing `d_nested_locations` buffer. ## Intentionally out of scope Deferred to subsequent PRs per the split plan in NVIDIA/cudf-spark#14069: - **3b.3** — nested string/bytes/enum-as-string, default values inside nested messages, proto2 required-field checks for nested children - **3b.4** — recursive nested STRUCT support, depth-aware grandchild parent location computation - **3b.5** — count/scan kernels for repeated fields inside nested messages - **3b.6** — `build_repeated_child_list_column` for repeated children ## Tests Seven new tests in `ProtobufTest.java` covering the top-level STRUCT scalar path: - `testNestedMessageInt32Child` - `testNestedMessageLastOneWins` - `testNestedMessageMultipleScalarChildren` - `testNestedMessageAbsentParentIsNull` - `testZeroLengthNestedMessageChildNullability` - `testChildlessNestedMessagePresence` - `testFailfastNestedRepeatedWrongWireType` ## Size ~740 lines added across 5 source files plus the test file; each source file stays under the ~300 LOC per-file target. --------- Signed-off-by: Haoyang Li <haoyangl@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
# Conflicts: # src/main/cpp/src/protobuf/protobuf.cu # src/main/cpp/src/protobuf/protobuf_builders.cu # src/main/cpp/src/protobuf/protobuf_kernels.cu # src/main/cpp/src/protobuf/protobuf_kernels.cuh # src/test/java/com/nvidia/spark/rapids/jni/ProtobufTest.java
# Conflicts: # src/main/cpp/benchmarks/CMakeLists.txt # src/main/cpp/src/protobuf/protobuf.cu # src/main/cpp/src/protobuf/protobuf_builders.cu # src/main/cpp/src/protobuf/protobuf_host_helpers.hpp # src/main/cpp/src/protobuf/protobuf_kernels.cu # src/main/cpp/src/protobuf/protobuf_kernels.cuh # src/test/java/com/nvidia/spark/rapids/jni/ProtobufTest.java
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com> # Conflicts: # src/main/cpp/src/protobuf/protobuf.cu # src/main/cpp/src/protobuf/protobuf_builders.cu # src/main/cpp/src/protobuf/protobuf_host_helpers.hpp # src/main/cpp/src/protobuf/protobuf_kernels.cu # src/main/cpp/src/protobuf/protobuf_kernels.cuh # src/main/cpp/src/protobuf/protobuf_types.cuh # src/test/java/com/nvidia/spark/rapids/jni/ProtobufTest.java
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
|
NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release. |
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
This PR adds a protocol buffer decoder with a large subset of Proto2 features, to support spark expression
from_protobuf.The code is ready to me but too large to review. I'm splitting it into small parts:
Part 0: #4373
PR Split Plan
This PR is being split into a chain of focused PRs to make review manageable. Each PR is independently buildable and testable, and each subsequent PR is largely additive.
Please check NVIDIA/cudf-spark#14069 for detailed pr split plan.
Summary
This PR adds a GPU-accelerated protobuf decoder that converts
LIST<INT8/UINT8>columns (one serialized protobuf message per row) into nested cuDFSTRUCTcolumns via JNI. This is the native kernel layer that powersfrom_protobuf()GPU acceleration in the spark-rapids plugin.The implementation spans ~6,500 lines of new CUDA/C++/Java code and ~3,900 lines of tests, organized into a clean four-file C++ architecture plus a validated Java schema API.
Key capabilities
int32,int64,uint32,uint64,sint32/sint64(zigzag),fixed32/sfixed32/fixed64/sfixed64,float,double,bool,string,bytesArrayType(StructType)— repeated nested messages with arbitrary child fieldsPerformance characteristics
FIELD_LOOKUP_TABLE_MAX = 4096)Architecture
File structure
Dependency graph
Multi-pass decode algorithm
The decoder processes each batch of messages through multiple GPU passes:
count_repeated_fields_kernel): One thread per row. Scans message bytes to count repeated field occurrences and record nested message locations. Handles both packed and unpacked repeated encoding.thrust::exclusive_scan): Prefix sum on repeated counts to compute output array offsets.scan_all_fields_kernel+scan_all_repeated_occurrences_kernel): Records exact byte locations (offset + length) for every target field. Last-one-wins semantics for duplicate scalar fields.MAX_NESTED_STRUCT_DECODE_DEPTH = 10.Flattened schema representation
The protobuf schema is represented as parallel arrays passed through JNI. Fields are ordered in pre-order traversal (parent before children):
fieldNumbers[]parentIndices[]depthLevels[]wireTypes[]outputTypeIds[]encodings[]isRepeated[]isRequired[]hasDefaultValue[]defaultInts/Floats/Bools/Strings[]enumValidValues[][]enumNames[][][]Example for
message Outer { int32 a = 1; Inner b = 2; } message Inner { int32 x = 1; string y = 2; }:Test coverage
107 JUnit tests in
ProtobufTest.java+ 13 tests inProtobufSchemaDescriptorTest.java, organized by feature:Benchmarks
8 NVBench benchmarks in
protobuf_decode.cu:BM_protobuf_flat_scalarsBM_protobuf_nestedBM_protobuf_repeatedBM_protobuf_wide_repeated_messageBM_protobuf_repeated_child_listsBM_protobuf_repeated_child_string_count_scanBM_protobuf_repeated_child_string_buildBM_protobuf_many_repeatedReview Guide
This PR is large (~12,000 lines total) but has a well-defined layered architecture. This guide provides a recommended reading order, key areas to focus on per file, and a mental model for understanding the code.
Recommended reading order
Read bottom-up from the API surface to the kernel internals:
protobuf.hppProtobufSchemaDescriptor.javaProtobuf.javaProtobufJni.cppprotobuf_common.cuh§1: typesfield_location,device_nested_field_descriptor, etc.protobuf_common.cuh§2: device helpersread_varint,skip_field,get_field_data_location,decode_tag,lookup_fieldprotobuf_common.cuh§3: LocationProvidersTopLevelLocationProvider,NestedLocationProvider, etc. — these abstract how extraction kernels compute byte offsetsprotobuf_common.cuh§4: template kernelsextract_varint_kernel,extract_fixed_kernel,extract_lengths_kernel,copy_varlen_data_kernel, batched variantsprotobuf_common.cuh§5: template host functionsextract_typed_column,build_repeated_scalar_column,extract_and_build_string_or_bytes_column,validate_enum_and_propagate_rowsprotobuf_kernels.cu§1: scanscan_all_fields_kernel— the core single-pass field scannerprotobuf_kernels.cu§2: count/scan repeatedcount_repeated_fields_kernel,scan_all_repeated_occurrences_kernel, shared__device__helpersprotobuf_kernels.cu§3: nestedscan_nested_message_fields_kernel,scan_repeated_message_children_kernel, compute kernelsprotobuf_kernels.cu§4: validationcheck_required_fields_kernel,validate_enum_values_kernel, enum-string kernelsprotobuf_builders.cu§1: utilitiesmake_null_column,make_empty_column_safe,make_null_list_column_with_childprotobuf_builders.cu§2: enum-stringmake_enum_string_lookup_tables,build_enum_string_column,build_repeated_enum_string_columnprotobuf_builders.cu§3: nested structbuild_nested_struct_column— most complex builder, recursive depth handlingprotobuf_builders.cu§4: repeated structbuild_repeated_struct_column,build_repeated_child_list_column— repeated-in-repeatedprotobuf.cudecode_protobuf_to_struct: orchestration, batched scalar extraction, PERMISSIVE null propagationProtobufTest.javaprotobuf_decode.cuTotal estimated review time: ~5-6 hours for a thorough review.
Key review areas by priority
P0: Correctness-critical
scan_all_fields_kernel(protobuf_kernels.cu): This is the most important kernel. One thread per row, scans all bytes of each message, records field locations. Key things to verify:cur < msg_end)atomicCAS(no races)row_has_invalid_datais set on parse errors so the row can be nullifiedcount_repeated_fields_kernel(protobuf_kernels.cu): Counts repeated field occurrences. Must correctly distinguish packed vs unpacked encoding. Packed detection:wire_type == WT_LENbutexpected_wire_type != WT_LEN.build_nested_struct_column(protobuf_builders.cu): Recursive builder for nested messages. Verify:MAX_NESTED_STRUCT_DECODE_DEPTH)is_repeatedchildren inside nested messages get proper LIST wrappingPERMISSIVE mode null propagation (protobuf.cu, lines ~1163-1190): Invalid enum values must nullify the entire struct row AND propagate to all descendants via
apply_parent_mask_to_row_aligned_column+propagate_nulls_to_descendants.JNI memory safety (ProtobufJni.cpp): Every
GetObjectArrayElement/GetByteArrayElements/GetIntArrayElementsmust have a matchingDeleteLocalRef/ReleaseXxxArrayElements. Verify no leaks in the enum_names triple-nested loop.P1: Robustness
validate_decode_context(protobuf.hpp): Validates all schema invariants before any GPU work. Check completeness — duplicate field numbers under same parent, encoding compatibility, enum metadata non-empty forENC_ENUM_STRING.ProtobufSchemaDescriptor.validate(Java): Mirrors C++ validation. Defensive copies in constructor, re-validation on deserialization.Varint parsing (protobuf_common.cuh
read_varint): 10th byte must only use lowest bit. Truncated/malformed varints must returnfalse.Wire type handling (
skip_field,get_wire_type_size): VerifyWT_SGROUPuses iterative handling with depth cap of 32 (not recursive).WT_EGROUPis rejected as standalone.P2: Performance
Batched scalar extraction (protobuf.cu, lines ~400-470): Fields are grouped into 12 categories by type and extracted with 2D kernel launches. Verify grouping logic covers all type+encoding combinations.
Field lookup tables (
build_field_lookup_table,build_index_lookup_tablein protobuf_common.cuh): O(1) field_number → index mapping when max field number ≤FIELD_LOOKUP_TABLE_MAX. Falls back to linear scan otherwise.String two-phase construction:
extract_lengths_kernel→make_offsets_child_column→copy_varlen_data_kernel. Verify no off-by-one in offset calculations.Things to watch for
uvector.size()afteruvector.release()in the same expression. The code caches sizes before releasing.num_rows + 1elements.count_repeated_in_nested_kernelandscan_repeated_in_nested_kernelhandle both packed and unpacked within nested message boundaries.cudf::logic_errorfor data errors: The code usescudf::logic_errorfor wire-format errors in strict mode. This is semantically imprecise (it conventionally signals API misuse), but functionally correct.Mapping review to test coverage
scan_all_fields_kernelcount/scan_repeatedbuild_nested_struct_columnvalidate_enum + enum-stringcheck_required_fields_kernel