Skip to content

Consolidate protobuf scan and decode interfaces #4836

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
thirtiseven wants to merge 12 commits into
base: main
Choose a base branch
from
Open

Consolidate protobuf scan and decode interfaces #4836

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
f0d1897
Consolidate protobuf scan and decode interfaces
thirtiseven Jul 17, 2026
48d80c8
add a regression tests
thirtiseven Jul 20, 2026
964ba8a
AssertUtils.assertStructColumnsAreEqual refactor
thirtiseven Jul 20, 2026
6fe0c3e
address comment
thirtiseven Jul 20, 2026
a33b516
address comments
thirtiseven Jul 23, 2026
dd80cb5
address comments
thirtiseven Jul 24, 2026
661c357
address comments
thirtiseven Jul 27, 2026
882e82d
style
thirtiseven Jul 27, 2026
65cbdf7
address comments
thirtiseven Jul 29, 2026
eaa2a8e
address comments
thirtiseven Jul 30, 2026
93f8e48
address comments
thirtiseven Aug 5, 2026
67907da
Merge branch 'main' into protobuf-refactor-consolidate-interfaces
igorpeshansky Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
610 changes: 161 additions & 449 deletions src/main/cpp/src/protobuf/protobuf.cu
View file
Open in desktop

Large diffs are not rendered by default.

10 changes: 0 additions & 10 deletions src/main/cpp/src/protobuf/protobuf.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,6 @@ enum class proto_encoding : int {
ENUM_STRING = 3,
};

CUDF_HOST_DEVICE constexpr int encoding_value(proto_encoding encoding)
{
return static_cast<int>(encoding);
}

constexpr int MAX_FIELD_NUMBER = (1 << 29) - 1;

