Skip to content

Commit 73f3086

Browse files
Lordwormsalamb
andauthored
Minor: Support protobuf serialization for Utf8View and BinaryView (#12165)
* Minor: Support protobuf serialization for ScalarValue::Utf8View and ScalarValue::BinaryView * Use correct protobuf type, add coverage for DataType serialization --------- Co-authored-by: Andrew Lamb <[email protected]>
1 parent ddbdf4b commit 73f3086

File tree

2 files changed

+29
-8
lines changed

2 files changed

+29
-8
lines changed

datafusion/proto-common/src/to_proto/mod.rs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,9 +153,11 @@ impl TryFrom<&DataType> for protobuf::arrow_type::ArrowTypeEnum {
153153
Self::Interval(protobuf::IntervalUnit::from(interval_unit) as i32)
154154
}
155155
DataType::Binary => Self::Binary(EmptyMessage {}),
156+
DataType::BinaryView => Self::BinaryView(EmptyMessage {}),
156157
DataType::FixedSizeBinary(size) => Self::FixedSizeBinary(*size),
157158
DataType::LargeBinary => Self::LargeBinary(EmptyMessage {}),
158159
DataType::Utf8 => Self::Utf8(EmptyMessage {}),
160+
DataType::Utf8View => Self::Utf8View(EmptyMessage {}),
159161
DataType::LargeUtf8 => Self::LargeUtf8(EmptyMessage {}),
160162
DataType::List(item_type) => Self::List(Box::new(protobuf::List {
161163
field_type: Some(Box::new(item_type.as_ref().try_into()?)),
@@ -210,7 +212,7 @@ impl TryFrom<&DataType> for protobuf::arrow_type::ArrowTypeEnum {
210212
"Proto serialization error: The RunEndEncoded data type is not yet supported".to_owned()
211213
))
212214
}
213-
DataType::Utf8View | DataType::BinaryView | DataType::ListView(_) | DataType::LargeListView(_) => {
215+
DataType::ListView(_) | DataType::LargeListView(_) => {
214216
return Err(Error::General(format!("Proto serialization error: {val} not yet supported")))
215217
}
216218
};

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1191,7 +1191,7 @@ impl LogicalExtensionCodec for UDFExtensionCodec {
11911191
}
11921192

11931193
#[test]
1194-
fn round_trip_scalar_values() {
1194+
fn round_trip_scalar_values_and_data_types() {
11951195
let should_pass: Vec<ScalarValue> = vec![
11961196
ScalarValue::Boolean(None),
11971197
ScalarValue::Float32(None),
@@ -1245,6 +1245,8 @@ fn round_trip_scalar_values() {
12451245
ScalarValue::UInt64(Some(0)),
12461246
ScalarValue::Utf8(Some(String::from("Test string "))),
12471247
ScalarValue::LargeUtf8(Some(String::from("Test Large utf8"))),
1248+
ScalarValue::Utf8View(Some(String::from("Test stringview"))),
1249+
ScalarValue::BinaryView(Some(b"binaryview".to_vec())),
12481250
ScalarValue::Date32(Some(0)),
12491251
ScalarValue::Date32(Some(i32::MAX)),
12501252
ScalarValue::Date32(None),
@@ -1471,21 +1473,38 @@ fn round_trip_scalar_values() {
14711473
ScalarValue::FixedSizeBinary(5, None),
14721474
];
14731475

1474-
for test_case in should_pass.into_iter() {
1475-
let proto: protobuf::ScalarValue = (&test_case)
1476-
.try_into()
1477-
.expect("failed conversion to protobuf");
1478-
1476+
// ScalarValue directly
1477+
for test_case in should_pass.iter() {
1478+
let proto: protobuf::ScalarValue =
1479+
test_case.try_into().expect("failed conversion to protobuf");
14791480
let roundtrip: ScalarValue = (&proto)
14801481
.try_into()
14811482
.expect("failed conversion from protobuf");
14821483

14831484
assert_eq!(
1484-
test_case, roundtrip,
1485+
test_case, &roundtrip,
14851486
"ScalarValue was not the same after round trip!\n\n\
14861487
Input: {test_case:?}\n\nRoundtrip: {roundtrip:?}"
14871488
);
14881489
}
1490+
1491+
// DataType conversion
1492+
for test_case in should_pass.iter() {
1493+
let dt = test_case.data_type();
1494+
1495+
let proto: protobuf::ArrowType = (&dt)
1496+
.try_into()
1497+
.expect("datatype failed conversion to protobuf");
1498+
let roundtrip: DataType = (&proto)
1499+
.try_into()
1500+
.expect("datatype failed conversion from protobuf");
1501+
1502+
assert_eq!(
1503+
dt, roundtrip,
1504+
"DataType was not the same after round trip!\n\n\
1505+
Input: {dt:?}\n\nRoundtrip: {roundtrip:?}"
1506+
);
1507+
}
14891508
}
14901509

14911510
// create a map array [{joe:1}, {blogs:2, foo:4}, {}, null] for testing

0 commit comments

Comments
 (0)