Skip to content

Commit 62331e2

Browse files
sdf-jklsdf-jkl
committed
Add arg_variant_kind_err helper
1 parent c705c2e commit 62331e2

5 files changed

Lines changed: 27 additions & 22 deletions

File tree

‎src/shared.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,17 @@ pub fn arg_shape_err(udf: &str, arg_index: u8, expected: &str, actual: &str) ->
300300
))
301301
}
302302

303+
/// Helper for invalid Variant kind in an argument.
304+
pub fn arg_variant_kind_err(
305+
udf: &str,
306+
arg_index: u8,
307+
expected_variant_kind: &str,
308+
) -> DataFusionError {
309+
DataFusionError::Execution(format!(
310+
"{udf} arg #{arg_index}: expected variant {expected_variant_kind}"
311+
))
312+
}
313+
303314
/// Helper for missing argument field metadata.
304315
pub fn arg_field_meta_missing_err(udf: &str, arg_index: u8) -> DataFusionError {
305316
DataFusionError::Execution(format!("{udf} arg #{arg_index} field metadata is missing"))

‎src/variant_list_delete.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@ use datafusion::{
1313
use parquet_variant::{Variant, VariantBuilder};
1414
use parquet_variant_compute::{VariantArray, VariantType};
1515

16-
use crate::shared::{arg_shape_err, args_count_err, ensure, try_parse_variant_scalar};
16+
use crate::shared::{
17+
arg_shape_err, arg_variant_kind_err, args_count_err, ensure, try_parse_variant_scalar,
18+
};
1719

1820
#[derive(Debug, Hash, PartialEq, Eq)]
1921
pub struct VariantListDelete {
@@ -41,7 +43,7 @@ fn try_parse_index_scalar(scalar: &ScalarValue) -> Result<usize> {
4143

4244
fn delete_list_element(variant_list: Variant, index: usize) -> Result<(Vec<u8>, Vec<u8>)> {
4345
let Variant::List(variant_list) = variant_list else {
44-
return exec_err!("expected variant list");
46+
return Err(arg_variant_kind_err("variant_list_delete", 1, "list"));
4547
};
4648

4749
if index >= variant_list.len() {

‎src/variant_list_insert.rs‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ use std::sync::Arc;
33
use arrow::array::StructArray;
44
use arrow_schema::{DataType, Field, Fields};
55
use datafusion::{
6-
common::exec_err,
76
error::{DataFusionError, Result},
87
logical_expr::{
98
ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
@@ -13,7 +12,9 @@ use datafusion::{
1312
use parquet_variant::{Variant, VariantBuilder};
1413
use parquet_variant_compute::{VariantArray, VariantType};
1514

16-
use crate::shared::{arg_shape_err, args_count_err, ensure, try_parse_variant_scalar};
15+
use crate::shared::{
16+
arg_shape_err, arg_variant_kind_err, args_count_err, ensure, try_parse_variant_scalar,
17+
};
1718

1819
#[derive(Debug, Hash, PartialEq, Eq)]
1920
pub struct VariantListInsert {
@@ -96,7 +97,6 @@ impl ScalarUDFImpl for VariantListInsert {
9697

9798
let out: StructArray = {
9899
let (m, v) = create_variant_list_with_new_elements(
99-
self.name(),
100100
variant_list,
101101
[variant_to_insert].into_iter(),
102102
)?;
@@ -121,7 +121,6 @@ impl ScalarUDFImpl for VariantListInsert {
121121
.map(|v| {
122122
v.map(|v| {
123123
create_variant_list_with_new_elements(
124-
self.name(),
125124
v,
126125
[variant_to_insert.clone()].into_iter(),
127126
)
@@ -160,7 +159,6 @@ impl ScalarUDFImpl for VariantListInsert {
160159
match (variant_list_to_update, element_to_append) {
161160
(Some(variant_list), Some(element_to_append)) => {
162161
let (m, v) = create_variant_list_with_new_elements(
163-
self.name(),
164162
variant_list,
165163
[element_to_append].into_iter(),
166164
)?;
@@ -211,12 +209,11 @@ impl ScalarUDFImpl for VariantListInsert {
211209
// note: I wonder if we can abstract this away
212210
// it would be good to profile and see if this pocket of code is slow
213211
fn create_variant_list_with_new_elements<'m, 'v>(
214-
udf_name: &str,
215212
variant_list: Variant,
216213
elements_to_insert: impl Iterator<Item = Variant<'m, 'v>>,
217214
) -> Result<(Vec<u8>, Vec<u8>)> {
218215
let Variant::List(variant_list) = variant_list else {
219-
return exec_err!("{udf_name} arg #1: expected variant list");
216+
return Err(arg_variant_kind_err("variant_list_insert", 1, "list"));
220217
};
221218

222219
// note: I wonder if we can abstract this away

‎src/variant_object_delete.rs‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ use std::sync::Arc;
33
use arrow::array::StructArray;
44
use arrow_schema::{DataType, Field, Fields};
55
use datafusion::{
6-
common::exec_err,
76
error::{DataFusionError, Result},
87
logical_expr::{
98
ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
@@ -14,8 +13,8 @@ use parquet_variant::{Variant, VariantBuilder};
1413
use parquet_variant_compute::{VariantArray, VariantType};
1514

1615
use crate::shared::{
17-
arg_null_error, arg_shape_err, args_count_err, ensure, try_parse_string_scalar,
18-
try_parse_variant_scalar,
16+
arg_null_error, arg_shape_err, arg_variant_kind_err, args_count_err, ensure,
17+
try_parse_string_scalar, try_parse_variant_scalar,
1918
};
2019

2120
#[derive(Debug, Hash, PartialEq, Eq)]
@@ -101,7 +100,7 @@ impl ScalarUDFImpl for VariantObjectDelete {
101100
let variant_object = try_parse_variant_scalar(scalar_variant_object_to_update)?;
102101
let variant_object = variant_object.value(0);
103102
let Variant::Object(variant_object) = variant_object else {
104-
return exec_err!("expected variant object");
103+
return Err(arg_variant_kind_err(self.name(), 1, "object"));
105104
};
106105

107106
let mut v = VariantBuilder::new();
@@ -131,7 +130,7 @@ impl ScalarUDFImpl for VariantObjectDelete {
131130
v_opt
132131
.map(|variant_object| {
133132
let Variant::Object(variant_object) = variant_object else {
134-
return exec_err!("expected variant object");
133+
return Err(arg_variant_kind_err(self.name(), 1, "object"));
135134
};
136135

137136
let mut v = VariantBuilder::new();

‎src/variant_object_insert.rs‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ use std::sync::Arc;
33
use arrow::array::StructArray;
44
use arrow_schema::{DataType, Field, Fields};
55
use datafusion::{
6-
common::exec_err,
76
error::{DataFusionError, Result},
87
logical_expr::{
98
ColumnarValue, ReturnFieldArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility,
@@ -14,8 +13,8 @@ use parquet_variant::{Variant, VariantBuilder};
1413
use parquet_variant_compute::{VariantArray, VariantType};
1514

1615
use crate::shared::{
17-
arg_null_error, arg_shape_err, args_count_err, ensure, try_parse_string_scalar,
18-
try_parse_variant_scalar,
16+
arg_null_error, arg_shape_err, arg_variant_kind_err, args_count_err, ensure,
17+
try_parse_string_scalar, try_parse_variant_scalar,
1918
};
2019

2120
#[derive(Debug, Hash, PartialEq, Eq)]
@@ -121,7 +120,7 @@ impl ScalarUDFImpl for VariantObjectInsert {
121120
let variant_object = try_parse_variant_scalar(scalar_variant_object_to_update)?;
122121
let variant_object = variant_object.value(0);
123122
let Variant::Object(variant_object) = variant_object else {
124-
return exec_err!("{} arg #1: expected variant object", self.name());
123+
return Err(arg_variant_kind_err(self.name(), 1, "object"));
125124
};
126125

127126
let mut v = VariantBuilder::new();
@@ -149,10 +148,7 @@ impl ScalarUDFImpl for VariantObjectInsert {
149148
v_opt
150149
.map(|variant_object| {
151150
let Variant::Object(variant_object) = variant_object else {
152-
return exec_err!(
153-
"{} arg #1: expected variant object",
154-
self.name()
155-
);
151+
return Err(arg_variant_kind_err(self.name(), 1, "object"));
156152
};
157153

158154
let mut v = VariantBuilder::new();

0 commit comments

Comments
 (0)