Skip to content

Consolidate protobuf scan and decode interfaces - #4836

Open
thirtiseven wants to merge 12 commits into
NVIDIA:mainfrom
thirtiseven:protobuf-refactor-consolidate-interfaces
Open

Consolidate protobuf scan and decode interfaces#4836
thirtiseven wants to merge 12 commits into
NVIDIA:mainfrom
thirtiseven:protobuf-refactor-consolidate-interfaces

Conversation

@thirtiseven

@thirtiseven thirtiseven commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR consolidates the protobuf scan and decode interfaces without changing the Java API.

  • make scan_message_field_locations<MismatchPolicy> the common message walker used by scan_all_fields_kernel, count_repeated_fields_kernel, scan_all_field_occurrences_in_message, and scan_nested_message_fields_kernel; callers provide descriptor lookup plus on_singular and on_repeated callbacks through message_scan_context
  • move scalar value handling into decode_varint_value and decode_fixed_value, which are now shared by extract_varint_kernel, extract_fixed_kernel, extract_varint_batched_kernel, and extract_fixed_batched_kernel
  • reuse extract_integer_into_buffers from both extract_typed_column and build_repeated_scalar_column
  • replace parallel decode and validation arguments with protobuf_field_decode_request, required_field_input_view, and protobuf_value_domain_view
  • pass enum data through enum_value_device_view, enum_domain_device_view, and enum_string_lookup_device_view
  • pass repeated_field_work directly through build_repeated_scalar_column, build_repeated_string_column, build_repeated_enum_string_column, and build_repeated_child_list_column; the object keeps schema_idx, total_count, LIST offsets, and the occurrences scratch buffer together
  • validate non-empty repeated work with validate_nonempty_repeated_field_work
  • replace numeric batched-scalar group indices with scalar_group
  • derive row counts in launch_scan_all_fields, launch_count_repeated_fields, and launch_scan_all_field_occurrences from d_in.size()
  • Refactor existing ProtobufTest cases to build expectedStruct and expectedValues columns and compare them with actualStruct via AssertUtils.assertStructColumnsAreEqual, replacing manual copyToHost() child traversal and spot checks.

The refactor removes 246 net lines from the six production files. It does not add protobuf field types or change the Java API.

Intentional behavior change

scan_nested_message_fields_kernel now calls scan_message_field_locations<wire_type_mismatch_policy::skip>. When a singular field inside a nested message has a wire type different from field_descriptor::expected_wire_type, the field is skipped and remains absent, while later sibling fields continue decoding.

This matches protobuf-java, which stores the mismatched field in the nested message's UnknownFieldSet, and Spark CPU, where ProtobufDataToCatalyst checks DynamicMessage.getUnknownFields only on the root message.

Top-level behavior is unchanged:

  • scan_all_fields_kernel uses wire_type_mismatch_policy::report_error
  • count_repeated_fields_kernel uses wire_type_mismatch_policy::report_error_and_skip for fields from nested.schema_lookup
  • nested singular and repeated children use wire_type_mismatch_policy::skip

Part of NVIDIA/cudf-spark#14069.

Testing

  • Native build
  • com.nvidia.spark.rapids.jni.ProtobufTest (154 tests)
  • added testNestedSingularWrongWireType_FailfastSkipsMismatchedField
  • added testNestedSingularWrongWireType_PermissiveSkipsMismatchedField
  • verified nested wrong-wire parsing with protobuf-java 3.19.6, 3.25.5, 4.29.3, and 4.33.0

@nvauto

nvauto commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

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>
@thirtiseven
thirtiseven marked this pull request as ready for review July 20, 2026 09:02
@greptile-apps

greptile-apps Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR consolidates the protobuf scan and decode interfaces by unifying six separate scanning kernels under a single scan_message_field_locations<MismatchPolicy> template, replacing fragmented batched_scalar_desc and dispatcher macros with a SCALAR_KINDS table, and bundling repeated-field state into repeated_field_work. The net effect is –246 lines across six production files with no Java API change.

  • Scan unification: scan_all_fields_kernel, count_repeated_fields_kernel, scan_all_field_occurrences_in_message, and scan_nested_message_fields_kernel now all delegate to scan_message_field_locations<MismatchPolicy>, which drives field traversal through caller-supplied on_singular/on_repeated callbacks and a compile-time mismatch policy.
  • Intentional behavior change: scan_nested_message_fields_kernel switches from report_error_and_abort to continue_silently for singular fields inside nested messages — a wire-type mismatch is silently skipped instead of aborting the scan, matching protobuf-java semantics. Two new tests (testNestedSingularWrongWireType_Failfast/Permissive) confirm both modes now produce the same output.
  • Scalar dispatch rewrite: The old macro-generated extract_varint_batched_kernel/extract_fixed_batched_kernel pair is replaced by a single extract_scalar_batched_kernel<T, DecodeFn> driven by dispatch_scalar_decoder and the compile-time SCALAR_KINDS table; decode_varint_value/decode_fixed_value become reusable per-element helpers.

Confidence Score: 5/5

Safe to merge — well-scoped refactoring with one intentional and tested behavioral change, 154 passing tests, and no Java API change.

