diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 26b7ee68aad..a392248b002 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -785,6 +785,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 08313ae21a7..ca207dec07d 100644 --- a/cpp/include/cudf/io/experimental/variant.hpp +++ b/cpp/include/cudf/io/experimental/variant.hpp @@ -8,14 +8,17 @@ #include #include #include +#include #include #include #include +#include #include #include #include +#include #include /** @@ -150,6 +153,39 @@ namespace io::parquet::experimental { cuda::stream_ref 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. 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 + * `@throws` std::invalid_argument if `column_names` contains duplicate entries + * `@throws` std::invalid_argument if any name in `column_names` contains '.' or '[' + */ +[[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 00000000000..97fb3202e73 --- /dev/null +++ b/cpp/src/io/parquet/experimental/variant_encode.cu @@ -0,0 +1,969 @@ +/* + * 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 +#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; } + + // Accumulate as a negative value to correctly represent INT64_MIN. + constexpr int64_t INT64_MIN_VAL = int64_t{-9223372036854775807LL - 1}; + 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'; + if (result < (INT64_MIN_VAL + d) / 10) { return cuda::std::nullopt; } + result = result * 10 - d; + ++i; + } + if (!negative) { + if (result == INT64_MIN_VAL) { return cuda::std::nullopt; } + return -result; + } + return 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') { + 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; + } + 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; +} + +// Validates `s` against the strict JSON number grammar (no leading zeros, digit required +// before/after '.', digit required after 'e'/'E', etc). Anything that fails this check is not +// a legal JSON number and must not be silently coerced into 0.0 by parse_float64. +__device__ bool is_valid_json_number(cudf::string_view s) +{ + auto const* data = s.data(); + auto const n = s.size_bytes(); + if (n == 0) { return false; } + + size_type i = 0; + if (data[i] == '-') { + ++i; + if (i >= n) { return false; } + } + if (data[i] == '0') { + ++i; + } else if (data[i] >= '1' && data[i] <= '9') { + ++i; + while (i < n && data[i] >= '0' && data[i] <= '9') { + ++i; + } + } else { + return false; + } + if (i < n && data[i] == '.') { + ++i; + if (i >= n || data[i] < '0' || data[i] > '9') { return false; } + while (i < n && data[i] >= '0' && data[i] <= '9') { + ++i; + } + } + if (i < n && (data[i] == 'e' || data[i] == 'E')) { + ++i; + if (i < n && (data[i] == '+' || data[i] == '-')) { ++i; } + if (i >= n || data[i] < '0' || data[i] > '9') { return false; } + while (i < n && data[i] >= '0' && data[i] <= '9') { + ++i; + } + } + return i == n; +} + +// Classification of a get_json_object result string, used to drive both size computation and +// encoding so the two stay in lock-step. +enum class scalar_kind : uint8_t { NUL, BOOL_TRUE, BOOL_FALSE, STRING, INT, FLOAT, INVALID }; + +// `raw` is the string returned by get_json_object with strip_quotes_from_single_strings=false. +// Anything that is not one of JSON null/true/false/string/number (in particular a nested +// object/array, or a malformed literal) is reported as INVALID rather than guessed at. +__device__ scalar_kind classify_scalar(cudf::string_view raw) +{ + if (raw.size_bytes() == 0) { return scalar_kind::INVALID; } + if (raw == cudf::string_view("null", 4)) { return scalar_kind::NUL; } + if (raw == cudf::string_view("true", 4)) { return scalar_kind::BOOL_TRUE; } + if (raw == cudf::string_view("false", 5)) { return scalar_kind::BOOL_FALSE; } + if (raw.data()[0] == '"') { + if (raw.size_bytes() >= 2 && raw.data()[raw.size_bytes() - 1] == '"') { + return scalar_kind::STRING; + } + return scalar_kind::INVALID; + } + if (!is_valid_json_number(raw)) { return scalar_kind::INVALID; } + return is_float_number(raw) ? scalar_kind::FLOAT : scalar_kind::INT; +} + +__device__ int hex_digit(char c) +{ + if (c >= '0' && c <= '9') { return c - '0'; } + if (c >= 'a' && c <= 'f') { return c - 'a' + 10; } + if (c >= 'A' && c <= 'F') { return c - 'A' + 10; } + return -1; +} + +// Encodes `cp` as UTF-8 into `out` (which must have room for up to 4 bytes) and returns the +// number of bytes written. +__device__ int encode_utf8(uint32_t cp, uint8_t* out) +{ + if (cp <= 0x7Fu) { + out[0] = static_cast(cp); + return 1; + } + if (cp <= 0x7FFu) { + out[0] = static_cast(0xC0u | (cp >> 6u)); + out[1] = static_cast(0x80u | (cp & 0x3Fu)); + return 2; + } + if (cp <= 0xFFFFu) { + out[0] = static_cast(0xE0u | (cp >> 12u)); + out[1] = static_cast(0x80u | ((cp >> 6u) & 0x3Fu)); + out[2] = static_cast(0x80u | (cp & 0x3Fu)); + return 3; + } + out[0] = static_cast(0xF0u | (cp >> 18u)); + out[1] = static_cast(0x80u | ((cp >> 12u) & 0x3Fu)); + out[2] = static_cast(0x80u | ((cp >> 6u) & 0x3Fu)); + out[3] = static_cast(0x80u | (cp & 0x3Fu)); + return 4; +} + +// Computes the number of UTF-8 bytes that JSON-unescaping `[s, s+len)` (the content between the +// quotes of a JSON string) would produce, without writing anything. Returns nullopt if the +// content contains an invalid or truncated escape sequence. +__device__ cuda::std::optional json_string_unescaped_size(char const* s, size_type len) +{ + size_type count = 0; + size_type i = 0; + while (i < len) { + char c = s[i]; + if (c != '\\') { + ++count; + ++i; + continue; + } + ++i; + if (i >= len) { return cuda::std::nullopt; } + char e = s[i]; + switch (e) { + case '"': + case '\\': + case '/': + case 'b': + case 'f': + case 'n': + case 'r': + case 't': + ++count; + ++i; + break; + case 'u': { + ++i; + if (i + 4 > len) { return cuda::std::nullopt; } + int cp = 0; + for (int k = 0; k < 4; ++k) { + int d = hex_digit(s[i + k]); + if (d < 0) { return cuda::std::nullopt; } + cp = (cp << 4) | d; + } + i += 4; + auto codepoint = static_cast(cp); + if (codepoint >= 0xD800u && codepoint <= 0xDBFFu) { + if (i + 6 > len || s[i] != '\\' || s[i + 1] != 'u') { return cuda::std::nullopt; } + int lo = 0; + for (int k = 0; k < 4; ++k) { + int d = hex_digit(s[i + 2 + k]); + if (d < 0) { return cuda::std::nullopt; } + lo = (lo << 4) | d; + } + if (lo < 0xDC00 || lo > 0xDFFF) { return cuda::std::nullopt; } + i += 6; + codepoint = + 0x10000u + ((codepoint - 0xD800u) << 10u) + (static_cast(lo) - 0xDC00u); + } else if (codepoint >= 0xDC00u && codepoint <= 0xDFFFu) { + return cuda::std::nullopt; // unpaired low surrogate + } + uint8_t buf[4]; + count += encode_utf8(codepoint, buf); + break; + } + default: return cuda::std::nullopt; + } + } + return count; +} + +// JSON-unescapes `[s, s+len)` into `out`. Assumes the content was already validated by +// json_string_unescaped_size (same grammar); returns the pointer past the last byte written. +__device__ uint8_t* write_json_string_unescaped(char const* s, size_type len, uint8_t* out) +{ + size_type i = 0; + while (i < len) { + char c = s[i]; + if (c != '\\') { + *out++ = static_cast(c); + ++i; + continue; + } + ++i; + char e = s[i]; + switch (e) { + case '"': + *out++ = static_cast('"'); + ++i; + break; + case '\\': + *out++ = static_cast('\\'); + ++i; + break; + case '/': + *out++ = static_cast('/'); + ++i; + break; + case 'b': + *out++ = static_cast('\b'); + ++i; + break; + case 'f': + *out++ = static_cast('\f'); + ++i; + break; + case 'n': + *out++ = static_cast('\n'); + ++i; + break; + case 'r': + *out++ = static_cast('\r'); + ++i; + break; + case 't': + *out++ = static_cast('\t'); + ++i; + break; + case 'u': { + ++i; + int cp = 0; + for (int k = 0; k < 4; ++k) { + cp = (cp << 4) | hex_digit(s[i + k]); + } + i += 4; + auto codepoint = static_cast(cp); + if (codepoint >= 0xD800u && codepoint <= 0xDBFFu) { + int lo = 0; + for (int k = 0; k < 4; ++k) { + lo = (lo << 4) | hex_digit(s[i + 2 + k]); + } + i += 6; + codepoint = + 0x10000u + ((codepoint - 0xD800u) << 10u) + (static_cast(lo) - 0xDC00u); + } + out += encode_utf8(codepoint, out); + break; + } + default: break; + } + } + return out; +} + +// Size in bytes of the VARIANT encoding for one JSON scalar value string. If `error_flag` is +// non-null, sets *error_flag (via atomicExch) and returns a harmless placeholder size when `raw` +// is not a legal scalar (nested object/array or malformed literal) or contains an invalid string +// escape; callers passing a non-null flag must check it after the kernel completes and fail the +// whole encode rather than trust the size. Callers that have already validated every row via a +// prior checked pass (see write_values_kernel) may pass nullptr, since INVALID cannot recur for +// the same input. +__device__ size_type encoded_field_size(cudf::string_view raw, int32_t* error_flag = nullptr) +{ + switch (classify_scalar(raw)) { + case scalar_kind::NUL: + case scalar_kind::BOOL_TRUE: + case scalar_kind::BOOL_FALSE: return 1; + case scalar_kind::STRING: { + size_type str_len = static_cast(raw.size_bytes()) - 2; + auto const unescaped_len = json_string_unescaped_size(raw.data() + 1, str_len); + if (!unescaped_len.has_value()) { + if (error_flag != nullptr) { atomicExch(error_flag, 1); } + return 1; + } + if (*unescaped_len <= 63) { return 1 + *unescaped_len; } // SHORT_STRING + return 1 + 4 + *unescaped_len; // LONG_STRING (1 hdr + 4-byte len + bytes) + } + case scalar_kind::FLOAT: + case scalar_kind::INT: return 9; // header + 8-byte double/int64 + case scalar_kind::INVALID: + default: + if (error_flag != nullptr) { atomicExch(error_flag, 1); } + return 1; + } +} + +// Write VARIANT bytes for one JSON scalar value string into `out`. +// Returns pointer past the last written byte. Assumes `raw` was already validated by a prior +// call to encoded_field_size that did not set the error flag. +__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); + }; + + switch (classify_scalar(raw)) { + case scalar_kind::NUL: { + *out++ = make_prim_header(primitive_type::NULLVAL); + return out; + } + case scalar_kind::BOOL_TRUE: { + *out++ = make_prim_header(primitive_type::BOOLEAN_TRUE); + return out; + } + case scalar_kind::BOOL_FALSE: { + *out++ = make_prim_header(primitive_type::BOOLEAN_FALSE); + return out; + } + case scalar_kind::STRING: { + auto const* str_start = raw.data() + 1; + size_type str_len = static_cast(raw.size_bytes()) - 2; + auto const unescaped_len = json_string_unescaped_size(str_start, str_len); + if (!unescaped_len.has_value()) { + // Should be unreachable: already validated by encoded_field_size. + *out++ = make_prim_header(primitive_type::NULLVAL); + return out; + } + auto const n = *unescaped_len; + if (n <= 63) { + *out++ = static_cast(basic_type::SHORT_STRING) | (static_cast(n) << 2u); + return write_json_string_unescaped(str_start, str_len, out); + } + // LONG_STRING + *out++ = make_prim_header(primitive_type::LONG_STRING); + uint32_t len32 = static_cast(n); + cuda::std::memcpy(out, &len32, 4); + out += 4; + return write_json_string_unescaped(str_start, str_len, out); + } + case scalar_kind::FLOAT: { + *out++ = make_prim_header(primitive_type::FLOAT64); + double val = parse_float64(raw); + cuda::std::memcpy(out, &val, sizeof(double)); + return out + sizeof(double); + } + case scalar_kind::INT: { + auto parsed = try_parse_int64(raw); + if (parsed.has_value()) { + *out++ = make_prim_header(primitive_type::INT64); + int64_t ival = *parsed; + cuda::std::memcpy(out, &ival, sizeof(int64_t)); + return out + sizeof(int64_t); + } + // Failed to parse as INT64 (out-of-range): fall back to FLOAT64. + *out++ = make_prim_header(primitive_type::FLOAT64); + double val = parse_float64(raw); + cuda::std::memcpy(out, &val, sizeof(double)); + return out + sizeof(double); + } + case scalar_kind::INVALID: + default: { + // Should be unreachable: already validated by encoded_field_size. + *out++ = make_prim_header(primitive_type::NULLVAL); + return out; + } + } +} + +// ─── 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, + int32_t* error_flag) +{ + 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; + int64_t 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), error_flag); + } + + // 1 (value_metadata) + NUM_ELEMENTS_SIZE + n_present*FIELD_ID_SIZE + // + (n_present+1)*FIELD_OFFSET_SIZE + values_bytes + // Accumulated as int64_t: a row's total blob size (dominated by values_bytes, which sums + // per-field encoded sizes that can each be up to ~4 GB) can exceed the int32 size_type range. + value_sizes[row] = int64_t{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, cudf::get_current_device_resource_ref()); + 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); + } + + 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, {}, stream, mr); + } + + // ── 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]; + }); + 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); + 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, cudf::get_current_device_resource_ref()); + + // ── 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) { + 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, cudf::get_current_device_resource_ref())); + } + + // ── 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, cudf::get_current_device_resource_ref()); + + // ── Input null mask ─────────────────────────────────────────────────────── + size_type const null_count = input.null_count(); + // copy_bitmask handles input.offset() and yields an offset-free mask. This is an internal + // scratch buffer (never adopted into the returned column), so it is allocated against the + // current device resource rather than the caller-supplied `mr`. + rmm::device_buffer owned_null_mask = + (null_count > 0) ? cudf::detail::copy_bitmask(input.null_mask(), + input.offset(), + input.offset() + num_rows, + stream, + cudf::get_current_device_resource_ref()) + : rmm::device_buffer{}; + bitmask_type const* input_null_mask = + (null_count > 0) ? static_cast(owned_null_mask.data()) : nullptr; + + // ── Compute per-row value blob sizes ───────────────────────────────────── + // Accumulated as int64_t: the sum of per-row blob sizes across the whole column can exceed + // the int32 size_type range even though the final list-column offsets must be int32. + rmm::device_uvector value_sizes( + num_rows, stream, cudf::get_current_device_resource_ref()); + { + cudf::detail::device_scalar error_flag( + 0, stream, cudf::get_current_device_resource_ref()); + 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, + error_flag.data()); + CUDF_CUDA_TRY(cudaGetLastError()); + CUDF_EXPECTS(error_flag.value(stream) == 0, + "encode_strings_to_variant encountered a JSON value that is not a supported " + "scalar (nested object/array), or a malformed literal or string escape; " + "VARIANT encoding requires scalar, non-nested field values only", + std::invalid_argument); + } + + // ── Prefix-sum to get per-row value offsets ─────────────────────────────── + // Scanned in int64_t first so overflow can be detected before truncating into the int32 + // offsets that the VARIANT value list column requires. + rmm::device_uvector value_offsets_wide( + num_rows + 1, stream, cudf::get_current_device_resource_ref()); + { + CUDF_CUDA_TRY(cudaMemsetAsync(value_offsets_wide.data(), 0, sizeof(int64_t), stream.value())); + thrust::inclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + value_sizes.begin(), + value_sizes.end(), + value_offsets_wide.begin() + 1); + } + + int64_t total_value_bytes_wide{}; + cudf::detail::cuda_memcpy(host_span{&total_value_bytes_wide, 1}, + device_span{value_offsets_wide.data() + num_rows, 1}, + stream); + CUDF_EXPECTS(total_value_bytes_wide <= std::numeric_limits::max(), + "encode_strings_to_variant output exceeds the VARIANT value list column " + "size limit (a LIST column's offsets are always 32-bit)", + std::overflow_error); + size_type const total_value_bytes = static_cast(total_value_bytes_wide); + + // Narrow the validated wide offsets down to the int32 offsets the list column needs. + rmm::device_uvector value_offsets(num_rows + 1, stream, mr); + thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + value_offsets_wide.begin(), + value_offsets_wide.end(), + value_offsets.begin(), + cuda::proclaim_return_type( + [] __device__(int64_t v) { return static_cast(v); })); + + // ── Allocate and write value blobs ──────────────────────────────────────── + 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 06c0462ebed..051459eda16 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 00000000000..c70d04f8f03 --- /dev/null +++ b/cpp/tests/io/experimental/variant_encode_test.cpp @@ -0,0 +1,501 @@ +/* + * 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 + +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); +} + +// Encode json_rows[offset .. offset+size) via a sliced strings_column_view. +std::unique_ptr encode_sliced(std::vector const& json_rows, + std::vector const& col_names, + cudf::size_type offset, + cudf::size_type size) +{ + cudf::test::strings_column_wrapper w(json_rows.begin(), json_rows.end()); + auto full_col = w.release(); + cudf::column_view sliced = cudf::slice(full_col->view(), {offset, offset + size})[0]; + cudf::strings_column_view scv{sliced}; + 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"); + cudf::test::fixed_width_column_wrapper expected{3.14}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*floats, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowFloatExponent) +{ + auto variant = encode({R"({"f":1.5e2})"}, {"f"}); + + auto floats = extract_float64(variant->view(), "$.f"); + cudf::test::fixed_width_column_wrapper expected{150.0}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*floats, expected); +} + +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, SingleRowShortStringBoundary) +{ + // 63 bytes is the maximum for the SHORT_STRING encoding path + std::string boundary_str(63, 'x'); + auto json = R"({"s":")" + boundary_str + R"("})"; + auto variant = encode({json}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{boundary_str}; + 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, SingleRowShortStringNonAscii) +{ + // "café" – é is U+00E9, encoded as 2 UTF-8 bytes (0xC3 0xA9), total 5 bytes + auto variant = encode({R"({"s":"café"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"café"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowLongStringNonAscii) +{ + // 35 × "é" (2 UTF-8 bytes each) = 70 bytes → LONG_STRING path + std::string non_ascii_long; + for (int i = 0; i < 35; ++i) { + non_ascii_long += "\xC3\xA9"; // UTF-8 for é + } + auto json = R"({"s":")" + non_ascii_long + R"("})"; + auto variant = encode({json}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{non_ascii_long}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowEscapedString) +{ + // The JSON string contains an escaped quote and an escaped backslash; the VARIANT payload + // must contain the unescaped bytes a"b\c, not the raw escape sequences. + auto variant = encode({R"({"s":"a\"b\\c"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"a\"b\\c"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowCommonEscapes) +{ + // \n \t \r \b \f \/ all unescape to single control/ASCII bytes. + auto variant = encode({R"({"s":"a\nb\tc\rd\be\ff\/g"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"a\nb\tc\rd\be\ff/g"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} +TEST_F(EncodeStringsToVariantTest, SingleRowUnicodeEscape) +{ + // \u00e9 is the JSON \uXXXX escape for U+00E9 (e-acute), which must be unescaped to its + // 2-byte UTF-8 encoding, not copied as the literal 6-character escape sequence. + auto variant = encode({R"({"s":"caf\u00e9"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"café"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowLiteralEmojiPassthrough) +{ + // A literal (non-escaped) 4-byte UTF-8 emoji is copied through unchanged. + auto variant = encode({R"({"s":"😀"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"\xF0\x9F\x98\x80"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, SingleRowSurrogatePairEscape) +{ + // U+1F600 (grinning face emoji) as a JSON \uD83D\uDE00 UTF-16 surrogate pair escape, which + // must be combined and unescaped to its 4-byte UTF-8 encoding. + auto variant = encode({R"({"s":"\uD83D\uDE00"})"}, {"s"}); + + auto strs = extract_string(variant->view(), "$.s"); + cudf::test::strings_column_wrapper expected{"\xF0\x9F\x98\x80"}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*strs, expected); +} + +TEST_F(EncodeStringsToVariantTest, UnpairedSurrogateEscapeRejected) +{ + // A high surrogate with no following low surrogate is not valid UTF-16/8 and must be rejected + // rather than silently mis-encoded. + EXPECT_THROW(encode({R"({"s":"\uD83D"})"}, {"s"}), std::invalid_argument); +} + +TEST_F(EncodeStringsToVariantTest, NestedObjectFieldValueRejected) +{ + // The encoder only supports scalar, non-nested field values. + EXPECT_THROW(encode({R"({"a":{"x":1}})"}, {"a"}), std::invalid_argument); +} + +TEST_F(EncodeStringsToVariantTest, NestedArrayFieldValueRejected) +{ + EXPECT_THROW(encode({R"({"a":[1,2,3]})"}, {"a"}), std::invalid_argument); +} + +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"); + cudf::test::fixed_width_column_wrapper expected({0}, {false}); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*ints, expected); +} + +TEST_F(EncodeStringsToVariantTest, SlicedInputView) +{ + // Build a 5-row column and encode only rows 1–3 via a sliced strings_column_view. + auto variant = + encode_sliced({R"({"a":0})", R"({"a":10})", R"({"a":20})", R"({"a":30})", R"({"a":40})"}, + {"a"}, + /*offset=*/1, + /*size=*/3); + + ASSERT_EQ(variant->size(), 3); + auto a_vals = extract_int64(variant->view(), "$.a"); + cudf::test::fixed_width_column_wrapper expected{10, 20, 30}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, expected); +} + +// ─── 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, MultipleRowsManyBlocks) +{ + // 512 rows ensures work is spread across multiple CUDA blocks (typically 256 threads each). + constexpr int N = 512; + std::vector json_rows; + json_rows.reserve(N); + for (int i = 0; i < N; ++i) { + json_rows.push_back(R"({"a":)" + std::to_string(i) + R"(})"); + } + + auto variant = encode(json_rows, {"a"}); + ASSERT_EQ(variant->size(), N); + + auto a_vals = extract_int64(variant->view(), "$.a"); + ASSERT_EQ(a_vals->size(), N); + EXPECT_EQ(a_vals->null_count(), 0); + + std::vector exp_vals(N); + for (int i = 0; i < N; ++i) { + exp_vals[i] = i; + } + cudf::test::fixed_width_column_wrapper expected(exp_vals.begin(), exp_vals.end()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*a_vals, expected); +} + +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); +} + +TEST_F(EncodeStringsToVariantTest, FieldCountLimit) +{ + std::vector names; + for (int i = 0; i < 255; ++i) { + names.push_back("f" + std::to_string(i)); + } + + auto variant = encode({R"({})"}, names); + EXPECT_EQ(variant->size(), 1); + + names.push_back("f255"); + EXPECT_THROW(encode({R"({})"}, names), std::invalid_argument); +}