This is an automated email from the ASF dual-hosted git repository.
jiayu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new 1193bb2 Add argument matcher or (#69)
1193bb2 is described below
commit 1193bb29ad88c198508d0e9fcbfddf0fab241940
Author: jp <[email protected]>
AuthorDate: Fri Sep 12 07:24:19 2025 -0700
Add argument matcher or (#69)
---
rust/sedona-functions/src/sd_format.rs | 2 +-
rust/sedona-functions/src/st_perimeter.rs | 4 +-
rust/sedona-functions/src/st_transform.rs | 9 ++--
rust/sedona-schema/src/matchers.rs | 70 +++++++++++++++++++++++++++----
4 files changed, 71 insertions(+), 14 deletions(-)
diff --git a/rust/sedona-functions/src/sd_format.rs
b/rust/sedona-functions/src/sd_format.rs
index aa13875..b1bb33a 100644
--- a/rust/sedona-functions/src/sd_format.rs
+++ b/rust/sedona-functions/src/sd_format.rs
@@ -81,7 +81,7 @@ impl SedonaScalarKernel for SDFormatDefault {
let matcher = ArgMatcher::new(
vec![
ArgMatcher::is_any(),
- ArgMatcher::is_optional(ArgMatcher::is_string()),
+ ArgMatcher::optional(ArgMatcher::is_string()),
],
formatted_type,
);
diff --git a/rust/sedona-functions/src/st_perimeter.rs
b/rust/sedona-functions/src/st_perimeter.rs
index 993136a..d6e0208 100644
--- a/rust/sedona-functions/src/st_perimeter.rs
+++ b/rust/sedona-functions/src/st_perimeter.rs
@@ -28,8 +28,8 @@ pub fn st_perimeter_udf() -> SedonaScalarUDF {
ArgMatcher::new(
vec![
ArgMatcher::is_geometry_or_geography(),
- ArgMatcher::is_optional(ArgMatcher::is_boolean()),
- ArgMatcher::is_optional(ArgMatcher::is_boolean()),
+ ArgMatcher::optional(ArgMatcher::is_boolean()),
+ ArgMatcher::optional(ArgMatcher::is_boolean()),
],
SedonaType::Arrow(DataType::Float64),
),
diff --git a/rust/sedona-functions/src/st_transform.rs
b/rust/sedona-functions/src/st_transform.rs
index 9030fcb..d028db6 100644
--- a/rust/sedona-functions/src/st_transform.rs
+++ b/rust/sedona-functions/src/st_transform.rs
@@ -32,9 +32,12 @@ pub fn st_transform_udf() -> SedonaScalarUDF {
ArgMatcher::new(
vec![
ArgMatcher::is_geometry_or_geography(),
- ArgMatcher::is_string(),
- ArgMatcher::is_optional(ArgMatcher::is_string()),
- ArgMatcher::is_optional(ArgMatcher::is_boolean()),
+ ArgMatcher::or(vec![ArgMatcher::is_string(),
ArgMatcher::is_numeric()]),
+ ArgMatcher::optional(ArgMatcher::or(vec![
+ ArgMatcher::is_string(),
+ ArgMatcher::is_numeric(),
+ ])),
+ ArgMatcher::optional(ArgMatcher::is_boolean()),
],
SedonaType::Wkb(Edges::Planar, None),
),
diff --git a/rust/sedona-schema/src/matchers.rs
b/rust/sedona-schema/src/matchers.rs
index 0f95cc5..ba725c8 100644
--- a/rust/sedona-schema/src/matchers.rs
+++ b/rust/sedona-schema/src/matchers.rs
@@ -85,12 +85,12 @@ impl ArgMatcher {
if arg == &&SedonaType::Arrow(DataType::Null) ||
matcher.match_type(arg) {
arg_iter.next(); // Consume the argument
continue; // Move to the next matcher
- } else if matcher.is_optional() {
+ } else if matcher.optional() {
continue; // Skip the optional matcher
} else {
return false; // Non-optional matcher failed
}
- } else if matcher.is_optional() {
+ } else if matcher.optional() {
continue; // Skip remaining optional matchers
} else {
return false; // Non-optional matcher failed with no arguments
left
@@ -179,11 +179,18 @@ impl ArgMatcher {
}
/// Matches any argument that is optional
- pub fn is_optional(
+ pub fn optional(
matcher: Arc<dyn TypeMatcher + Send + Sync>,
) -> Arc<dyn TypeMatcher + Send + Sync> {
Arc::new(OptionalMatcher { inner: matcher })
}
+
+ /// Matches if any of the given matchers match
+ pub fn or(
+ matchers: Vec<Arc<dyn TypeMatcher + Send + Sync>>,
+ ) -> Arc<dyn TypeMatcher + Send + Sync> {
+ Arc::new(OrMatcher { matchers })
+ }
}
/// A TypeMatcher is a predicate on a [SedonaType]
@@ -198,7 +205,7 @@ pub trait TypeMatcher: Debug {
fn match_type(&self, arg: &SedonaType) -> bool;
/// If this argument is optional, return true
- fn is_optional(&self) -> bool {
+ fn optional(&self) -> bool {
false
}
@@ -244,7 +251,7 @@ impl TypeMatcher for OptionalMatcher {
self.inner.match_type(arg)
}
- fn is_optional(&self) -> bool {
+ fn optional(&self) -> bool {
true
}
@@ -253,6 +260,21 @@ impl TypeMatcher for OptionalMatcher {
}
}
+#[derive(Debug)]
+struct OrMatcher {
+ matchers: Vec<Arc<dyn TypeMatcher + Send + Sync>>,
+}
+
+impl TypeMatcher for OrMatcher {
+ fn match_type(&self, arg: &SedonaType) -> bool {
+ self.matchers.iter().any(|m| m.match_type(arg))
+ }
+
+ fn type_if_null(&self) -> Option<SedonaType> {
+ None
+ }
+}
+
#[derive(Debug)]
struct IsGeometryOrGeography {}
@@ -446,8 +468,8 @@ mod tests {
let matcher = ArgMatcher::new(
vec![
ArgMatcher::is_geometry(),
- ArgMatcher::is_optional(ArgMatcher::is_boolean()),
- ArgMatcher::is_optional(ArgMatcher::is_numeric()),
+ ArgMatcher::optional(ArgMatcher::is_boolean()),
+ ArgMatcher::optional(ArgMatcher::is_numeric()),
],
SedonaType::Arrow(DataType::Null),
);
@@ -486,6 +508,38 @@ mod tests {
]));
}
+ #[test]
+ fn or_matcher() {
+ let matcher = ArgMatcher::new(
+ vec![
+ ArgMatcher::is_geometry(),
+ ArgMatcher::or(vec![ArgMatcher::is_boolean(),
ArgMatcher::is_numeric()]),
+ ],
+ SedonaType::Arrow(DataType::Null),
+ );
+
+ // Matches first arg
+ assert!(matcher.matches(&[WKB_GEOMETRY,
SedonaType::Arrow(DataType::Boolean),]));
+
+ // Matches second arg
+ assert!(matcher.matches(&[WKB_GEOMETRY,
SedonaType::Arrow(DataType::Int32)]));
+
+ // No match when second arg is incorrect type
+ assert!(!matcher.matches(&[WKB_GEOMETRY, WKB_GEOMETRY]));
+
+ // No match when first arg is incorrect type
+ assert!(!matcher.matches(&[
+ SedonaType::Arrow(DataType::Boolean),
+ SedonaType::Arrow(DataType::Boolean)
+ ]));
+
+ // Return type if null
+ assert_eq!(
+ ArgMatcher::or(vec![ArgMatcher::is_boolean(),
ArgMatcher::is_numeric()]).type_if_null(),
+ None
+ );
+ }
+
#[test]
fn arg_matcher_matches_null() {
for type_matcher in [
@@ -498,7 +552,7 @@ mod tests {
ArgMatcher::is_string(),
ArgMatcher::is_binary(),
ArgMatcher::is_boolean(),
- ArgMatcher::is_optional(ArgMatcher::is_numeric()),
+ ArgMatcher::optional(ArgMatcher::is_numeric()),
] {
let matcher = ArgMatcher::new(vec![type_matcher],
SedonaType::Arrow(DataType::Null));
assert!(matcher.matches(&[SedonaType::Arrow(DataType::Null)]));