The changes consolidate previously duplicated scan/decode paths without altering observable outputs for valid inputs. The one intentional change (silently skipping a mismatched singular field inside a nested message instead of aborting) is explicitly documented, matches protobuf-java behavior, and is covered by two new test cases. Kernel argument ordering and stride indexing were verified to be internally consistent across all call sites.

Files Needing Attention: No files require special attention. The output_index field defaulting to -1 in field_descriptor is only exercised in the unified count-scan path where it is always explicitly initialized.

Important Files Changed

Filename Overview
src/main/cpp/src/protobuf/protobuf_kernels.cu Core behavioral file: four scan kernels unified under scan_message_field_locations; intentional continue_silently change for nested singulars; enum kernels use typed views.
src/main/cpp/src/protobuf/protobuf_kernels.cuh New decode_varint_value/decode_fixed_value helpers, SCALAR_KINDS table, dispatch_scalar_decoder template, and unified extract_scalar_batched_kernel replace old macro-generated kernel pairs.
src/main/cpp/src/protobuf/protobuf_types.cuh Removed obsolete device view types; added scalar_value_input/output, enum_*_device_view typed views, output_index field to field_descriptor, and 17 device_layout_compatible static_asserts.
src/main/cpp/src/protobuf/protobuf_host_helpers.hpp New aggregated parameter structs (protobuf_field_decode_request, repeated_field_work_bundle), validate_nonempty_repeated_field_work, make_repeated_field_work_bundle, and make_offsets_column added.
src/main/cpp/src/protobuf/protobuf.cu Count-scan unified with output_indices; macro-based scalar batch dispatch replaced by SCALAR_KINDS fold; INT32 enum fallback detection simplified; new enum metadata validation.
src/main/cpp/src/protobuf/protobuf_builders.cu build_protobuf_field_values_column moved to kernels.cuh; repeated builders validate via validate_nonempty_repeated_field_work; enum_string_lookup_tables gains typed .view() method.
src/main/cpp/src/protobuf/protobuf_device_helpers.cuh get_wire_type_size, skip_field, get_field_data_location, and proto_tag::wire_type all migrated from int to proto_wire_type enum for type safety.
src/main/java/com/nvidia/spark/rapids/jni/ProtobufSchemaDescriptor.java validateEnumMetadata gains outputTypeId parameter with new validation: enum metadata now requires INT32/DEFAULT or STRING/ENUM_STRING encoding.
src/test/java/com/nvidia/spark/rapids/jni/ProtobufTest.java 154 tests refactored to AssertUtils.assertStructColumnsAreEqual; two new wrong-wire-type tests confirm continue_silently behavioral change; enumMetadata calls use new String varargs form.

Reviews (10): Last reviewed commit: "Merge branch 'main' into protobuf-refact..." | Re-trigger Greptile

Comment thread src/test/java/com/nvidia/spark/rapids/jni/ProtobufTest.java
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Comment thread src/main/cpp/src/protobuf/protobuf_types.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_types.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_host_helpers.hpp Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Comment thread src/test/java/com/nvidia/spark/rapids/jni/ProtobufSchemaDescriptorBuilder.java Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_host_helpers.hpp Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
field_location location) {
field_locations[f] = location;
auto record_singular = [&](int f, field_location location) {
field_locations[fields.lookup.data[f].output_index] = location;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

That's not quite what I asked. Have we tried benchmarks to confirm that this is actually a performance issue either way? Your new code is simpler, so readability-wise it's an improvement regardless (though we might want to make the same change to scan_nested_message_fields_kernel and clean up the unused output_index assignments in make_field_descriptors when output_indices is empty).

Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_builders.cu Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
Comment thread src/main/cpp/src/protobuf/protobuf_types.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment on lines 681 to 684
if (!field.enum_valid_values.empty()) {
validate_enum_and_propagate_rows(
out, valid, field.enum_valid_values, decode_ctx, {num_items, top_row_indices}, stream);
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

[Really optional] This is actually the only difference between the INT32 case and extract_and_build_scalar_field_column<int32_t>. I wish we could add this as an if constexpr (std::is_same_v<T, int32_t>) block into extract_and_build_scalar_field_column, and replace the INT32 case as well, but even if we inline extract_and_build_scalar_column, as suggested in #4836 (comment), we'd still need to propagate top_row_indices, so let's defer and think on how we could consolidate it in a future PR.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Agreed and let's defer this.

Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
igorpeshansky
igorpeshansky previously approved these changes Aug 4, 2026

@igorpeshansky igorpeshansky left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM :shipit: with some minor remaining cleanup items.

Comment thread src/main/cpp/src/protobuf/protobuf.cu Outdated
Comment thread src/main/cpp/src/protobuf/protobuf.cu
Comment thread src/main/cpp/src/protobuf/protobuf_kernels.cuh Outdated
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
@thirtiseven

Copy link
Copy Markdown
Collaborator Author

build

@igorpeshansky

Copy link
Copy Markdown
Collaborator

build

@igorpeshansky igorpeshansky left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM :shipit:

Builds were failing due to #4932, so needed to sync past 0235e25. Should work now.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants