Skip to content

Commit 6e231f2

Browse files
adriangbadriangb
authored
Remove use of deprecated DataFusion methods (#81)
1 parent 7cecdd1 commit 6e231f2

12 files changed

+68
-37
lines changed

‎benches/main.rs

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,29 @@
11
use codspeed_criterion_compat::{criterion_group, criterion_main, Bencher, Criterion};
22

3-
use datafusion::common::ScalarValue;
3+
use datafusion::arrow::datatypes::DataType;
44
use datafusion::logical_expr::ColumnarValue;
5+
use datafusion::{common::ScalarValue, logical_expr::ScalarFunctionArgs};
56
use datafusion_functions_json::udfs::{json_contains_udf, json_get_str_udf};
67

78
fn bench_json_contains(b: &mut Bencher) {
89
let json_contains = json_contains_udf();
9-
let args = &[
10+
let args = vec![
1011
ColumnarValue::Scalar(ScalarValue::Utf8(Some(
1112
r#"{"a": {"aa": "x", "ab: "y"}, "b": []}"#.to_string(),
1213
))),
1314
ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))),
1415
ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))),
1516
];
1617

17-
b.iter(|| json_contains.invoke_batch(args, 1).unwrap());
18+
b.iter(|| {
19+
json_contains
20+
.invoke_with_args(ScalarFunctionArgs {
21+
args: args.clone(),
22+
number_rows: 1,
23+
return_type: &DataType::Boolean,
24+
})
25+
.unwrap()
26+
});
1827
}
1928

2029
fn bench_json_get_str(b: &mut Bencher) {
@@ -27,7 +36,15 @@ fn bench_json_get_str(b: &mut Bencher) {
2736
ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))),
2837
];
2938

30-
b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
39+
b.iter(|| {
40+
json_get_str
41+
.invoke_with_args(ScalarFunctionArgs {
42+
args: args.to_vec(),
43+
number_rows: 1,
44+
return_type: &DataType::Utf8,
45+
})
46+
.unwrap()
47+
});
3148
}
3249

3350
fn criterion_benchmark(c: &mut Criterion) {

‎src/json_as_text.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44
use datafusion::arrow::array::{ArrayRef, StringArray, StringBuilder};
55
use datafusion::arrow::datatypes::DataType;
66
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
88
use jiter::Peek;
99

1010
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
@@ -49,8 +49,8 @@ impl ScalarUDFImpl for JsonAsText {
4949
return_type_check(arg_types, self.name(), DataType::Utf8)
5050
}
5151

52-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
53-
invoke::<StringArray>(args, jiter_json_as_text)
52+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
53+
invoke::<StringArray>(&args.args, jiter_json_as_text)
5454
}
5555

5656
fn aliases(&self) -> &[String] {

‎src/json_contains.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use datafusion::arrow::array::BooleanBuilder;
55
use datafusion::arrow::datatypes::DataType;
66
use datafusion::common::arrow::array::{ArrayRef, BooleanArray};
77
use datafusion::common::{plan_err, Result, ScalarValue};
8-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
8+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
99

1010
use crate::common::{invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
1111
use crate::common_macros::make_udf_function;
@@ -53,8 +53,8 @@ impl ScalarUDFImpl for JsonContains {
5353
}
5454
}
5555

56-
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
57-
invoke::<BooleanArray>(args, jiter_json_contains)
56+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
57+
invoke::<BooleanArray>(&args.args, jiter_json_contains)
5858
}
5959

6060
fn aliases(&self) -> &[String] {

‎src/json_get.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use datafusion::arrow::array::ArrayRef;
55
use datafusion::arrow::array::UnionArray;
66
use datafusion::arrow::datatypes::DataType;
77
use datafusion::common::Result as DataFusionResult;
8-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
8+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
99
use datafusion::scalar::ScalarValue;
1010
use jiter::{Jiter, NumberAny, NumberInt, Peek};
1111

@@ -55,8 +55,8 @@ impl ScalarUDFImpl for JsonGet {
5555
return_type_check(arg_types, self.name(), JsonUnion::data_type())
5656
}
5757

58-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
59-
invoke::<JsonUnion>(args, jiter_json_get_union)
58+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
59+
invoke::<JsonUnion>(&args.args, jiter_json_get_union)
6060
}
6161

6262
fn aliases(&self) -> &[String] {

‎src/json_get_bool.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::any::Any;
33
use datafusion::arrow::array::BooleanArray;
44
use datafusion::arrow::datatypes::DataType;
55
use datafusion::common::Result as DataFusionResult;
6-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
6+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
77
use jiter::Peek;
88

99
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
@@ -48,8 +48,8 @@ impl ScalarUDFImpl for JsonGetBool {
4848
return_type_check(arg_types, self.name(), DataType::Boolean).map(|_| DataType::Boolean)
4949
}
5050

51-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
52-
invoke::<BooleanArray>(args, jiter_json_get_bool)
51+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
52+
invoke::<BooleanArray>(&args.args, jiter_json_get_bool)
5353
}
5454

5555
fn aliases(&self) -> &[String] {

‎src/json_get_float.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44
use datafusion::arrow::array::{ArrayRef, Float64Array, Float64Builder};
55
use datafusion::arrow::datatypes::DataType;
66
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
88
use jiter::{NumberAny, Peek};
99

1010
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
@@ -49,8 +49,8 @@ impl ScalarUDFImpl for JsonGetFloat {
4949
return_type_check(arg_types, self.name(), DataType::Float64)
5050
}
5151

52-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
53-
invoke::<Float64Array>(args, jiter_json_get_float)
52+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
53+
invoke::<Float64Array>(&args.args, jiter_json_get_float)
5454
}
5555

