Skip to content

Commit feff9ce

Browse files
liamzwbaoliamzwbao
committed
[Variant]: Implement DataType::FixedSizeList support for cast_to_variant kernel
1 parent fc1ef41 commit feff9ce

File tree

2 files changed

+116
-10
lines changed

2 files changed

+116
-10
lines changed

‎parquet-variant-compute/src/arrow_to_variant.rs‎

Lines changed: 22 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@
1717

1818
use crate::type_conversion::{decimal_to_variant_decimal, CastOptions};
1919
use arrow::array::{
20-
Array, AsArray, GenericBinaryArray, GenericListArray, GenericListViewArray, GenericStringArray,
21-
OffsetSizeTrait, PrimitiveArray,
20+
Array, AsArray, FixedSizeListArray, GenericBinaryArray, GenericListArray, GenericListViewArray,
21+
GenericStringArray, OffsetSizeTrait, PrimitiveArray,
2222
};
2323
use arrow::compute::kernels::cast;
2424
use arrow::datatypes::{
@@ -82,6 +82,7 @@ pub(crate) enum ArrowToVariantRowBuilder<'a> {
8282
LargeList(ListArrowToVariantBuilder<'a, GenericListArray<i64>>),
8383
ListView(ListArrowToVariantBuilder<'a, GenericListViewArray<i32>>),
8484
LargeListView(ListArrowToVariantBuilder<'a, GenericListViewArray<i64>>),
85+
FixedSizeList(ListArrowToVariantBuilder<'a, FixedSizeListArray>),
8586
Struct(StructArrowToVariantBuilder<'a>),
8687
Map(MapArrowToVariantBuilder<'a>),
8788
Union(UnionArrowToVariantBuilder<'a>),
@@ -138,6 +139,7 @@ impl<'a> ArrowToVariantRowBuilder<'a> {
138139
LargeList(b) => b.append_row(builder, index),
139140
ListView(b) => b.append_row(builder, index),
140141
LargeListView(b) => b.append_row(builder, index),
142+
FixedSizeList(b) => b.append_row(builder, index),
141143
Struct(b) => b.append_row(builder, index),
142144
Map(b) => b.append_row(builder, index),
143145
Union(b) => b.append_row(builder, index),
@@ -255,6 +257,10 @@ pub(crate) fn make_arrow_to_variant_row_builder<'a>(
255257
array.as_list_view(),
256258
options,
257259
)?),
260+
DataType::FixedSizeList(_, _) => FixedSizeList(ListArrowToVariantBuilder::new(
261+
array.as_fixed_size_list(),
262+
options,
263+
)?),
258264
DataType::Struct(_) => Struct(StructArrowToVariantBuilder::new(
259265
array.as_struct(),
260266
options,
@@ -281,11 +287,6 @@ pub(crate) fn make_arrow_to_variant_row_builder<'a>(
281287
)));
282288
}
283289
},
284-
dt => {
285-
return Err(ArrowError::CastError(format!(
286-
"Unsupported data type for casting to Variant: {dt}",
287-
)));
288-
}
289290
};
290291
Ok(builder)
291292
}
@@ -523,7 +524,8 @@ impl NullArrowToVariantBuilder {
523524
}
524525
}
525526

526-
/// Generic list builder for List, LargeList, ListView, and LargeListView types
527+
/// Generic list builder for ListLikeArray types including List, LargeList, ListView, LargeListView,
528+
/// and FixedSizeList
527529
pub(crate) struct ListArrowToVariantBuilder<'a, L: ListLikeArray> {
528530
list_array: &'a L,
529531
values_builder: Box<ArrowToVariantRowBuilder<'a>>,
@@ -599,6 +601,18 @@ impl<O: OffsetSizeTrait> ListLikeArray for GenericListViewArray<O> {
599601
}
600602
}
601603

604+
impl ListLikeArray for FixedSizeListArray {
605+
fn values(&self) -> &dyn Array {
606+
self.values()
607+
}
608+
609+
fn element_range(&self, index: usize) -> Range<usize> {
610+
let value_length = self.value_length().as_usize();
611+
let offset = index * value_length;
612+
offset..(offset + value_length)
613+
}
614+
}
615+
602616
/// Struct builder for StructArray
603617
pub(crate) struct StructArrowToVariantBuilder<'a> {
604618
struct_array: &'a arrow::array::StructArray,

‎parquet-variant-compute/src/cast_to_variant.rs‎

