Skip to content

Commit ed98690

Browse files
irenjjirenjj
authored
chore: Migrate Regex function to invoke_with_args (#14728)
* chore" Migrate Regex function to invoke_with_args * fix * fix issues
1 parent 45d9820 commit ed98690

File tree

4 files changed

+81
-75
lines changed

4 files changed

+81
-75
lines changed

‎datafusion/functions/src/regex/regexpcount.rs

+68-66
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,12 @@ impl ScalarUDFImpl for RegexpCountFunc {
108108
Ok(Int64)
109109
}
110110

111-
fn invoke_batch(
111+
fn invoke_with_args(
112112
&self,
113-
args: &[ColumnarValue],
114-
_number_rows: usize,
113+
args: datafusion_expr::ScalarFunctionArgs,
115114
) -> Result<ColumnarValue> {
115+
let args = &args.args;
116+
116117
let len = args
117118
.iter()
118119
.fold(Option::<usize>::None, |acc, arg| match arg {
@@ -618,6 +619,7 @@ fn count_matches(
618619
mod tests {
619620
use super::*;
620621
use arrow::array::{GenericStringArray, StringViewArray};
622+
use datafusion_expr::ScalarFunctionArgs;
621623

622624
#[test]
623625
fn test_regexp_count() {
@@ -655,11 +657,11 @@ mod tests {
655657
let v_sv = ScalarValue::Utf8(Some(v.to_string()));
656658
let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
657659
let expected = expected.get(pos).cloned();
658-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
659-
let re = RegexpCountFunc::new().invoke_batch(
660-
&[ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
661-
1,
662-
);
660+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
661+
args: vec![ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
662+
number_rows: 2,
663+
return_type: &Int64,
664+
});
663665
match re {
664666
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
665667
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -670,11 +672,11 @@ mod tests {
670672
// largeutf8
671673
let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
672674
let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
673-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
674-
let re = RegexpCountFunc::new().invoke_batch(
675-
&[ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
676-
1,
677-
);
675+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
676+
args: vec![ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
677+
number_rows: 2,
678+
return_type: &Int64,
679+
});
678680
match re {
679681
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
680682
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -685,11 +687,11 @@ mod tests {
685687
// utf8view
686688
let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
687689
let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
688-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
689-
let re = RegexpCountFunc::new().invoke_batch(
690-
&[ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
691-
1,
692-
);
690+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
691+
args: vec![ColumnarValue::Scalar(v_sv), ColumnarValue::Scalar(regex_sv)],
692+
number_rows: 2,
693+
return_type: &Int64,
694+
});
693695
match re {
694696
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
695697
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -711,15 +713,15 @@ mod tests {
711713
let regex_sv = ScalarValue::Utf8(Some(regex.to_string()));
712714
let start_sv = ScalarValue::Int64(Some(start));
713715
let expected = expected.get(pos).cloned();
714-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
715-
let re = RegexpCountFunc::new().invoke_batch(
716-
&[
716+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
717+
args: vec![
717718
ColumnarValue::Scalar(v_sv),
718719
ColumnarValue::Scalar(regex_sv),
719720
ColumnarValue::Scalar(start_sv.clone()),
720721
],
721-
1,
722-
);
722+
number_rows: 3,
723+
return_type: &Int64,
724+
});
723725
match re {
724726
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
725727
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -730,15 +732,15 @@ mod tests {
730732
// largeutf8
731733
let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
732734
let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
733-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
734-
let re = RegexpCountFunc::new().invoke_batch(
735-
&[
735+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
736+
args: vec![
736737
ColumnarValue::Scalar(v_sv),
737738
ColumnarValue::Scalar(regex_sv),
738739
ColumnarValue::Scalar(start_sv.clone()),
739740
],
740-
1,
741-
);
741+
number_rows: 3,
742+
return_type: &Int64,
743+
});
742744
match re {
743745
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
744746
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -749,15 +751,15 @@ mod tests {
749751
// utf8view
750752
let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
751753
let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
752-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
753-
let re = RegexpCountFunc::new().invoke_batch(
754-
&[
754+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
755+
args: vec![
755756
ColumnarValue::Scalar(v_sv),
756757
ColumnarValue::Scalar(regex_sv),
757-
ColumnarValue::Scalar(start_sv),
758+
ColumnarValue::Scalar(start_sv.clone()),
758759
],
759-
1,
760-
);
760+
number_rows: 3,
761+
return_type: &Int64,
762+
});
761763
match re {
762764
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
763765
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -781,16 +783,16 @@ mod tests {
781783
let start_sv = ScalarValue::Int64(Some(start));
782784
let flags_sv = ScalarValue::Utf8(Some(flags.to_string()));
783785
let expected = expected.get(pos).cloned();
784-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
785-
let re = RegexpCountFunc::new().invoke_batch(
786-
&[
786+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
787+
args: vec![
787788
ColumnarValue::Scalar(v_sv),
788789
ColumnarValue::Scalar(regex_sv),
789790
ColumnarValue::Scalar(start_sv.clone()),
790791
ColumnarValue::Scalar(flags_sv.clone()),
791792
],
792-
1,
793-
);
793+
number_rows: 4,
794+
return_type: &Int64,
795+
});
794796
match re {
795797
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
796798
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -802,16 +804,16 @@ mod tests {
802804
let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
803805
let regex_sv = ScalarValue::LargeUtf8(Some(regex.to_string()));
804806
let flags_sv = ScalarValue::LargeUtf8(Some(flags.to_string()));
805-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
806-
let re = RegexpCountFunc::new().invoke_batch(
807-
&[
807+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
808+
args: vec![
808809
ColumnarValue::Scalar(v_sv),
809810
ColumnarValue::Scalar(regex_sv),
810811
ColumnarValue::Scalar(start_sv.clone()),
811812
ColumnarValue::Scalar(flags_sv.clone()),
812813
],
813-
1,
814-
);
814+
number_rows: 4,
815+
return_type: &Int64,
816+
});
815817
match re {
816818
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
817819
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -823,16 +825,16 @@ mod tests {
823825
let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
824826
let regex_sv = ScalarValue::Utf8View(Some(regex.to_string()));
825827
let flags_sv = ScalarValue::Utf8View(Some(flags.to_string()));
826-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
827-
let re = RegexpCountFunc::new().invoke_batch(
828-
&[
828+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
829+
args: vec![
829830
ColumnarValue::Scalar(v_sv),
830831
ColumnarValue::Scalar(regex_sv),
831-
ColumnarValue::Scalar(start_sv),
832+
ColumnarValue::Scalar(start_sv.clone()),
832833
ColumnarValue::Scalar(flags_sv.clone()),
833834
],
834-
1,
835-
);
835+
number_rows: 4,
836+
return_type: &Int64,
837+
});
836838
match re {
837839
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
838840
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -905,16 +907,16 @@ mod tests {
905907
let start_sv = ScalarValue::Int64(Some(start));
906908
let flags_sv = ScalarValue::Utf8(flags.get(pos).map(|f| f.to_string()));
907909
let expected = expected.get(pos).cloned();
908-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
909-
let re = RegexpCountFunc::new().invoke_batch(
910-
&[
910+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
911+
args: vec![
911912
ColumnarValue::Scalar(v_sv),
912913
ColumnarValue::Scalar(regex_sv),
913914
ColumnarValue::Scalar(start_sv.clone()),
914915
ColumnarValue::Scalar(flags_sv.clone()),
915916
],
916-
1,
917-
);
917+
number_rows: 4,
918+
return_type: &Int64,
919+
});
918920
match re {
919921
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
920922
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -926,16 +928,16 @@ mod tests {
926928
let v_sv = ScalarValue::LargeUtf8(Some(v.to_string()));
927929
let regex_sv = ScalarValue::LargeUtf8(regex.get(pos).map(|s| s.to_string()));
928930
let flags_sv = ScalarValue::LargeUtf8(flags.get(pos).map(|f| f.to_string()));
929-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
930-
let re = RegexpCountFunc::new().invoke_batch(
931-
&[
931+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
932+
args: vec![
932933
ColumnarValue::Scalar(v_sv),
933934
ColumnarValue::Scalar(regex_sv),
934935
ColumnarValue::Scalar(start_sv.clone()),
935936
ColumnarValue::Scalar(flags_sv.clone()),
936937
],
937-
1,
938-
);
938+
number_rows: 4,
939+
return_type: &Int64,
940+
});
939941
match re {
940942
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
941943
assert_eq!(v, expected, "regexp_count scalar test failed");
@@ -947,16 +949,16 @@ mod tests {
947949
let v_sv = ScalarValue::Utf8View(Some(v.to_string()));
948950
let regex_sv = ScalarValue::Utf8View(regex.get(pos).map(|s| s.to_string()));
949951
let flags_sv = ScalarValue::Utf8View(flags.get(pos).map(|f| f.to_string()));
950-
#[allow(deprecated)] // TODO: migrate to invoke_with_args
951-
let re = RegexpCountFunc::new().invoke_batch(
952-
&[
952+
let re = RegexpCountFunc::new().invoke_with_args(ScalarFunctionArgs {
953+
args: vec![
953954
ColumnarValue::Scalar(v_sv),
954955
ColumnarValue::Scalar(regex_sv),
955-
ColumnarValue::Scalar(start_sv),
956+
ColumnarValue::Scalar(start_sv.clone()),
956957
ColumnarValue::Scalar(flags_sv.clone()),
957958
],
958-
1,
959-
);
959+
number_rows: 4,
960+
return_type: &Int64,
961+
});
960962
match re {
961963
Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => {
962964
assert_eq!(v, expected, "regexp_count scalar test failed");

‎datafusion/functions/src/regex/regexplike.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -110,11 +110,12 @@ impl ScalarUDFImpl for RegexpLikeFunc {
110110
})
111111
}
112112

113-
fn invoke_batch(
113+
fn invoke_with_args(
114114
&self,
115-
args: &[ColumnarValue],
116-
_number_rows: usize,
115+
args: datafusion_expr::ScalarFunctionArgs,
117116
) -> Result<ColumnarValue> {
117+
let args = &args.args;
118+
118119
let len = args
119120
.iter()
120121
.fold(Option::<usize>::None, |acc, arg| match arg {

‎datafusion/functions/src/regex/regexpmatch.rs

+4-3
Original file line numberDiff line numberDiff line change
@@ -118,11 +118,12 @@ impl ScalarUDFImpl for RegexpMatchFunc {
118118
other => DataType::List(Arc::new(Field::new_list_field(other.clone(), true))),
119119
})
120120
}
121-
fn invoke_batch(
121+
122+
fn invoke_with_args(
122123
&self,
123-
args: &[ColumnarValue],
124-
_number_rows: usize,
124+
args: datafusion_expr::ScalarFunctionArgs,
125125
) -> Result<ColumnarValue> {
126+
let args = &args.args;
126127
let len = args
127128
.iter()
128129
.fold(Option::<usize>::None, |acc, arg| match arg {

‎datafusion/functions/src/regex/regexpreplace.rs

+5-3
Original file line numberDiff line numberDiff line change
@@ -147,11 +147,13 @@ impl ScalarUDFImpl for RegexpReplaceFunc {
147147
}
148148
})
149149
}
150-
fn invoke_batch(
150+
151+
fn invoke_with_args(
151152
&self,
152-
args: &[ColumnarValue],
153-
_number_rows: usize,
153+
args: datafusion_expr::ScalarFunctionArgs,
154154
) -> Result<ColumnarValue> {
155+
let args = &args.args;
156+
155157
let len = args
156158
.iter()
157159
.fold(Option::<usize>::None, |acc, arg| match arg {

0 commit comments

Comments
 (0)