Skip to content

Commit 6009223

Browse files
adriangbadriangb
committed
Merge branch 'update-methods' into sorted-json
2 parents 1dc59bd + a739fd7 commit 6009223

File tree

12 files changed

+95
-40
lines changed

12 files changed

+95
-40
lines changed

‎benches/main.rs‎

Lines changed: 48 additions & 7 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_top_level_sorted_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 make_json_negative_testcase() -> String {
@@ -46,7 +63,15 @@ fn bench_json_get_str_negative(b: &mut Bencher) {
4663
ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), // lexicographically less than "b1"
4764
];
4865

49-
b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
66+
b.iter(|| {
67+
json_get_str
68+
.invoke_with_args(ScalarFunctionArgs {
69+
args: args.to_vec(),
70+
number_rows: 1,
71+
return_type: &DataType::Utf8,
72+
})
73+
.unwrap()
74+
});
5075
}
5176

5277
fn bench_json_get_str_sorted(b: &mut Bencher) {
@@ -59,7 +84,15 @@ fn bench_json_get_str_sorted(b: &mut Bencher) {
5984
ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))),
6085
];
6186

62-
b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
87+
b.iter(|| {
88+
json_get_str
89+
.invoke_with_args(ScalarFunctionArgs {
90+
args: args.to_vec(),
91+
number_rows: 1,
92+
return_type: &DataType::Utf8,
93+
})
94+
.unwrap()
95+
});
6396
}
6497

6598
fn bench_json_get_str_sorted_negative(b: &mut Bencher) {
@@ -69,7 +102,15 @@ fn bench_json_get_str_sorted_negative(b: &mut Bencher) {
69102
ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), // lexicographically less than "b1"
70103
];
71104

72-
b.iter(|| json_get_str.invoke_batch(args, 1).unwrap());
105+
b.iter(|| {
106+
json_get_str
107+
.invoke_with_args(ScalarFunctionArgs {
108+
args: args.to_vec(),
109+
number_rows: 1,
110+
return_type: &DataType::Utf8,
111+
})
112+
.unwrap()
113+
});
73114
}
74115

75116
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::{
@@ -70,8 +70,8 @@ impl ScalarUDFImpl for JsonAsText {
7070
return_type_check(arg_types, self.name(), DataType::Utf8)
7171
}
7272

73-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
74-
invoke::<StringArray>(args, |json, path| jiter_json_as_text(json, path, self.sorted))
73+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
74+
invoke::<StringArray>(&args.args, |args, path| jiter_json_as_text(args, path, self.sorted))
7575
}
7676

7777
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;
@@ -72,8 +72,8 @@ impl ScalarUDFImpl for JsonContains {
7272
}
7373
}
7474

75-
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
76-
invoke::<BooleanArray>(args, |json, path| jiter_json_contains(json, path, self.sorted))
75+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
76+
invoke::<BooleanArray>(&args.args, |json, path| jiter_json_contains(json, path, self.sorted))
7777
}
7878

7979
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

@@ -73,8 +73,8 @@ impl ScalarUDFImpl for JsonGet {
7373
return_type_check(arg_types, self.name(), JsonUnion::data_type())
7474
}
7575

76-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
77-
invoke::<JsonUnion>(args, |json, path| jiter_json_get_union(json, path, self.sorted))
76+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
77+
invoke::<JsonUnion>(&args.args, |json, path| jiter_json_get_union(json, path, self.sorted))
7878
}
7979

8080
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, Sortedness};
@@ -67,8 +67,8 @@ impl ScalarUDFImpl for JsonGetBool {
6767
return_type_check(arg_types, self.name(), DataType::Boolean).map(|_| DataType::Boolean)
6868
}
6969

70-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
71-
invoke::<BooleanArray>(args, |json, path| jiter_json_get_bool(json, path, self.sorted))
70+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
71+
invoke::<BooleanArray>(&args.args, |json, path| jiter_json_get_bool(json, path, self.sorted))
7272
}
7373

7474
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::{
@@ -70,8 +70,8 @@ impl ScalarUDFImpl for JsonGetFloat {
7070
return_type_check(arg_types, self.name(), DataType::Float64)
7171
}
7272

73-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
74-
invoke::<Float64Array>(args, |json_data, path| {
73+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
74+
invoke::<Float64Array>(&args.args, |json_data, path| {
7575
jiter_json_get_float(json_data, path, self.sorted)
7676
})
7777
}

‎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::{
@@ -70,8 +70,8 @@ impl ScalarUDFImpl for JsonGetInt {
7070
return_type_check(arg_types, self.name(), DataType::Int64)
7171
}
7272

73-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
74-
invoke::<Int64Array>(args, |json, path| jiter_json_get_int(json, path, self.sorted))
73+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
74+
invoke::<Int64Array>(&args.args, |json, path| jiter_json_get_int(json, path, self.sorted))
7575
}
7676

7777
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, Sortedness};
99
use crate::common_macros::make_udf_function;
@@ -66,8 +66,8 @@ impl ScalarUDFImpl for JsonGetJson {
6666
return_type_check(arg_types, self.name(), DataType::Utf8)
6767
}
6868

69-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
70-
invoke::<StringArray>(args, |json, path| jiter_json_get_json(json, path, self.sorted))
69+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
70+
invoke::<StringArray>(&args.args, |json, path| jiter_json_get_json(json, path, self.sorted))
7171
}
7272

7373
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, Sortedness};
@@ -67,8 +67,8 @@ impl ScalarUDFImpl for JsonGetStr {
6767
return_type_check(arg_types, self.name(), DataType::Utf8)
6868
}
6969

70-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
71-
invoke::<StringArray>(args, |json, path| jiter_json_get_str(json, path, self.sorted))
70+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
71+
invoke::<StringArray>(&args.args, |json, path| jiter_json_get_str(json, path, self.sorted))
7272
}
7373

7474
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::{
@@ -73,8 +73,8 @@ impl ScalarUDFImpl for JsonLength {
7373
return_type_check(arg_types, self.name(), DataType::UInt64)
7474
}
7575

76-
fn invoke(&self, args: &[ColumnarValue]) -> DataFusionResult<ColumnarValue> {
77-
invoke::<UInt64Array>(args, |json, path| jiter_json_length(json, path, self.sorted))
76+
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> {
77+
invoke::<UInt64Array>(&args.args, |json, path| jiter_json_length(json, path, self.sorted))
7878
}
7979

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

0 commit comments

Comments
 (0)