enum class proto_wire_type : int {
Expand All @@ -53,11 +48,6 @@ enum class proto_wire_type : int {
I32BIT = 5,
};

CUDF_HOST_DEVICE constexpr int wire_type_value(proto_wire_type wire_type)
{
return static_cast<int>(wire_type);
}

constexpr int MAX_NESTING_DEPTH = 10;

struct nested_field_descriptor {
Expand Down
420 changes: 109 additions & 311 deletions src/main/cpp/src/protobuf/protobuf_builders.cu
View file
Open in desktop

Large diffs are not rendered by default.

47 changes: 25 additions & 22 deletions src/main/cpp/src/protobuf/protobuf_device_helpers.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,10 +82,10 @@ void set_error_once_async(protobuf_error* error_flag,
protobuf_error error,
rmm::cuda_stream_view stream);

__device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t const* end)
__device__ inline int get_wire_type_size(proto_wire_type wt, uint8_t const* cur, uint8_t const* end)
{
switch (wt) {
case wire_type_value(proto_wire_type::VARINT): {
case proto_wire_type::VARINT: {
// Need to scan to find the end of varint
int count = 0;
while (cur < end && count < MAX_VARINT_BYTES) {
Expand All @@ -94,15 +94,15 @@ __device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t con
}
return -1; // Invalid varint
}
case wire_type_value(proto_wire_type::I64BIT):
case proto_wire_type::I64BIT:
// Check if there's enough data for 8 bytes
if (end - cur < 8) return -1;
return 8;
case wire_type_value(proto_wire_type::I32BIT):
case proto_wire_type::I32BIT:
// Check if there's enough data for 4 bytes
if (end - cur < 4) return -1;
return 4;
case wire_type_value(proto_wire_type::LEN): {
case proto_wire_type::LEN: {
uint64_t len;
int n;
if (!read_varint(cur, end, len, n)) return -1;
Expand All @@ -112,7 +112,7 @@ __device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t con
}
return n + static_cast<int>(len);
}
case wire_type_value(proto_wire_type::SGROUP): {
case proto_wire_type::SGROUP: {
auto const* start = cur;
int depth = 1;
while (cur < end && depth > 0) {
Expand All @@ -121,24 +121,24 @@ __device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t con
if (!read_varint(cur, end, key, key_bytes)) return -1;
cur += key_bytes;

int inner_wt = static_cast<int>(key & 0x7);
if (inner_wt == wire_type_value(proto_wire_type::EGROUP)) {
auto const inner_wt = static_cast<proto_wire_type>(key & 0x7);
if (inner_wt == proto_wire_type::EGROUP) {
--depth;
if (depth == 0) { return static_cast<int>(cur - start); }
} else if (inner_wt == wire_type_value(proto_wire_type::SGROUP)) {
} else if (inner_wt == proto_wire_type::SGROUP) {
if (++depth > 32) return -1;
} else {
int inner_size = -1;
switch (inner_wt) {
case wire_type_value(proto_wire_type::VARINT): {
case proto_wire_type::VARINT: {
uint64_t dummy;
int vbytes;
if (!read_varint(cur, end, dummy, vbytes)) return -1;
inner_size = vbytes;
break;
}
case wire_type_value(proto_wire_type::I64BIT): inner_size = 8; break;
case wire_type_value(proto_wire_type::LEN): {
case proto_wire_type::I64BIT: inner_size = 8; break;
case proto_wire_type::LEN: {
uint64_t len;
int len_bytes;
if (!read_varint(cur, end, len, len_bytes)) return -1;
Expand All @@ -148,7 +148,7 @@ __device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t con
inner_size = len_bytes + static_cast<int>(len);
break;
}
case wire_type_value(proto_wire_type::I32BIT): inner_size = 4; break;
case proto_wire_type::I32BIT: inner_size = 4; break;
default: return -1;
}
if (inner_size < 0 || cur + inner_size > end) return -1;
Expand All @@ -157,21 +157,21 @@ __device__ inline int get_wire_type_size(int wt, uint8_t const* cur, uint8_t con
}
return -1;
}
case wire_type_value(proto_wire_type::EGROUP): return 0;
case proto_wire_type::EGROUP: return 0;
default: return -1;
}
}

__device__ inline bool skip_field(uint8_t const* cur,
uint8_t const* end,
int wt,
proto_wire_type wt,
uint8_t const*& out_cur)
{
// A bare end-group is only valid while a start-group payload is being parsed recursively inside
// get_wire_type_size(wire_type_value(proto_wire_type::SGROUP)).
// get_wire_type_size(proto_wire_type::SGROUP).
// The scan/count kernels should never accept it as a standalone field because Spark CPU treats
// unmatched end-groups as malformed protobuf.
if (wt == wire_type_value(proto_wire_type::EGROUP)) { return false; }
if (wt == proto_wire_type::EGROUP) { return false; }

int size = get_wire_type_size(wt, cur, end);
if (size < 0) return false;
Expand All @@ -185,10 +185,13 @@ __device__ inline bool skip_field(uint8_t const* cur,
* Get the data offset and length for a field at current position.
* Returns true on success, false on error.
*/
__device__ inline bool get_field_data_location(
uint8_t const* cur, uint8_t const* end, int wt, int32_t& data_offset, int32_t& data_length)
__device__ inline bool get_field_data_location(uint8_t const* cur,
uint8_t const* end,
proto_wire_type wt,
int32_t& data_offset,
int32_t& data_length)
{
if (wt == wire_type_value(proto_wire_type::LEN)) {
if (wt == proto_wire_type::LEN) {
// For length-delimited, read the length prefix
uint64_t len;
int len_bytes;
Expand Down Expand Up @@ -246,7 +249,7 @@ __device__ inline bool check_message_bounds(T start,

struct proto_tag {
int field_number;
int wire_type;
proto_wire_type wire_type;
};

__device__ inline bool decode_tag(uint8_t const*& cur,
Expand All @@ -268,7 +271,7 @@ __device__ inline bool decode_tag(uint8_t const*& cur,
return false;
}
tag.field_number = static_cast<int>(fn);
tag.wire_type = static_cast<int>(key & 0x7);
tag.wire_type = static_cast<proto_wire_type>(key & 0x7);
return true;
}

Expand Down
Loading