Skip to content

Commit eaa2a8e

Browse files
thirtiseventhirtiseven
committed
address comments
Signed-off-by: Haoyang Li <haoyangl@nvidia.com>
1 parent 65cbdf7 commit eaa2a8e

5 files changed

Lines changed: 101 additions & 119 deletions

File tree

‎src/main/cpp/src/protobuf/protobuf.cu‎

Lines changed: 23 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -311,16 +311,16 @@ void validate_decode_context(protobuf_decode_context const& context)
311311
std::to_string(i),
312312
std::invalid_argument);
313313

314-
auto const has_enum_metadata =
315-
!(context.enum_valid_values[i].empty() && context.enum_names[i].empty());
314+
auto const has_enum_metadata = !context.enum_valid_values[i].empty();
316315
auto const is_numeric_enum =
317316
type.id() == cudf::type_id::INT32 && field.encoding == proto_encoding::DEFAULT;
318317
auto const is_string_enum =
319318
type.id() == cudf::type_id::STRING && field.encoding == proto_encoding::ENUM_STRING;
320319
CUDF_EXPECTS(!has_enum_metadata || is_numeric_enum || is_string_enum,
321320
"protobuf decode context: enum metadata requires INT32/DEFAULT or "
322321
"STRING/ENUM_STRING at field " +
323-
std::to_string(i),
322+
std::to_string(i) + " (output_type=" + cudf::type_to_name(type) +
323+
", encoding=" + std::to_string(static_cast<int>(field.encoding)) + ")",
324324
std::invalid_argument);
325325