Lines changed: 94 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,8 @@ mod tests {
9191
ArrayRef, BinaryArray, BooleanArray, Date32Array, Date64Array, Decimal128Array,
9292
Decimal256Array, Decimal32Array, Decimal64Array, DictionaryArray, DurationMicrosecondArray,
9393
DurationMillisecondArray, DurationNanosecondArray, DurationSecondArray,
94-
FixedSizeBinaryBuilder, Float16Array, Float32Array, Float64Array, GenericByteBuilder,
95-
GenericByteViewBuilder, Int16Array, Int32Array, Int64Array, Int8Array,
94+
FixedSizeBinaryBuilder, FixedSizeListBuilder, Float16Array, Float32Array, Float64Array,
95+
GenericByteBuilder, GenericByteViewBuilder, Int16Array, Int32Array, Int64Array, Int8Array,
9696
IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, LargeListArray,
9797
LargeListViewBuilder, LargeStringArray, ListArray, ListViewBuilder, MapArray, NullArray,
9898
StringArray, StringRunBuilder, StringViewArray, StructArray, Time32MillisecondArray,
@@ -1407,6 +1407,98 @@ mod tests {
14071407
);
14081408
}
14091409

1410+
#[test]
1411+
fn test_cast_to_variant_fixed_size_list() {
1412+
let mut builder = FixedSizeListBuilder::new(Int32Array::builder(0), 2);
1413+
builder.values().append_value(0);
1414+
builder.values().append_value(1);
1415+
builder.append(true); // First list: [0, 1]
1416+
1417+
builder.values().append_null();
1418+
builder.values().append_value(3);
1419+
builder.append(true); // Second list: [null, 3]
1420+
1421+
builder.values().append_value(4);
1422+
builder.values().append_null();
1423+
builder.append(false); // Third list: null
1424+
1425+
builder.values().append_nulls(2);
1426+
builder.append(true); // Last list: [null, null]
1427+
1428+
let fixed_size_list_array = builder.finish();
1429+
1430+
// Expected values
1431+
let (metadata, value) = {
1432+
let mut builder = VariantBuilder::new();
1433+
let mut list = builder.new_list();
1434+
list.append_value(0i32);
1435+
list.append_value(1i32);
1436+
list.finish();
1437+
builder.finish()
1438+
};
1439+
let variant0 = Variant::new(&metadata, &value);
1440+
1441+
let (metadata, value) = {
1442+
let mut builder = VariantBuilder::new();
1443+
let mut list = builder.new_list();
1444+
list.append_null();
1445+
list.append_value(3i32);
1446+
list.finish();
1447+
builder.finish()
1448+
};
1449+
let variant1 = Variant::new(&metadata, &value);
1450+
1451+
let (metadata, value) = {
1452+
let mut builder = VariantBuilder::new();
1453+
let mut list = builder.new_list();
1454+
list.append_null();
1455+
list.append_null();
1456+
list.finish();
1457+
builder.finish()
1458+
};
1459+
let variant3 = Variant::new(&metadata, &value);
1460+
1461+
run_test(
1462+
Arc::new(fixed_size_list_array),
1463+
vec![Some(variant0), Some(variant1), None, Some(variant3)],
1464+
);
1465+
}
1466+
1467+
#[test]
1468+
fn test_cast_to_variant_sliced_fixed_size_list() {
1469+
// Create a FixedSizeListArray with size 2
1470+
let mut builder = FixedSizeListBuilder::new(Int64Array::builder(0), 2);
1471+
builder.values().append_value(0);
1472+
builder.values().append_value(1);
1473+
builder.append(true); // First list: [0, 1]
1474+
1475+
builder.values().append_null();
1476+
builder.values().append_value(3);
1477+
builder.append(true); // Second list: [null, 3]
1478+
1479+
builder.values().append_value(4);
1480+
builder.values().append_null();
1481+
builder.append(false); // Third list: null
1482+
1483+
let fixed_size_list_array = builder.finish();
1484+
1485+
// Expected value for slice(1, 2) - should get the second and third elements
1486+
let (metadata, value) = {
1487+
let mut builder = VariantBuilder::new();
1488+
let mut list = builder.new_list();
1489+
list.append_null();
1490+
list.append_value(3i64);
1491+
list.finish();
1492+
builder.finish()
1493+
};
1494+
let variant = Variant::new(&metadata, &value);
1495+
1496+
run_test(
1497+
Arc::new(fixed_size_list_array.slice(1, 2)),
1498+
vec![Some(variant), None],
1499+
);
1500+
}
1501+
14101502
#[test]
14111503
fn test_cast_to_variant_struct() {
14121504
// Test a simple struct with two fields: id (int64) and age (int32)

0 commit comments

Comments
 (0)