5656
fn aliases(&self) -> &[String] {

‎src/json_get_int.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44
use datafusion::arrow::array::{ArrayRef, Int64Array, Int64Builder};
55
use datafusion::arrow::datatypes::DataType;
66
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
88
use jiter::{NumberInt, Peek};
99

1010
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
@@ -49,8 +49,8 @@ impl ScalarUDFImpl for JsonGetInt {
4949
return_type_check(arg_types, self.name(), DataType::Int64)
5050
}
5151

52-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
53-
invoke::<Int64Array>(args, jiter_json_get_int)
52+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
53+
invoke::<Int64Array>(&args.args, jiter_json_get_int)
5454
}
5555

5656
fn aliases(&self) -> &[String] {

‎src/json_get_json.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::any::Any;
33
use datafusion::arrow::array::StringArray;
44
use datafusion::arrow::datatypes::DataType;
55
use datafusion::common::Result as DataFusionResult;
6-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
6+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
77

88
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
99
use crate::common_macros::make_udf_function;
@@ -47,8 +47,8 @@ impl ScalarUDFImpl for JsonGetJson {
4747
return_type_check(arg_types, self.name(), DataType::Utf8)
4848
}
4949

50-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
51-
invoke::<StringArray>(args, jiter_json_get_json)
50+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
51+
invoke::<StringArray>(&args.args, jiter_json_get_json)
5252
}
5353

5454
fn aliases(&self) -> &[String] {

‎src/json_get_str.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::any::Any;
33
use datafusion::arrow::array::StringArray;
44
use datafusion::arrow::datatypes::DataType;
55
use datafusion::common::Result as DataFusionResult;
6-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
6+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
77
use jiter::Peek;
88

99
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, JsonPath};
@@ -48,8 +48,8 @@ impl ScalarUDFImpl for JsonGetStr {
4848
return_type_check(arg_types, self.name(), DataType::Utf8)
4949
}
5050

51-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
52-
invoke::<StringArray>(args, jiter_json_get_str)
51+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
52+
invoke::<StringArray>(&args.args, jiter_json_get_str)
5353
}
5454

5555
fn aliases(&self) -> &[String] {

‎src/json_length.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44
use datafusion::arrow::array::{ArrayRef, UInt64Array, UInt64Builder};
55
use datafusion::arrow::datatypes::DataType;
66
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
88
use jiter::Peek;
99

1010
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
@@ -49,8 +49,8 @@ impl ScalarUDFImpl for JsonLength {
4949
return_type_check(arg_types, self.name(), DataType::UInt64)
5050
}
5151

52-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
53-
invoke::<UInt64Array>(args, jiter_json_length)
52+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
53+
invoke::<UInt64Array>(&args.args, jiter_json_length)
5454
}
5555

5656
fn aliases(&self) -> &[String] {

‎src/json_object_keys.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use std::sync::Arc;
44
use datafusion::arrow::array::{ArrayRef, ListBuilder, StringBuilder};
55
use datafusion::arrow::datatypes::{DataType, Field};
66
use datafusion::common::{Result as DataFusionResult, ScalarValue};
7-
use datafusion::logical_expr::{ColumnarValue, ScalarUDFImpl, Signature, Volatility};
7+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility};
88
use jiter::Peek;
99

1010
use crate::common::{get_err, invoke, jiter_json_find, return_type_check, GetError, InvokeResult, JsonPath};
@@ -53,8 +53,8 @@ impl ScalarUDFImpl for JsonObjectKeys {
5353
)
5454
}
5555

56-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
57-
invoke::<BuildListArray>(args, jiter_json_object_keys)
56+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
57+
invoke::<BuildListArray>(&args.args, jiter_json_object_keys)
5858
}
5959

6060
fn aliases(&self) -> &[String] {

‎tests/main.rs

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use datafusion::arrow::datatypes::{Field, Int64Type, Int8Type, Schema};
55
use datafusion::arrow::{array::StringDictionaryBuilder, datatypes::DataType};
66
use datafusion::assert_batches_eq;
77
use datafusion::common::ScalarValue;
8-
use datafusion::logical_expr::ColumnarValue;
8+
use datafusion::logical_expr::{ColumnarValue, ScalarFunctionArgs};
99
use datafusion::prelude::SessionContext;
1010
use datafusion_functions_json::udfs::json_get_str_udf;
1111
use utils::{create_context, display_val, logical_plan, run_query, run_query_dict, run_query_large, run_query_params};
@@ -500,7 +500,14 @@ fn test_json_get_utf8() {
500500
ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))),
501501
];
502502

503-
let ColumnarValue::Scalar(sv) = json_get_str.invoke_batch(args, 1).unwrap() else {
503+
let ColumnarValue::Scalar(sv) = json_get_str
504+
.invoke_with_args(ScalarFunctionArgs {
505+
args: args.to_vec(),
506+
number_rows: 1,
507+
return_type: &DataType::Utf8,
508+
})
509+
.unwrap()
510+
else {
504511
panic!("expected scalar")
505512
};
506513

@@ -518,7 +525,14 @@ fn test_json_get_large_utf8() {
518525
ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("aa".to_string()))),
519526
];
520527

521-
let ColumnarValue::Scalar(sv) = json_get_str.invoke_batch(args, 1).unwrap() else {
528+
let ColumnarValue::Scalar(sv) = json_get_str
529+
.invoke_with_args(ScalarFunctionArgs {
530+
args: args.to_vec(),
531+
number_rows: 1,
532+
return_type: &DataType::LargeUtf8,
533+
})
534+
.unwrap()
535+
else {
522536
panic!("expected scalar")
523537
};
524538

0 commit comments

Comments
 (0)