326326
if (field.encoding == proto_encoding::ENUM_STRING) {
@@ -573,13 +573,13 @@ std::unique_ptr<cudf::column> decode_protobuf_to_struct(cudf::column_view const&
573573
// Batched scalar extraction: group non-special fixed-width fields by extraction
574574
// category and extract all fields of each category with a single 2D kernel launch.
575575
{
576-
static constexpr auto fallback = scalar_kinds.size();
577-
std::array<std::vector<int>, scalar_kinds.size() + 1> group_lists;
576+
static constexpr auto fallback = SCALAR_KINDS.size();
577+
std::array<std::vector<int>, SCALAR_KINDS.size() + 1> group_lists;
578578
// Implied constexpr (https://en.cppreference.com/cpp/language/lambda).
579579
auto find_group = [](cudf::type_id type, proto_encoding encoding) {
580580
auto const decode = get_scalar_decode_kind(type, encoding);
581-
auto const it = std::ranges::find(scalar_kinds, scalar_kind{type, decode});
582-
return static_cast<size_t>(it - scalar_kinds.begin());
581+
auto const it = std::ranges::find(SCALAR_KINDS, scalar_kind{type, decode});
582+
return static_cast<size_t>(it - SCALAR_KINDS.begin());
583583
};
584584

585585
for (int i = 0; i < num_scalar; i++) {
@@ -638,18 +638,16 @@ std::unique_ptr<cudf::column> decode_protobuf_to_struct(cudf::column_view const&
638638
}
639639
};
640640

641-
auto launch_index = [&](auto index_constant) {
642-
constexpr auto index = decltype(index_constant)::value;
643-
using T = std::conditional_t<scalar_kinds[index].type == cudf::type_id::BOOL8,
644-
uint8_t,
645-
cudf::id_to_type<scalar_kinds[index].type>>;
646-
dispatch_scalar_decoder<T>(scalar_kinds[index].decode, [&]<auto DecodeFn>() {
647-
launch_decoder.template operator()<T, DecodeFn>(group_lists[index]);
641+
auto launch_index = [&]<std::size_t I, std::size_t = I>() {
642+
constexpr auto type = SCALAR_KINDS[I].type;
643+
using T = std::conditional_t<type == cudf::type_id::BOOL8, uint8_t, cudf::id_to_type<type>>;
644+
dispatch_scalar_decoder<T, SCALAR_KINDS[I].decode>([&]<auto DecodeFn>() {
645+
launch_decoder.template operator()<T, DecodeFn>(group_lists[I]);
648646
});
649647
};
650648
[&]<size_t... I>(std::index_sequence<I...>) {
651-
(launch_index(std::integral_constant<size_t, I>{}), ...);
652-
}(std::make_index_sequence<scalar_kinds.size()>{});
649+
(launch_index.template operator()<I>(), ...);
650+
}(std::make_index_sequence<SCALAR_KINDS.size()>{});
653651

654652
// Per-field fallback (INT32 with enum, etc.)
655653
for (int i : group_lists[fallback]) {
@@ -691,17 +689,15 @@ std::unique_ptr<cudf::column> decode_protobuf_to_struct(cudf::column_view const&
691689
// Process repeated fields (three-phase: offsets → combined scan → build columns)
692690
if (num_repeated > 0) {
693691
// Phases A and B: build LIST offsets and allocate occurrence buffers.
694-
auto rep_work = make_repeated_field_work_bundle(
695-
std::views::iota(0, num_repeated),
696-
num_repeated,
697-
[&](int repeated_position) { return repeated_field_indices[repeated_position]; },
698-
d_repeated_info.data(),
699-
num_rows,
700-
schema_context,
701-
"Top-level repeated field",
702-
stream,
703-
mr,
704-
scratch_mr);
692+
auto rep_work = make_repeated_field_work_bundle(std::views::iota(0, num_repeated),
693+
repeated_field_indices,
694+
d_repeated_info.data(),
695+
num_rows,
696+
schema_context,
697+
"Top-level repeated field",
698+
stream,
699+
mr,
700+
scratch_mr);
705701

706702
if (!rep_work.scan_descriptors.empty()) {
707703
auto scan_bundle =

‎src/main/cpp/src/protobuf/protobuf_builders.cu‎

Lines changed: 9 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -538,17 +538,15 @@ std::unique_ptr<cudf::column> build_nested_struct_column(
538538
decode_ctx,
539539
stream);
540540

541-
auto repeated_work = make_repeated_field_work_bundle(
542-
repeated_child_positions,
543-
num_child_fields,
544-
[&](int child_position) { return child_field_indices[child_position]; },
545-
d_repeated_info.data(),
546-
input.num_rows,
547-
schema,
548-
"Repeated nested-field",
549-
stream,
550-
mr,
551-
scratch_mr);
541+
auto repeated_work = make_repeated_field_work_bundle(repeated_child_positions,
542+
child_field_indices,
543+
d_repeated_info.data(),
544+
input.num_rows,
545+
schema,
546+
"Repeated nested-field",
547+
stream,
548+
mr,
549+
scratch_mr);
552550

553551
if (!repeated_work.scan_descriptors.empty()) {
554552
auto scan_bundle =

‎src/main/cpp/src/protobuf/protobuf_host_helpers.hpp‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -215,11 +215,10 @@ inline std::unique_ptr<cudf::column> make_offsets_column(cudf::size_type num_row
215215
0);
216216
}
217217

218-
template <typename PositionRange, typename SchemaIndexFn>
218+
template <typename PositionRange>
219219
inline repeated_field_work_bundle make_repeated_field_work_bundle(
220220
PositionRange const& field_positions,
221-
int num_fields,
222-
SchemaIndexFn get_schema_index,
221+
std::vector<int> const& schema_indices,
223222
field_occurrence_count const* repeated_info,
224223
int num_rows,
225224
protobuf_schema const& schema,
@@ -228,7 +227,7 @@ inline repeated_field_work_bundle make_repeated_field_work_bundle(
228227
rmm::device_async_resource_ref output_mr,
229228
rmm::device_async_resource_ref scratch_mr)
230229
{
231-
CUDF_EXPECTS(num_fields >= 0, std::string{__func__} + ": field count must be non-negative");
230+
auto const num_fields = static_cast<int>(schema_indices.size());
232231
repeated_field_work_bundle result{
233232
std::vector<std::optional<repeated_field_work>>(num_fields),
234233
cudf::detail::make_pinned_vector_async<field_occurrence_scan_desc>(0, stream)};
@@ -239,7 +238,7 @@ inline repeated_field_work_bundle make_repeated_field_work_bundle(
239238
std::string{__func__} + ": field position is out of bounds");
240239
CUDF_EXPECTS(repeated_info != nullptr,
241240
std::string{__func__} + ": repeated count buffer must be non-null");
242-
auto const schema_idx = get_schema_index(field_position);
241+
auto const schema_idx = schema_indices[field_position];
243242
// The source is row-major; expose one field's counts without materializing another buffer.
244243
auto counts_begin = thrust::make_transform_iterator(
245244
thrust::make_counting_iterator<int>(0),

‎src/main/cpp/src/protobuf/protobuf_kernels.cuh‎

Lines changed: 64 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -240,37 +240,37 @@ __device__ inline void decode_fixed_value(scalar_value_input input,
240240
if (output.valid) output.valid[index] = true;
241241
}
242242

243-
enum class scalar_decode_kind : uint8_t { fixed, varint, zigzag };
243+
enum class scalar_decode_kind : uint8_t { FIXED, VARINT, ZIGZAG };
244244

245245
struct scalar_kind {
246246
cudf::type_id type;
247247
scalar_decode_kind decode;
248248
bool operator==(scalar_kind const&) const = default;
249249
};
250250

251-
inline constexpr auto scalar_kinds = std::to_array<scalar_kind>({
252-
{cudf::type_id::INT32, scalar_decode_kind::varint},
253-
{cudf::type_id::UINT32, scalar_decode_kind::varint},
254-
{cudf::type_id::INT64, scalar_decode_kind::varint},
255-
{cudf::type_id::UINT64, scalar_decode_kind::varint},
256-
{cudf::type_id::BOOL8, scalar_decode_kind::varint},
257-
{cudf::type_id::INT32, scalar_decode_kind::zigzag},
258-
{cudf::type_id::INT64, scalar_decode_kind::zigzag},
259-
{cudf::type_id::FLOAT32, scalar_decode_kind::fixed},
260-
{cudf::type_id::FLOAT64, scalar_decode_kind::fixed},
261-
{cudf::type_id::INT32, scalar_decode_kind::fixed},
262-
{cudf::type_id::UINT32, scalar_decode_kind::fixed},
263-
{cudf::type_id::INT64, scalar_decode_kind::fixed},
264-
{cudf::type_id::UINT64, scalar_decode_kind::fixed},
251+
inline constexpr auto SCALAR_KINDS = std::to_array<scalar_kind>({
252+
{cudf::type_id::INT32, scalar_decode_kind::VARINT},
253+
{cudf::type_id::UINT32, scalar_decode_kind::VARINT},
254+
{cudf::type_id::INT64, scalar_decode_kind::VARINT},
255+
{cudf::type_id::UINT64, scalar_decode_kind::VARINT},
256+
{cudf::type_id::BOOL8, scalar_decode_kind::VARINT},
257+
{cudf::type_id::INT32, scalar_decode_kind::ZIGZAG},
258+
{cudf::type_id::INT64, scalar_decode_kind::ZIGZAG},
259+
{cudf::type_id::FLOAT32, scalar_decode_kind::FIXED},
260+
{cudf::type_id::FLOAT64, scalar_decode_kind::FIXED},
261+
{cudf::type_id::INT32, scalar_decode_kind::FIXED},
262+
{cudf::type_id::UINT32, scalar_decode_kind::FIXED},
263+
{cudf::type_id::INT64, scalar_decode_kind::FIXED},
264+
{cudf::type_id::UINT64, scalar_decode_kind::FIXED},
265265
});
266266

267267
constexpr scalar_decode_kind get_scalar_decode_kind(cudf::type_id type, proto_encoding encoding)
268268
{
269269
using enum cudf::type_id;
270270
using enum proto_encoding;
271-
return type == FLOAT32 || type == FLOAT64 || encoding == FIXED ? scalar_decode_kind::fixed
272-
: encoding == ZIGZAG ? scalar_decode_kind::zigzag
273-
: scalar_decode_kind::varint;
271+
return type == FLOAT32 || type == FLOAT64 || encoding == FIXED ? scalar_decode_kind::FIXED
272+
: encoding == ZIGZAG ? scalar_decode_kind::ZIGZAG
273+
: scalar_decode_kind::VARINT;
274274
}
275275

276276
template <typename T>
@@ -279,52 +279,61 @@ inline scalar_decode_kind get_scalar_decode_kind(proto_encoding encoding)
279279
if constexpr (std::is_floating_point_v<T>) {
280280
CUDF_EXPECTS(encoding == proto_encoding::DEFAULT || encoding == proto_encoding::FIXED,
281281
"Floating-point protobuf extraction requires default or fixed encoding");
282-
return scalar_decode_kind::fixed;
283282
} else if (encoding == proto_encoding::FIXED) {
284-
if constexpr (sizeof(T) == 4 || sizeof(T) == 8) {
285-
return scalar_decode_kind::fixed;
286-
} else {
287-
CUDF_FAIL("Fixed-width protobuf extraction requires a 32-bit or 64-bit output type");
288-
}
283+
CUDF_EXPECTS(sizeof(T) == 4 || sizeof(T) == 8,
284+
"Fixed-width protobuf extraction requires a 32-bit or 64-bit output type");
289285
} else if constexpr (std::is_signed_v<T>) {
290286
CUDF_EXPECTS(encoding == proto_encoding::DEFAULT || encoding == proto_encoding::ZIGZAG,
291287
"Signed varint protobuf extraction requires default or zigzag encoding");
292-
return encoding == proto_encoding::ZIGZAG ? scalar_decode_kind::zigzag
293-
: scalar_decode_kind::varint;
294-
} else {
288+
} else if constexpr (std::is_integral_v<T>) {
295289
CUDF_EXPECTS(encoding == proto_encoding::DEFAULT,
296290
"Unsigned varint protobuf extraction requires default encoding");
297-
return scalar_decode_kind::varint;
291+
} else {
292+
CUDF_FAIL("Varint protobuf extraction requires an integral output type");
293+
}
294+
return get_scalar_decode_kind(
295+
std::is_floating_point_v<T> ? cudf::type_id::FLOAT32 : cudf::type_id::INT32, encoding);
296+
}
297+
298+
template <typename T, scalar_decode_kind Decode, typename F>
299+
constexpr void dispatch_scalar_decoder(F&& f)
300+
{
301+
if constexpr (Decode == scalar_decode_kind::FIXED) {
302+
static_assert(sizeof(T) == 4 || sizeof(T) == 8);
303+
std::forward<F>(f).template operator()<decode_fixed_value<T>>();
304+
} else if constexpr (Decode == scalar_decode_kind::VARINT) {
305+
static_assert(std::is_integral_v<T>);
306+
std::forward<F>(f).template operator()<decode_varint_value<T, false>>();
307+
} else if constexpr (Decode == scalar_decode_kind::ZIGZAG) {
308+
static_assert(std::is_integral_v<T> && std::is_signed_v<T>);
309+
std::forward<F>(f).template operator()<decode_varint_value<T, true>>();
310+
} else {
311+
static_assert(Decode == scalar_decode_kind::FIXED || Decode == scalar_decode_kind::VARINT ||
312+
Decode == scalar_decode_kind::ZIGZAG);
298313
}
299314
}
300315

301316
template <typename T, typename F>
302317
inline void dispatch_scalar_decoder(scalar_decode_kind decode, F&& f)
303318
{
304319
switch (decode) {
305-
case scalar_decode_kind::fixed:
320+
case scalar_decode_kind::FIXED:
306321
if constexpr (sizeof(T) == 4 || sizeof(T) == 8) {
307-
f.template operator()<decode_fixed_value<T>>();
308-
} else {
309-
CUDF_FAIL("Fixed-width protobuf extraction requires a 32-bit or 64-bit output type");
322+
return dispatch_scalar_decoder<T, scalar_decode_kind::FIXED>(std::forward<F>(f));
310323
}
311324
break;
312-
case scalar_decode_kind::varint:
325+
case scalar_decode_kind::VARINT:
313326
if constexpr (std::is_integral_v<T>) {
314-
f.template operator()<decode_varint_value<T, false>>();
315-
} else {
316-
CUDF_FAIL("Varint protobuf extraction requires an integral output type");
327+
return dispatch_scalar_decoder<T, scalar_decode_kind::VARINT>(std::forward<F>(f));
317328
}
318329
break;
319-
case scalar_decode_kind::zigzag:
330+
case scalar_decode_kind::ZIGZAG:
320331
if constexpr (std::is_integral_v<T> && std::is_signed_v<T>) {
321-
f.template operator()<decode_varint_value<T, true>>();
322-
} else {
323-
CUDF_FAIL("Zigzag protobuf extraction requires a signed integral output type");
332+
return dispatch_scalar_decoder<T, scalar_decode_kind::ZIGZAG>(std::forward<F>(f));
324333
}
325334
break;
326-
default: CUDF_FAIL("Unknown protobuf scalar decode kind");
327335
}
336+
CUDF_UNREACHABLE("Invalid protobuf scalar decode kind/type combination");
328337
}
329338

330339
template <typename OutputType, auto DecodeFn, typename LocationProvider>
@@ -463,23 +472,6 @@ inline std::pair<rmm::device_buffer, cudf::size_type> make_null_mask_from_valid(
463472
return cudf::detail::valid_if(begin, end, pred, stream, mr);
464473
}
465474

466-
template <typename T, typename LaunchFn>
467-
std::unique_ptr<cudf::column> extract_and_build_scalar_column(cudf::data_type dt,
468-
int num_rows,
469-
LaunchFn&& launch_extract,
470-
rmm::cuda_stream_view stream,
471-
rmm::device_async_resource_ref mr)
472-
{
473-
rmm::device_uvector<T> out(num_rows, stream, mr);
474-
rmm::device_uvector<bool> valid(num_rows, stream, mr);
475-
if (num_rows == 0) {
476-
return std::make_unique<cudf::column>(dt, 0, out.release(), rmm::device_buffer{}, 0);
477-
}
478-
launch_extract(out.data(), valid.data());
479-
auto [mask, null_count] = make_null_mask_from_valid(valid, num_rows, stream, mr);
480-
return std::make_unique<cudf::column>(dt, num_rows, out.release(), std::move(mask), null_count);
481-
}
482-
483475
template <typename T, typename LocationProvider>
484476
inline void extract_scalar_into_buffers(uint8_t const* message_data,
485477
LocationProvider const& loc_provider,
@@ -521,21 +513,20 @@ std::unique_ptr<cudf::column> extract_and_build_scalar_field_column(
521513
rmm::cuda_stream_view stream,
522514
rmm::device_async_resource_ref mr)
523515
{
524-
return extract_and_build_scalar_column<T>(
525-
field.output_type,
516+
if (num_rows == 0) { return cudf::make_empty_column(field.output_type); }
517+
rmm::device_uvector<T> out(num_rows, stream, mr);
518+
rmm::device_uvector<bool> valid(num_rows, stream, mr);
519+
extract_scalar_into_buffers<T, LocationProvider>(
520+
message_data,
521+
loc_provider,
526522
num_rows,
527-
[&](T* out_ptr, bool* valid_ptr) {
528-
extract_scalar_into_buffers<T, LocationProvider>(
529-
message_data,
530-
loc_provider,
531-
num_rows,
532-
field.schema.encoding,
533-
make_scalar_decode_options<T>(field),
534-
{out_ptr, valid_ptr, decode_ctx.error->data()},
535-
stream);
536-
},
537-
stream,
538-
mr);
523+
field.schema.encoding,
524+
make_scalar_decode_options<T>(field),
525+
{out.data(), valid.data(), decode_ctx.error->data()},
526+
stream);
527+
auto [mask, null_count] = make_null_mask_from_valid(valid, num_rows, stream, mr);
528+
return std::make_unique<cudf::column>(
529+
field.output_type, num_rows, out.release(), std::move(mask), null_count);
539530
}
540531

541532
template <typename LocationProvider, typename ValidityFn>
@@ -665,9 +656,7 @@ inline std::unique_ptr<cudf::column> extract_typed_column(protobuf_field_decode_
665656
return extract_and_build_scalar_field_column<uint8_t>(
666657
field, message_data, loc_provider, num_items, decode_ctx, stream, mr);
667658
case cudf::type_id::INT32: {
668-
if (num_items == 0) {
669-
return std::make_unique<cudf::column>(dt, 0, rmm::device_buffer{}, rmm::device_buffer{}, 0);
670-
}
659+
if (num_items == 0) { return cudf::make_empty_column(dt); }
671660
rmm::device_uvector<int32_t> out(num_items, stream, mr);
672661
rmm::device_uvector<bool> valid(num_items, stream, mr);
673662
extract_scalar_into_buffers<int32_t, LocationProvider>(

‎src/main/cpp/src/protobuf/protobuf_types.cuh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,7 @@ struct field_descriptor {
9595
int field_number; // Protobuf field number
9696
proto_wire_type expected_wire_type; // Expected wire type for this field
9797
bool is_repeated; // Repeated children are scanned via count/scan kernels
98-
int output_index = -1; // Column within the matching output buffer, or -1 when unused
98+
int output_index = -1; // Matching output column, or -1 when unused
9999
};
100100

101101
/**

0 commit comments

Comments
 (0)