Skip to content

Commit 62770b6

Browse files
liamzwbaoalamb
andauthored
[Variant]: Implement DataType::Dictionary support for cast_to_variant kernel (#8173)
# Which issue does this PR close? - Closes #8062 # Rationale for this change # What changes are included in this PR? Implement `DataType::Dictionary` in `cast_to_variant` # Are these changes tested? Yes # Are there any user-facing changes? New cast type supported Co-authored-by: Andrew Lamb <[email protected]>
1 parent 653ca78 commit 62770b6

File tree

1 file changed

+64
-6
lines changed

1 file changed

+64
-6
lines changed

parquet-variant-compute/src/cast_to_variant.rs

Lines changed: 64 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -502,6 +502,27 @@ pub fn cast_to_variant(input: &dyn Array) -> Result<VariantArray, ArrowError> {
502502
builder
503503
);
504504
}
505+
DataType::Dictionary(_, _) => {
506+
let dict_array = input.as_any_dictionary();
507+
let values_variant_array = cast_to_variant(dict_array.values().as_ref())?;
508+
let normalized_keys = dict_array.normalized_keys();
509+
let keys = dict_array.keys();
510+
511+
for (i, key_idx) in normalized_keys.iter().enumerate() {
512+
if keys.is_null(i) {
513+
builder.append_null();
514+
continue;
515+
}
516+
517+
if values_variant_array.is_null(*key_idx) {
518+
builder.append_null();
519+
continue;
520+
}
521+
522+
let value = values_variant_array.value(*key_idx);
523+
builder.append_variant(value);
524+
}
525+
}
505526
dt => {
506527
return Err(ArrowError::CastError(format!(
507528
"Unsupported data type for casting to Variant: {dt:?}",
@@ -520,12 +541,12 @@ mod tests {
520541
use super::*;
521542
use arrow::array::{
522543
ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
523-
Decimal256Array, Decimal32Array, Decimal64Array, FixedSizeBinaryBuilder, Float16Array,
524-
Float32Array, Float64Array, GenericByteBuilder, GenericByteViewBuilder, Int16Array,
525-
Int32Array, Int64Array, Int8Array, IntervalYearMonthArray, LargeStringArray, NullArray,
526-
StringArray, StringViewArray, StructArray, Time32MillisecondArray, Time32SecondArray,
527-
Time64MicrosecondArray, Time64NanosecondArray, UInt16Array, UInt32Array, UInt64Array,
528-
UInt8Array,
544+
Decimal256Array, Decimal32Array, Decimal64Array, DictionaryArray, FixedSizeBinaryBuilder,
545+
Float16Array, Float32Array, Float64Array, GenericByteBuilder, GenericByteViewBuilder,
546+
Int16Array, Int32Array, Int64Array, Int8Array, IntervalYearMonthArray, LargeStringArray,
547+
NullArray, StringArray, StringViewArray, StructArray, Time32MillisecondArray,
548+
Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, UInt16Array, UInt32Array,
549+
UInt64Array, UInt8Array,
529550
};
530551
use arrow::buffer::NullBuffer;
531552
use arrow_schema::{Field, Fields};
@@ -1826,6 +1847,43 @@ mod tests {
18261847
);
18271848
}
18281849

1850+
#[test]
1851+
fn test_cast_to_variant_dictionary() {
1852+
let values = StringArray::from(vec!["apple", "banana", "cherry", "date"]);
1853+
let keys = Int32Array::from(vec![Some(0), Some(1), None, Some(2), Some(0), Some(3)]);
1854+
let dict_array = DictionaryArray::<Int32Type>::try_new(keys, Arc::new(values)).unwrap();
1855+
1856+
run_test(
1857+
Arc::new(dict_array),
1858+
vec![
1859+
Some(Variant::from("apple")),
1860+
Some(Variant::from("banana")),
1861+
None,
1862+
Some(Variant::from("cherry")),
1863+
Some(Variant::from("apple")),
1864+
Some(Variant::from("date")),
1865+
],
1866+
);
1867+
}
1868+
1869+
#[test]
1870+
fn test_cast_to_variant_dictionary_with_nulls() {
1871+
// Test dictionary with null values in the values array
1872+
let values = StringArray::from(vec![Some("a"), None, Some("c")]);
1873+
let keys = Int8Array::from(vec![Some(0), Some(1), Some(2), Some(0)]);
1874+
let dict_array = DictionaryArray::<Int8Type>::try_new(keys, Arc::new(values)).unwrap();
1875+
1876+
run_test(
1877+
Arc::new(dict_array),
1878+
vec![
1879+
Some(Variant::from("a")),
1880+
None, // key 1 points to null value
1881+
Some(Variant::from("c")),
1882+
Some(Variant::from("a")),
1883+
],
1884+
);
1885+
}
1886+
18291887
/// Converts the given `Array` to a `VariantArray` and tests the conversion
18301888
/// against the expected values. It also tests the handling of nulls by
18311889
/// setting one element to null and verifying the output.

0 commit comments

Comments
 (0)