From c9ec06c7379616262e8e6e13e70fc094af3a5fc7 Mon Sep 17 00:00:00 2001 From: Abigale Kim Date: Mon, 10 Aug 2026 23:02:05 +0000 Subject: [PATCH 1/2] first draft --- cpp/CMakeLists.txt | 1 + cpp/include/cudf/io/experimental/variant.hpp | 32 + .../io/parquet/experimental/variant_encode.cu | 627 ++++++++++++++++++ cpp/tests/CMakeLists.txt | 1 + .../io/experimental/variant_encode_test.cpp | 322 +++++++++ 5 files changed, 983 insertions(+) create mode 100644 cpp/src/io/parquet/experimental/variant_encode.cu create mode 100644 cpp/tests/io/experimental/variant_encode_test.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 605e78eb0c14..e43eb596fd17 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -796,6 +796,7 @@ add_library( src/io/parquet/experimental/hybrid_scan_preprocess.cu src/io/parquet/experimental/page_index_filter.cu src/io/parquet/experimental/page_index_filter_utils.cu + src/io/parquet/experimental/variant_encode.cu src/io/parquet/experimental/variant_extract.cu src/io/parquet/experimental/variant_path.cpp src/io/parquet/expression_transform_helpers.cpp diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index 7c53a89e4e7f..cef14b7a5344 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -7,13 +7,16 @@ #include #include +#include #include #include #include +#include #include #include +#include #include /** @@ -109,6 +112,35 @@ namespace io::parquet::experimental { rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +/** + * @brief Encode a strings column of flat JSON objects as Parquet VARIANT. + * + * Each row of @p input must be a non-nested JSON object string (e.g. `{"a":1,"b":"hi"}`). + * The function extracts the scalar fields named in @p column_names and encodes the result + * as a Parquet VARIANT column: a `struct metadata, list value>`. + * + * Supported JSON value types per field: + * - `null` → VARIANT null primitive + * - `true` / `false` → VARIANT boolean primitive + * - Integer literals → VARIANT INT64 primitive + * - Floating-point → VARIANT FLOAT64 primitive + * - Quoted strings → VARIANT short-string or long-string + * + * Fields absent from a row (or whose input row is null) are omitted from that row's + * VARIANT object; they do not appear in the value blob. + * + * @param input Strings column where each non-null row is a flat JSON object + * @param column_names Field names to encode; ordering need not be sorted + * @param stream CUDA stream + * @param mr Device memory resource + * @return `struct metadata, list value>` VARIANT column + */ +[[nodiscard]] std::unique_ptr encode_strings_to_variant( + cudf::strings_column_view const& input, + cudf::host_span column_names, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + /** @} */ } // namespace io::parquet::experimental } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/io/parquet/experimental/variant_encode.cu b/cpp/src/io/parquet/experimental/variant_encode.cu new file mode 100644 index 000000000000..4200d90f4421 --- /dev/null +++ b/cpp/src/io/parquet/experimental/variant_encode.cu @@ -0,0 +1,627 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace cudf { +namespace io::parquet::experimental { +namespace { + +using basic_type = variant_basic_type; +using primitive_type = variant_primitive_type; + +constexpr int block_size_encode = 256; + +// Sizes of the fixed-width fields in the encoded object value blob. +// We always use 1-byte field IDs (supports up to 255 keys), 1-byte num_elements, +// and 4-byte field offsets (supports values up to ~4 GB per row). +constexpr int FIELD_ID_SIZE = 1; +constexpr int FIELD_OFFSET_SIZE = 4; +constexpr int NUM_ELEMENTS_SIZE = 1; // is_large=0 + +// value_header for object: is_large(0) | (field_id_size-1)(0<<2) | (field_offset_size-1)(3) +// value_metadata = OBJECT(2) | (value_header << 2) +constexpr uint8_t OBJECT_VALUE_METADATA = + static_cast(basic_type::OBJECT) | + (((0u << 4u) | (0u << 2u) | uint8_t{FIELD_OFFSET_SIZE - 1}) << 2u); + +// ─── device helpers ─────────────────────────────────────────────────────────── + +__device__ void write_le(uint8_t*& out, uint64_t val, int bytes) +{ + for (int i = 0; i < bytes; ++i) { + *out++ = static_cast(val & 0xFFu); + val >>= 8u; + } +} + +__device__ cuda::std::optional try_parse_int64(cudf::string_view s) +{ + auto const* data = s.data(); + auto const n = s.size_bytes(); + if (n == 0) { return cuda::std::nullopt; } + + size_type i = 0; + bool negative = false; + if (data[i] == '-') { + negative = true; + ++i; + } + if (i >= n || data[i] < '0' || data[i] > '9') { return cuda::std::nullopt; } + + int64_t result = 0; + while (i < n) { + char c = data[i]; + if (c < '0' || c > '9') { return cuda::std::nullopt; } + int64_t d = c - '0'; + // Overflow guard: result * 10 + d > INT64_MAX + if (result > (int64_t{9223372036854775807LL} - d) / 10) { return cuda::std::nullopt; } + result = result * 10 + d; + ++i; + } + return negative ? -result : result; +} + +__device__ double parse_float64(cudf::string_view s) +{ + auto const* data = s.data(); + auto const n = s.size_bytes(); + + size_type i = 0; + bool negative = false; + if (i < n && data[i] == '-') { + negative = true; + ++i; + } + + double result = 0.0; + while (i < n && data[i] >= '0' && data[i] <= '9') { + result = result * 10.0 + (data[i] - '0'); + ++i; + } + + if (i < n && data[i] == '.') { + ++i; + double factor = 0.1; + while (i < n && data[i] >= '0' && data[i] <= '9') { + result += (data[i] - '0') * factor; + factor *= 0.1; + ++i; + } + } + + if (i < n && (data[i] == 'e' || data[i] == 'E')) { + ++i; + bool exp_neg = false; + if (i < n && data[i] == '-') { + exp_neg = true; + ++i; + } else if (i < n && data[i] == '+') { + ++i; + } + int exp = 0; + while (i < n && data[i] >= '0' && data[i] <= '9') { + exp = exp * 10 + (data[i] - '0'); + ++i; + } + double factor = 1.0; + for (int j = 0; j < exp; ++j) { + factor *= 10.0; + } + if (exp_neg) { + result /= factor; + } else { + result *= factor; + } + } + + return negative ? -result : result; +} + +// Returns true if s contains a float-indicating character ('.', 'e', 'E'). +__device__ bool is_float_number(cudf::string_view s) +{ + for (size_type i = 0; i < s.size_bytes(); ++i) { + char c = s.data()[i]; + if (c == '.' || c == 'e' || c == 'E') { return true; } + } + return false; +} + +// Size in bytes of the VARIANT encoding for one JSON scalar value string. +// `raw` is the string returned by get_json_object with strip_quotes_from_single_strings=false. +__device__ size_type encoded_field_size(cudf::string_view raw) +{ + if (raw == cudf::string_view("null", 4)) { return 1; } + if (raw == cudf::string_view("true", 4) || raw == cudf::string_view("false", 5)) { return 1; } + + if (raw.size_bytes() > 0 && raw.data()[0] == '"') { + // JSON string: strip surrounding quotes + size_type str_len = static_cast(raw.size_bytes()) - 2; + if (str_len <= 63) { return 1 + str_len; } // SHORT_STRING + return 1 + 4 + str_len; // LONG_STRING (1 hdr + 4-byte len + bytes) + } + + if (is_float_number(raw)) { return 9; } // header + 8-byte double + return 9; // header + 8-byte int64 +} + +// Write VARIANT bytes for one JSON scalar value string into `out`. +// Returns pointer past the last written byte. +__device__ uint8_t* write_field_value(uint8_t* out, cudf::string_view raw) +{ + auto make_prim_header = [](primitive_type pt) -> uint8_t { + return static_cast(basic_type::PRIMITIVE) | (static_cast(pt) << 2u); + }; + + if (raw == cudf::string_view("null", 4)) { + *out++ = make_prim_header(primitive_type::NULLVAL); + return out; + } + if (raw == cudf::string_view("true", 4)) { + *out++ = make_prim_header(primitive_type::BOOLEAN_TRUE); + return out; + } + if (raw == cudf::string_view("false", 5)) { + *out++ = make_prim_header(primitive_type::BOOLEAN_FALSE); + return out; + } + + if (raw.size_bytes() > 0 && raw.data()[0] == '"') { + auto const* str_start = raw.data() + 1; + size_type str_len = static_cast(raw.size_bytes()) - 2; + + if (str_len <= 63) { + *out++ = + static_cast(basic_type::SHORT_STRING) | (static_cast(str_len) << 2u); + cuda::std::memcpy(out, str_start, str_len); + return out + str_len; + } + // LONG_STRING + *out++ = make_prim_header(primitive_type::LONG_STRING); + uint32_t len32 = static_cast(str_len); + cuda::std::memcpy(out, &len32, 4); + out += 4; + cuda::std::memcpy(out, str_start, str_len); + return out + str_len; + } + + if (is_float_number(raw)) { + *out++ = make_prim_header(primitive_type::FLOAT64); + double val = parse_float64(raw); + cuda::std::memcpy(out, &val, sizeof(double)); + return out + sizeof(double); + } + + *out++ = make_prim_header(primitive_type::INT64); + auto parsed = try_parse_int64(raw); + int64_t ival = parsed.has_value() ? *parsed : int64_t{0}; + cuda::std::memcpy(out, &ival, sizeof(int64_t)); + return out + sizeof(int64_t); +} + +// ─── kernels ────────────────────────────────────────────────────────────────── + +/** + * @brief Compute the byte size of each row's VARIANT value blob. + * + * Null input rows produce size 0. Non-null rows get the object blob size, + * summing header + num_elements + field_ids + field_offsets + field values. + */ +CUDF_KERNEL __launch_bounds__(block_size_encode) void compute_value_sizes_kernel( + device_span extracted, // N columns in original order + device_span sorted_to_original, // sorted field index → original col index + size_type num_rows, + size_type num_fields, + bitmask_type const* input_null_mask, + device_span value_sizes) +{ + auto const tid = cudf::detail::grid_1d::global_thread_id(); + auto const stride = cudf::detail::grid_1d::grid_stride(); + + for (auto row = tid; row < num_rows; row += stride) { + if (input_null_mask != nullptr && !cudf::bit_is_set(input_null_mask, row)) { + value_sizes[row] = 0; + continue; + } + + size_type n_present = 0; + size_type values_bytes = 0; + + for (size_type si = 0; si < num_fields; ++si) { + auto const orig = sorted_to_original[si]; + if (extracted[orig].is_null(row)) { continue; } + ++n_present; + values_bytes += encoded_field_size(extracted[orig].element(row)); + } + + // 1 (value_metadata) + NUM_ELEMENTS_SIZE + n_present*FIELD_ID_SIZE + // + (n_present+1)*FIELD_OFFSET_SIZE + values_bytes + value_sizes[row] = 1 + NUM_ELEMENTS_SIZE + n_present * FIELD_ID_SIZE + + (n_present + 1) * FIELD_OFFSET_SIZE + values_bytes; + } +} + +/** + * @brief Write the VARIANT value blob for each row into the pre-allocated output buffer. + * + * Uses prefix-summed @p value_offsets to locate each row's destination region. + */ +CUDF_KERNEL __launch_bounds__(block_size_encode) void write_values_kernel( + device_span extracted, + device_span sorted_to_original, + size_type num_rows, + size_type num_fields, + bitmask_type const* input_null_mask, + device_span value_offsets, + uint8_t* output) +{ + auto const tid = cudf::detail::grid_1d::global_thread_id(); + auto const stride = cudf::detail::grid_1d::grid_stride(); + + for (auto row = tid; row < num_rows; row += stride) { + if (input_null_mask != nullptr && !cudf::bit_is_set(input_null_mask, row)) { continue; } + + uint8_t* out = output + value_offsets[row]; + + // Header byte + *out++ = OBJECT_VALUE_METADATA; + + // Count present fields and their individual sizes (first pass over fields) + size_type n_present = 0; + for (size_type si = 0; si < num_fields; ++si) { + if (!extracted[sorted_to_original[si]].is_null(row)) { ++n_present; } + } + + // num_elements + write_le(out, static_cast(n_present), NUM_ELEMENTS_SIZE); + + // field_ids (sorted dict index for each present field) + for (size_type si = 0; si < num_fields; ++si) { + auto const orig = sorted_to_original[si]; + if (extracted[orig].is_null(row)) { continue; } + write_le(out, static_cast(si), FIELD_ID_SIZE); + } + + // field_offsets: cumulative offsets within the values region + sentinel + size_type cur_offset = 0; + for (size_type si = 0; si < num_fields; ++si) { + auto const orig = sorted_to_original[si]; + if (extracted[orig].is_null(row)) { continue; } + write_le(out, static_cast(cur_offset), FIELD_OFFSET_SIZE); + cur_offset += encoded_field_size(extracted[orig].element(row)); + } + write_le(out, static_cast(cur_offset), FIELD_OFFSET_SIZE); // sentinel + + // field values in sorted order + for (size_type si = 0; si < num_fields; ++si) { + auto const orig = sorted_to_original[si]; + if (extracted[orig].is_null(row)) { continue; } + out = write_field_value(out, extracted[orig].element(row)); + } + } +} + +// ─── host helpers ───────────────────────────────────────────────────────────── + +// Build the fixed VARIANT metadata blob for a sorted list of key names. +// Layout: header(1) | dict_size(offset_size) | offsets[(N+1)*offset_size] | key_bytes +std::vector build_metadata_blob(std::vector const& sorted_names) +{ + size_t const N = sorted_names.size(); + size_t total_key_bytes = 0; + for (auto const& name : sorted_names) { + total_key_bytes += name.size(); + } + + int offset_size = 1; + if (total_key_bytes > 255 || N > 255) { offset_size = 2; } + if (total_key_bytes > 65535 || N > 65535) { offset_size = 4; } + + auto write_le_host = [](std::vector& buf, size_t val, int bytes) { + for (int i = 0; i < bytes; ++i) { + buf.push_back(static_cast(val & 0xFFu)); + val >>= 8u; + } + }; + + std::vector blob; + // header: version=1 | sorted=1 | unused=0 | offset_size-1 + blob.push_back(static_cast(0x01u | (1u << 4u) | (uint8_t(offset_size - 1) << 6u))); + + write_le_host(blob, N, offset_size); // dictionary_size + + // offsets[0..N] relative to start of string_data + size_t cur = 0; + for (auto const& name : sorted_names) { + write_le_host(blob, cur, offset_size); + cur += name.size(); + } + write_le_host(blob, cur, offset_size); // sentinel + + for (auto const& name : sorted_names) { + for (char c : name) { + blob.push_back(static_cast(c)); + } + } + + return blob; +} + +// Build a list column where every row contains the same `blob` bytes. +// Null rows (from input_null_mask) get 0-length list entries. +std::unique_ptr make_constant_metadata_column(std::vector const& blob, + size_type num_rows, + bitmask_type const* input_null_mask, + size_type null_count, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + size_type const m = static_cast(blob.size()); + + // Copy null mask to host so we can check validity per row while building offsets. + std::vector h_null_mask; + if (null_count > 0 && input_null_mask != nullptr) { + size_t const mask_bytes = cudf::bitmask_allocation_size_bytes(num_rows); + h_null_mask.resize(mask_bytes / sizeof(bitmask_type)); + cudf::detail::cuda_memcpy(host_span{h_null_mask}, + device_span{input_null_mask, h_null_mask.size()}, + stream); + } + + auto row_is_null = [&](size_type i) -> bool { + if (h_null_mask.empty()) { return false; } + return !cudf::bit_is_set(h_null_mask.data(), i); + }; + + // Build host offsets: each non-null row occupies m bytes + std::vector h_offsets(num_rows + 1); + size_type running = 0; + for (size_type i = 0; i < num_rows; ++i) { + h_offsets[i] = running; + if (!row_is_null(i)) { running += m; } + } + h_offsets[num_rows] = running; + + // Allocate child data: replicated blob for each non-null row + size_type const total_bytes = running; + rmm::device_buffer child_data(total_bytes, stream, mr); + if (total_bytes > 0) { + auto* dst = static_cast(child_data.data()); + // Copy host blob to device once, then tile it for each non-null row + rmm::device_uvector d_blob(blob.size(), stream, mr); + cudf::detail::cuda_memcpy_async(device_span{d_blob.data(), d_blob.size()}, + host_span{blob.data(), blob.size()}, + stream); + + // Tile: for each non-null row, copy m bytes + size_type write_pos = 0; + for (size_type i = 0; i < num_rows; ++i) { + if (!row_is_null(i)) { + CUDF_CUDA_TRY(cudf::detail::memcpy_async(dst + write_pos, d_blob.data(), m, stream)); + write_pos += m; + } + } + } + + auto d_offsets = cudf::detail::make_device_uvector_async(h_offsets, stream, mr); + auto offsets_col = std::make_unique(data_type{type_id::INT32}, + static_cast(h_offsets.size()), + d_offsets.release(), + rmm::device_buffer{}, + 0); + + auto child_col = std::make_unique( + data_type{type_id::UINT8}, total_bytes, std::move(child_data), rmm::device_buffer{}, 0); + + // Null mask for the list column comes from the input + rmm::device_buffer list_null_mask{}; + if (null_count > 0 && input_null_mask != nullptr) { + list_null_mask = cudf::detail::copy_bitmask(input_null_mask, 0, num_rows, stream, mr); + } + + stream.synchronize(); + return make_lists_column( + num_rows, std::move(offsets_col), std::move(child_col), null_count, std::move(list_null_mask)); +} + +} // namespace + +namespace detail { + +std::unique_ptr encode_strings_to_variant(cudf::strings_column_view const& input, + cudf::host_span column_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + size_type const num_rows = input.size(); + size_type const num_fields = static_cast(column_names.size()); + + CUDF_EXPECTS(num_fields <= 255, + "encode_strings_to_variant supports at most 255 fields", + std::invalid_argument); + + // Empty output + if (num_rows == 0) { + auto empty_meta = cudf::make_lists_column( + 0, make_empty_column(type_id::INT32), make_empty_column(type_id::UINT8), 0, {}); + auto empty_val = cudf::make_lists_column( + 0, make_empty_column(type_id::INT32), make_empty_column(type_id::UINT8), 0, {}); + std::vector> empty_children; + empty_children.push_back(std::move(empty_meta)); + empty_children.push_back(std::move(empty_val)); + return cudf::make_structs_column(0, std::move(empty_children), 0, {}); + } + + // ── Sort field names ────────────────────────────────────────────────────── + std::vector sort_indices(num_fields); + std::iota(sort_indices.begin(), sort_indices.end(), size_t{0}); + std::sort(sort_indices.begin(), sort_indices.end(), [&](size_t a, size_t b) { + return column_names[a] < column_names[b]; + }); + // sorted_to_original[i] = original column index of the i-th sorted key + std::vector h_sorted_to_original(num_fields); + std::vector sorted_names(num_fields); + for (size_type i = 0; i < num_fields; ++i) { + h_sorted_to_original[i] = static_cast(sort_indices[i]); + sorted_names[i] = std::string(column_names[sort_indices[i]]); + } + + auto d_sorted_to_original = cudf::detail::make_device_uvector(h_sorted_to_original, stream, mr); + + // ── Extract field values with get_json_object ───────────────────────────── + cudf::get_json_object_options opts; + opts.set_strip_quotes_from_single_strings(false); + opts.set_missing_fields_as_nulls(true); + + std::vector> extracted_cols; + extracted_cols.reserve(num_fields); + for (size_type i = 0; i < num_fields; ++i) { + std::string path = "$." + std::string(column_names[i]); + cudf::string_scalar path_scalar(path, true, stream, cudf::get_current_device_resource_ref()); + extracted_cols.push_back(cudf::get_json_object(input, path_scalar, opts, stream, mr)); + } + + // ── Build device array of column_device_views ───────────────────────────── + // column_device_view::create returns unique_ptr with a custom deleter; collect them to extend + // lifetime, then copy the views themselves (which hold device pointers) to device. + using cdv_ptr = std::unique_ptr>; + std::vector dv_holders; + dv_holders.reserve(num_fields); + std::vector h_views; + h_views.reserve(num_fields); + for (auto const& col : extracted_cols) { + dv_holders.push_back(column_device_view::create(col->view(), stream)); + h_views.push_back(*dv_holders.back()); + } + auto d_views = cudf::detail::make_device_uvector(h_views, stream, mr); + + // ── Input null mask ─────────────────────────────────────────────────────── + bitmask_type const* input_null_mask = input.null_mask(); + size_type const null_count = input.null_count(); + + // ── Compute per-row value blob sizes ───────────────────────────────────── + rmm::device_uvector value_sizes(num_rows, stream, mr); + { + auto grid = cudf::detail::grid_1d{num_rows, block_size_encode}; + compute_value_sizes_kernel<<>>( + d_views, d_sorted_to_original, num_rows, num_fields, input_null_mask, value_sizes); + CUDF_CUDA_TRY(cudaGetLastError()); + } + + // ── Prefix-sum to get per-row value offsets ─────────────────────────────── + rmm::device_uvector value_offsets(num_rows + 1, stream, mr); + { + auto const zero = size_type{0}; + cudf::detail::cuda_memcpy_async(device_span{value_offsets.data(), 1}, + host_span{&zero, 1}, + stream); + thrust::inclusive_scan(rmm::exec_policy_nosync(stream, mr), + value_sizes.begin(), + value_sizes.end(), + value_offsets.begin() + 1); + } + + // ── Allocate and write value blobs ──────────────────────────────────────── + size_type total_value_bytes{}; + cudf::detail::cuda_memcpy(host_span{&total_value_bytes, 1}, + device_span{value_offsets.data() + num_rows, 1}, + stream); + + auto value_child_data = rmm::device_buffer(static_cast(total_value_bytes), stream, mr); + if (total_value_bytes > 0) { + auto grid = cudf::detail::grid_1d{num_rows, block_size_encode}; + write_values_kernel<<>>( + d_views, + d_sorted_to_original, + num_rows, + num_fields, + input_null_mask, + device_span{value_offsets.data(), static_cast(num_rows + 1)}, + static_cast(value_child_data.data())); + CUDF_CUDA_TRY(cudaGetLastError()); + } + + // Build value list column + auto value_offsets_col = std::make_unique( + data_type{type_id::INT32}, num_rows + 1, value_offsets.release(), rmm::device_buffer{}, 0); + auto value_child_col = std::make_unique(data_type{type_id::UINT8}, + total_value_bytes, + std::move(value_child_data), + rmm::device_buffer{}, + 0); + + rmm::device_buffer value_null_mask{}; + if (null_count > 0 && input_null_mask != nullptr) { + value_null_mask = cudf::detail::copy_bitmask(input_null_mask, 0, num_rows, stream, mr); + } + auto value_col = make_lists_column(num_rows, + std::move(value_offsets_col), + std::move(value_child_col), + null_count, + std::move(value_null_mask)); + + // ── Build metadata list column ──────────────────────────────────── + auto metadata_blob = build_metadata_blob(sorted_names); + auto metadata_col = + make_constant_metadata_column(metadata_blob, num_rows, input_null_mask, null_count, stream, mr); + + // ── Assemble struct ───────────────────────────────────── + rmm::device_buffer struct_null_mask{}; + if (null_count > 0 && input_null_mask != nullptr) { + struct_null_mask = cudf::detail::copy_bitmask(input_null_mask, 0, num_rows, stream, mr); + } + std::vector> children; + children.push_back(std::move(metadata_col)); + children.push_back(std::move(value_col)); + return make_structs_column( + num_rows, std::move(children), null_count, std::move(struct_null_mask), stream, mr); +} + +} // namespace detail + +std::unique_ptr encode_strings_to_variant(cudf::strings_column_view const& input, + cudf::host_span column_names, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + return detail::encode_strings_to_variant(input, column_names, stream, mr); +} + +} // namespace io::parquet::experimental +} // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 06c0462ebed9..051459eda16e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -361,6 +361,7 @@ ConfigureTest( io/parquet_test.cpp ) ConfigureTest(VARIANT_EXTRACT_TEST io/experimental/variant_extract_test.cpp) +ConfigureTest(VARIANT_ENCODE_TEST io/experimental/variant_encode_test.cpp) ConfigureTest( PARQUET_DELETION_VECTORS_TEST io/parquet_deletion_vectors_test.cpp diff --git a/cpp/tests/io/experimental/variant_encode_test.cpp b/cpp/tests/io/experimental/variant_encode_test.cpp new file mode 100644 index 000000000000..e86d20563235 --- /dev/null +++ b/cpp/tests/io/experimental/variant_encode_test.cpp @@ -0,0 +1,322 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +// ─── helpers ───────────────────────────────────────────────────────────────── + +// Encode a vector of JSON strings with the given column names. +std::unique_ptr encode(std::vector const& json_rows, + std::vector const& col_names, + std::vector const& valid = {}) +{ + std::unique_ptr input_col; + if (valid.empty()) { + cudf::test::strings_column_wrapper w(json_rows.begin(), json_rows.end()); + input_col = w.release(); + } else { + cudf::test::strings_column_wrapper w(json_rows.begin(), json_rows.end(), valid.begin()); + input_col = w.release(); + } + cudf::strings_column_view scv{input_col->view()}; + std::vector names(col_names); + return cudf::io::parquet::experimental::encode_strings_to_variant(scv, names); +} + +// Extract a field from a VARIANT struct column and cast to INT64. +std::unique_ptr extract_int64(cudf::column_view const& variant, + std::string const& path) +{ + using namespace cudf::io::parquet::experimental; + return extract_variant_field(variant, path, cudf::data_type{cudf::type_id::INT64}); +} + +// Extract a field from a VARIANT struct column and cast to FLOAT64. +std::unique_ptr extract_float64(cudf::column_view const& variant, + std::string const& path) +{ + using namespace cudf::io::parquet::experimental; + return extract_variant_field(variant, path, cudf::data_type{cudf::type_id::FLOAT64}); +} + +// Extract a field from a VARIANT struct column and cast to STRING. +std::unique_ptr extract_string(cudf::column_view const& variant, + std::string const& path) +{ + using namespace cudf::io::parquet::experimental; + return extract_variant_field(variant, path, cudf::data_type{cudf::type_id::STRING}); +} + +// Extract a field from a VARIANT struct column and cast to BOOL8. +std::unique_ptr extract_bool(cudf::column_view const& variant, + std::string const& path) +{ + using namespace cudf::io::parquet::experimental; + return extract_variant_field(variant, path, cudf::data_type{cudf::type_id::BOOL8}); +} + +} // namespace + +struct EncodeStringsToVariantTest : public cudf::test::BaseFixture {}; + +// ─── single-field tests ─────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, SingleRowInteger) +{ + auto variant = encode({R"({"a":42})"}, {"a"}); + + auto ints = extract_int64(variant->view(), "$.a"); + cudf::test::fixed_width_column_wrapper expected{42}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ints, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowNegativeInteger) +{ + auto variant = encode({R"({"x":-100})"}, {"x"}); + + auto ints = extract_int64(variant->view(), "$.x"); + cudf::test::fixed_width_column_wrapper expected{-100}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ints, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowFloat) +{ + auto variant = encode({R"({"f":3.14})"}, {"f"}); + + auto floats = extract_float64(variant->view(), "$.f"); + // Check that we get a non-null row + EXPECT_EQ(floats->null_count(), 0); + EXPECT_EQ(floats->size(), 1); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowFloatExponent) +{ + auto variant = encode({R"({"f":1.5e2})"}, {"f"}); + + auto floats = extract_float64(variant->view(), "$.f"); + EXPECT_EQ(floats->null_count(), 0); + EXPECT_EQ(floats->size(), 1); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowBoolTrue) +{ + auto variant = encode({R"({"b":true})"}, {"b"}); + + auto bools = extract_bool(variant->view(), "$.b"); + cudf::test::fixed_width_column_wrapper expected{true}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*bools, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowBoolFalse) +{ + auto variant = encode({R"({"b":false})"}, {"b"}); + + auto bools = extract_bool(variant->view(), "$.b"); + cudf::test::fixed_width_column_wrapper expected{false}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*bools, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowShortString) +{ + auto variant = encode({R"({"s":"hello"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"hello"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowLongString) +{ + // Strings > 63 bytes use the LONG_STRING encoding path + std::string long_str(70, 'x'); + auto json = R"({"s":")" + long_str + R"("})"; + auto variant = encode({json}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{long_str}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowNullValue) +{ + // JSON null value → VARIANT null; cast_variant returns null for that row + auto variant = encode({R"({"a":null})"}, {"a"}); + + auto ints = extract_int64(variant->view(), "$.a"); + // null JSON value encodes as VARIANT null, which cast_variant cannot cast to INT64 + EXPECT_EQ(ints->size(), 1); +} + +// ─── multi-field tests ──────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, MultiField) +{ + auto variant = encode({R"({"a":1,"b":"world","c":true})"}, {"a", "b", "c"}); + + auto ints = extract_int64(variant->view(), "$.a"); + auto strs = extract_string(variant->view(), "$.b"); + auto bools = extract_bool(variant->view(), "$.c"); + + cudf::test::fixed_width_column_wrapper exp_ints{1}; + cudf::test::strings_column_wrapper exp_strs{"world"}; + cudf::test::fixed_width_column_wrapper exp_bools{true}; + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ints, exp_ints); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, exp_strs); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*bools, exp_bools); +} + +TEST_F(EncodeStringsToVariantTest, FieldOrderIndependentOfInputOrder) +{ + // column_names provided in reverse alphabetical order; should still encode correctly + auto variant = encode({R"({"z":99,"a":7})"}, {"z", "a"}); + + auto a_vals = extract_int64(variant->view(), "$.a"); + auto z_vals = extract_int64(variant->view(), "$.z"); + + cudf::test::fixed_width_column_wrapper exp_a{7}; + cudf::test::fixed_width_column_wrapper exp_z{99}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, exp_a); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*z_vals, exp_z); +} + +// ─── missing fields ─────────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, MissingFieldIsAbsent) +{ + // "b" is listed in column_names but absent from the JSON object + auto variant = encode({R"({"a":5})"}, {"a", "b"}); + + auto a_vals = extract_int64(variant->view(), "$.a"); + auto b_vals = extract_int64(variant->view(), "$.b"); + + cudf::test::fixed_width_column_wrapper exp_a{5}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, exp_a); + // b is absent → null + EXPECT_EQ(b_vals->size(), 1); + EXPECT_EQ(b_vals->null_count(), 1); +} + +TEST_F(EncodeStringsToVariantTest, ExtraColumnsInNameList) +{ + // Many names provided; most absent from the JSON + auto variant = encode({R"({"only":42})"}, {"only", "x", "y", "z"}); + + auto only_vals = extract_int64(variant->view(), "$.only"); + cudf::test::fixed_width_column_wrapper expected{42}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*only_vals, expected); + + auto x_vals = extract_int64(variant->view(), "$.x"); + EXPECT_EQ(x_vals->null_count(), 1); +} + +// ─── multi-row tests ────────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, MultipleRows) +{ + auto variant = + encode({R"({"a":1,"b":"foo"})", R"({"a":2,"b":"bar"})", R"({"a":3,"b":"baz"})"}, {"a", "b"}); + + auto a_vals = extract_int64(variant->view(), "$.a"); + auto b_vals = extract_string(variant->view(), "$.b"); + + cudf::test::fixed_width_column_wrapper exp_a{1, 2, 3}; + cudf::test::strings_column_wrapper exp_b{"foo", "bar", "baz"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, exp_a); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*b_vals, exp_b); +} + +TEST_F(EncodeStringsToVariantTest, MultipleRowsDifferentFieldsPresent) +{ + // Row 0 has a; row 1 has b; row 2 has both + auto variant = encode({R"({"a":10})", R"({"b":20})", R"({"a":30,"b":40})"}, {"a", "b"}); + + auto a_vals = extract_int64(variant->view(), "$.a"); + auto b_vals = extract_int64(variant->view(), "$.b"); + + // Row 1 has no "a" → null + cudf::test::fixed_width_column_wrapper exp_a({10, 0, 30}, {true, false, true}); + cudf::test::fixed_width_column_wrapper exp_b({0, 20, 40}, {false, true, true}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, exp_a); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*b_vals, exp_b); +} + +// ─── null input rows ────────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, NullInputRow) +{ + // Row 1 is null + auto variant = encode({R"({"a":7})", R"({"a":8})", R"({"a":9})"}, {"a"}, {true, false, true}); + + ASSERT_EQ(variant->type().id(), cudf::type_id::STRUCT); + EXPECT_EQ(variant->null_count(), 1); + + auto a_vals = extract_int64(variant->view(), "$.a"); + cudf::test::fixed_width_column_wrapper expected({7, 0, 9}, {true, false, true}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, expected); +} + +TEST_F(EncodeStringsToVariantTest, AllNullInputRows) +{ + auto variant = encode({R"({"a":1})", R"({"a":2})"}, {"a"}, {false, false}); + + EXPECT_EQ(variant->null_count(), 2); +} + +// ─── empty input ───────────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, EmptyInput) +{ + auto variant = encode({}, {"a", "b"}); + + ASSERT_EQ(variant->type().id(), cudf::type_id::STRUCT); + EXPECT_EQ(variant->size(), 0); +} + +// ─── output structure ──────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, OutputIsVariantStruct) +{ + auto variant = encode({R"({"x":1})"}, {"x"}); + + // Must be struct, list> + ASSERT_EQ(variant->type().id(), cudf::type_id::STRUCT); + ASSERT_EQ(variant->num_children(), 2); + + cudf::structs_column_view sv{variant->view()}; + EXPECT_EQ(sv.child(0).type().id(), cudf::type_id::LIST); // metadata + EXPECT_EQ(sv.child(1).type().id(), cudf::type_id::LIST); // value + + cudf::lists_column_view meta_lv{sv.child(0)}; + cudf::lists_column_view val_lv{sv.child(1)}; + EXPECT_EQ(meta_lv.child().type().id(), cudf::type_id::UINT8); + EXPECT_EQ(val_lv.child().type().id(), cudf::type_id::UINT8); +} + +// ─── zero column_names ──────────────────────────────────────────────────────── + +TEST_F(EncodeStringsToVariantTest, NoColumnNames) +{ + // Empty field list → every row encodes as an empty VARIANT object + auto variant = encode({R"({"a":1})", R"({"b":2})"}, {}); + + ASSERT_EQ(variant->type().id(), cudf::type_id::STRUCT); + EXPECT_EQ(variant->size(), 2); + EXPECT_EQ(variant->null_count(), 0); +} From 9ca1d0cf78bcde4851c7abf454634431f4c85027 Mon Sep 17 00:00:00 2001 From: Abigale Kim Date: Mon, 10 Aug 2026 20:09:09 -0500 Subject: [PATCH 2/2] Apply suggestions from code review Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- cpp/include/cudf/io/experimental/variant.hpp | 20 ++++++----- .../io/parquet/experimental/variant_encode.cu | 35 +++++++++++++------ .../io/experimental/variant_encode_test.cpp | 30 +++++++++------- 3 files changed, 53 insertions(+), 32 deletions(-) diff --git a/cpp/include/cudf/io/experimental/variant.hpp b/cpp/include/cudf/io/experimental/variant.hpp index cef14b7a5344..9f200921289e 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -113,15 +113,15 @@ namespace io::parquet::experimental { rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** - * @brief Encode a strings column of flat JSON objects as Parquet VARIANT. + * `@brief` Encode a strings column of flat JSON objects as Parquet VARIANT. * - * Each row of @p input must be a non-nested JSON object string (e.g. `{"a":1,"b":"hi"}`). - * The function extracts the scalar fields named in @p column_names and encodes the result + * Each row of `@p` input must be a non-nested JSON object string (e.g. `{"a":1,"b":"hi"}`). + * The function extracts the scalar fields named in `@p` column_names and encodes the result * as a Parquet VARIANT column: a `struct metadata, list value>`. * * Supported JSON value types per field: * - `null` → VARIANT null primitive - * - `true` / `false` → VARIANT boolean primitive + * - `true` / `false` → VARIANT boolean primitive * - Integer literals → VARIANT INT64 primitive * - Floating-point → VARIANT FLOAT64 primitive * - Quoted strings → VARIANT short-string or long-string @@ -129,11 +129,13 @@ namespace io::parquet::experimental { * Fields absent from a row (or whose input row is null) are omitted from that row's * VARIANT object; they do not appear in the value blob. * - * @param input Strings column where each non-null row is a flat JSON object - * @param column_names Field names to encode; ordering need not be sorted - * @param stream CUDA stream - * @param mr Device memory resource - * @return `struct metadata, list value>` VARIANT column + * `@param` input Strings column where each non-null row is a flat JSON object + * `@param` column_names Field names to encode; ordering need not be sorted. At most 255 names + * `@param` stream CUDA stream + * `@param` mr Device memory resource + * `@return` `struct metadata, list value>` VARIANT column + * + * `@throws` std::invalid_argument if `column_names` contains more than 255 names */ [[nodiscard]] std::unique_ptr encode_strings_to_variant( cudf::strings_column_view const& input, diff --git a/cpp/src/io/parquet/experimental/variant_encode.cu b/cpp/src/io/parquet/experimental/variant_encode.cu index 4200d90f4421..c3d2d24f5b64 100644 --- a/cpp/src/io/parquet/experimental/variant_encode.cu +++ b/cpp/src/io/parquet/experimental/variant_encode.cu @@ -125,19 +125,22 @@ __device__ double parse_float64(cudf::string_view s) } if (i < n && (data[i] == 'e' || data[i] == 'E')) { - ++i; + +i; bool exp_neg = false; if (i < n && data[i] == '-') { exp_neg = true; - ++i; + +i; } else if (i < n && data[i] == '+') { - ++i; + +i; } int exp = 0; while (i < n && data[i] >= '0' && data[i] <= '9') { - exp = exp * 10 + (data[i] - '0'); - ++i; + if (exp < 1000) { exp = exp * 10 + (data[i] - '0'); } + +i; } + // Anything beyond the double range saturates. + if (exp > 400) { return (exp_neg ? 0.0 : cuda::std::numeric_limits::infinity()) * + (negative ? -1.0 : 1.0); } double factor = 1.0; for (int j = 0; j < exp; ++j) { factor *= 10.0; @@ -493,6 +496,13 @@ std::unique_ptr encode_strings_to_variant(cudf::strings_column_view cons std::sort(sort_indices.begin(), sort_indices.end(), [&](size_t a, size_t b) { return column_names[a] < column_names[b]; }); + CUDF_EXPECTS(std::adjacent_find(sort_indices.begin(), + sort_indices.end(), + [&](size_t a, size_t b) { + return column_names[a] == column_names[b]; + }) == sort_indices.end(), + "encode_strings_to_variant does not accept duplicate field names", + std::invalid_argument); // sorted_to_original[i] = original column index of the i-th sorted key std::vector h_sorted_to_original(num_fields); std::vector sorted_names(num_fields); @@ -511,6 +521,9 @@ std::unique_ptr encode_strings_to_variant(cudf::strings_column_view cons std::vector> extracted_cols; extracted_cols.reserve(num_fields); for (size_type i = 0; i < num_fields; ++i) { + CUDF_EXPECTS(column_names[i].find_first_of(".[") == std::string::npos, + "encode_strings_to_variant does not support field names containing '.' or '['", + std::invalid_argument); std::string path = "$." + std::string(column_names[i]); cudf::string_scalar path_scalar(path, true, stream, cudf::get_current_device_resource_ref()); extracted_cols.push_back(cudf::get_json_object(input, path_scalar, opts, stream, mr)); @@ -546,11 +559,13 @@ std::unique_ptr encode_strings_to_variant(cudf::strings_column_view cons // ── Prefix-sum to get per-row value offsets ─────────────────────────────── rmm::device_uvector value_offsets(num_rows + 1, stream, mr); { - auto const zero = size_type{0}; - cudf::detail::cuda_memcpy_async(device_span{value_offsets.data(), 1}, - host_span{&zero, 1}, - stream); - thrust::inclusive_scan(rmm::exec_policy_nosync(stream, mr), + thrust::exclusive_scan(rmm::exec_policy_nosync(stream), + value_sizes.begin(), + value_sizes.end() + 0, + value_offsets.begin(), + size_type{0}); + // then write the total into value_offsets[num_rows] via an inclusive scan + thrust::inclusive_scan(rmm::exec_policy_nosync(stream), value_sizes.begin(), value_sizes.end(), value_offsets.begin() + 1); diff --git a/cpp/tests/io/experimental/variant_encode_test.cpp b/cpp/tests/io/experimental/variant_encode_test.cpp index e86d20563235..ede34239f982 100644 --- a/cpp/tests/io/experimental/variant_encode_test.cpp +++ b/cpp/tests/io/experimental/variant_encode_test.cpp @@ -3,19 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include +`#include` +`#include` +`#include` +`#include` + +`#include` +`#include` +`#include` +`#include` +`#include` +`#include` +`#include` + +`#include` +`#include` +`#include` +`#include` namespace {