From 49a67d2a45eb8db2c4ddcb860f7e55ad50bac4c3 Mon Sep 17 00:00:00 2001 From: Niranjan Artal Date: Wed, 19 Aug 2026 14:08:35 -0700 Subject: [PATCH 1/2] Expose Variant logical type IDs in Java --- .../ai/rapids/cudf/VariantLogicalType.java | 65 ++++++ .../java/ai/rapids/cudf/VariantUtils.java | 21 ++ java/src/main/native/src/VariantUtilsJni.cpp | 17 ++ .../java/ai/rapids/cudf/VariantUtilsTest.java | 189 ++++++++++++++++++ 4 files changed, 292 insertions(+) create mode 100644 java/src/main/java/ai/rapids/cudf/VariantLogicalType.java diff --git a/java/src/main/java/ai/rapids/cudf/VariantLogicalType.java b/java/src/main/java/ai/rapids/cudf/VariantLogicalType.java new file mode 100644 index 000000000000..b1a1697eb069 --- /dev/null +++ b/java/src/main/java/ai/rapids/cudf/VariantLogicalType.java @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package ai.rapids.cudf; + +/** + * Logical type identifiers returned by {@link VariantUtils#getVariantTypeId(ColumnView)}. + * + *

These values mirror the explicitly assigned numeric values of the experimental C++ + * {@code cudf::io::parquet::experimental::variant_logical_type} enum. The IDs are part of the + * Java API contract and must not be derived from the Java enum ordinal. + */ +@Experimental +public enum VariantLogicalType { + OBJECT(0), + ARRAY(1), + NULL_VALUE(2), + BOOLEAN(3), + LONG_VALUE(4), + STRING(5), + DOUBLE_VALUE(6), + DECIMAL(7), + DATE(8), + TIMESTAMP(9), + TIMESTAMP_NTZ(10), + FLOAT_VALUE(11), + BINARY(12), + UUID(13), + TIME_NTZ(14); + + private static final VariantLogicalType[] TYPES = VariantLogicalType.values(); + + private final int nativeId; + + VariantLogicalType(int nativeId) { + this.nativeId = nativeId; + } + + /** + * Get the value stored in the {@code UINT8} result column. + * + * @return the native logical type ID + */ + public int getNativeId() { + return nativeId; + } + + /** + * Find the named logical type for a native ID. + * + * @param nativeId ID returned in a valid result row + * @return the corresponding logical type + * @throws IllegalArgumentException if {@code nativeId} is not recognized + */ + public static VariantLogicalType fromNative(int nativeId) { + for (VariantLogicalType type : TYPES) { + if (type.nativeId == nativeId) { + return type; + } + } + throw new IllegalArgumentException("Unknown Variant logical type ID: " + nativeId); + } +} diff --git a/java/src/main/java/ai/rapids/cudf/VariantUtils.java b/java/src/main/java/ai/rapids/cudf/VariantUtils.java index 4c8e11d9df32..17b5d9c70876 100644 --- a/java/src/main/java/ai/rapids/cudf/VariantUtils.java +++ b/java/src/main/java/ai/rapids/cudf/VariantUtils.java @@ -48,6 +48,25 @@ public static ColumnVector getVariantFieldValue(ColumnView variantStruct, String return new ColumnVector(getVariantFieldValue(variantStruct.getNativeView(), path)); } + /** + * Return the logical type ID of each raw Variant-encoded value. + * + *

The input must be a LIST<UINT8> column. The result is a UINT8 column containing the + * IDs defined by {@link VariantLogicalType}. A result row is null when the input row is null, + * the value blob is empty, or its header is unrecognized. An encoded Variant null is represented + * by a valid {@link VariantLogicalType#NULL_VALUE} row. Only the header byte is inspected, so a + * recognized header is classified even if the remaining payload is truncated. + * + *

This API mirrors an experimental libcudf API and is subject to change. + * + * @param valueBytes LIST<UINT8> column of raw Variant-encoded values + * @return owning UINT8 column of logical type IDs + */ + public static ColumnVector getVariantTypeId(ColumnView valueBytes) { + Objects.requireNonNull(valueBytes, "valueBytes"); + return new ColumnVector(getVariantTypeId(valueBytes.getNativeView())); + } + /** * Decode raw Variant-encoded value bytes into {@code targetType}. Supported target types are * {@link DType#STRING}, {@link DType#INT8}, {@link DType#INT16}, {@link DType#INT32}, and @@ -76,6 +95,8 @@ public static ColumnVector extractVariantField( private static native long getVariantFieldValue(long variantStructHandle, String path); + private static native long getVariantTypeId(long valueBytesHandle); + private static native long castVariantValue(long valueBytesHandle, int cudfTypeId); private static native long extractVariantField( diff --git a/java/src/main/native/src/VariantUtilsJni.cpp b/java/src/main/native/src/VariantUtilsJni.cpp index d427da7b6bd3..7c848fb7c36d 100644 --- a/java/src/main/native/src/VariantUtilsJni.cpp +++ b/java/src/main/native/src/VariantUtilsJni.cpp @@ -31,6 +31,23 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_getVariantFieldValue( JNI_CATCH(env, 0); } +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_getVariantTypeId( + JNIEnv* env, jclass, jlong value_bytes_handle) +{ + JNI_NULL_CHECK(env, value_bytes_handle, "value bytes column is null", 0); + JNI_TRY + { + cudf::jni::auto_set_device(env); + auto const& value_bytes = *reinterpret_cast(value_bytes_handle); + return cudf::jni::release_as_jlong( + cudf::io::parquet::experimental::get_variant_type_id( + value_bytes, + cudf::get_default_stream(), + cudf::get_current_device_resource_ref())); + } + JNI_CATCH(env, 0); +} + JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_castVariantValue(JNIEnv* env, jclass, jlong value_bytes_handle, diff --git a/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java b/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java index d9c9633805ed..871357780e9b 100644 --- a/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java +++ b/java/src/test/java/ai/rapids/cudf/VariantUtilsTest.java @@ -21,6 +21,7 @@ import java.util.Map; import static ai.rapids.cudf.AssertUtils.assertColumnsAreEqual; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; public class VariantUtilsTest extends CudfTestBase { @@ -252,6 +253,18 @@ private static int header(int kind, int value) { private static int simple(int type) { return header(0x00, type); } + + private static List primitiveHeader(int type) { + return bytes(simple(type)); + } + + private static List shortStringHeader(int length) { + return bytes(header(0x01, length)); + } + + private static List containerHeader(boolean array) { + return bytes(header(array ? 0x03 : 0x02, 0)); + } } // Test data is constructed according to Apache Parquet's Variant encoding spec: @@ -389,6 +402,182 @@ void getThenCastFieldValue() { } } + @Test + void variantLogicalTypeIdsMatchNativeValues() { + VariantLogicalType[] types = { + VariantLogicalType.OBJECT, + VariantLogicalType.ARRAY, + VariantLogicalType.NULL_VALUE, + VariantLogicalType.BOOLEAN, + VariantLogicalType.LONG_VALUE, + VariantLogicalType.STRING, + VariantLogicalType.DOUBLE_VALUE, + VariantLogicalType.DECIMAL, + VariantLogicalType.DATE, + VariantLogicalType.TIMESTAMP, + VariantLogicalType.TIMESTAMP_NTZ, + VariantLogicalType.FLOAT_VALUE, + VariantLogicalType.BINARY, + VariantLogicalType.UUID, + VariantLogicalType.TIME_NTZ + }; + + for (int nativeId = 0; nativeId < types.length; nativeId++) { + assertEquals(nativeId, types[nativeId].getNativeId()); + assertEquals(types[nativeId], VariantLogicalType.fromNative(nativeId)); + } + assertThrows(IllegalArgumentException.class, () -> VariantLogicalType.fromNative(-1)); + assertThrows(IllegalArgumentException.class, () -> VariantLogicalType.fromNative(15)); + } + + @Test + void getVariantTypeIdCoversLogicalTypesAndPhysicalAliases() { + // get_variant_type_id intentionally classifies only the header byte. Supplying header-only + // blobs here both pins every physical-to-logical mapping and verifies that truncated payloads + // with recognized headers remain classifiable. + try (ColumnVector values = ColumnVector.fromLists( + BINARY_TYPE, + VariantEncoder.containerHeader(false), // OBJECT + VariantEncoder.containerHeader(true), // ARRAY + VariantEncoder.primitiveHeader(0), // NULL_VALUE + VariantEncoder.primitiveHeader(1), // BOOLEAN_TRUE + VariantEncoder.primitiveHeader(2), // BOOLEAN_FALSE + VariantEncoder.primitiveHeader(3), // INT8 + VariantEncoder.primitiveHeader(4), // INT16 + VariantEncoder.primitiveHeader(5), // INT32 + VariantEncoder.primitiveHeader(6), // INT64 + VariantEncoder.shortStringHeader(0), // SHORT_STRING + VariantEncoder.primitiveHeader(16), // LONG_STRING + VariantEncoder.primitiveHeader(7), // FLOAT64 + VariantEncoder.primitiveHeader(8), // DECIMAL4 + VariantEncoder.primitiveHeader(9), // DECIMAL8 + VariantEncoder.primitiveHeader(10), // DECIMAL16 + VariantEncoder.primitiveHeader(11), // DATE + VariantEncoder.primitiveHeader(12), // TIMESTAMP_MICROS + VariantEncoder.primitiveHeader(18), // TIMESTAMP_NANOS + VariantEncoder.primitiveHeader(13), // TIMESTAMP_NTZ_MICROS + VariantEncoder.primitiveHeader(19), // TIMESTAMP_NTZ_NANOS + VariantEncoder.primitiveHeader(14), // FLOAT32 + VariantEncoder.primitiveHeader(15), // BINARY + VariantEncoder.primitiveHeader(20), // UUID + VariantEncoder.primitiveHeader(17)); // TIME_NTZ_MICROS + ColumnVector result = VariantUtils.getVariantTypeId(values); + ColumnVector expected = ColumnVector.fromUnsignedBytes( + (byte) VariantLogicalType.OBJECT.getNativeId(), + (byte) VariantLogicalType.ARRAY.getNativeId(), + (byte) VariantLogicalType.NULL_VALUE.getNativeId(), + (byte) VariantLogicalType.BOOLEAN.getNativeId(), + (byte) VariantLogicalType.BOOLEAN.getNativeId(), + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + (byte) VariantLogicalType.STRING.getNativeId(), + (byte) VariantLogicalType.STRING.getNativeId(), + (byte) VariantLogicalType.DOUBLE_VALUE.getNativeId(), + (byte) VariantLogicalType.DECIMAL.getNativeId(), + (byte) VariantLogicalType.DECIMAL.getNativeId(), + (byte) VariantLogicalType.DECIMAL.getNativeId(), + (byte) VariantLogicalType.DATE.getNativeId(), + (byte) VariantLogicalType.TIMESTAMP.getNativeId(), + (byte) VariantLogicalType.TIMESTAMP.getNativeId(), + (byte) VariantLogicalType.TIMESTAMP_NTZ.getNativeId(), + (byte) VariantLogicalType.TIMESTAMP_NTZ.getNativeId(), + (byte) VariantLogicalType.FLOAT_VALUE.getNativeId(), + (byte) VariantLogicalType.BINARY.getNativeId(), + (byte) VariantLogicalType.UUID.getNativeId(), + (byte) VariantLogicalType.TIME_NTZ.getNativeId())) { + assertEquals(DType.UINT8, result.getType()); + assertColumnsAreEqual(expected, result); + } + } + + @Test + void getVariantTypeIdNullAndInvalidHeaderBehavior() { + try (ColumnVector values = ColumnVector.fromLists( + BINARY_TYPE, + VariantEncoder.primitiveHeader(0), + null, + bytes(), + bytes(0xfc), + VariantEncoder.primitiveHeader(6)); + ColumnVector result = VariantUtils.getVariantTypeId(values); + ColumnVector expected = ColumnVector.fromBoxedUnsignedBytes( + (byte) VariantLogicalType.NULL_VALUE.getNativeId(), + null, + null, + null, + (byte) VariantLogicalType.LONG_VALUE.getNativeId())) { + assertColumnsAreEqual(expected, result); + } + } + + @Test + void getVariantTypeIdAcceptsExtractedValues() { + try (ColumnVector variant = makeXyzVariantColumn(); + ColumnVector values = VariantUtils.getVariantFieldValue(variant, "x"); + ColumnVector result = VariantUtils.getVariantTypeId(values); + ColumnVector expected = ColumnVector.fromBoxedUnsignedBytes( + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + (byte) VariantLogicalType.LONG_VALUE.getNativeId(), + null)) { + assertColumnsAreEqual(expected, result); + } + } + + @Test + void getVariantTypeIdSupportsSlicedInput() { + try (ColumnVector values = ColumnVector.fromLists( + BINARY_TYPE, + VariantEncoder.primitiveHeader(0), + VariantEncoder.containerHeader(true), + VariantEncoder.primitiveHeader(14), + VariantEncoder.shortStringHeader(0), + VariantEncoder.containerHeader(false)); + ColumnVector slice = values.subVector(1, 4); + ColumnVector result = VariantUtils.getVariantTypeId(slice); + ColumnVector expected = ColumnVector.fromUnsignedBytes( + (byte) VariantLogicalType.ARRAY.getNativeId(), + (byte) VariantLogicalType.FLOAT_VALUE.getNativeId(), + (byte) VariantLogicalType.STRING.getNativeId())) { + assertColumnsAreEqual(expected, result); + } + } + + @Test + void getVariantTypeIdHandlesEmptyAndAllNullInput() { + try (ColumnVector empty = ColumnVector.fromLists(BINARY_TYPE); + ColumnVector emptyResult = VariantUtils.getVariantTypeId(empty); + ColumnVector expectedEmpty = ColumnVector.fromUnsignedBytes(); + ColumnVector allNull = ColumnVector.fromLists( + BINARY_TYPE, (List) null, (List) null); + ColumnVector allNullResult = VariantUtils.getVariantTypeId(allNull); + ColumnVector expectedAllNull = ColumnVector.fromBoxedUnsignedBytes(null, null)) { + assertEquals(DType.UINT8, emptyResult.getType()); + assertColumnsAreEqual(expectedEmpty, emptyResult); + assertColumnsAreEqual(expectedAllNull, allNullResult); + } + } + + @Test + void getVariantTypeIdRejectsInvalidInput() { + ListType listOfInt = new ListType(true, new BasicType(false, DType.INT32)); + try (ColumnVector notAList = ColumnVector.fromInts(1); + ColumnVector wrongChildType = ColumnVector.fromLists(listOfInt, Arrays.asList(1, 2))) { + assertThrows(CudfException.class, () -> VariantUtils.getVariantTypeId(notAList)); + assertThrows(CudfException.class, () -> VariantUtils.getVariantTypeId(wrongChildType)); + } + assertThrows(NullPointerException.class, () -> VariantUtils.getVariantTypeId(null)); + } + + @Test + void emptyInputUnsupportedDirectCastThrows() { + try (ColumnVector empty = ColumnVector.fromLists(BINARY_TYPE)) { + assertThrows(IllegalArgumentException.class, + () -> VariantUtils.castVariantValue(empty, DType.UINT32)); + } + } + @Test void emptyInputProducesEmptyOutput() { try (ColumnVector variant = ColumnVector.fromStructs(VARIANT_TYPE); From a6f3d5c7df0c699eb196d5599f0c211a00c3a134 Mon Sep 17 00:00:00 2001 From: Niranjan Artal Date: Wed, 19 Aug 2026 17:15:13 -0700 Subject: [PATCH 2/2] address clang format --- java/src/main/native/src/VariantUtilsJni.cpp | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/java/src/main/native/src/VariantUtilsJni.cpp b/java/src/main/native/src/VariantUtilsJni.cpp index 7c848fb7c36d..40e250e00795 100644 --- a/java/src/main/native/src/VariantUtilsJni.cpp +++ b/java/src/main/native/src/VariantUtilsJni.cpp @@ -31,19 +31,17 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_getVariantFieldValue( JNI_CATCH(env, 0); } -JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_getVariantTypeId( - JNIEnv* env, jclass, jlong value_bytes_handle) +JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_VariantUtils_getVariantTypeId(JNIEnv* env, + jclass, + jlong value_bytes_handle) { JNI_NULL_CHECK(env, value_bytes_handle, "value bytes column is null", 0); JNI_TRY { cudf::jni::auto_set_device(env); auto const& value_bytes = *reinterpret_cast(value_bytes_handle); - return cudf::jni::release_as_jlong( - cudf::io::parquet::experimental::get_variant_type_id( - value_bytes, - cudf::get_default_stream(), - cudf::get_current_device_resource_ref())); + return cudf::jni::release_as_jlong(cudf::io::parquet::experimental::get_variant_type_id( + value_bytes, cudf::get_default_stream(), cudf::get_current_device_resource_ref())); } JNI_CATCH(env, 0); }